Enable multiple llms

sia/__main__.py

@@ -11,18 +11,13 @@ from .openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
 from .web_agent import WebAgent
-from .web_agent import WebAgent
 from .web_io_buffer import WebIOBuffer
-from .web_io_buffer import WebIOBuffer
-from .web_socket_manager import WebSocketManager
+from .web.api import Api
+from .web.static import Static
+from .web.websockts import Websockets
 from .working_memory import WorkingMemory
 from .xml_validator import XMLValidator
 
-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 Main:
     @classmethod
     async def create(cls, config: Config):
@@ -32,68 +27,65 @@ class Main:
         self._system_prompt = self._config.system_prompt.read_text()
         self._action_schema = self._config.action_schema.read_text()
 
-        match self._config.llm_engine:
-            case "local":
-                self._llm = LocalLlmEngine(
-                    self._config.model,
-                    self._config.temperature,
-                    self._config.token_limit
-                )
-            case "hf":
-                self._llm = HfLlmEngine(
-                    self._config.model,
-                    self._config.temperature,
-                    self._config.api_token,
-                )
-            case "openai":
-                self._llm = OpenAILlmEngine(
-                    self._config.model,
-                    self._config.temperature,
-                    self._config.token_limit,
-                    self._config.api_token,
-                )
-            case "mistral":
-                self._llm = MistralLlmEngine(
-                    self._config.model,
-                    self._config.temperature,
-                    self._config.api_token,
-                    self._config.token_limit
-                )
-            case _:
-                raise ValueError(f"Invalid LLM engine: {self._config.llm_engine}")
+        # Initialize LLM engines based on config
+        self._llms = {}
+        
+        if config.local_enabled:
+            self._llms['local'] = LocalLlmEngine(
+                config.local_model,
+                config.local_temperature,
+                config.local_token_limit
+            )
+            
+        if config.openai_enabled:
+            self._llms['openai'] = OpenAILlmEngine(
+                config.openai_model,
+                config.openai_temperature,
+                config.openai_token_limit,
+                config.openai_api_key,
+            )
+            
+        if config.hf_enabled:
+            self._llms['hf'] = HfLlmEngine(
+                config.hf_model,
+                config.hf_temperature,
+                config.hf_api_key,
+            )
+            
+        if config.mistral_enabled:
+            self._llms['mistral'] = MistralLlmEngine(
+                config.mistral_model,
+                config.mistral_temperature,
+                config.mistral_token_limit,
+                config.mistral_api_key,
+            )
+
+        if not self._llms:
+            raise ValueError("No LLM engines enabled in configuration")
+
         self._io_buffer = WebIOBuffer()
         self._agent = WebAgent(
             system_prompt=self._system_prompt,
             action_schema=self._action_schema,
             working_memory=WorkingMemory(),
             metrics=SystemMetrics(),
-            llm=self._llm,
+            llms=self._llms,
             validator=XMLValidator(self._action_schema),
             parser=ResponseParser(self._io_buffer),
             iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
         )
-        self._ws_manager = await WebSocketManager.create(
-            self._agent,
-            self._io_buffer
-        )
 
         self._app = web.Application()
-        self._init_routes()
+        self._api = Api(self._app, self._agent, self._io_buffer)
+        self._websockets = Websockets(self._app, self._agent, self._io_buffer)
+        self._static = Static(self._app, self._config)
+
         return self
 
     @property
     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._config.static_files, show_index=False)
