Implemented webserver
This commit is contained in:
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
|
||||
)
|
||||
Reference in New Issue
Block a user