Implemented webserver

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)
+        +add_state_change_handler(handler Callable) void
-        +get_current_state() WebAgentState
+        +add_response_change_handler(handler Callable) void
-        +proceed() void
+        +approve_context() void
-        +add_state_change_handler(handler func) 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)
+        +WebServer(agent WebAgent, io_buffer WebIOBuffer, host str, port int)
-        -broadcast_state_change() void
+        +start() void
+        +stop() void
+        +broadcast(message_type str, data str, **kwargs) void
+    }
+
+    class ClientMessage {
+        <<enumeration>>
+        APPROVE_CONTEXT
+        APPROVE_RESPONSE
+        MODIFY_RESPONSE
+        SEND_INPUT
+    }
+
+    class ServerMessage {
+        <<enumeration>>
+        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

requirements.txt

@@ -1,4 +1,5 @@
 accelerate
+aiohttp
 bs4
 torch
 transformers

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
+        )

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 = "<reasoning>test</reasoning>"
+        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)

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 = "<test>context</test>"
+            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"], "<reasoning>test</reasoning>")
+            
+            # 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("<invalid>")
+            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)