Enable multiple llms

sia/web/api.py

@@ -0,0 +1,97 @@
+from aiohttp import web
+import json
+import asyncio
+
+from ..web_agent import WebAgent
+from ..web_io_buffer import WebIOBuffer
+from .util import wrap_async
+
+class Api:
+    def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer):
+        self._app = app
+        self._agent = agent
+        self._io_buffer = io_buffer
+        
+        self._init_routes()
+        
+    def _init_routes(self):
+        """Initialize REST API and WebSocket routes."""
+        self._app.router.add_post("/api/inference/{llm}", self._run_inference)
+        self._app.router.add_post("/api/approve/{llm}", self._approve_response)
+        self._app.router.add_post("/api/context", self._modify_context)
+        self._app.router.add_post("/api/input", self._send_input)
+        self._app.router.add_post("/api/clear", self._clear_output)
+        self._app.router.add_get("/api/output/{llm}", self._get_output)
+        self._app.router.add_get("/api/llms", self._get_llms)
+        
+    @property
+    def app(self):
+        return self._app
+
+    async def _run_inference(self, request: web.Request) -> web.Response:
+        """Start inference on specified LLM."""
+        llm_name = request.match_info["llm"]
+        try:
+            await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name)
+            return web.Response(status=200)
+        except (ValueError, RuntimeError) as e:
+            return web.Response(status=400, text=str(e))
+
+    async def _approve_response(self, request: web.Request) -> web.Response:
+        """Approve response from specified LLM."""
+        llm_name = request.match_info["llm"]
+        data = await request.json()
+        response = data.get("response")
+        if not response:
+            return web.Response(status=400, text="Missing response in request body")
+        try:
+            self._agent.approve_response(llm_name, response)
+            return web.Response(status=200)
+        except ValueError as e:
+            return web.Response(status=400, text=str(e))
+
+    async def _modify_context(self, request: web.Request) -> web.Response:
+        """Modify the current context."""
+        data = await request.json()
+        context = data.get("context")
+        if not context:
+            return web.Response(status=400, text="Missing context in request body")
+        self._agent.modify_context(context)
+        return web.Response(status=200)
+
+    async def _send_input(self, request: web.Request) -> web.Response:
+        """Send input to the IO buffer."""
+        data = await request.json()
+        input_text = data.get("input")
+        if not input_text:
+            return web.Response(status=400, text="Missing input in request body")
+        self._io_buffer.append_stdin(input_text)
+        return web.Response(status=200)
+
+    async def _clear_output(self, request: web.Request) -> web.Response:
+        """Clear the stdout buffer."""
+        self._io_buffer.clear_stdout()
+        return web.Response(status=200)
+
+    async def _get_output(self, request: web.Request) -> web.Response:
+        """Get complete output for specified LLM."""
+        llm_name = request.match_info["llm"]
+        try:
+            output = self._agent.get_output(llm_name)
+            return web.Response(
+                text=json.dumps({"output": output}),
+                content_type="application/json"
+            )
+        except ValueError as e:
+            return web.Response(status=400, text=str(e))
+
+    async def _get_llms(self, request: web.Request) -> web.Response:
+        """Get all LLMs and their current states."""
+        states = self._agent.llms
+        return web.Response(
+            text=json.dumps({
+                name: state.name 
+                for name, state in states.items()
+            }),
+            content_type="application/json"
+        )

sia/web/context_websocket.py

@@ -0,0 +1,55 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+
+from .util import wrap_async
+from ..web_agent import WebAgent
+
+class ContextWebSocket:
+    """
+    WebSocket handler for context changes.
+    Broadcasts context updates to all connected clients.
+    """
+    
+    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: str, generated: bool):
+        """Handle context changes from the WebAgent."""
+        await self._broadcast_message({
+            "context": context,
+            "generated": generated
+        })
+
+    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+        
+        self._clients.add(ws)
+        
+        try:
+            # Send initial context
+            await ws.send_json({
+                "context": self._agent.context,
+                "generated": True
+            })
+                
+            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

