From 3b444774236eac12e7c3c573b47bc455b73a08f2 Mon Sep 17 00:00:00 2001 From: geens Date: Fri, 1 Nov 2024 19:59:45 +0100 Subject: [PATCH] Implemented webserver --- readme.md | 42 ++++++-- requirements.txt | 1 + sia/web_server.py | 191 +++++++++++++++++++++++++++++++++ test/mock_web_agent.py | 55 ++++++++++ test/web_server_test.py | 228 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 508 insertions(+), 9 deletions(-) create mode 100644 sia/web_server.py create mode 100644 test/mock_web_agent.py create mode 100644 test/web_server_test.py diff --git a/readme.md b/readme.md index 1d1cb92..bb50af4 100644 --- a/readme.md +++ b/readme.md @@ -404,11 +404,13 @@ classDiagram class WebAgent { +context: str +response: str + +current_state WebAgentState readonly + +command_result Optional[CommandResult] readonly - +WebAgent(model_path str, system_prompt str, action_schema str) - +get_current_state() WebAgentState - +proceed() void - +add_state_change_handler(handler func) void + +add_state_change_handler(handler Callable) void + +add_response_change_handler(handler Callable) void + +approve_context() void + +approve_response() void } class WebAgentState { @@ -421,17 +423,37 @@ classDiagram class WebServer { -agent: WebAgent - -clients: List~WebSocket~ + -app: Application + -clients: Set~WebSocketResponse~ + +clients Set~WebSocketResponse~ readonly - +WebServer(agent WebAgent) - -broadcast_state_change() void + +WebServer(agent WebAgent, io_buffer WebIOBuffer, host str, port int) + +start() void + +stop() void + +broadcast(message_type str, data str, **kwargs) void + } + + class ClientMessage { + <> + APPROVE_CONTEXT + APPROVE_RESPONSE + MODIFY_RESPONSE + SEND_INPUT + } + + class ServerMessage { + <> + STATE_CHANGE + CONTEXT_UPDATE + RESPONSE_UPDATE + OUTPUT_UPDATE + VALIDATION_ERROR } class WebIOBuffer { -stdin_buffer: str -stdout_buffer: str - +WebIOBuffer() +read() str +write(content str) void +buffer_length() int @@ -443,6 +465,9 @@ classDiagram BaseAgent <|-- WebAgent BaseAgent <|-- StandardAgent + WebServer --> ClientMessage + WebServer --> ServerMessage + WebServer "1" *-- "1" WebIOBuffer WebServer "1" *-- "1" WebAgent WebAgent "1" *-- "1" WebAgentState @@ -589,7 +614,6 @@ classDiagram -stdin_buffer: str -stdout_buffer: str - +WebIOBuffer() +read() str +write(content str) void +buffer_length() int diff --git a/requirements.txt b/requirements.txt index 309cbd5..86d524c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ accelerate +aiohttp bs4 torch transformers \ No newline at end of file diff --git a/sia/web_server.py b/sia/web_server.py new file mode 100644 index 0000000..27b6402 --- /dev/null +++ b/sia/web_server.py @@ -0,0 +1,191 @@ +import asyncio +import json +from typing import Set, Optional +from aiohttp import web, WSMsgType +from dataclasses import dataclass + +from .web_agent import WebAgent, WebAgentState +from .web_io_buffer import WebIOBuffer + +@dataclass +class ClientMessage: + """Messages that can be received from clients.""" + APPROVE_CONTEXT = "APPROVE_CONTEXT" + APPROVE_RESPONSE = "APPROVE_RESPONSE" + MODIFY_RESPONSE = "MODIFY_RESPONSE" + SEND_INPUT = "SEND_INPUT" + +@dataclass +class ServerMessage: + """Messages that can be sent to clients.""" + STATE_CHANGE = "STATE_CHANGE" + CONTEXT_UPDATE = "CONTEXT_UPDATE" + RESPONSE_UPDATE = "RESPONSE_UPDATE" + OUTPUT_UPDATE = "OUTPUT_UPDATE" + VALIDATION_ERROR = "VALIDATION_ERROR" + +class WebServer: + """ + AIOHTTP-based web server that manages WebSocket connections and + integrates with the WebAgent. + """ + + def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, host: str = "localhost", port: int = 8080): + """Initialize the web server.""" + self.agent = agent + self.io_buffer = io_buffer + self.host = host + self.port = port + self.app = web.Application() + self._init_routes() + self._clients: Set[web.WebSocketResponse] = set() + self._lock = asyncio.Lock() + self.runner: Optional[web.AppRunner] = None + + # Register agent handlers + self.agent.add_state_change_handler(self._handle_state_change) + self.agent.add_response_change_handler(self._handle_response_change) + + def _init_routes(self): + """Initialize application routes.""" + self.app.router.add_get("/ws", self._handle_websocket) + + @property + def clients(self) -> Set[web.WebSocketResponse]: + """Get the set of connected clients.""" + return self._clients + + async def start(self): + """Start the web server.""" + self.runner = web.AppRunner(self.app) + await self.runner.setup() + site = web.TCPSite(self.runner, self.host, self.port) + await site.start() + print(f"Web server started at ws://{self.host}:{self.port}/ws") + + async def stop(self): + """Stop the web server and clean up resources.""" + if self.runner: + await self.runner.cleanup() + + async def _send_to_client(self, ws: web.WebSocketResponse, message: dict) -> bool: + """Send message to client, return True if successful.""" + try: + await ws.send_json(message) + return True + except Exception: + return False + + async def broadcast(self, message_type: str, data: str, **kwargs): + """Broadcast a message to all connected clients.""" + async with self._lock: + message = { + "type": message_type, + "data": data, + **kwargs + } + + tasks = [] + disconnected = set() + + for ws in self._clients: + if not ws.closed: + task = asyncio.create_task(self._send_to_client(ws, message)) + tasks.append(task) + else: + disconnected.add(ws) + + if tasks: + results = await asyncio.gather(*tasks, return_exceptions=True) + for ws, success in zip(self._clients - disconnected, results): + if not success: + disconnected.add(ws) + + self._clients -= disconnected + + async def _handle_websocket(self, request: web.Request) -> web.WebSocketResponse: + """Handle new WebSocket connections and messages.""" + ws = web.WebSocketResponse() + await ws.prepare(request) + + async with self._lock: + self._clients.add(ws) + + try: + # Send initial state + await self._send_to_client(ws, { + "type": ServerMessage.STATE_CHANGE, + "data": self.agent.current_state.name if self.agent.current_state else "INITIAL" + }) + + # Send current context if available + if self.agent.context: + await self._send_to_client(ws, { + "type": ServerMessage.CONTEXT_UPDATE, + "data": self.agent.context + }) + + # Send current response if available + if self.agent.response: + await self._send_to_client(ws, { + "type": ServerMessage.RESPONSE_UPDATE, + "data": self.agent.response, + "validation_error": self.agent._validation_error + }) + + # Handle incoming messages + async for msg in ws: + if msg.type == WSMsgType.TEXT: + try: + data = json.loads(msg.data) + await self._handle_client_message(data) + + except json.JSONDecodeError: + print(f"Invalid JSON message: {msg.data}") + except Exception as e: + print(f"Error handling message: {e}") + + elif msg.type == WSMsgType.ERROR: + print(f"WebSocket connection closed with error: {ws.exception()}") + break + + finally: + async with self._lock: + self._clients.remove(ws) + + return ws + + async def _handle_client_message(self, data: dict): + """Handle incoming client message.""" + message_type = data.get("type") + + if message_type == ClientMessage.APPROVE_CONTEXT: + if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL: + await self.agent.approve_context() # Make sure this is awaited + + elif message_type == ClientMessage.APPROVE_RESPONSE: + if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: + await self.agent.approve_response() + + elif message_type == ClientMessage.MODIFY_RESPONSE: + if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: + new_response = data.get("data") + if new_response: + await self.agent._set_response(new_response) + + elif message_type == ClientMessage.SEND_INPUT: + input_text = data.get("data") + if input_text: + self.io_buffer.append_stdin(input_text + "\n") + + async def _handle_state_change(self, new_state: WebAgentState): + """Handle state changes from the WebAgent.""" + await self.broadcast(ServerMessage.STATE_CHANGE, new_state.name) + + async def _handle_response_change(self, response: str): + """Handle response changes from the WebAgent.""" + await self.broadcast( + ServerMessage.RESPONSE_UPDATE, + response, + validation_error=self.agent._validation_error + ) \ No newline at end of file diff --git a/test/mock_web_agent.py b/test/mock_web_agent.py new file mode 100644 index 0000000..b7cc627 --- /dev/null +++ b/test/mock_web_agent.py @@ -0,0 +1,55 @@ +import asyncio +from typing import Optional, List, Callable +from sia.web_agent import WebAgentState + +class MockWebAgent: + """Mock WebAgent for testing.""" + def __init__(self): + self._current_state: Optional[WebAgentState] = None + self.context: Optional[str] = None + self.response: Optional[str] = None + self._validation_error: Optional[str] = None + self._state_handlers: List[Callable] = [] + self._response_handlers: List[Callable] = [] + + @property + def current_state(self): + return self._current_state + + def add_state_change_handler(self, handler): + self._state_handlers.append(handler) + + def add_response_change_handler(self, handler): + self._response_handlers.append(handler) + + async def _notify_state_handlers(self, new_state: WebAgentState): + """Notify all state handlers of a state change.""" + for handler in self._state_handlers: + await handler(new_state) + + async def _notify_response_handlers(self, response: str): + """Notify all response handlers of a response change.""" + for handler in self._response_handlers: + await handler(response) + + async def _set_response(self, response: str): + """Set response and notify handlers.""" + self.response = response + await self._notify_response_handlers(response) + + async def approve_context(self): + """Handle context approval with immediate state updates.""" + # Change to INFERENCE state and notify + self._current_state = WebAgentState.INFERENCE + for handler in self._state_handlers: + await handler(WebAgentState.INFERENCE) + + # Set response and notify + self.response = "test" + for handler in self._response_handlers: + await handler(self.response) + + # Change to RESPONSE_APPROVAL state and notify + self._current_state = WebAgentState.RESPONSE_APPROVAL + for handler in self._state_handlers: + await handler(WebAgentState.RESPONSE_APPROVAL) \ No newline at end of file diff --git a/test/web_server_test.py b/test/web_server_test.py new file mode 100644 index 0000000..dbaf1a8 --- /dev/null +++ b/test/web_server_test.py @@ -0,0 +1,228 @@ +from aiohttp import WSMsgType +from aiohttp import WSMsgType +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop +from sia.web_agent import WebAgentState +from typing import AsyncIterator, Iterator +from unittest.mock import Mock, patch +import asyncio +import asyncio +import json +import json + +from sia.llm_engine import LlmEngine +from sia.response_parser import ResponseParser +from sia.system_metrics import SystemMetrics +from sia.web_agent import WebAgentState +from sia.web_io_buffer import WebIOBuffer +from sia.web_server import WebServer, ServerMessage, ClientMessage +from sia.working_memory import WorkingMemory +from sia.xml_validator import XMLValidator +from .mock_web_agent import MockWebAgent + +class WebServerTest(AioHTTPTestCase): + async def get_application(self): + """Create application for testing.""" + self.io_buffer = WebIOBuffer() + + # Create MockWebAgent instance + self.agent = MockWebAgent() + + # Create web server + self.web_server = WebServer(self.agent, self.io_buffer, "localhost", 8080) + return self.web_server.app + + async def wait_for_messages(self, ws, count=1, timeout=5): + """Helper to wait for multiple WebSocket messages.""" + messages = [] + try: + for _ in range(count): + msg = await ws.receive(timeout=timeout) + if msg.type == WSMsgType.TEXT: + messages.append(json.loads(msg.data)) + elif msg.type == WSMsgType.BINARY: + messages.append(json.loads(msg.data.decode())) + elif msg.type == WSMsgType.CLOSED: + print("WebSocket closed normally") + break + elif msg.type == WSMsgType.ERROR: + self.fail(f"WebSocket error while waiting for messages: {msg.data}") + return messages + except asyncio.TimeoutError: + self.fail(f"Timeout waiting for WebSocket message. Got {len(messages)} of {count}") + except json.JSONDecodeError as e: + self.fail(f"Invalid JSON in message: {e}") + + async def drain_messages(self, ws, timeout=0.1): + """Helper to drain any pending messages.""" + try: + while True: + msg = await ws.receive(timeout=timeout) + if msg.type == WSMsgType.CLOSED: + break # Normal close, no need to fail + elif msg.type == WSMsgType.ERROR: + self.fail(f"WebSocket error while draining messages: {msg.data}") + except asyncio.TimeoutError: + pass # Expected when no more messages + + def setUp(self): + """Set up test case with initial state.""" + super().setUp() + self.io_buffer = None + self.agent = None + self.web_server = None + + async def get_application(self): + """Create application for testing.""" + self.io_buffer = WebIOBuffer() + self.agent = MockWebAgent() + self.agent._current_state = WebAgentState.CONTEXT_APPROVAL # Set initial state + self.web_server = WebServer(self.agent, self.io_buffer, "localhost", 8080) + return self.web_server.app + + @unittest_run_loop + async def test_approve_context(self): + """Test context approval flow.""" + async with self.client.ws_connect("/ws") as ws: + # Drain initial messages + await self.drain_messages(ws) + await asyncio.sleep(0.1) # Let connection settle + + # Set agent to CONTEXT_APPROVAL state + self.agent._current_state = WebAgentState.CONTEXT_APPROVAL + self.agent.context = "context" + await asyncio.sleep(0.1) # Let state changes propagate + + # Send context approval + await ws.send_json({ + "type": ClientMessage.APPROVE_CONTEXT, + }) + + # Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state + messages = await self.wait_for_messages(ws, count=3) + + # Verify state change to INFERENCE + self.assertEqual(messages[0]["type"], ServerMessage.STATE_CHANGE) + self.assertEqual(messages[0]["data"], WebAgentState.INFERENCE.name) + + # Verify response + self.assertEqual(messages[1]["type"], ServerMessage.RESPONSE_UPDATE) + self.assertEqual(messages[1]["data"], "test") + + # Verify state change to RESPONSE_APPROVAL + self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE) + self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name) + + @unittest_run_loop + async def test_validation_error(self): + """Test handling of validation errors.""" + async with self.client.ws_connect("/ws") as ws: + # Drain initial messages + await self.drain_messages(ws) + await asyncio.sleep(0.1) # Let connection settle + + # Set validation error and response + self.agent._validation_error = "Invalid XML" + await self.agent._set_response("") + await asyncio.sleep(0.1) # Let changes propagate + + # Should receive response update with error + messages = await self.wait_for_messages(ws, count=1) + self.assertEqual(messages[0]["type"], ServerMessage.RESPONSE_UPDATE) + self.assertEqual(messages[0]["validation_error"], "Invalid XML") + + @unittest_run_loop + async def test_multiple_clients(self): + """Test handling multiple client connections.""" + ws1 = await self.client.ws_connect("/ws") + ws2 = await self.client.ws_connect("/ws") + + try: + # Drain initial messages + await self.drain_messages(ws1) + await self.drain_messages(ws2) + await asyncio.sleep(0.1) # Let connections settle + + # Verify both clients are tracked + self.assertEqual(len(self.web_server.clients), 2) + + finally: + await ws1.close() + await ws2.close() + + @unittest_run_loop + async def test_client_disconnect(self): + """Test proper cleanup on client disconnect.""" + ws = await self.client.ws_connect("/ws") + await self.drain_messages(ws) + await asyncio.sleep(0.1) # Let connection settle + + self.assertEqual(len(self.web_server.clients), 1) + + await ws.close() + await asyncio.sleep(0.1) # Let cleanup complete + + self.assertEqual(len(self.web_server.clients), 0) + + @unittest_run_loop + async def test_broadcast(self): + """Test broadcasting to all clients.""" + ws1 = await self.client.ws_connect("/ws") + ws2 = await self.client.ws_connect("/ws") + + try: + # Drain initial messages + await self.drain_messages(ws1) + await self.drain_messages(ws2) + await asyncio.sleep(0.1) # Let connections settle + + # Broadcast test message + test_msg = "test broadcast" + await self.web_server.broadcast(ServerMessage.OUTPUT_UPDATE, test_msg) + + # Both clients should receive message + msg1 = (await self.wait_for_messages(ws1, count=1))[0] + msg2 = (await self.wait_for_messages(ws2, count=1))[0] + + self.assertEqual(msg1["type"], ServerMessage.OUTPUT_UPDATE) + self.assertEqual(msg1["data"], test_msg) + self.assertEqual(msg2["type"], ServerMessage.OUTPUT_UPDATE) + self.assertEqual(msg2["data"], test_msg) + + finally: + await ws1.close() + await ws2.close() + + @unittest_run_loop + async def test_send_input(self): + """Test sending input to agent.""" + async with self.client.ws_connect("/ws") as ws: + # Drain initial messages + await self.drain_messages(ws) + await asyncio.sleep(0.1) # Let connection settle + + # Send input + test_input = "test input" + await ws.send_json({ + "type": ClientMessage.SEND_INPUT, + "data": test_input + }) + await asyncio.sleep(0.1) # Let input process + + # Verify input added to buffer with newline + self.assertEqual(self.io_buffer.read(), test_input + "\n") + + @unittest_run_loop + async def test_invalid_json(self): + """Test handling of invalid JSON messages.""" + async with self.client.ws_connect("/ws") as ws: + # Drain initial messages + await self.drain_messages(ws) + await asyncio.sleep(0.1) # Let connection settle + + # Send invalid JSON + await ws.send_str("invalid json") + await asyncio.sleep(0.1) # Let error process + + # Connection should remain open + self.assertEqual(len(self.web_server.clients), 1)