57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
from aiohttp import web, WSMsgType
|
|
from typing import Dict, Set
|
|
|
|
from .util import wrap_async
|
|
from ..web_agent import WebAgent, LlmState
|
|
|
|
class LlmWebSocket:
|
|
"""
|
|
WebSocket handler for LLM state changes.
|
|
Broadcasts state updates to all connected clients.
|
|
"""
|
|
|
|
def __init__(self, agent: WebAgent):
|
|
self._agent = agent
|
|
self._clients: Set[web.WebSocketResponse] = set()
|
|
self._agent.add_llm_change_handler(wrap_async(self._handle_state_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
|
|
|
|
async def _handle_state_change(self, llm_name: str, new_state: LlmState):
|
|
"""Handle state changes from the WebAgent."""
|
|
await self._broadcast_message({
|
|
"llm": llm_name,
|
|
"state": new_state.name
|
|
})
|
|
|
|
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
|
|
"""Handle new WebSocket connections."""
|
|
ws = web.WebSocketResponse(heartbeat=30)
|
|
await ws.prepare(request)
|
|
|
|
try:
|
|
# Send initial states for all LLMs
|
|
states = self._agent.llms
|
|
for llm_name, state in states.items():
|
|
await ws.send_json({
|
|
"llm": llm_name,
|
|
"state": state.name
|
|
})
|
|
|
|
self._clients.add(ws)
|
|
|
|
async for msg in ws:
|
|
if msg.type == WSMsgType.ERROR:
|
|
print(f"WebSocket connection closed with error: {ws.exception()}")
|
|
finally:
|
|
self._clients.remove(ws)
|
|
|
|
return ws |