sia/web/llm_websocket.py

@@ -0,0 +1,57 @@
+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

sia/web/static.py

@@ -0,0 +1,32 @@
+from aiohttp import web
+import mimetypes
+
+from ..config import Config
+
+mimetypes.add_type("application/javascript", ".js")
+mimetypes.add_type("application/javascript", ".jsx")
+mimetypes.add_type("text/javascript", ".js")
+mimetypes.add_type("text/javascript", ".jsx")
+
+class Static:
+    def __init__(self, app: web.Application, config: Config):
+        self._config = config
+
+        app.router.add_get("/", self._serve_index)
+        app.router.add_static("/static/", self._config.static_files, show_index=False)
+        app.router.add_static("/assets/", self._config.static_files / "assets", show_index=False)
+        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._config.static_files / "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"
+        )

sia/web/stdout_websocket.py

@@ -0,0 +1,54 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+import asyncio
+
+from .util import wrap_async
+from ..web_io_buffer import WebIOBuffer
+
+class StdoutWebSocket:
+    """
+    WebSocket handler for stdout changes.
+    Broadcasts stdout updates to all connected clients.
+    """
+    
+    def __init__(self, io_buffer: WebIOBuffer):
+        self._io_buffer = io_buffer
+        self._clients: Set[web.WebSocketResponse] = set()
+        self._io_buffer.add_stdout_change_handler(wrap_async(self._handle_stdout_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_stdout_change(self, output: str):
+        """Handle stdout changes from the WebIOBuffer."""
+        await self._broadcast_message({
+            "output": output
+        })
+
+    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+        
+        self._clients.add(ws)
+        
+        try:
+            # Send initial stdout content
+            await ws.send_json({
+                "output": self._io_buffer.get_stdout()
+            })
+                
+            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

sia/web/token_websocket.py

@@ -0,0 +1,51 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+import asyncio
+import json
+
+from .util import wrap_async
+from ..web_agent import WebAgent
+
+class TokenWebSocket:
+    """
+    WebSocket handler for LLM token streaming.
+    Broadcasts new tokens to all connected clients.
+    """
+    
+    def __init__(self, agent: WebAgent):
+        self._agent = agent
+        self._clients: Set[web.WebSocketResponse] = set()
+        self._agent.add_token_handler(wrap_async(self._handle_new_token))
+        
+    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_new_token(self, llm_name: str, token: str):
+        """Handle new tokens from the WebAgent."""
+        await self._broadcast_message({
+            "llm": llm_name,
+            "token": token
+        })
+
+    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+        
+        self._clients.add(ws)
+        
+        try:
+            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

sia/web/util.py

@@ -0,0 +1,8 @@
+import asyncio
+
+def wrap_async(coro_func):
+    """Wraps an async callback to be safely called from another thread."""
+    loop = asyncio.get_event_loop()
+    def wrapper(*args, **kwargs):
+        loop.call_soon_threadsafe(lambda: asyncio.create_task(coro_func(*args, **kwargs)))
+    return wrapper

sia/web/websockts.py

@@ -0,0 +1,20 @@
+from aiohttp import web
+
+from ..web_agent import WebAgent
+from ..web_io_buffer import WebIOBuffer
+from .llm_websocket import LlmWebSocket
+from .context_websocket import ContextWebSocket 
+from .token_websocket import TokenWebSocket
+from .stdout_websocket import StdoutWebSocket
+
+class Websockets:
+    def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer):
+        self._llm_ws = LlmWebSocket(agent)
+        self._context_ws = ContextWebSocket(agent)
+        self._token_ws = TokenWebSocket(agent)
+        self._stdout_ws = StdoutWebSocket(io_buffer)
+
+        app.router.add_get("/ws/llm", self._llm_ws.handle_connection)
+        app.router.add_get("/ws/context", self._context_ws.handle_connection)
+        app.router.add_get("/ws/token", self._token_ws.handle_connection)
+        app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection)