diff --git a/Dockerfile b/Dockerfile index dc612af..a865d42 100644 --- a/Dockerfile +++ b/Dockerfile @@ -11,7 +11,7 @@ FROM requirements AS sia-test COPY ./ /root/sia/ WORKDIR /root/sia/ RUN mkdir -p /root/model -CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py", "-v"] +CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py"] FROM node:20-alpine AS web-test WORKDIR /app diff --git a/readme.md b/readme.md index 8071e4c..252dfc2 100644 --- a/readme.md +++ b/readme.md @@ -414,34 +414,32 @@ classDiagram class WebAgent { +context: str - +response: str + +response: str readonly +current_state WebAgentState readonly +command_result Optional[CommandResult] readonly + +validation_error Optional[str] readonly +add_state_change_handler(handler Callable) void +add_response_change_handler(handler Callable) void +approve_context() void + +set__response(response str) void +approve_response() void } class WebAgentState { + <> UPDATE CONTEXT_APPROVAL INFERENCE RESPONSE_APPROVAL + STOPPED } - class WebServer { - -agent: WebAgent - -app: Application - -clients: Set~WebSocketResponse~ - +clients Set~WebSocketResponse~ readonly + class WebSocketManager { + -web_sockets: Set~WebSocket~ - +WebServer(agent WebAgent, io_buffer WebIOBuffer, host str, port int) - +start() void - +stop() void - +broadcast(message_type str, data str, **kwargs) void + +WebServer(agent WebAgent, io_buffer WebIOBuffer, static_files path, host str, port int) } class ClientMessage { diff --git a/sia/__main__.py b/sia/__main__.py index 3da918f..e8d2fc7 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -1,65 +1,105 @@ -import asyncio +from concurrent.futures import ThreadPoolExecutor +from aiohttp import web from pathlib import Path +import asyncio +import mimetypes +import time from .llm_engine import LlmEngine from .system_metrics import SystemMetrics from .web_agent import WebAgent from .web_io_buffer import WebIOBuffer -from .web_server import WebServer +from .web_socket_manager import WebSocketManager from .working_memory import WorkingMemory from .xml_validator import XMLValidator from .response_parser import ResponseParser +from .web_agent import WebAgent +from .web_io_buffer import WebIOBuffer -async def run_web_server(): - """Run the web server with default configuration.""" - base_dir = Path(__file__).parent.parent +mimetypes.add_type("application/javascript", ".js") +mimetypes.add_type("application/javascript", ".jsx") +mimetypes.add_type("text/javascript", ".js") +mimetypes.add_type("text/javascript", ".jsx") - # Load system prompt and action schema - system_prompt = (base_dir / "system_prompt.md").read_text() - action_schema = (base_dir / "action_schema.xsd").read_text() +class TestLLM: + def infer(self, prompt: str, context: str): + yield "" + time.sleep(2) + yield "test reasoning" + time.sleep(2) + yield "" - # Initialize core components - llm = LlmEngine("/root/model") - io_buffer = WebIOBuffer() - agent = WebAgent( - system_prompt=system_prompt, - action_schema=action_schema, - working_memory=WorkingMemory(), - system_metrics=SystemMetrics(), - llm=llm, - validator=XMLValidator(action_schema), - parser=ResponseParser(io_buffer) - ) +class Main: + def __init__(self): + self._base_dir = Path(__file__).parent.parent + self._system_prompt = (self._base_dir / "system_prompt.md").read_text() + self._action_schema = (self._base_dir / "action_schema.xsd").read_text() + self._static_dir = self._base_dir / "static" - # Start the web server - server = WebServer(agent, io_buffer, "0.0.0.0", 8080) - await server.start() + self._llm = LlmEngine("/root/model") + #self._llm = TestLLM() + self._io_buffer = WebIOBuffer() + self._agent = WebAgent( + system_prompt=self._system_prompt, + action_schema=self._action_schema, + working_memory=WorkingMemory(), + system_metrics=SystemMetrics(), + llm=self._llm, + validator=XMLValidator(self._action_schema), + parser=ResponseParser(self._io_buffer) + ) + self._ws_manager = WebSocketManager( + self._agent, + self._io_buffer + ) - try: - while True: - if server._response_tasks: - done, _ = await asyncio.wait( - server._response_tasks, - timeout=0.1, - return_when=asyncio.FIRST_COMPLETED - ) - for task in done: - try: - await task - except asyncio.CancelledError: - print("Task was cancelled") - except Exception as e: - print(f"Task error: {e}") - server._response_tasks.discard(task) # Use discard to prevent KeyError - await asyncio.sleep(0.1) - except KeyboardInterrupt: - print("\nShutting down...") - finally: - await server.stop() + self._app = web.Application() + self._init_routes() -def main(): - """Main entry point for the SIA application.""" - asyncio.run(run_web_server()) + async def app(self): + return self._app + + def _init_routes(self): + """Initialize application routes.""" + self._app.middlewares.append(self._cors_middleware) + self._app.router.add_get("/ws", self._ws_manager.handle_websocket) + self._app.router.add_get("/", self._serve_index) + self._app.router.add_static("/static/", self._static_dir, show_index=False) + self._app.router.add_static("/assets/", self._static_dir / "assets", show_index=False) + self._app.router.add_get("/{path:.*}", self._serve_index) + + async def _serve_index(self, request: web.Request) -> web.Response: + """Serve the React application HTML for any unmatched routes.""" + index_path = self._static_dir / "index.html" + if not index_path.exists(): + raise web.HTTPNotFound() + + with open(index_path, "r") as f: + html_content = f.read() + + return web.Response( + text=html_content, + content_type="text/html" + ) + + @web.middleware + async def _cors_middleware(self, request: web.Request, handler): + """Handle CORS headers.""" + if request.method == "OPTIONS": + response = web.Response() + else: + response = await handler(request) + + response.headers.update({ + 'Access-Control-Allow-Origin': '*', + 'Access-Control-Allow-Methods': 'GET, POST, OPTIONS', + 'Access-Control-Allow-Headers': 'Content-Type', + 'Access-Control-Max-Age': '86400', + }) + return response if __name__ == "__main__": - main() + main = Main() + app = asyncio.run(main.app()) + print("Web server started at http://localhost:8080") + web.run_app(app, port=8080) \ No newline at end of file diff --git a/sia/web_agent.py b/sia/web_agent.py index e63245e..02794e9 100644 --- a/sia/web_agent.py +++ b/sia/web_agent.py @@ -24,6 +24,7 @@ class WebAgentState(Enum): CONTEXT_APPROVAL = auto() INFERENCE = auto() RESPONSE_APPROVAL = auto() + STOPPED = auto() class WebAgent(BaseAgent): """ @@ -45,12 +46,12 @@ class WebAgent(BaseAgent): Initialize web agent with required components. """ super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser) - self.response = "" + self._response = "" self._system_prompt = system_prompt self._current_state = WebAgentState.CONTEXT_APPROVAL self._validation_error: Optional[str] = None self._state_change_handlers: List[Callable[[WebAgentState], None]] = [] - self._response_change_handlers: List[Callable[[str], None]] = [] + self._response_change_handlers: List[Callable[[str, str], None]] = [] self._command_result: Optional[CommandResult] = None self.context = self._compile_context() @@ -63,6 +64,16 @@ class WebAgent(BaseAgent): def command_result(self) -> Optional[CommandResult]: """Get the result of the last command execution.""" return self._command_result + + @property + def validation_error(self) -> Optional[str]: + """Get the validation error, if any.""" + return self._validation_error + + @property + def response(self) -> str: + """Get the current response.""" + return self._response def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None: """ @@ -74,7 +85,7 @@ class WebAgent(BaseAgent): if handler not in self._state_change_handlers: self._state_change_handlers.append(handler) - def add_response_change_handler(self, handler: Callable[[str], None]) -> None: + def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None: """ Add a callback for response changes. @@ -89,20 +100,13 @@ class WebAgent(BaseAgent): for handler in self._state_change_handlers: handler(self._current_state) - def _notify_response_change(self) -> None: - """Notify all handlers of response change.""" - for handler in self._response_change_handlers: - handler(self.response) - - def _set_response(self, response: str) -> None: + def set_response(self, response: str) -> None: """Set response and notify handlers.""" - self.response = response - self._validation_error = self._validator.validate(response) + self._response = response + validation_error = self._validator.validate(response) + self._validation_error = validation_error for handler in self._response_change_handlers: - try: - handler(response) - except Exception as e: - print(f"Error in response handler: {e}") + handler(response, validation_error) def _set_state(self, new_state: WebAgentState) -> None: """ @@ -114,39 +118,22 @@ class WebAgent(BaseAgent): self._current_state = new_state self._notify_state_change() - async def approve_context(self) -> None: + def approve_context(self) -> None: if self._current_state != WebAgentState.CONTEXT_APPROVAL: error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})" raise Exception(error_msg) self._set_state(WebAgentState.INFERENCE) - self._set_response("") + self.set_response("") response_token_iter = self._llm.infer(self._system_prompt, self.context) response = "" - try: - for token in response_token_iter: - response += token - await self._set_response(response) - print(f"{token}") - except Exception as e: - print(f"Error during inference: {e}") - self._set_state(WebAgentState.CONTEXT_APPROVAL) - return - - await self._set_response(response) + for token in response_token_iter: + response += token + self.set_response(response) + print(f"{token}") self._set_state(WebAgentState.RESPONSE_APPROVAL) - async def _set_response(self, response: str) -> None: - """Set response and notify handlers asynchronously.""" - self.response = response - self._validation_error = self._validator.validate(response) - for handler in self._response_change_handlers: - try: - await handler(response) - except Exception as e: - print(f"Error in response handler: {e}") - def approve_response(self) -> None: if self._current_state != WebAgentState.RESPONSE_APPROVAL: raise Exception("Not in RESPONSE_APPROVAL state") @@ -157,10 +144,12 @@ class WebAgent(BaseAgent): if isinstance(parse_result, Command): result = parse_result.execute(self._working_memory) self._command_result = result + if result.should_stop: + self._set_state(WebAgentState.STOPPED) + return else: self._working_memory.add_entry(parse_result) self._working_memory.update() self.context = self._compile_context() - self._set_state(WebAgentState.CONTEXT_APPROVAL) diff --git a/sia/web_io_buffer.py b/sia/web_io_buffer.py index 9af92fe..4399ac0 100644 --- a/sia/web_io_buffer.py +++ b/sia/web_io_buffer.py @@ -10,30 +10,16 @@ class WebIOBuffer(IOBuffer): def __init__(self): self._stdin_buffer = "" self._stdout_buffer = "" - self._lock = Lock() - - def append_stdin(self, content: str) -> None: - """Thread-safe append to stdin buffer.""" - if not content: - return - - with self._lock: - self._stdin_buffer += content def read(self) -> str: """Thread-safe read from stdin buffer.""" - with self._lock: - content = self._stdin_buffer - self._stdin_buffer = "" - return content + content = self._stdin_buffer + self._stdin_buffer = "" + return content def write(self, content: str) -> None: """Thread-safe write to stdout buffer.""" - if not content: - return - - with self._lock: - self._stdout_buffer += content + self._stdout_buffer += content def buffer_length(self) -> int: """ @@ -51,20 +37,12 @@ class WebIOBuffer(IOBuffer): Args: content: String content to append to the stdin buffer """ - if content: - self._stdin_buffer += content + self._stdin_buffer += content def get_stdout(self) -> str: """Thread-safe get stdout content.""" - with self._lock: - return self._stdout_buffer + return self._stdout_buffer def clear_stdout(self) -> None: """Thread-safe clear stdout buffer.""" - with self._lock: - self._stdout_buffer = "" - - def buffer_length(self) -> int: - """Thread-safe get stdin buffer length.""" - with self._lock: - return len(self._stdin_buffer) + self._stdout_buffer = "" \ No newline at end of file diff --git a/sia/web_server.py b/sia/web_server.py deleted file mode 100644 index bb19e73..0000000 --- a/sia/web_server.py +++ /dev/null @@ -1,317 +0,0 @@ -from aiohttp import web, WSMsgType -from pathlib import Path -from typing import Set, Optional -import asyncio -import json -from dataclasses import dataclass -import mimetypes -from functools import partial - -from .web_agent import WebAgent, WebAgentState -from .web_io_buffer import WebIOBuffer - -mimetypes.add_type("application/javascript", ".js") -mimetypes.add_type("application/javascript", ".jsx") -mimetypes.add_type("text/javascript", ".js") -mimetypes.add_type("text/javascript", ".jsx") - -@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. Serves static files from a web directory. - """ - def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, - host: str = "0.0.0.0", port: int = 8080, - static_dir: Optional[Path] = None): - """Initialize the web server.""" - self.agent = agent - self.io_buffer = io_buffer - self.host = host - self.port = port - self.app = web.Application() - self.static_dir = static_dir or Path(__file__).parent.parent / "static" - self._init_routes() - self._clients: Set[web.WebSocketResponse] = set() - self._lock = asyncio.Lock() - self._response_tasks: List[asyncio.Task] = [] # Initialize tasks list - self.runner: Optional[web.AppRunner] = None - - # Register wrapped agent handlers - self.agent.add_state_change_handler(self._wrap_async_handler(self._handle_state_change)) - self.agent.add_response_change_handler(self._wrap_async_handler(self._handle_response_change)) - - def _wrap_async_handler(self, coro_func): - """Wrap an async handler function to be called synchronously.""" - def wrapper(*args, **kwargs): - loop = asyncio.get_event_loop() - task = loop.create_task(coro_func(*args, **kwargs)) - self._response_tasks.append(task) - self._response_tasks = [t for t in self._response_tasks if not t.done()] - return task - return wrapper - - def _init_routes(self): - """Initialize application routes.""" - self.app.middlewares.append(self._cors_middleware) - self.app.router.add_get("/ws", self._handle_websocket) - - if self.static_dir.exists(): - self.app.router.add_get("/", self._serve_index) - self.app.router.add_static("/static/", self.static_dir, show_index=False) - self.app.router.add_static("/assets/", self.static_dir / "assets", show_index=False) - self.app.router.add_get("/{path:.*}", self._serve_index) - - 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: - self.agent.approve_context() - - elif message_type == ClientMessage.APPROVE_RESPONSE: - if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: - 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: - self.agent._set_response(new_response) - - elif message_type == ClientMessage.SEND_INPUT: - input_text = data.get("data") - if input_text: - if not input_text.endswith('\n'): - input_text += '\n' - - self.io_buffer.append_stdin(input_text) - - current_length = self.io_buffer.buffer_length() - return - else: - print(f"Agent in wrong state: {self.agent.current_state}") - - async def _serve_index(self, request: web.Request) -> web.Response: - """Serve the React application HTML for any unmatched routes.""" - index_path = self.static_dir / "index.html" - if not index_path.exists(): - raise web.HTTPNotFound() - - with open(index_path, "r") as f: - html_content = f.read() - - return web.Response( - text=html_content, - content_type="text/html" - ) - - @web.middleware - async def _cors_middleware(self, request: web.Request, handler): - """Handle CORS headers.""" - if request.method == "OPTIONS": - response = web.Response() - else: - response = await handler(request) - - response.headers.update({ - 'Access-Control-Allow-Origin': '*', - 'Access-Control-Allow-Methods': 'GET, POST, OPTIONS', - 'Access-Control-Allow-Headers': 'Content-Type', - 'Access-Control-Max-Age': '86400', - }) - return response - - @property - def clients(self) -> Set[web.WebSocketResponse]: - """Get the set of connected clients.""" - return self._clients - - async def start(self): - """Start the web server.""" - # Verify static directory exists - if not self.static_dir.exists(): - raise FileNotFoundError(f"Static directory not found: {self.static_dir}") - - 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 http://{self.host}:{self.port}") - print(f"Serving static files from: {self.static_dir}") - print(f"WebSocket endpoint 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.""" - if not ws.closed: - try: - await ws.send_json(message) - return True - except Exception as e: - print(f"Error sending to client: {e}") - return False - 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((ws, task)) - else: - disconnected.add(ws) - - if tasks: - for ws, task in tasks: - try: - success = await task - if not success: - disconnected.add(ws) - except Exception: - 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(heartbeat=30) - 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): - message_type = data.get("type") - - if message_type == ClientMessage.APPROVE_CONTEXT: - if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL: - await self.agent.approve_context() - elif message_type == ClientMessage.APPROVE_RESPONSE: - if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: - self.agent.approve_response() # Now synchronous - - elif message_type == ClientMessage.MODIFY_RESPONSE: - if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: - new_response = data.get("data") - if new_response: - self.agent._set_response(new_response) - - elif message_type == ClientMessage.SEND_INPUT: - input_text = data.get("data") - if input_text: - self.agent.set_input(input_text) - self.io_buffer.append_stdin(input_text + "\n") - else: - print(f"Agent in wrong state: {self.agent.current_state}") - - 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 - ) - - 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: - try: - await self._send_to_client(ws, message) - except Exception as e: - print(f"Error sending to client: {e}") - disconnected.add(ws) - else: - disconnected.add(ws) - - self._clients -= disconnected \ No newline at end of file diff --git a/sia/web_socket_manager.py b/sia/web_socket_manager.py new file mode 100644 index 0000000..bd7f3b8 --- /dev/null +++ b/sia/web_socket_manager.py @@ -0,0 +1,114 @@ +from aiohttp import web, WSMsgType +from typing import Any, Dict, Set +import asyncio +import json +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" + +class WebSocketManager: + """ + Manages WebSocket connections and integrates with the WebAgent + """ + def __init__(self, + agent: WebAgent, + io_buffer: WebIOBuffer): + """Initialize the web server.""" + self.agent = agent + self.io_buffer = io_buffer + self._clients: Set[web.WebSocketResponse] = set() + + self.agent.add_state_change_handler(self._handle_state_change) + self.agent.add_response_change_handler(self._handle_response_change) + + async def broadcast_message(self, message: Dict): + """Broadcast message to all connected clients.""" + disconnected = set() + for ws in self._clients: + try: + await ws.send_json(message) + except ConnectionResetError: + disconnected.add(ws) + self._clients -= disconnected + + def _handle_state_change(self, new_state: WebAgentState): + """Handle state changes from the WebAgent.""" + asyncio.run(self.broadcast_message({ + "type": ServerMessage.STATE_CHANGE, + "state": new_state.name, + })) + + def _handle_response_change(self, response: str, validation_error: str): + """Handle response changes from the WebAgent.""" + asyncio.run(self.broadcast_message({ + "type": ServerMessage.RESPONSE_UPDATE, + "response": response, + "validation_error": validation_error + })) + + async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse: + """Handle new WebSocket connections and messages.""" + ws = web.WebSocketResponse(heartbeat=30) + await ws.prepare(request) + + self._clients.add(ws) + + try: + await ws.send_json({ + "type": ServerMessage.STATE_CHANGE, + "state": self.agent.current_state.name + }) + await ws.send_json({ + "type": ServerMessage.CONTEXT_UPDATE, + "context": self.agent.context + }) + await ws.send_json({ + "type": ServerMessage.RESPONSE_UPDATE, + "response": self.agent.response, + "validation_error": self.agent._validation_error + }) + async for msg in ws: + if msg.type == WSMsgType.TEXT: + #try: + data = json.loads(msg.data) + await self._handle_client_message(request, data) + #except Exception as e: + # print(f"Error handling incoming websocket message: {msg}\n{e}") + elif msg.type == WSMsgType.ERROR: + print(f"WebSocket connection closed with error: {ws.exception()}") + finally: + self._clients.remove(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")) \ No newline at end of file diff --git a/test/mock_web_agent.py b/test/mock_web_agent.py deleted file mode 100644 index 17fca99..0000000 --- a/test/mock_web_agent.py +++ /dev/null @@ -1,53 +0,0 @@ -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) - - def set_input(self, input_text): - pass - - def _notify_state_handlers(self, new_state: WebAgentState): - """Notify all state handlers of a state change.""" - for handler in self._state_handlers: - handler(new_state) - - def _notify_response_handlers(self, response: str): - """Notify all response handlers of a response change.""" - for handler in self._response_handlers: - handler(response) - - def _set_response(self, response: str): - """Set response and notify handlers.""" - self.response = response - self._notify_response_handlers(response) - - def approve_context(self): - """Handle context approval with immediate state updates.""" - # Change to INFERENCE state and notify - self._current_state = WebAgentState.INFERENCE - self._notify_state_handlers(WebAgentState.INFERENCE) - - # Set response and notify - self._set_response("test") - - # Change to RESPONSE_APPROVAL state and notify - self._current_state = WebAgentState.RESPONSE_APPROVAL - self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL) diff --git a/test/web_agent_test.py b/test/web_agent_test.py index e0ecf14..1c8524d 100644 --- a/test/web_agent_test.py +++ b/test/web_agent_test.py @@ -88,7 +88,7 @@ class WebAgentTest(unittest.TestCase): """Track state changes for verification.""" self.state_changes.append(state) - def response_change_handler(self, response: str): + def response_change_handler(self, response: str, validation_error: str): """Track response changes for verification.""" self.response_changes.append(response) @@ -97,7 +97,7 @@ class WebAgentTest(unittest.TestCase): self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL) self.assertEqual(self.agent.response, "") self.assertIsNotNone(self.agent.context) - self.assertIsNone(self.agent._validation_error) + self.assertIsNone(self.agent.validation_error) self.assertEqual(len(self.agent._state_change_handlers), 0) self.assertEqual(len(self.agent._response_change_handlers), 0) @@ -133,7 +133,7 @@ class WebAgentTest(unittest.TestCase): self.agent.approve_response() self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception)) - async def test_context_approval_flow(self): + def test_context_approval_flow(self): """Test complete context approval flow with state transitions.""" # Register handlers self.agent.add_state_change_handler(self.state_change_handler) @@ -145,27 +145,16 @@ class WebAgentTest(unittest.TestCase): yield "test reasoning" self.mock_llm.infer.side_effect = mock_infer - # Set initial state - self.agent._current_state = WebAgentState.CONTEXT_APPROVAL - # Approve context self.agent.approve_context() - # Wait for response tasks to complete - await asyncio.gather(*self.web_server._response_tasks) - # Verify state transitions expected_states = [ WebAgentState.INFERENCE, WebAgentState.RESPONSE_APPROVAL ] self.assertEqual(self.state_changes, expected_states) - - # Updated test: Should now expect streaming updates - self.assertEqual(len(self.response_changes), len(self.received_messages)) - for i, msg in enumerate(self.received_messages): - self.assertEqual(msg["type"], "RESPONSE_UPDATE") - self.assertEqual(msg["data"], self.response_changes[i]) + self.assertEqual(self.response_changes[-1], "test reasoning") def test_response_approval_flow_command(self): """Test response approval flow with command execution.""" @@ -206,7 +195,7 @@ class WebAgentTest(unittest.TestCase): # Set initial state and response self.agent._current_state = WebAgentState.RESPONSE_APPROVAL - self.agent.response = "test reasoning" + self.agent.set_response("test reasoning") # Approve response self.agent.approve_response() @@ -224,22 +213,13 @@ class WebAgentTest(unittest.TestCase): ] self.assertEqual(self.state_changes, expected_states) - async def test_response_validation(self): + def test_response_validation(self): """Test response validation during response setting.""" - # Register handlers self.agent.add_response_change_handler(self.response_change_handler) - - # Mock validator to fail error_message = "Invalid XML" self.mock_validator.validate.return_value = error_message - - # Set response using async method - await self.agent._set_response("") - - # Verify validation error stored - self.assertEqual(self.agent._validation_error, error_message) - - # Verify response updated + self.agent.set_response("") + self.assertEqual(self.agent.validation_error, error_message) self.assertEqual(self.response_changes, [""]) def tearDown(self): diff --git a/test/web_server_test.py b/test/web_server_test.py deleted file mode 100644 index ea3a870..0000000 --- a/test/web_server_test.py +++ /dev/null @@ -1,251 +0,0 @@ -from aiohttp import WSMsgType -from aiohttp import WSMsgType -from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop -from pathlib import Path -from sia.web_agent import WebAgentState -from unittest.mock import patch -import asyncio -import asyncio -import time -import json -import json -import shutil -import tempfile - -from sia.web_agent import WebAgentState -from sia.web_io_buffer import WebIOBuffer -from sia.web_server import WebServer, ServerMessage, ClientMessage -from .mock_web_agent import MockWebAgent - -class WebServerTest(AioHTTPTestCase): - async def get_application(self): - """Create application for testing.""" - self.io_buffer = WebIOBuffer() - - self.agent = MockWebAgent() - - # Create temporary directory for static files - self.temp_dir = Path(tempfile.mkdtemp()) - self.static_dir = self.temp_dir / "static" - self.static_dir.mkdir(exist_ok=True) - - # Create web server with temp static directory - self.agent._current_state = WebAgentState.CONTEXT_APPROVAL - self.web_server = WebServer( - self.agent, - self.io_buffer, - "localhost", - 8080, - static_dir=self.static_dir - ) - return self.web_server.app - - def tearDown(self): - """Clean up temporary directory.""" - if hasattr(self, 'temp_dir'): - shutil.rmtree(self.temp_dir) - - 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: - 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, - }) - await asyncio.sleep(0.1) # Let state changes propagate - - # 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" - self.agent._set_response("") # Now synchronous - 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() - - async def wait_for_buffer_content(self, timeout=1.0): - start = time.time() - while time.time() - start < timeout: - if self.io_buffer.buffer_length() > 0: - return True - await asyncio.sleep(0.01) - return False - - async def test_send_input(self): - async with self.client.ws_connect("/ws") as ws: - await self.drain_messages(ws) - await asyncio.sleep(0.1) - - test_input = "test input" - await ws.send_json({ - "type": ClientMessage.SEND_INPUT, - "data": test_input - }) - - # Wait for buffer to have content - self.assertTrue(await self.wait_for_buffer_content()) - - content = self.io_buffer.read() - self.assertEqual( - content, - test_input + "\n", - f"Expected '{test_input}\\n' but got '{content}'" - ) - - @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) diff --git a/test/web_socket_manager_test.py b/test/web_socket_manager_test.py new file mode 100644 index 0000000..b6f0a59 --- /dev/null +++ b/test/web_socket_manager_test.py @@ -0,0 +1,179 @@ +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop +from aiohttp.web import AppRunner +from unittest.mock import Mock, patch, call +import asyncio +import unittest +import xml.etree.ElementTree as ET + +from sia.web_agent import WebAgent, WebAgentState +from sia.web_io_buffer import WebIOBuffer +from sia.web_socket_manager import WebSocketManager, ClientMessage, ServerMessage +from sia.working_memory import WorkingMemory +from sia.system_metrics import SystemMetrics +from sia.llm_engine import LlmEngine +from sia.xml_validator import XMLValidator +from sia.response_parser import ResponseParser + +class WebSocketManagerTest(AioHTTPTestCase): + async def get_application(self): + """Create test application with WebSocket handler""" + # Create real components where beneficial for testing + self.io_buffer = WebIOBuffer() + self.working_memory = WorkingMemory() + + # Create mock metrics with proper generate_context implementation + self.mock_metrics = Mock(spec=SystemMetrics) + metrics_elem = ET.Element("context") + metrics_elem.attrib = { + "time": "2024-10-31T12:00:00Z", + "cpu": "10", + "gpu": "20", + "memory_used": "1000", + "memory_total": "2000", + "disk_used": "5000", + "disk_total": "10000", + "context": "0", + "stdin": "0" + } + self.mock_metrics.generate_context.return_value = metrics_elem + + # Create minimal mocks for other components + self.mock_llm = Mock(spec=LlmEngine) + def mock_infer(prompt: str, context: str): + yield "test reasoning" + self.mock_llm.infer.side_effect = mock_infer + self.mock_validator = Mock(spec=XMLValidator) + self.mock_validator.validate.return_value = None + + # Create parser with real IO buffer + self.parser = ResponseParser(self.io_buffer) + + # Create agent with mix of real/mock components + self.agent = WebAgent( + system_prompt="test prompt", + action_schema="test schema", + working_memory=self.working_memory, + system_metrics=self.mock_metrics, + llm=self.mock_llm, + validator=self.mock_validator, + parser=self.parser + ) + + # Create WebSocketManager + self.ws_manager = WebSocketManager(self.agent, self.io_buffer) + + # Create application + app = web.Application() + app.router.add_get('/ws', self.ws_manager.handle_websocket) + return app + + @unittest_run_loop + async def test_initial_connection(self): + """Test initial WebSocket connection and state messages""" + async with self.client.ws_connect('/ws') as ws: + # Verify initial state message + msg = await ws.receive_json() + self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE) + self.assertEqual(msg['state'], WebAgentState.CONTEXT_APPROVAL.name) + + # Verify initial context message + msg = await ws.receive_json() + self.assertEqual(msg['type'], ServerMessage.CONTEXT_UPDATE) + self.assertIsNotNone(msg['context']) + + # Verify initial response message + msg = await ws.receive_json() + self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE) + self.assertEqual(msg['response'], '') + self.assertIsNone(msg['validation_error']) + + @unittest_run_loop + async def test_approve_context(self): + """Test context approval flow""" + async with self.client.ws_connect('/ws') as ws: + for _ in range(3): + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + await ws.send_json({ + 'type': ClientMessage.APPROVE_CONTEXT + }) + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE) + self.assertEqual(msg['state'], WebAgentState.INFERENCE.name) + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE) + self.assertEqual(msg['response'], "") + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE) + self.assertEqual(msg['response'], "test reasoning") + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE) + self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name) + + @unittest_run_loop + async def test_send_input(self): + """Test sending input through WebSocket""" + test_input = "test input" + + async with self.client.ws_connect('/ws') as ws: + for _ in range(3): + await ws.receive_json() + + # Send input + await ws.send_json({ + 'type': ClientMessage.SEND_INPUT, + 'input': test_input + }) + + # Delay for processing + await asyncio.sleep(0.5) + + # Verify input was added to buffer + self.assertEqual(self.io_buffer.read(), test_input) + + @unittest_run_loop + async def test_modify_response(self): + """Test modifying response in RESPONSE_APPROVAL state""" + # Set agent to RESPONSE_APPROVAL state + async with self.client.ws_connect('/ws') as ws: + for _ in range(3): + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + await ws.send_json({ + 'type': ClientMessage.APPROVE_CONTEXT + }) + for _ in range(4): + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + modified_response = "modified test" + await ws.send_json({ + 'type': ClientMessage.MODIFY_RESPONSE, + 'response': modified_response + }) + msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE) + self.assertEqual(msg['response'], modified_response) + + @unittest_run_loop + async def test_multiple_clients(self): + """Test broadcasting to multiple WebSocket clients""" + # Connect first client + async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2: + for _ in range(3): + await asyncio.wait_for(ws1.receive_json(), timeout=0.5) + await asyncio.wait_for(ws2.receive_json(), timeout=0.5) + await ws1.send_json({ + 'type': ClientMessage.APPROVE_CONTEXT + }) + msg1 = await ws1.receive_json() + msg2 = await ws2.receive_json() + self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE) + self.assertEqual(msg1['state'], WebAgentState.INFERENCE.name) + self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE) + self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name) + + def tearDown(self): + """Clean up resources""" + super().tearDown() + self.io_buffer = None + self.working_memory = None + self.agent = None + self.ws_manager = None