diff --git a/sia/web/api.py b/sia/web/api.py index af8042e..a21a004 100644 --- a/sia/web/api.py +++ b/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)) \ No newline at end of file + 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" + ) \ No newline at end of file diff --git a/sia/web/context_info_websocket.py b/sia/web/context_info_websocket.py new file mode 100644 index 0000000..1d0f279 --- /dev/null +++ b/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 diff --git a/sia/web/websockets.py b/sia/web/websockets.py index 6e79d02..8ddd5a8 100644 --- a/sia/web/websockets.py +++ b/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) \ No newline at end of file diff --git a/sia/web_agent.py b/sia/web_agent.py index 0421318..87c325b 100644 --- a/sia/web_agent.py +++ b/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]) \ No newline at end of file + 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) \ No newline at end of file diff --git a/web/src/components/Layout/MobileControls.tsx b/web/src/components/Layout/MobileControls.tsx index 005e995..1fe7206 100644 --- a/web/src/components/Layout/MobileControls.tsx +++ b/web/src/components/Layout/MobileControls.tsx @@ -5,17 +5,17 @@ const MobileControls: React.FC = () => { return (

Agent Controls

- + - +
- +
- +
diff --git a/web/src/components/Layout/Ribbon.tsx b/web/src/components/Layout/Ribbon.tsx index 6743b83..b5bf900 100644 --- a/web/src/components/Layout/Ribbon.tsx +++ b/web/src/components/Layout/Ribbon.tsx @@ -4,9 +4,9 @@ import { AgentSection, LLMSection, AutoApproverSection, MemorySection } from '.. const Ribbon: React.FC = () => { const screenSize = useScreenSize(); - + if (screenSize === 'mobile') return null; - + return (
diff --git a/web/src/components/Memory/MemoryEditor.tsx b/web/src/components/Memory/MemoryEditor.tsx index eb7377a..0f9b567 100644 --- a/web/src/components/Memory/MemoryEditor.tsx +++ b/web/src/components/Memory/MemoryEditor.tsx @@ -1,7 +1,8 @@ import React, { useState } from 'react'; import { useWebSocketContext } from '../../contexts/WebSocketContext'; import { api, EntryTypes, getInitialEntryData } from '../../services/api'; -import { +import { + ContextInfoEntry, ReasoningEntry, SingleEntry, RepeatEntry, @@ -253,7 +254,9 @@ const MemoryEditor: React.FC = ({ className = '' }) => {
)} - + + + {memory.length === 0 ? (

No memory entries

diff --git a/web/src/components/Memory/entries/ContextInfoEntry.tsx b/web/src/components/Memory/entries/ContextInfoEntry.tsx new file mode 100644 index 0000000..adc5bac --- /dev/null +++ b/web/src/components/Memory/entries/ContextInfoEntry.tsx @@ -0,0 +1,30 @@ +import React from 'react'; +import { useWebSocketContext } from '../../../contexts/WebSocketContext'; + +const ContextInfoEntry: React.FC = () => { + const { contextInfo } = useWebSocketContext(); + + if (!contextInfo) { + return null; + } + + return ( +
+
+ Context Info +
+
+
+ {Object.entries(contextInfo).map(([key, value]) => ( +
+ {key}: + {String(value)} +
+ ))} +
+
+
+ ); +}; + +export default ContextInfoEntry; diff --git a/web/src/components/Memory/entries/index.ts b/web/src/components/Memory/entries/index.ts index a0e8bf5..3bcd42d 100644 --- a/web/src/components/Memory/entries/index.ts +++ b/web/src/components/Memory/entries/index.ts @@ -1,5 +1,6 @@ export { default as BaseEntry } from './BaseEntry'; export { default as BackgroundEntry } from './BackgroundEntry'; +export { default as ContextInfoEntry } from './ContextInfoEntry'; export { default as ParseErrorEntry } from './ParseErrorEntry'; export { default as ReadEntry } from './ReadEntry'; export { default as ReasoningEntry } from './ReasoningEntry'; diff --git a/web/src/contexts/WebSocketContext.tsx b/web/src/contexts/WebSocketContext.tsx index b06bbc9..6933e07 100644 --- a/web/src/contexts/WebSocketContext.tsx +++ b/web/src/contexts/WebSocketContext.tsx @@ -1,6 +1,12 @@ import React, { createContext, useContext, useEffect, useRef, useState } from 'react'; import { AutoApproverConfig, ChatMessage } from '../types'; +interface ContextInfo { + [key: string]: string | number; + token_count: number; + token_limit: number; +} + interface WebSocketContextType { state: { agentState: 'IDLE' | 'INFERENCE' | 'PROCESSING_RESPONSE'; @@ -17,6 +23,7 @@ interface WebSocketContextType { llms: {name: string}[]; setActiveLLM: (llm: string) => Promise; connectionStatus: Record; + contextInfo: ContextInfo | null; } const defaultState: WebSocketContextType = { @@ -40,6 +47,7 @@ const defaultState: WebSocketContextType = { llms: [], setActiveLLM: async () => {}, connectionStatus: {}, + contextInfo: null, }; const WebSocketContext = createContext(defaultState); @@ -56,6 +64,7 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child const [autoApprover, setAutoApprover] = useState(defaultState.autoApprover); const [llms, setLLMs] = useState<{name: string}[]>([]); const [connectionStatus, setConnectionStatus] = useState>({}); + const [contextInfo, setContextInfo] = useState(null); const sockets = useRef<{[key: string]: WebSocket | null}>({}); @@ -101,7 +110,7 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child }, []); useEffect(() => { - const endpoints = ['ws/state', 'ws/response', 'ws/memory', 'ws/auto_approver', 'ws/chat']; + const endpoints = ['ws/state', 'ws/response', 'ws/memory', 'ws/auto_approver', 'ws/chat', 'ws/context_info']; const newConnectionStatus: Record = {}; const createWebSocket = (endpoint: string) => { @@ -157,20 +166,22 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child setMessages(prevMessages => { const updatedIds = data.messages.map((m: ChatMessage) => m.id); const filteredMessages = prevMessages.filter(m => !updatedIds.includes(m.id)); - + return [...filteredMessages, ...data.messages]; }); } - + if (data.deleted_message_ids) { - setMessages(prevMessages => + setMessages(prevMessages => prevMessages.filter(m => !data.deleted_message_ids.includes(m.id)) ); } - + if (data.last_read_timestamp) { setLastReadTimestamp(data.last_read_timestamp); } + } else if (endpoint === 'ws/context_info') { + setContextInfo(data); } } catch (error) { console.error(`Failed to parse message from ${endpoint}:`, error); @@ -232,7 +243,8 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child isConnected, llms, setActiveLLM, - connectionStatus + connectionStatus, + contextInfo }; return (