Implemented webserver
This commit is contained in:
42
readme.md
42
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 {
|
||||
<<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
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
accelerate
|
||||
aiohttp
|
||||
bs4
|
||||
torch
|
||||
transformers
|
||||
191
sia/web_server.py
Normal file
191
sia/web_server.py
Normal file
@@ -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
|
||||
)
|
||||
55
test/mock_web_agent.py
Normal file
55
test/mock_web_agent.py
Normal file
@@ -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)
|
||||
228
test/web_server_test.py
Normal file
228
test/web_server_test.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user