-        self._app.router.add_static("/assets/", self._config.static_files / "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._config.static_files / "index.html"
@@ -108,22 +100,6 @@ class Main:
             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__":
     loop = asyncio.new_event_loop()
     config = Config()

sia/auto_approver.py

@@ -29,7 +29,7 @@ class AutoApprover:
         self._response_thread: Thread | None = None
         self._config_change_handlers: list[ConfigChangeHandler] = []
         
-        self.agent.add_state_change_handler(self._handle_state_change)
+        self.agent.add_llm_change_handler(self._handle_state_change)
 
     @property
     def config(self) -> AutoApproverConfig:

sia/base_agent.py

@@ -22,7 +22,6 @@ class BaseAgent(ABC):
             action_schema: str, 
             working_memory: WorkingMemory,
             metrics: SystemMetrics,
-            llm: LlmEngine,
             validator: XMLValidator,
             parser: ResponseParser,
         ):
@@ -33,7 +32,6 @@ class BaseAgent(ABC):
         self._action_schema = action_schema
         self._working_memory = working_memory
         self._metrics = metrics
-        self._llm = llm
         self._validator = validator
         self._parser = parser
         
@@ -47,7 +45,7 @@ class BaseAgent(ABC):
         """Get the system prompt."""
         return f"{self._system_prompt}\n{self._action_schema}"
             
