Separate and tested web_socket_manager
This commit is contained in:
136
sia/__main__.py
136
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 "<reasoning>"
|
||||
time.sleep(2)
|
||||
yield "test reasoning"
|
||||
time.sleep(2)
|
||||
yield "</reasoning>"
|
||||
|
||||
# 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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = ""
|
||||
@@ -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
|
||||
114
sia/web_socket_manager.py
Normal file
114
sia/web_socket_manager.py
Normal file
@@ -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"))
|
||||
Reference in New Issue
Block a user