Add context info in ui

sia/web/api.py

@@ -58,6 +58,8 @@ class Api:
         self._app.router.add_post("/api/memory/entry/{id}/update", self._update_entry)
         self._app.router.add_post("/api/memory/load_iteration", self._load_iteration)
 
+        self._app.router.add_get("/api/context_info", self._get_context_info)
+
     async def _get_llms(self, request: web.Request) -> web.Response:
         return web.Response(
             text=json.dumps(
@@ -293,4 +295,11 @@ class Api:
                 
             return web.Response(status=200)
         except Exception as e:
-            return web.Response(status=400, text=str(e))
+            return web.Response(status=400, text=str(e))
+
+    async def _get_context_info(self, request: web.Request) -> web.Response:
+        """Get current context info."""
+        return web.Response(
+            text=json.dumps(self._agent.context_info),
+            content_type="application/json"
+        )

sia/web/context_info_websocket.py

@@ -0,0 +1,48 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+
+from .util import wrap_async
+from ..web_agent import WebAgent
+
+class ContextInfoWebSocket:
+    """
+    WebSocket handler for context info updates.
+    Broadcasts context info to all connected clients when it changes.
+    """
+
+    def __init__(self, agent: WebAgent):
+        self._agent = agent
+        self._clients: Set[web.WebSocketResponse] = set()
+        self._agent.add_context_change_handler(wrap_async(self._handle_context_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_context_change(self, context_info: Dict):
+        """Handle context changes from the agent."""
+        await self._broadcast_message(context_info)
+
+    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 context info
+            await ws.send_json(self._agent.context_info)
+            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.discard(ws)
+
+        return ws

sia/web/websockets.py

@@ -6,6 +6,7 @@ from ..web_agent import WebAgent
 from ..working_memory import WorkingMemory
 from .auto_approver_websocket import AutoApproverWebSocket
 from .chat_websocket import ChatWebSocket
+from .context_info_websocket import ContextInfoWebSocket
 from .memory_websocket import MemoryWebSocket
 from .response_websocket import ResponseWebSocket
 from .state_websocket import StateWebSocket
@@ -14,12 +15,14 @@ class Websockets:
     def __init__(self, app: web.Application, agent: WebAgent, io_buffer: ChatIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
         self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
         self._chat_ws = ChatWebSocket(io_buffer)
+        self._context_info_ws = ContextInfoWebSocket(agent)
         self._memory_ws = MemoryWebSocket(working_memory)
         self._response_ws = ResponseWebSocket(agent)
         self._state_ws = StateWebSocket(agent)
 
         app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)
         app.router.add_get("/ws/chat", self._chat_ws.handle_connection)
+        app.router.add_get("/ws/context_info", self._context_info_ws.handle_connection)
         app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
         app.router.add_get("/ws/response", self._response_ws.handle_connection)
         app.router.add_get("/ws/state", self._state_ws.handle_connection)

sia/web_agent.py

@@ -45,6 +45,7 @@ class WebAgent(BaseAgent):
         
         self._state_change_handlers: List[Callable[[AgentState], None]] = []
         self._selected_llm_change_handlers: List[Callable[[str], None]] = []
+        self._context_change_handlers: List[Callable[[Dict], None]] = []
 
         self._update_compiled_context()
         self._working_memory.add_change_handler(self._update_compiled_context)
@@ -81,6 +82,13 @@ class WebAgent(BaseAgent):
     def add_state_change_handler(self, handler: Callable[[AgentState], None]) -> None:
         self._state_change_handlers.append(handler)
 
+    @property
+    def context_info(self) -> Dict:
+        return dict(self._compiled_context.attrib)
+
+    def add_context_change_handler(self, handler: Callable[[Dict], None]) -> None:
+        self._context_change_handlers.append(handler)
+
     def run_inference(self) -> None:
         """Start inference on specified LLM"""
         if self._state != AgentState.IDLE:
@@ -132,4 +140,15 @@ class WebAgent(BaseAgent):
             handler(state)
     
     def _update_compiled_context(self):
-        self._compiled_context = self._compile_context(self._llms[self.active_llm])
+        self._compiled_context = self._compile_context(self._llms[self.active_llm])
+        self._notify_context_change()
+
+    def _notify_context_change(self):
+        """Notify all context change handlers."""
+        for handler in self._context_change_handlers:
+            handler(self.context_info)
+
+    def _notify_selected_llm_change(self):
+        """Notify all selected LLM change handlers."""
+        for handler in self._selected_llm_change_handlers:
+            handler(self._selected_llm)