-    def _compile_context(self) -> str:
+    def _compile_context(self, llmEngine: LlmEngine) -> str:
         """
         Compile the current context for LLM inference.
         Includes system metrics and working memory entries.
@@ -76,8 +74,8 @@ class BaseAgent(ABC):
         context_str = pretty_print_element(context)
         
         # Calculate token usage percentage
-        token_count = self._llm.token_count(self.system_prompt, context_str)
-        token_limit = self._llm.token_limit()
+        token_count = llmEngine.token_count(self.system_prompt, context_str)
+        token_limit = llmEngine.token_limit()
         context_usage = (float(token_count) / float(token_limit)) * 100.0
         
         # Update context usage metric

sia/config.py

@@ -7,20 +7,11 @@ import os
 
 @dataclass
 class Config:
-    """
-    Configuration class that handles both command line and environment variables.
-    
-    Command line arguments take precedence over environment variables.
-    Environment variables serve as defaults that can be overridden via CLI.
-    """
-
     def __init__(self):
-        """
-        Create configuration from command line arguments and environment variables.
-        Required arguments must be provided either via CLI or environment variables.
-        """
         load_dotenv()
         parser = argparse.ArgumentParser(description='SIA - Self Improving Agent')
+        
+        # Core configuration
         parser.add_argument(
             '--system-prompt',
             type=Path,
@@ -39,6 +30,8 @@ class Config:
             default=os.getenv('SIA_ITERATIONS_DIR', 'iterations'),
             help='Path to the directory for storing iterations (default: iterations, env: SIA_ITERATIONS_DIR)'
         )
+
+        # Web server configuration
         parser.add_argument(
             '--server',
             action='store_true',
@@ -63,108 +56,239 @@ class Config:
             default=self._parse_optional_path('SIA_STATIC_FILES', './static/'),
             help='Path to static web files (default: ./static/, env: SIA_STATIC_FILES)'
         )
+
+        # Local LLM configuration
         parser.add_argument(
-            '--llm-engine',
-            type=str,
-            default=os.getenv('SIA_LLM_ENGINE', 'local'),
-            help='LLM engine (default: local, env: SIA_LLM_ENGINE)'
+            '--local-enable',
+            action='store_true',
+            default=self._parse_bool_env('SIA_LOCAL_ENABLED', False),
+            help='Enable local LLM engine (env: SIA_LOCAL_ENABLED)'
         )
         parser.add_argument(
-            '--api-token',
+            '--local-model',
             type=str,
-            default=os.getenv('SIA_API_TOKEN'),
-            help='API access token (env: SIA_API_TOKEN)'
+            default=os.getenv('SIA_LOCAL_MODEL', '/root/model/'),
+            help='Path to local model directory (default: /root/model/, env: SIA_LOCAL_MODEL)'
         )
         parser.add_argument(
-            '--model',
-            type=str,
-            default=os.getenv('SIA_MODEL', '/root/model/'),
-            help='Path to the model directory (default: /root/model/, env: SIA_MODEL)'
-        )
-        parser.add_argument(
-            '--temperature',
+            '--local-temperature',
             type=float,
-            default=float(os.getenv('SIA_TEMPERATURE', '0.7')),
-            help='LLM temperature parameter (default: 0.7, env: SIA_TEMPERATURE)'
+            default=float(os.getenv('SIA_LOCAL_TEMPERATURE', '0.7')),
+            help='Local LLM temperature (default: 0.7, env: SIA_LOCAL_TEMPERATURE)'
         )
         parser.add_argument(
-            '--token-limit',
+            '--local-token-limit',
             type=int,
-            default=os.getenv('SIA_TOKEN_LIMIT'),
-            help='Token limit for the LLM (env: SIA_TOKEN_LIMIT)'
+            default=int(os.getenv('SIA_LOCAL_TOKEN_LIMIT', '2048')),
+            help='Local LLM token limit (env: SIA_LOCAL_TOKEN_LIMIT)'
         )
+
+        # OpenAI configuration
+        parser.add_argument(
+            '--openai-enable',
+            action='store_true',
+            default=self._parse_bool_env('SIA_OPENAI_ENABLED', False),
+            help='Enable OpenAI LLM engine (env: SIA_OPENAI_ENABLED)'
+        )
+        parser.add_argument(
+            '--openai-model',
+            type=str,
+            default=os.getenv('SIA_OPENAI_MODEL', 'gpt-3.5-turbo'),
+            help='OpenAI model name (default: gpt-3.5-turbo, env: SIA_OPENAI_MODEL)'
+        )
+        parser.add_argument(
+            '--openai-temperature',
+            type=float,
+            default=float(os.getenv('SIA_OPENAI_TEMPERATURE', '0.7')),
+            help='OpenAI temperature (default: 0.7, env: SIA_OPENAI_TEMPERATURE)'
+        )
+        parser.add_argument(
+            '--openai-token-limit',
+            type=int,
+            default=int(os.getenv('SIA_OPENAI_TOKEN_LIMIT', '4096')),
+            help='OpenAI token limit (env: SIA_OPENAI_TOKEN_LIMIT)'
+        )
+        parser.add_argument(
+            '--openai-api-key',
+            type=str,
+            default=os.getenv('SIA_OPENAI_API_KEY'),
+            help='OpenAI API key (env: SIA_OPENAI_API_KEY)'
+        )
+
+        # Hugging Face configuration
+        parser.add_argument(
+            '--hf-enable',
+            action='store_true',
+            default=self._parse_bool_env('SIA_HF_ENABLED', False),
+            help='Enable Hugging Face LLM engine (env: SIA_HF_ENABLED)'
+        )
+        parser.add_argument(
+            '--hf-model',
+            type=str,
+            default=os.getenv('SIA_HF_MODEL'),
+            help='Hugging Face model name (env: SIA_HF_MODEL)'
+        )
+        parser.add_argument(
+            '--hf-temperature',
+            type=float,
+            default=float(os.getenv('SIA_HF_TEMPERATURE', '0.7')),
+            help='Hugging Face temperature (default: 0.7, env: SIA_HF_TEMPERATURE)'
+        )
+        parser.add_argument(
+            '--hf-api-key',
+            type=str,
+            default=os.getenv('SIA_HF_API_KEY'),
+            help='Hugging Face API key (env: SIA_HF_API_KEY)'
+        )
+
+        # Mistral configuration
+        parser.add_argument(
+            '--mistral-enable',
+            action='store_true',
+            default=self._parse_bool_env('SIA_MISTRAL_ENABLED', False),
+            help='Enable Mistral LLM engine (env: SIA_MISTRAL_ENABLED)'
+        )
+        parser.add_argument(
+            '--mistral-model',
+            type=str,
+            default=os.getenv('SIA_MISTRAL_MODEL'),
+            help='Mistral model name (env: SIA_MISTRAL_MODEL)'
+        )
+        parser.add_argument(
+            '--mistral-temperature',
+            type=float,
+            default=float(os.getenv('SIA_MISTRAL_TEMPERATURE', '0.7')),
+            help='Mistral temperature (default: 0.7, env: SIA_MISTRAL_TEMPERATURE)'
+        )
+        parser.add_argument(
+            '--mistral-token-limit',
+            type=int,
+            default=int(os.getenv('SIA_MISTRAL_TOKEN_LIMIT', '4096')),
+            help='Mistral token limit (env: SIA_MISTRAL_TOKEN_LIMIT)'
+        )
+        parser.add_argument(
+            '--mistral-api-key',
+            type=str,
+            default=os.getenv('SIA_MISTRAL_API_KEY'),
+            help='Mistral API key (env: SIA_MISTRAL_API_KEY)'
+        )
+
         self.args = parser.parse_args()
-    
+
     def _parse_bool_env(self, env_var: str, default: bool) -> bool:
-        """Parse boolean environment variable."""
         val = os.getenv(env_var)
         if val is None:
             return default
         return val.lower() in ('true', '1', 'yes', 'on')
     
     def _parse_optional_path(self, env_var: str, default: Optional[Path]) -> Optional[Path]:
-        """Parse optional Path environment variable."""
         val = os.getenv(env_var)
         if val is None:
             return default
         return Path(val)
 
+    # Core properties
     @property
     def system_prompt(self) -> Path:
-        """Path to the system prompt file."""
         return self.args.system_prompt
     
     @property
     def action_schema(self) -> Path:
-        """Path to the action schema file."""
         return self.args.action_schema
     
     @property
     def iterations_dir(self) -> Path:
-        """Path to the directory for storing iterations."""
         return self.args.iterations_dir
 
+    # Server properties
     @property
     def server(self) -> bool:
-        """Enable web server for debugging and human feedback."""
         return self.args.server
 
     @property
     def host(self) -> str:
-        """Web server host."""
         return self.args.host
     
     @property
     def port(self) -> int:
-        """Web server port."""
         return self.args.port
     
     @property
     def static_files(self) -> Path:
-        """Path to static web files."""
         return self.args.static_files
 
+    # Local LLM properties
     @property
-    def llm_engine(self) -> str:
-        """LLM engine."""
-        return self.args.llm_engine
-    
-    @property
-    def api_token(self) -> Optional[str]:
-        """API access token."""
-        return self.args.api_token
+    def local_enabled(self) -> bool:
+        return self.args.local_enable
 
     @property
-    def model(self) -> str:
-        """Path to the model directory."""
-        return self.args.model
+    def local_model(self) -> str:
+        return self.args.local_model
 
     @property
-    def temperature(self) -> float:
-        """LLM temperature parameter."""
-        return self.args.temperature
-    
+    def local_temperature(self) -> float:
+        return self.args.local_temperature
+
     @property
-    def token_limit(self) -> int:
-        """Token limit for the LLM."""
-        return self.args.token_limit
+    def local_token_limit(self) -> int:
+        return self.args.local_token_limit
+
+    # OpenAI properties
+    @property
+    def openai_enabled(self) -> bool:
+        return self.args.openai_enable
+
+    @property
+    def openai_model(self) -> str:
+        return self.args.openai_model
+
+    @property
+    def openai_temperature(self) -> float:
+        return self.args.openai_temperature
+
+    @property
+    def openai_token_limit(self) -> int:
+        return self.args.openai_token_limit
+
+    @property
+    def openai_api_key(self) -> Optional[str]:
+        return self.args.openai_api_key
+
+    # Hugging Face properties
+    @property
+    def hf_enabled(self) -> bool:
+        return self.args.hf_enable
+
+    @property
+    def hf_model(self) -> str:
+        return self.args.hf_model
+
+    @property
+    def hf_temperature(self) -> float:
+        return self.args.hf_temperature
+
+    @property
+    def hf_api_key(self) -> Optional[str]:
+        return self.args.hf_api_key
+
+    # Mistral properties
+    @property
+    def mistral_enabled(self) -> bool:
+        return self.args.mistral_enable
+
+    @property
+    def mistral_model(self) -> str:
+        return self.args.mistral_model
+
+    @property
+    def mistral_temperature(self) -> float:
+        return self.args.mistral_temperature
+
+    @property
+    def mistral_token_limit(self) -> int:
+        return self.args.mistral_token_limit
+
+    @property
+    def mistral_api_key(self) -> Optional[str]:
+        return self.args.mistral_api_key

sia/mistral_llm_engine.py

@@ -13,13 +13,13 @@ class MistralLlmEngine(LlmEngine):
             self,
             model: str,
             temperature: float,
+            token_limit: int,
             api_key: str,
-            token_limit: int
         ):
         self._model = model
         self._temperature = temperature
-        self._api_key = api_key
         self._token_limit = token_limit
+        self._api_key = api_key
         self._client = Mistral(api_key=api_key)
         self._tokenizer = MistralTokenizer.v3()
 

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)

sia/web_agent.py

@@ -1,6 +1,7 @@
 from enum import Enum, auto
-from threading import Thread, Lock
-from typing import Callable, List, Optional
+from threading import Lock
+from typing import Callable, Dict, List, Optional
+from collections import defaultdict
 
 from .base_agent import BaseAgent
 from .command import Command
@@ -12,200 +13,148 @@ from .system_metrics import SystemMetrics
 from .working_memory import WorkingMemory
 from .xml_validator import XMLValidator
 
-class WebAgentState(Enum):
-    """
-    States for the web agent state machine.
-    """
-    UPDATE = auto()
-    """Updating system metrics and working memory entries"""
-    CONTEXT_APPROVAL = auto()
-    """Waiting for human approval of context"""
+class LlmState(Enum):
+    NO_OUTPUT = auto()
     INFERENCE = auto()
-    """Processing context through LLM"""
-    RESPONSE_APPROVAL = auto()
-    """Waiting for human approval of LLM response"""
-    STOPPED = auto()
+    OUTPUT = auto()
 
 class WebAgent(BaseAgent):
-    """
-    Agent implementation for interactive web interface.
-    
-    Uses a state machine to allow human intervention between steps.
-    Broadcasts state changes to registered handlers.
-    """
-    
     def __init__(
             self,
             system_prompt: str,
             action_schema: str, 
             working_memory: WorkingMemory,
             metrics: SystemMetrics,
-            llm: LlmEngine,
+            llms: Dict[str, LlmEngine],
             validator: XMLValidator,
             parser: ResponseParser,
             iteration_logger: IterationLogger,
         ):
-        """
-        Initialize web agent with required components.
-        """
         super().__init__(
             system_prompt,
             action_schema,
             working_memory,
-            metrics, llm,
+            metrics,
             validator,
             parser
         )
-        self._response = ""
-        self._state = WebAgentState.CONTEXT_APPROVAL
-        self._validation_error: Optional[str] = None
-        self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
-        self._context_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()
-        self._state_lock = Lock()
+        self._llms = llms
         self._iteration_logger = iteration_logger
+        self._llm_states: Dict[str, LlmState] = {name: LlmState.NO_OUTPUT for name in llms}
+        self._llm_outputs: Dict[str, str] = defaultdict(str)
+        self._validation_error: Optional[str] = None
+        self._command_result: Optional[CommandResult] = None
+        self._context = self._compile_context(next(iter(self._llms.values())))
+        
+        # Locks
+        self._llm_lock = Lock()
+        self._output_lock = Lock()
+        
+        # Event handlers
+        self._llm_change_handlers: List[Callable[[str, LlmState], None]] = []
+        self._token_handlers: List[Callable[[str, str], None]] = []
+        self._context_change_handlers: List[Callable[[str, bool], None]] = []
 
     @property
-    def state(self) -> WebAgentState:
-        """Get the current state of the agent."""
-        return self._state
-    
+    def llms(self) -> Dict[str, LlmState]:
+        """Get current state of all LLMs"""
+        with self._llm_lock:
+            return self._llm_states.copy()
+
+    @property
+    def context(self) -> str:
+        return self._context
+
     @property
     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
-    
-    @property
-    def context(self) -> str:
-        """Get the current context."""
-        return self._context
-        
-    def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
-        """
-        Add a callback for state changes.
-        
-        Args:
-            handler: Function to call with new state
-        """
-        if handler not in self._state_change_handlers:
-            self._state_change_handlers.append(handler)
-    
-    def add_context_change_handler(self, handler: Callable[[str], None]) -> None:
-        """
-        Add a callback for context changes.
 
