Basic web loop works
This commit is contained in:
@@ -34,8 +34,10 @@ class TestLLM:
|
|||||||
yield "</reasoning>"
|
yield "</reasoning>"
|
||||||
|
|
||||||
class Main:
|
class Main:
|
||||||
def __init__(self):
|
@classmethod
|
||||||
self._config = Config()
|
async def create(cls, config: Config):
|
||||||
|
self = cls()
|
||||||
|
self._config = config
|
||||||
|
|
||||||
self._system_prompt = self._config.system_prompt.read_text()
|
self._system_prompt = self._config.system_prompt.read_text()
|
||||||
self._action_schema = self._config.action_schema.read_text()
|
self._action_schema = self._config.action_schema.read_text()
|
||||||
@@ -62,13 +64,14 @@ class Main:
|
|||||||
validator=XMLValidator(self._action_schema),
|
validator=XMLValidator(self._action_schema),
|
||||||
parser=ResponseParser(self._io_buffer)
|
parser=ResponseParser(self._io_buffer)
|
||||||
)
|
)
|
||||||
self._ws_manager = WebSocketManager(
|
self._ws_manager = await WebSocketManager.create(
|
||||||
self._agent,
|
self._agent,
|
||||||
self._io_buffer
|
self._io_buffer
|
||||||
)
|
)
|
||||||
|
|
||||||
self._app = web.Application()
|
self._app = web.Application()
|
||||||
self._init_routes()
|
self._init_routes()
|
||||||
|
return self
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def app(self):
|
def app(self):
|
||||||
@@ -114,6 +117,8 @@ class Main:
|
|||||||
return response
|
return response
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main = Main()
|
loop = asyncio.new_event_loop()
|
||||||
print("Web server started at http://localhost:8080")
|
config = Config()
|
||||||
web.run_app(main.app, port=8080)
|
main = loop.run_until_complete(Main.create(config))
|
||||||
|
print(f"Web server started at http://{config.host}:{config.port}")
|
||||||
|
web.run_app(main.app, loop=loop, host=config.host, port=config.port)
|
||||||
@@ -42,8 +42,8 @@ class Config:
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
'--host',
|
'--host',
|
||||||
type=str,
|
type=str,
|
||||||
default=os.getenv('SIA_SERVER_HOST', 'localhost'),
|
default=os.getenv('SIA_SERVER_HOST', '0.0.0.0'),
|
||||||
help='Web server host (default: localhost, env: SIA_SERVER_HOST)'
|
help='Web server host (default: 0.0.0.0, env: SIA_SERVER_HOST)'
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
'--port',
|
'--port',
|
||||||
|
|||||||
137
sia/web_agent.py
137
sia/web_agent.py
@@ -1,4 +1,5 @@
|
|||||||
from enum import Enum, auto
|
from enum import Enum, auto
|
||||||
|
from threading import Thread, Lock
|
||||||
from typing import Callable, List, Optional
|
from typing import Callable, List, Optional
|
||||||
|
|
||||||
from .base_agent import BaseAgent
|
from .base_agent import BaseAgent
|
||||||
@@ -13,17 +14,15 @@ from .xml_validator import XMLValidator
|
|||||||
class WebAgentState(Enum):
|
class WebAgentState(Enum):
|
||||||
"""
|
"""
|
||||||
States for the web agent state machine.
|
States for the web agent state machine.
|
||||||
|
|
||||||
States:
|
|
||||||
UPDATE: Updating system metrics and working memory entries
|
|
||||||
CONTEXT_APPROVAL: Waiting for human approval of context
|
|
||||||
INFERENCE: Processing context through LLM
|
|
||||||
RESPONSE_APPROVAL: Waiting for human approval of LLM response
|
|
||||||
"""
|
"""
|
||||||
UPDATE = auto()
|
UPDATE = auto()
|
||||||
|
"""Updating system metrics and working memory entries"""
|
||||||
CONTEXT_APPROVAL = auto()
|
CONTEXT_APPROVAL = auto()
|
||||||
|
"""Waiting for human approval of context"""
|
||||||
INFERENCE = auto()
|
INFERENCE = auto()
|
||||||
|
"""Processing context through LLM"""
|
||||||
RESPONSE_APPROVAL = auto()
|
RESPONSE_APPROVAL = auto()
|
||||||
|
"""Waiting for human approval of LLM response"""
|
||||||
STOPPED = auto()
|
STOPPED = auto()
|
||||||
|
|
||||||
class WebAgent(BaseAgent):
|
class WebAgent(BaseAgent):
|
||||||
@@ -48,17 +47,19 @@ class WebAgent(BaseAgent):
|
|||||||
super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
|
super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
|
||||||
self._response = ""
|
self._response = ""
|
||||||
self._system_prompt = system_prompt
|
self._system_prompt = system_prompt
|
||||||
self._current_state = WebAgentState.CONTEXT_APPROVAL
|
self._state = WebAgentState.CONTEXT_APPROVAL
|
||||||
self._validation_error: Optional[str] = None
|
self._validation_error: Optional[str] = None
|
||||||
self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
|
self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
|
||||||
|
self._context_change_handlers: List[Callable[[str], None]] = []
|
||||||
self._response_change_handlers: List[Callable[[str, str], None]] = []
|
self._response_change_handlers: List[Callable[[str, str], None]] = []
|
||||||
self._command_result: Optional[CommandResult] = None
|
self._command_result: Optional[CommandResult] = None
|
||||||
self.context = self._compile_context()
|
self._context = self._compile_context()
|
||||||
|
self._state_lock = Lock()
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def current_state(self) -> WebAgentState:
|
def state(self) -> WebAgentState:
|
||||||
"""Get the current state of the agent."""
|
"""Get the current state of the agent."""
|
||||||
return self._current_state
|
return self._state
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def command_result(self) -> Optional[CommandResult]:
|
def command_result(self) -> Optional[CommandResult]:
|
||||||
@@ -74,6 +75,11 @@ class WebAgent(BaseAgent):
|
|||||||
def response(self) -> str:
|
def response(self) -> str:
|
||||||
"""Get the current response."""
|
"""Get the current response."""
|
||||||
return self._response
|
return self._response
|
||||||
|
|
||||||
|
@property
|
||||||
|
def context(self) -> str:
|
||||||
|
"""Get the current context."""
|
||||||
|
return self._context
|
||||||
|
|
||||||
def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
|
def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
|
||||||
"""
|
"""
|
||||||
@@ -85,6 +91,16 @@ class WebAgent(BaseAgent):
|
|||||||
if handler not in self._state_change_handlers:
|
if handler not in self._state_change_handlers:
|
||||||
self._state_change_handlers.append(handler)
|
self._state_change_handlers.append(handler)
|
||||||
|
|
||||||
|
def add_context_change_handler(self, handler: Callable[[str], None]) -> None:
|
||||||
|
"""
|
||||||
|
Add a callback for context changes.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
handler: Function to call with new context
|
||||||
|
"""
|
||||||
|
if handler not in self._context_change_handlers:
|
||||||
|
self._context_change_handlers.append(handler)
|
||||||
|
|
||||||
def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
|
def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
|
||||||
"""
|
"""
|
||||||
Add a callback for response changes.
|
Add a callback for response changes.
|
||||||
@@ -94,53 +110,54 @@ class WebAgent(BaseAgent):
|
|||||||
"""
|
"""
|
||||||
if handler not in self._response_change_handlers:
|
if handler not in self._response_change_handlers:
|
||||||
self._response_change_handlers.append(handler)
|
self._response_change_handlers.append(handler)
|
||||||
|
|
||||||
def _notify_state_change(self) -> None:
|
|
||||||
"""Notify all handlers of state change."""
|
|
||||||
for handler in self._state_change_handlers:
|
|
||||||
handler(self._current_state)
|
|
||||||
|
|
||||||
def set_response(self, response: str) -> None:
|
def modify_context(self, context: str) -> None:
|
||||||
"""Set response and notify handlers."""
|
with self._state_lock:
|
||||||
self._response = response
|
if self._state != WebAgentState.CONTEXT_APPROVAL:
|
||||||
validation_error = self._validator.validate(response)
|
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
|
||||||
self._validation_error = validation_error
|
raise Exception(error_msg)
|
||||||
for handler in self._response_change_handlers:
|
Thread(target=self._set_context, args=(context,)).start()
|
||||||
handler(response, validation_error)
|
|
||||||
|
def modify_response(self, response: str) -> None:
|
||||||
def _set_state(self, new_state: WebAgentState) -> None:
|
started = False
|
||||||
"""
|
try:
|
||||||
Set the current state and notify handlers.
|
self._state_lock.acquire()
|
||||||
|
if self._state != WebAgentState.RESPONSE_APPROVAL:
|
||||||
Args:
|
error_msg = f"Not in RESPONSE_APPROVAL state (current: {self._state})"
|
||||||
new_state: New state to set
|
raise Exception(error_msg)
|
||||||
"""
|
Thread(target=self._set_response, args=(response, self._state_lock)).start()
|
||||||
self._current_state = new_state
|
started = True
|
||||||
self._notify_state_change()
|
finally:
|
||||||
|
if not started:
|
||||||
|
self._state_lock.release()
|
||||||
|
|
||||||
def approve_context(self) -> None:
|
def approve_context(self) -> None:
|
||||||
if self._current_state != WebAgentState.CONTEXT_APPROVAL:
|
with self._state_lock:
|
||||||
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})"
|
if self._state != WebAgentState.CONTEXT_APPROVAL:
|
||||||
raise Exception(error_msg)
|
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
|
||||||
|
raise Exception(error_msg)
|
||||||
|
self._set_state(WebAgentState.INFERENCE)
|
||||||
|
Thread(target=self._approve_context_thread).start()
|
||||||
|
|
||||||
|
def approve_response(self) -> None:
|
||||||
|
with self._state_lock:
|
||||||
|
if self._state != WebAgentState.RESPONSE_APPROVAL:
|
||||||
|
raise Exception("Not in RESPONSE_APPROVAL state")
|
||||||
|
self._set_state(WebAgentState.UPDATE)
|
||||||
|
Thread(target=self._approve_response_thread).start()
|
||||||
|
|
||||||
self._set_state(WebAgentState.INFERENCE)
|
def _approve_context_thread(self) -> None:
|
||||||
self.set_response("")
|
self._set_response("")
|
||||||
|
|
||||||
response_token_iter = self._llm.infer(f"{self._system_prompt}\n{self._action_schema}", self.context)
|
response_token_iter = self._llm.infer(f"{self._system_prompt}\n{self._action_schema}", self.context)
|
||||||
response = ""
|
response = ""
|
||||||
for token in response_token_iter:
|
for token in response_token_iter:
|
||||||
response += token
|
response += token
|
||||||
self.set_response(response)
|
self._set_response(response)
|
||||||
print(f"{token}")
|
print(f"{token}")
|
||||||
self._set_state(WebAgentState.RESPONSE_APPROVAL)
|
self._set_state(WebAgentState.RESPONSE_APPROVAL)
|
||||||
|
|
||||||
def approve_response(self) -> None:
|
def _approve_response_thread(self) -> None:
|
||||||
if self._current_state != WebAgentState.RESPONSE_APPROVAL:
|
parse_result = self._parser.parse(self._response)
|
||||||
raise Exception("Not in RESPONSE_APPROVAL state")
|
|
||||||
|
|
||||||
self._set_state(WebAgentState.UPDATE)
|
|
||||||
|
|
||||||
parse_result = self._parser.parse(self.response)
|
|
||||||
if isinstance(parse_result, Command):
|
if isinstance(parse_result, Command):
|
||||||
result = parse_result.execute(self._working_memory)
|
result = parse_result.execute(self._working_memory)
|
||||||
self._command_result = result
|
self._command_result = result
|
||||||
@@ -151,5 +168,29 @@ class WebAgent(BaseAgent):
|
|||||||
self._working_memory.add_entry(parse_result)
|
self._working_memory.add_entry(parse_result)
|
||||||
|
|
||||||
self._working_memory.update()
|
self._working_memory.update()
|
||||||
self.context = self._compile_context()
|
self._set_context(self._compile_context())
|
||||||
self._set_state(WebAgentState.CONTEXT_APPROVAL)
|
self._set_state(WebAgentState.CONTEXT_APPROVAL)
|
||||||
|
|
||||||
|
def _set_state(self, state) -> None:
|
||||||
|
"""Notify all handlers of state change."""
|
||||||
|
self._state = state
|
||||||
|
for handler in self._state_change_handlers:
|
||||||
|
handler(self._state)
|
||||||
|
|
||||||
|
def _set_context(self, context) -> None:
|
||||||
|
"""Notify all handlers of context change."""
|
||||||
|
self._context = context
|
||||||
|
for handler in self._context_change_handlers:
|
||||||
|
handler(context)
|
||||||
|
|
||||||
|
def _set_response(self, response: str, lock: Optional[Lock] = None) -> None:
|
||||||
|
"""Set response and notify handlers."""
|
||||||
|
try:
|
||||||
|
self._response = response
|
||||||
|
validation_error = self._validator.validate(response)
|
||||||
|
self._validation_error = validation_error
|
||||||
|
for handler in self._response_change_handlers:
|
||||||
|
handler(response, validation_error)
|
||||||
|
finally:
|
||||||
|
if lock is not None:
|
||||||
|
lock.release()
|
||||||
@@ -27,18 +27,23 @@ class WebSocketManager:
|
|||||||
"""
|
"""
|
||||||
Manages WebSocket connections and integrates with the WebAgent
|
Manages WebSocket connections and integrates with the WebAgent
|
||||||
"""
|
"""
|
||||||
def __init__(self,
|
@classmethod
|
||||||
agent: WebAgent,
|
async def create(
|
||||||
io_buffer: WebIOBuffer):
|
cls,
|
||||||
|
agent: WebAgent,
|
||||||
|
io_buffer: WebIOBuffer):
|
||||||
"""Initialize the web server."""
|
"""Initialize the web server."""
|
||||||
|
self = cls()
|
||||||
self.agent = agent
|
self.agent = agent
|
||||||
self.io_buffer = io_buffer
|
self.io_buffer = io_buffer
|
||||||
self._clients: Set[web.WebSocketResponse] = set()
|
self._clients: Set[web.WebSocketResponse] = set()
|
||||||
|
|
||||||
self.agent.add_state_change_handler(self._handle_state_change)
|
self.agent.add_state_change_handler(self._wrap_async(self._handle_state_change))
|
||||||
self.agent.add_response_change_handler(self._handle_response_change)
|
self.agent.add_context_change_handler(self._wrap_async(self._handle_context_change))
|
||||||
|
self.agent.add_response_change_handler(self._wrap_async(self._handle_response_change))
|
||||||
|
return self
|
||||||
|
|
||||||
async def broadcast_message(self, message: Dict):
|
async def _broadcast_message(self, message: Dict):
|
||||||
"""Broadcast message to all connected clients."""
|
"""Broadcast message to all connected clients."""
|
||||||
disconnected = set()
|
disconnected = set()
|
||||||
for ws in self._clients:
|
for ws in self._clients:
|
||||||
@@ -47,21 +52,63 @@ class WebSocketManager:
|
|||||||
except ConnectionResetError:
|
except ConnectionResetError:
|
||||||
disconnected.add(ws)
|
disconnected.add(ws)
|
||||||
self._clients -= disconnected
|
self._clients -= disconnected
|
||||||
|
|
||||||
|
def _wrap_async(self, coro_func):
|
||||||
|
"""
|
||||||
|
Wraps an async callback to be safely called from another thread.
|
||||||
|
|
||||||
def _handle_state_change(self, new_state: WebAgentState):
|
Args:
|
||||||
|
coro_func: The async function to wrap
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A sync function that schedules the coroutine on the event loop
|
||||||
|
"""
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
def wrapper(*args, **kwargs):
|
||||||
|
loop.call_soon_threadsafe(lambda: asyncio.create_task(coro_func(*args, **kwargs)))
|
||||||
|
return wrapper
|
||||||
|
|
||||||
|
async def _handle_state_change(self, new_state: WebAgentState):
|
||||||
"""Handle state changes from the WebAgent."""
|
"""Handle state changes from the WebAgent."""
|
||||||
asyncio.run(self.broadcast_message({
|
await self._broadcast_message({
|
||||||
"type": ServerMessage.STATE_CHANGE,
|
"type": ServerMessage.STATE_CHANGE,
|
||||||
"state": new_state.name,
|
"state": new_state.name,
|
||||||
}))
|
})
|
||||||
|
|
||||||
|
async def _handle_context_change(self, context: str):
|
||||||
|
"""Handle context changes from the WebAgent."""
|
||||||
|
await self._broadcast_message({
|
||||||
|
"type": ServerMessage.CONTEXT_UPDATE,
|
||||||
|
"context": context
|
||||||
|
})
|
||||||
|
|
||||||
def _handle_response_change(self, response: str, validation_error: str):
|
async def _handle_response_change(self, response: str, validation_error: str):
|
||||||
"""Handle response changes from the WebAgent."""
|
"""Handle response changes from the WebAgent."""
|
||||||
asyncio.run(self.broadcast_message({
|
await self._broadcast_message({
|
||||||
"type": ServerMessage.RESPONSE_UPDATE,
|
"type": ServerMessage.RESPONSE_UPDATE,
|
||||||
"response": response,
|
"response": response,
|
||||||
"validation_error": validation_error
|
"validation_error": validation_error
|
||||||
}))
|
})
|
||||||
|
|
||||||
|
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
|
||||||
|
"""Handle incoming client message."""
|
||||||
|
message_type = data.get("type")
|
||||||
|
|
||||||
|
if message_type == ClientMessage.APPROVE_CONTEXT:
|
||||||
|
if self.agent.state == WebAgentState.CONTEXT_APPROVAL:
|
||||||
|
if data.get("context") is not None:
|
||||||
|
self.agent.modify_context(data.get("context"))
|
||||||
|
self.agent.approve_context()
|
||||||
|
elif message_type == ClientMessage.APPROVE_RESPONSE:
|
||||||
|
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
||||||
|
if data.get("response") is not None:
|
||||||
|
self.agent.modify_response(data.get("response"))
|
||||||
|
self.agent.approve_response()
|
||||||
|
elif message_type == ClientMessage.MODIFY_RESPONSE:
|
||||||
|
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
||||||
|
self.agent.modify_response(data.get("response"))
|
||||||
|
elif message_type == ClientMessage.SEND_INPUT:
|
||||||
|
self.io_buffer.append_stdin(data.get("input"))
|
||||||
|
|
||||||
async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
|
async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
|
||||||
"""Handle new WebSocket connections and messages."""
|
"""Handle new WebSocket connections and messages."""
|
||||||
@@ -73,7 +120,7 @@ class WebSocketManager:
|
|||||||
try:
|
try:
|
||||||
await ws.send_json({
|
await ws.send_json({
|
||||||
"type": ServerMessage.STATE_CHANGE,
|
"type": ServerMessage.STATE_CHANGE,
|
||||||
"state": self.agent.current_state.name
|
"state": self.agent.state.name
|
||||||
})
|
})
|
||||||
await ws.send_json({
|
await ws.send_json({
|
||||||
"type": ServerMessage.CONTEXT_UPDATE,
|
"type": ServerMessage.CONTEXT_UPDATE,
|
||||||
@@ -92,20 +139,4 @@ class WebSocketManager:
|
|||||||
print(f"WebSocket connection closed with error: {ws.exception()}")
|
print(f"WebSocket connection closed with error: {ws.exception()}")
|
||||||
finally:
|
finally:
|
||||||
self._clients.remove(ws)
|
self._clients.remove(ws)
|
||||||
return ws
|
return ws
|
||||||
|
|
||||||
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
|
|
||||||
"""Handle incoming client message."""
|
|
||||||
message_type = data.get("type")
|
|
||||||
|
|
||||||
if message_type == ClientMessage.APPROVE_CONTEXT:
|
|
||||||
if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
|
|
||||||
await asyncio.to_thread(self.agent.approve_context)
|
|
||||||
elif message_type == ClientMessage.APPROVE_RESPONSE:
|
|
||||||
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
|
|
||||||
await asyncio.to_thread(self.agent.approve_response)
|
|
||||||
elif message_type == ClientMessage.MODIFY_RESPONSE:
|
|
||||||
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
|
|
||||||
await asyncio.to_thread(self.agent.set_response, data.get("response"))
|
|
||||||
elif message_type == ClientMessage.SEND_INPUT:
|
|
||||||
self.io_buffer.append_stdin(data.get("input"))
|
|
||||||
@@ -94,7 +94,7 @@ class WebAgentTest(unittest.TestCase):
|
|||||||
|
|
||||||
def test_initialization(self):
|
def test_initialization(self):
|
||||||
"""Test initial state of web agent."""
|
"""Test initial state of web agent."""
|
||||||
self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL)
|
self.assertEqual(self.agent.state, WebAgentState.CONTEXT_APPROVAL)
|
||||||
self.assertEqual(self.agent.response, "")
|
self.assertEqual(self.agent.response, "")
|
||||||
self.assertIsNotNone(self.agent.context)
|
self.assertIsNotNone(self.agent.context)
|
||||||
self.assertIsNone(self.agent.validation_error)
|
self.assertIsNone(self.agent.validation_error)
|
||||||
@@ -121,7 +121,7 @@ class WebAgentTest(unittest.TestCase):
|
|||||||
def test_approve_context_state_error(self):
|
def test_approve_context_state_error(self):
|
||||||
"""Test approving context in wrong state."""
|
"""Test approving context in wrong state."""
|
||||||
# Force agent into UPDATE state
|
# Force agent into UPDATE state
|
||||||
self.agent._current_state = WebAgentState.UPDATE
|
self.agent._state = WebAgentState.UPDATE
|
||||||
|
|
||||||
with self.assertRaises(Exception) as context:
|
with self.assertRaises(Exception) as context:
|
||||||
asyncio.run(self.agent.approve_context())
|
asyncio.run(self.agent.approve_context())
|
||||||
@@ -169,7 +169,7 @@ class WebAgentTest(unittest.TestCase):
|
|||||||
self.parser.parse = Mock(return_value=mock_command)
|
self.parser.parse = Mock(return_value=mock_command)
|
||||||
|
|
||||||
# Set initial state and response
|
# Set initial state and response
|
||||||
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
|
self.agent._state = WebAgentState.RESPONSE_APPROVAL
|
||||||
self.agent._response = "<stop/>"
|
self.agent._response = "<stop/>"
|
||||||
|
|
||||||
# Approve response
|
# Approve response
|
||||||
@@ -194,8 +194,8 @@ class WebAgentTest(unittest.TestCase):
|
|||||||
self.agent.add_state_change_handler(self.state_change_handler)
|
self.agent.add_state_change_handler(self.state_change_handler)
|
||||||
|
|
||||||
# Set initial state and response
|
# Set initial state and response
|
||||||
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
|
self.agent._state = WebAgentState.RESPONSE_APPROVAL
|
||||||
self.agent.set_response("<reasoning>test reasoning</reasoning>")
|
self.agent._set_response("<reasoning>test reasoning</reasoning>")
|
||||||
|
|
||||||
# Approve response
|
# Approve response
|
||||||
self.agent.approve_response()
|
self.agent.approve_response()
|
||||||
@@ -218,7 +218,7 @@ class WebAgentTest(unittest.TestCase):
|
|||||||
self.agent.add_response_change_handler(self.response_change_handler)
|
self.agent.add_response_change_handler(self.response_change_handler)
|
||||||
error_message = "Invalid XML"
|
error_message = "Invalid XML"
|
||||||
self.mock_validator.validate.return_value = error_message
|
self.mock_validator.validate.return_value = error_message
|
||||||
self.agent.set_response("<invalid>")
|
self.agent._set_response("<invalid>")
|
||||||
self.assertEqual(self.agent.validation_error, error_message)
|
self.assertEqual(self.agent.validation_error, error_message)
|
||||||
self.assertEqual(self.response_changes, ["<invalid>"])
|
self.assertEqual(self.response_changes, ["<invalid>"])
|
||||||
|
|
||||||
|
|||||||
@@ -61,7 +61,7 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Create WebSocketManager
|
# Create WebSocketManager
|
||||||
self.ws_manager = WebSocketManager(self.agent, self.io_buffer)
|
self.ws_manager = await WebSocketManager.create(self.agent, self.io_buffer)
|
||||||
|
|
||||||
# Create application
|
# Create application
|
||||||
app = web.Application()
|
app = web.Application()
|
||||||
|
|||||||
@@ -24,35 +24,41 @@ const App = () => {
|
|||||||
const [showSidebar, setShowSidebar] = useState(false);
|
const [showSidebar, setShowSidebar] = useState(false);
|
||||||
const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
|
const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
|
||||||
|
|
||||||
const { wsState, sendMessage, setMessageHandler } = useWebSocket('ws://localhost:8080/ws');
|
const { wsState, sendMessage, addMessageHandler } = useWebSocket('ws://localhost:8080/ws');
|
||||||
|
|
||||||
// Set up WebSocket message handlers
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
setMessageHandler(ServerMessageType.STATE_CHANGE, (message) => {
|
addMessageHandler((message) => {
|
||||||
setAgentState(message.state);
|
switch(message.type) {
|
||||||
|
case ServerMessageType.STATE_CHANGE:
|
||||||
|
setAgentState(message.state);
|
||||||
|
break;
|
||||||
|
case ServerMessageType.CONTEXT_UPDATE:
|
||||||
|
setOriginalContext(message.context);
|
||||||
|
setModifiedContext(message.context);
|
||||||
|
break;
|
||||||
|
case ServerMessageType.RESPONSE_UPDATE:
|
||||||
|
setOriginalResponse(message.response);
|
||||||
|
setModifiedResponse(message.response);
|
||||||
|
setValidationError(message.validation_error);
|
||||||
|
break;
|
||||||
|
case ServerMessageType.OUTPUT_UPDATE:
|
||||||
|
setOutput(prev => prev + message.data);
|
||||||
|
break;
|
||||||
|
}
|
||||||
});
|
});
|
||||||
|
}, [addMessageHandler]);
|
||||||
setMessageHandler(ServerMessageType.CONTEXT_UPDATE, (message) => {
|
|
||||||
setOriginalContext(message.context);
|
|
||||||
setModifiedContext(message.context);
|
|
||||||
});
|
|
||||||
|
|
||||||
setMessageHandler(ServerMessageType.RESPONSE_UPDATE, (message) => {
|
|
||||||
setOriginalResponse(message.response);
|
|
||||||
setModifiedResponse(message.response);
|
|
||||||
setValidationError(message.validation_error);
|
|
||||||
});
|
|
||||||
|
|
||||||
setMessageHandler(ServerMessageType.OUTPUT_UPDATE, (message) => {
|
|
||||||
setOutput(prev => prev + message.data);
|
|
||||||
});
|
|
||||||
}, [setMessageHandler]);
|
|
||||||
|
|
||||||
const handleApprove = () => {
|
const handleApprove = () => {
|
||||||
if (agentState === AgentState.CONTEXT_APPROVAL) {
|
if (agentState === AgentState.CONTEXT_APPROVAL) {
|
||||||
sendMessage(ClientMessageType.APPROVE_CONTEXT);
|
sendMessage({
|
||||||
|
"type": ClientMessageType.APPROVE_CONTEXT,
|
||||||
|
"context": modifiedContext
|
||||||
|
});
|
||||||
} else if (agentState === AgentState.RESPONSE_APPROVAL) {
|
} else if (agentState === AgentState.RESPONSE_APPROVAL) {
|
||||||
sendMessage(ClientMessageType.APPROVE_RESPONSE, modifiedResponse);
|
sendMessage({
|
||||||
|
"type": ClientMessageType.APPROVE_RESPONSE,
|
||||||
|
"response": modifiedResponse
|
||||||
|
});
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -63,15 +69,28 @@ const App = () => {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
if (agentState === AgentState.APPROVE_RESPONSE) {
|
||||||
|
sendMessage({
|
||||||
|
"type": ClientMessageType.MODIFY_RESPONSE,
|
||||||
|
"response": modifiedResponse
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}, [modifiedResponse]);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
switch (agentState) {
|
switch (agentState) {
|
||||||
|
case AgentState.UPDATE: // fall through
|
||||||
case AgentState.CONTEXT_APPROVAL:
|
case AgentState.CONTEXT_APPROVAL:
|
||||||
setActiveTab(Tabs.CONTEXT);
|
setActiveTab(Tabs.CONTEXT);
|
||||||
break;
|
break;
|
||||||
|
case AgentState.INFERENCE: // fall through
|
||||||
case AgentState.RESPONSE_APPROVAL:
|
case AgentState.RESPONSE_APPROVAL:
|
||||||
setActiveTab(Tabs.RESPONSE);
|
setActiveTab(Tabs.RESPONSE);
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
setShowDiff(false);
|
||||||
|
setShowSidebar(false);
|
||||||
}, [agentState]);
|
}, [agentState]);
|
||||||
|
|
||||||
return (
|
return (
|
||||||
|
|||||||
@@ -19,7 +19,6 @@ export const Tabs = {
|
|||||||
};
|
};
|
||||||
|
|
||||||
export const ClientMessageType = {
|
export const ClientMessageType = {
|
||||||
MODIFY_CONTEXT: 'MODIFY_CONTEXT',
|
|
||||||
APPROVE_CONTEXT: 'APPROVE_CONTEXT',
|
APPROVE_CONTEXT: 'APPROVE_CONTEXT',
|
||||||
MODIFY_RESPONSE: 'MODIFY_RESPONSE',
|
MODIFY_RESPONSE: 'MODIFY_RESPONSE',
|
||||||
APPROVE_RESPONSE: 'APPROVE_RESPONSE',
|
APPROVE_RESPONSE: 'APPROVE_RESPONSE',
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ import { WebSocketState, MessageType } from '../constants';
|
|||||||
export const useWebSocket = (url) => {
|
export const useWebSocket = (url) => {
|
||||||
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
|
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
|
||||||
const wsRef = useRef(null);
|
const wsRef = useRef(null);
|
||||||
const messageHandlersRef = useRef({});
|
const handlersRef = useRef([]);
|
||||||
|
|
||||||
const connect = useCallback(() => {
|
const connect = useCallback(() => {
|
||||||
setWsState(WebSocketState.CONNECTING);
|
setWsState(WebSocketState.CONNECTING);
|
||||||
@@ -27,12 +27,7 @@ export const useWebSocket = (url) => {
|
|||||||
ws.onmessage = (event) => {
|
ws.onmessage = (event) => {
|
||||||
try {
|
try {
|
||||||
const message = JSON.parse(event.data);
|
const message = JSON.parse(event.data);
|
||||||
const handler = messageHandlersRef.current[message.type];
|
handlersRef.current.forEach(handler => handler(message));
|
||||||
if (handler) {
|
|
||||||
handler(message);
|
|
||||||
} else {
|
|
||||||
console.warn('No handler for message type:', message.type);
|
|
||||||
}
|
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error('Error handling WebSocket message:', error);
|
console.error('Error handling WebSocket message:', error);
|
||||||
}
|
}
|
||||||
@@ -48,25 +43,21 @@ export const useWebSocket = (url) => {
|
|||||||
};
|
};
|
||||||
}, [connect]);
|
}, [connect]);
|
||||||
|
|
||||||
const sendMessage = useCallback((type, data = null) => {
|
const sendMessage = useCallback((data) => {
|
||||||
if (wsRef.current?.readyState === WebSocket.OPEN) {
|
if (wsRef.current?.readyState === WebSocket.OPEN) {
|
||||||
const message = {
|
wsRef.current.send(JSON.stringify(data));
|
||||||
type,
|
|
||||||
data: data || undefined
|
|
||||||
};
|
|
||||||
wsRef.current.send(JSON.stringify(message));
|
|
||||||
} else {
|
} else {
|
||||||
console.error('WebSocket not connected, state:', wsRef.current?.readyState);
|
console.error('WebSocket not connected, state:', wsRef.current?.readyState);
|
||||||
}
|
}
|
||||||
}, []);
|
}, []);
|
||||||
|
|
||||||
const setMessageHandler = useCallback((type, handler) => {
|
const addMessageHandler = useCallback((handler) => {
|
||||||
messageHandlersRef.current[type] = handler;
|
handlersRef.current.push(handler);
|
||||||
}, []);
|
}, []);
|
||||||
|
|
||||||
return {
|
return {
|
||||||
wsState,
|
wsState,
|
||||||
sendMessage,
|
sendMessage,
|
||||||
setMessageHandler
|
addMessageHandler
|
||||||
};
|
};
|
||||||
};
|
};
|
||||||
Reference in New Issue
Block a user