-        Args:
-            handler: Function to call with new context
-        """
+    def add_llm_change_handler(self, handler: Callable[[str, LlmState], None]) -> None:
+        """Add handler for LLM state changes"""
+        if handler not in self._llm_change_handlers:
+            self._llm_change_handlers.append(handler)
+
+    def add_token_handler(self, handler: Callable[[str, str], None]) -> None:
+        """Add handler for new tokens"""
+        if handler not in self._token_handlers:
+            self._token_handlers.append(handler)
+
+    def add_context_change_handler(self, handler: Callable[[str, bool], None]) -> None:
+        """Add handler for context changes"""
         if handler not in self._context_change_handlers:
             self._context_change_handlers.append(handler)
-    
-    def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
-        """
-        Add a callback for response changes.
 
-        Args:
-            handler: Function to call with new response
-        """
-        if handler not in self._response_change_handlers:
-            self._response_change_handlers.append(handler)
-    
-    def modify_context(self, context: str) -> None:
-        with self._state_lock:
-            if self._state != WebAgentState.CONTEXT_APPROVAL:
-                error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
-                raise Exception(error_msg)
-        Thread(target=self._set_context, args=(context,)).start()
-    
-    def modify_response(self, response: str) -> None:
-        started = False
-        try:
-            self._state_lock.acquire()
-            if self._state != WebAgentState.RESPONSE_APPROVAL:
-                error_msg = f"Not in RESPONSE_APPROVAL state (current: {self._state})"
-                raise Exception(error_msg)
-            Thread(target=self._set_response, args=(response, self._state_lock)).start()
-            started = True
-        finally:
-            if not started:
-                self._state_lock.release()
+    def modify_context(self, context: str, generated: bool = False) -> None:
+        """Update context and reset all LLM states"""
+        with self._llm_lock:
+            self._context = context
+            self._llm_outputs.clear()
+            for llm_name in self._llms:
+                self._set_llm_state(llm_name, LlmState.NO_OUTPUT)
             
-    def approve_context(self) -> None:
-        with self._state_lock:
-            if self._state != WebAgentState.CONTEXT_APPROVAL:
-                error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
-                raise Exception(error_msg)
-            self._set_state(WebAgentState.INFERENCE)
-        Thread(target=self._approve_context_thread).start()
+            for handler in self._context_change_handlers:
+                handler(context, generated)
+
+    def run_inference(self, llm_name: str) -> None:
+        """Start inference on specified LLM"""
+        if llm_name not in self._llms:
+            raise ValueError(f"Unknown LLM: {llm_name}")
+
+        with self._llm_lock:
+            if self._llm_states[llm_name] != LlmState.NO_OUTPUT:
+                raise RuntimeError(f"LLM {llm_name} is not ready for inference")
+            self._set_llm_state(llm_name, LlmState.INFERENCE)
+
+        llm = self._llms[llm_name]
+        response_token_iter = llm.infer(self.system_prompt, self.context)
+        
+        with self._output_lock:
+            self._llm_outputs[llm_name] = ""
 
-    def approve_response(self) -> None:
-        with self._state_lock:
-            if self._state != WebAgentState.RESPONSE_APPROVAL:
-                raise Exception("Not in RESPONSE_APPROVAL state")
-            self._set_state(WebAgentState.UPDATE)
-            Thread(target=self._approve_response_thread).start()
-                
-    def _approve_context_thread(self) -> None:
-        self._set_response("")
-        response_token_iter = self._llm.infer(self.system_prompt, self.context)
-        response = ""
         for token in response_token_iter:
-            response += token
-            self._set_response(response)
-            print(f"{token}", end='')
-        self._set_state(WebAgentState.RESPONSE_APPROVAL)
-        print()
-    
-    def _approve_response_thread(self) -> None:
-        self._iteration_logger.log_iteration(self._context, self._response)
-        parse_result = self._parser.parse(self._response)
+            with self._output_lock:
+                self._llm_outputs[llm_name] += token
+            
+            for handler in self._token_handlers:
+                handler(llm_name, token)
+
+        with self._llm_lock:
+            self._set_llm_state(llm_name, LlmState.OUTPUT)
+
+    def get_output(self, llm_name: str) -> str:
+        """Get complete output for specified LLM"""
+        if llm_name not in self._llms:
+            raise ValueError(f"Unknown LLM: {llm_name}")
+            
+        with self._output_lock:
+            return self._llm_outputs[llm_name]
+
+    def approve_response(self, llm_name: str, response: str) -> None:
+        """Process approved response from specified LLM"""
+        if llm_name not in self._llms:
+            raise ValueError(f"Unknown LLM: {llm_name}")
+
+        self._iteration_logger.log_iteration(self._context, response)
+        
+        parse_result = self._parser.parse(response)
         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
-            self._working_memory.update()
+            if not result.should_stop:
+                self._working_memory.update()
         else:
             parse_result.update()
             self._working_memory.update()
             self._working_memory.add_entry(parse_result)
-        
-        self._set_context(self._compile_context())
-        self._set_state(WebAgentState.CONTEXT_APPROVAL)
             
-    def _set_state(self, state) -> None:
-        """Notify all handlers of state change."""
-        self._state = state
-        for handler in self._state_change_handlers:
-            handler(self._state)
-    
-    def _set_context(self, context) -> None:
-        """Notify all handlers of context change."""
-        self._context = context
-        for handler in self._context_change_handlers:
-            handler(context)
-    
-    def _set_response(self, response: str, lock: Optional[Lock] = None) -> None:
-        """Set response and notify handlers."""
-        try:
-            self._response = response
-            validation_error = self._validator.validate(response)
-            self._validation_error = validation_error
-            for handler in self._response_change_handlers:
-                handler(response, validation_error)
-        finally:
-            if lock is not None:
-                lock.release()
+        self.modify_context(self._compile_context(self._llms[llm_name]), True)
+
+    def _set_llm_state(self, llm_name: str, state: LlmState) -> None:
+        """Update LLM state and notify handlers"""
+        self._llm_states[llm_name] = state
+        for handler in self._llm_change_handlers:
+            handler(llm_name, state)

sia/web_socket_manager.py

@@ -43,7 +43,7 @@ class WebSocketManager:
         self._clients: Set[web.WebSocketResponse] = set()
         self.auto_approver = AutoApprover(agent)
 
-        self.agent.add_state_change_handler(self._wrap_async(self._handle_state_change))
+        self.agent.add_llm_change_handler(self._wrap_async(self._handle_state_change))
         self.agent.add_context_change_handler(self._wrap_async(self._handle_context_change))
         self.agent.add_response_change_handler(self._wrap_async(self._handle_response_change))
         self.io_buffer.add_stdout_change_handler(self._wrap_async(self._handle_stdout_change))