Enable multiple llms

2024-11-22 15:05:54 +01:00
parent bbb88f39e8
commit 8766a945c0
28 changed files with 972 additions and 525 deletions
--- a/run.sh
+++ b/run.sh
@@ -1,5 +1,7 @@
 #!/bin/bash

+set -e
+
 docker build \
    --tag sia \
    .
--- a/sia/main.py
+++ b/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()
--- a/sia/auto_approver.py
+++ b/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:
--- a/sia/base_agent.py
+++ b/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
--- a/sia/config.py
+++ b/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
+    def local_enabled(self) -> bool:
+        return self.args.local_enable

    @property
-    def api_token(self) -> Optional[str]:
-        """API access token."""
-        return self.args.api_token
+    def local_model(self) -> str:
+        return self.args.local_model

    @property
-    def model(self) -> str:
-        """Path to the model directory."""
-        return self.args.model
+    def local_temperature(self) -> float:
+        return self.args.local_temperature

    @property
-    def temperature(self) -> float:
-        """LLM temperature parameter."""
-        return self.args.temperature
+    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 token_limit(self) -> int:
-        """Token limit for the LLM."""
-        return self.args.token_limit
+    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
--- a/sia/mistral_llm_engine.py
+++ b/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()

--- a/sia/web/api.py
+++ b/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"
+        )
--- a/sia/web/context_websocket.py
+++ b/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
--- a/sia/web/llm_websocket.py
+++ b/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
--- a/sia/web/static.py
+++ b/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"
+        )
--- a/sia/web/stdout_websocket.py
+++ b/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
--- a/sia/web/token_websocket.py
+++ b/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
--- a/sia/web/util.py
+++ b/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
--- a/sia/web/websockts.py
+++ b/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)
--- a/sia/web_agent.py
+++ b/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
+    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)

-    @property
-    def context(self) -> str:
-        """Get the current context."""
-        return self._context
+    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_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_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.
+    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)
            
-        Args:
-            handler: Function to call with new response
-        """
-        if handler not in self._response_change_handlers:
-            self._response_change_handlers.append(handler)
+            for handler in self._context_change_handlers:
+                handler(context, generated)

-    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 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}")

-    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()
+        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)

-    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()
+        llm = self._llms[llm_name]
+        response_token_iter = llm.infer(self.system_prompt, self.context)
        
-    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()
+        with self._output_lock:
+            self._llm_outputs[llm_name] = ""

-    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()
+            with self._output_lock:
+                self._llm_outputs[llm_name] += token
            
-    def _approve_response_thread(self) -> None:
-        self._iteration_logger.log_iteration(self._context, self._response)
-        parse_result = self._parser.parse(self._response)
+            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)
+        self.modify_context(self._compile_context(self._llms[llm_name]), True)

-    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()
+    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)
--- a/sia/web_socket_manager.py
+++ b/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))
--- a/test/web_agent_test.py
+++ b/test/web_agent_test.py
@@ -95,22 +95,22 @@ class WebAgentTest(unittest.TestCase):
        self.assertEqual(self.agent.response, "")
        self.assertIsNotNone(self.agent.context)
        self.assertIsNone(self.agent.validation_error)
-        self.assertEqual(len(self.agent._state_change_handlers), 0)
+        self.assertEqual(len(self.agent._llm_change_handlers), 0)
        self.assertEqual(len(self.agent._response_change_handlers), 0)
        
    def test_handler_registration(self):
        """Test adding state and response change handlers."""
-        self.agent.add_state_change_handler(self.state_change_handler)
+        self.agent.add_llm_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
-        self.assertEqual(len(self.agent._state_change_handlers), 1)
+        self.assertEqual(len(self.agent._llm_change_handlers), 1)
        self.assertEqual(len(self.agent._response_change_handlers), 1)
        
        # Adding same handler twice should not duplicate
-        self.agent.add_state_change_handler(self.state_change_handler)
+        self.agent.add_llm_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
-        self.assertEqual(len(self.agent._state_change_handlers), 1)
+        self.assertEqual(len(self.agent._llm_change_handlers), 1)
        self.assertEqual(len(self.agent._response_change_handlers), 1)
        
    def test_approve_context_state_error(self):
@@ -129,7 +129,7 @@ class WebAgentTest(unittest.TestCase):
        
    def test_context_approval_flow(self):
        """Test complete context approval flow with state transitions."""
-        self.agent.add_state_change_handler(self.state_change_handler)
+        self.agent.add_llm_change_handler(self.state_change_handler)
        self.agent.add_response_change_handler(self.response_change_handler)
        
        self.agent.approve_context()
@@ -143,7 +143,7 @@ class WebAgentTest(unittest.TestCase):
        
    def test_response_approval_flow_command(self):
        """Test response approval flow with command execution."""
-        self.agent.add_state_change_handler(self.state_change_handler)
+        self.agent.add_llm_change_handler(self.state_change_handler)
        
        command_result = CommandResult.success()
        mock_command = MockCommand(command_result)
@@ -165,7 +165,7 @@ class WebAgentTest(unittest.TestCase):
        
    def test_response_approval_flow_entry(self):
        """Test response approval flow with entry creation."""
-        self.agent.add_state_change_handler(self.state_change_handler)
+        self.agent.add_llm_change_handler(self.state_change_handler)
        
        self.agent._state = WebAgentState.RESPONSE_APPROVAL
        self.agent._set_response("<reasoning>test reasoning</reasoning>")
--- a/web/package.json
+++ b/web/package.json
@@ -14,6 +14,7 @@
    "@monaco-editor/react": "^4.6.0",
    "@radix-ui/react-alert-dialog": "^1.0.5",
    "@radix-ui/react-scroll-area": "^1.0.5",
+    "@radix-ui/react-select": "^2.0.0",
    "@radix-ui/react-tabs": "^1.0.4",
    "class-variance-authority": "^0.7.0",
    "clsx": "^2.1.0",
--- a/web/src/components/App.jsx
+++ b/web/src/components/App.jsx
@@ -1,127 +1,187 @@
 import React, { useState, useEffect } from 'react';
-import { Header } from '@/components/Header';
+import { Header, Tabs } from '@/components/Header';
 import { Sidebar } from '@/components/Sidebar';
 import { FloatingButton } from '@/components/FloatingButton';
 import { ContentEditor } from '@/components/editors/ContentEditor';
 import { StandardEditor } from '@/components/editors/StandardEditor';
-import { useWebSocket } from '@/hooks/useWebSocket';
-import { AgentState, Tabs, ClientMessageType, ServerMessageType } from '@/constants';
+import { useWebSocket, WebSocketState } from '@/hooks/useWebSocket';

+export const LlmState = {
+  NO_OUTPUT: 'NO_OUTPUT',
+  INFERENCE: 'INFERENCE', 
+  OUTPUT: 'OUTPUT'
+};

 const App = () => {
  // Editor content state
-  const [originalContext, setOriginalContext] = useState('');
+  const [generatedContext, setGeneratedContext] = useState('');
  const [modifiedContext, setModifiedContext] = useState('');
-  const [originalResponse, setOriginalResponse] = useState('');
-  const [modifiedResponse, setModifiedResponse] = useState('');
-  const [validationError, setValidationError] = useState(null);
+  const [contextDirty, setContextDirty] = useState(false);
+  const [generatedResponses, setGeneratedResponses] = useState({});
+  const [modifiedResponses, setModifiedResponses] = useState({});
  const [input, setInput] = useState('');
  const [output, setOutput] = useState('');
  
+  // LLM state
+  const [llms, setLlms] = useState({});
+  const [activeLlm, setActiveLlm] = useState(null);
+  
  // UI state
  const [activeTab, setActiveTab] = useState(Tabs.CONTEXT);
  const [showDiff, setShowDiff] = useState(false);
  const [showSidebar, setShowSidebar] = useState(false);
-  const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
-  const [autoApproverConfig, setAutoApproverConfig] = useState({
-    context_enabled: false,
-    response_enabled: false,
-    context_timeout: 5.0,
-    response_timeout: 10.0
-  });

-  const { wsState, sendMessage, addMessageHandler } = useWebSocket(`${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws`);
+  // WebSocket connections
+  const wsRoot = `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws`
+  const llmWs = useWebSocket(`${wsRoot}/llm`);
+  const contextWs = useWebSocket(`${wsRoot}/context`);
+  const tokenWs = useWebSocket(`${wsRoot}/token`);
+  const stdoutWs = useWebSocket(`${wsRoot}/stdout`);
+  const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
+
+  // Handle llm state changes
  useEffect(() => {
-    addMessageHandler((message) => {
-      switch(message.type) {
-        case ServerMessageType.STATE_CHANGE:
-          setAgentState(message.state);
-          break;
-        case ServerMessageType.CONTEXT_UPDATE:
-          setOriginalContext(message.context);
-          setModifiedContext(message.context);
-          break;
-        case ServerMessageType.RESPONSE_UPDATE:
-          setOriginalResponse(message.response);
-          setModifiedResponse(message.response);
-          setValidationError(message.validation_error);
-          break;
-        case ServerMessageType.OUTPUT_UPDATE:
-          setOutput(message.output);
-          break;
-        case ServerMessageType.AUTO_APPROVER_CONFIG:
-          setAutoApproverConfig(message.config);
-          break;
+    llmWs.addMessageHandler((data) => {
+      setLlms(prevLlms => ({
+        ...prevLlms,
+        [data.llm]: data.state
+      }));
+      if (data.state === LlmState.NO_OUTPUT) {
+        setGeneratedResponses(prev => {
+          const updated = { ...prev };
+          delete updated[data.llm];
+          return updated;
+        });
+        setModifiedResponses(prev => {
+          const updated = { ...prev };
+          delete updated[data.llm];
+          return updated;
+        });
+      }
+      if (activeLlm === null) {
+        setActiveLlm(data.llm);
      } 
    });
-  }, [addMessageHandler]);
+  }, []);

-  const handleAutoApproverConfigChange = (config) => {
-    sendMessage({
-      type: ClientMessageType.AUTO_APPROVER_CONFIG,
-      config: config
+  // Handle context changes
+  useEffect(() => {
+    contextWs.addMessageHandler((data) => {
+      setContextDirty(false);
+      setModifiedContext(data.context);
+      if (data.generated) {
+        setGeneratedContext(data.context);
+      }
    });
+  }, []);
+
+  // Handle incoming tokens
+  useEffect(() => {
+    tokenWs.addMessageHandler((data) => {
+      setGeneratedResponses(prev => ({
+        ...prev,
+        [data.llm]: (prev[data.llm] || '') + data.token
+      }));
+      setModifiedResponses(prev => ({
+        ...prev,
+        [data.llm]: (prev[data.llm] || '') + data.token
+      }));
+    });
+  }, []);
+
+  // Handle stdout changes
+  useEffect(() => {
+    stdoutWs.addMessageHandler((data) => {
+      setOutput(data.output);
+    });
+  }, []);
+
+  const handleNextLlm = () => {
+    const llmNames = Object.keys(llms);
+    const currentIndex = llmNames.indexOf(activeLlm);
+    const nextIndex = (currentIndex + 1) % llmNames.length;
+    setActiveLlm(llmNames[nextIndex]);
  };

  const handleApprove = () => {
-    if (agentState === AgentState.CONTEXT_APPROVAL) {
-      sendMessage({
-        "type": ClientMessageType.APPROVE_CONTEXT,
-        "context": modifiedContext
+    const state = llms[activeLlm];
+    if (state === LlmState.NO_OUTPUT) {
+      if (contextDirty) {
+        fetch('/api/context', {
+          method: 'POST',
+          headers: { 'Content-Type': 'application/json' },
+          body: JSON.stringify({ context: modifiedContext })
+        });
+      }
+      fetch(`/api/inference/${activeLlm}`, {
+        method: 'POST',
      });
-    } else if (agentState === AgentState.RESPONSE_APPROVAL) {
-      sendMessage({
-        "type": ClientMessageType.APPROVE_RESPONSE,
-        "response": modifiedResponse
+    } else if (state === LlmState.OUTPUT) {
+      fetch(`/api/approve/${activeLlm}`, {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ response: modifiedResponses[activeLlm] })
      });
    }
  };

  const handleSendInput = () => {
    if (input.trim()) {
-      sendMessage({
-        "type": ClientMessageType.SEND_INPUT,
-        "input": input
+      fetch('/api/input', {
+        method: 'POST',
+        headers: { 'Content-Type': 'application/json' },
+        body: JSON.stringify({ input: input })
      });
      setInput('');
    }
  };

  const handleClearOutput = () => {
-    sendMessage({
-      "type": ClientMessageType.CLEAR_OUTPUT
-    });
+    fetch('/api/clear', { method: 'POST' });
  };

-  useEffect(() => {
-    if (agentState === AgentState.APPROVE_RESPONSE) {
-      sendMessage({
-        "type": ClientMessageType.MODIFY_RESPONSE,
-        "response": modifiedResponse
-      });
-    }
-  }, [modifiedResponse]);
+  const handleContextEdit = (context) => {
+    setContextDirty(true);
+    setGeneratedResponses({});
+    setModifiedResponses({});
+    const resetLlms = {};
+    for (const llm of Object.keys(llms)) {
+      resetLlms[llm] = LlmState.NO_OUTPUT;
+    }    setLlms(resetLlms);
+    setModifiedContext(context);
+  }

  useEffect(() => {
-    switch (agentState) {
-      case AgentState.UPDATE: // fall through
-      case AgentState.CONTEXT_APPROVAL:
+    const state = llms[activeLlm];
+    switch (state) {
+      case LlmState.NO_OUTPUT:
        setActiveTab(Tabs.CONTEXT);
+        setShowDiff(false);
+        setShowSidebar(false);
        break;
-      case AgentState.INFERENCE: // fall through
-      case AgentState.RESPONSE_APPROVAL:
+      default:
        setActiveTab(Tabs.RESPONSE);
+        setShowDiff(false);
+        setShowSidebar(false);
        break;
    }
-    setShowDiff(false);
-    setShowSidebar(false);
-  }, [agentState]);
+  }, [llms[activeLlm]]);
+
+  useEffect(() => {
+    if (llmWs.wsState === WebSocketState.CONNECTED && contextWs.wsState === WebSocketState.CONNECTED && tokenWs.wsState === WebSocketState.CONNECTED && stdoutWs.wsState === WebSocketState.CONNECTED) {
+      setWsState(WebSocketState.CONNECTED);
+    } else if (llmWs.wsState === WebSocketState.CONNECTING || contextWs.wsState === WebSocketState.CONNECTING || tokenWs.wsState === WebSocketState.CONNECTING || stdoutWs.wsState === WebSocketState.CONNECTING) {
+      setWsState(WebSocketState.CONNECTING);
+    } else {
+      setWsState(WebSocketState.DISCONNECTED);
+    }
+  }, [llmWs.wsState, contextWs.wsState, tokenWs.wsState, stdoutWs.wsState]);

  return (
    <div className="h-screen bg-gray-100 flex flex-col overflow-hidden">
      <Header
        wsState={wsState}
-        agentState={agentState}
+        agentState={llms[activeLlm]}
        activeTab={activeTab}
        setActiveTab={setActiveTab}
      />
@@ -131,19 +191,22 @@ const App = () => {
          {activeTab === Tabs.CONTEXT && (
            <ContentEditor
              showDiff={showDiff}
-              originalContent={originalContext}
+              originalContent={generatedContext}
              modifiedContent={modifiedContext}
-              onChange={value => setModifiedContext(value)}
+              onChange={handleContextEdit}
              validationError={null}
            />
          )}
          {activeTab === Tabs.RESPONSE && (
            <ContentEditor
              showDiff={showDiff}
-              originalContent={originalResponse}
-              modifiedContent={modifiedResponse}
-              onChange={value => setModifiedResponse(value)}
-              validationError={validationError}
+              originalContent={generatedResponses[activeLlm] || ''}
+              modifiedContent={modifiedResponses[activeLlm] || ''}
+              onChange={value => setModifiedResponses(prev => ({
+                ...prev,
+                [activeLlm]: value
+              }))}
+              validationError={null}
            />
          )}
          {activeTab === Tabs.INPUT && (
@@ -164,17 +227,18 @@ const App = () => {

        <Sidebar
          showSidebar={showSidebar}
-          activeTab={activeTab}
-          agentState={agentState}
-          validationError={validationError}
+          llms={llms}
+          activeLlm={activeLlm}
+          llmState={llms[activeLlm]}
          showDiff={showDiff}
          input={input}
+          output={output}
          onApprove={handleApprove}
          onToggleDiff={() => setShowDiff(!showDiff)}
          onSendInput={handleSendInput}
          onClearOutput={handleClearOutput}
-          autoApproverConfig={autoApproverConfig}
-          onAutoApproverConfigChange={handleAutoApproverConfigChange}
+          onLlmChange={setActiveLlm}
+          onNextLlm={handleNextLlm}
        />

        <FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
--- a/web/src/components/Header.jsx
+++ b/web/src/components/Header.jsx
@@ -1,5 +1,12 @@
 import React from 'react';
-import { WebSocketState, AgentState, Tabs } from '../constants';
+import { WebSocketState } from '../hooks/useWebSocket';
+
+export const Tabs = {
+  CONTEXT: 'CONTEXT',
+  RESPONSE: 'RESPONSE',
+  INPUT: 'INPUT',
+  OUTPUT: 'OUTPUT'
+};

 export const Header = ({ 
  wsState, 
--- a/web/src/components/Sidebar.jsx
+++ b/web/src/components/Sidebar.jsx
@@ -1,71 +1,50 @@
 import React from 'react';
 import { Button } from '@/components/ui/button';
-import { Switch } from '@/components/ui/switch';
-import { Input } from '@/components/ui/input';
-import { AgentState, Tabs } from '../constants';
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
+import { LlmState } from '@/components/App';

 export const Sidebar = ({
  showSidebar,
-  agentState,
+  llms,
+  activeLlm,
+  llmState,
  showDiff,
  input,
-  autoApproverConfig,
+  output,
  onApprove,
  onToggleDiff,
  onSendInput,
  onClearOutput,
-  onAutoApproverConfigChange,
+  onLlmChange,
+  onNextLlm,
 }) => (
  <aside className={`
    fixed top-0 right-0 h-screen w-64 bg-white shadow-lg 
    transition-transform duration-200 ease-in-out z-50
-    md:sticky md:h-[calc(100vh-4rem)] md:top-16
+    md:sticky md:h-[calc(100vh-4rem)]
    ${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
  `}>
-    <div className="p-4 space-y-4 md:pt-0">
-      <div className="space-y-2">
-        <div className="flex items-center justify-between">
-          <label>Auto-approve Context</label>
-          <Switch 
-            checked={autoApproverConfig.context_enabled}
-            onCheckedChange={(checked) => onAutoApproverConfigChange({
-              ...autoApproverConfig,
-              context_enabled: checked
-            })}
-          />
-        </div>
-        <Input
-          type="number"
-          value={autoApproverConfig.context_timeout}
-          onChange={(e) => onAutoApproverConfigChange({
-            ...autoApproverConfig,
-            context_timeout: parseFloat(e.target.value)
-          })}
-          disabled={autoApproverConfig.context_enabled}
-        />
-      </div>
+    <div className="p-4 space-y-4 overflow-y-auto h-full">
+      <Select value={activeLlm} onValueChange={onLlmChange}>
+        <SelectTrigger>
+          <SelectValue placeholder="Select LLM" />
+        </SelectTrigger>
+        <SelectContent>
+          {Object.keys(llms).map(llm => (
+            <SelectItem key={llm} value={llm}>
+              {llm} ({llms[llm]})
+            </SelectItem>
+          ))}
+        </SelectContent>
+      </Select>

-      <div className="space-y-2">
-        <div className="flex items-center justify-between">
-          <label>Auto-approve Response</label>
-          <Switch
-            checked={autoApproverConfig.response_enabled}
-            onCheckedChange={(checked) => onAutoApproverConfigChange({
-              ...autoApproverConfig,
-              response_enabled: checked
-            })}
-          />
-        </div>
-        <Input
-          type="number"
-          value={autoApproverConfig.response_timeout}
-          onChange={(e) => onAutoApproverConfigChange({
-            ...autoApproverConfig,
-            response_timeout: parseFloat(e.target.value)
-          })}
-          disabled={autoApproverConfig.response_enabled}
-        />
-      </div>
+      <Button
+        variant="outline"
+        onClick={onNextLlm}
+        className="w-full"
+      >
+        Next LLM
+      </Button>
      
      <Button
        variant="outline"
@@ -77,10 +56,14 @@ export const Sidebar = ({

      <Button 
        onClick={onApprove}
-        disabled={(agentState !== AgentState.CONTEXT_APPROVAL) && (agentState !== AgentState.RESPONSE_APPROVAL)}
+        disabled={llmState === LlmState.INFERENCE}
        className="w-full"
      >
-        Approve
+        {
+          llmState === LlmState.NO_OUTPUT ? 'Inference' :
+          llmState === LlmState.OUTPUT ? 'Approve' :
+          'In Progress'
+        }
      </Button>

      <Button 
@@ -93,7 +76,7 @@ export const Sidebar = ({

      <Button 
        onClick={onClearOutput}
-        variant="outline"
+        disabled={!output.trim()}
        className="w-full"
      >
        Clear Output
--- a/web/src/components/editors/ContextEditor.jsx
+++ b/web/src/components/editors/ContextEditor.jsx
@@ -1,15 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const ContextEditor = ({
-  originalContent,
-  modifiedContent,
-  showOriginal,
-  onChange,
-}) => (
-  <BaseEditor
-    content={showOriginal ? originalContent : modifiedContent}
-    onChange={onChange}
-    readOnly={showOriginal}
-  />
-);
--- a/web/src/components/editors/InputEditor.jsx
+++ b/web/src/components/editors/InputEditor.jsx
@@ -1,10 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const InputEditor = ({ content, onChange }) => (
-  <BaseEditor
-    content={content}
-    onChange={onChange}
-    language="plaintext"
-  />
-);
--- a/web/src/components/editors/OutputEditor.jsx
+++ b/web/src/components/editors/OutputEditor.jsx
@@ -1,10 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const OutputEditor = ({ content }) => (
-  <BaseEditor
-    content={content}
-    readOnly={true}
-    language="plaintext"
-  />
-);
--- a/web/src/components/editors/ResponseEditor.jsx
+++ b/web/src/components/editors/ResponseEditor.jsx
@@ -1,17 +0,0 @@
-import React from 'react';
-import { BaseEditor } from './BaseEditor';
-
-export const ResponseEditor = ({
-  originalContent,
-  modifiedContent,
-  showOriginal,
-  onChange,
-  validationError,
-}) => (
-  <BaseEditor
-    content={showOriginal ? originalContent : modifiedContent}
-    onChange={onChange}
-    readOnly={showOriginal}
-    validationError={validationError}
-  />
-);
--- a/web/src/components/ui/select.jsx
+++ b/web/src/components/ui/select.jsx
@@ -0,0 +1,52 @@
+import * as React from "react"
+import * as SelectPrimitive from "@radix-ui/react-select"
+import { Check, ChevronDown } from "lucide-react"
+
+const Select = SelectPrimitive.Root
+
+const SelectTrigger = React.forwardRef(({ className, children, ...props }, ref) => (
+  <SelectPrimitive.Trigger
+    ref={ref}
+    className="flex h-10 w-full items-center justify-between rounded-md border border-input bg-background px-3 py-2 text-sm ring-offset-background placeholder:text-muted-foreground focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50"
+    {...props}>
+    {children}
+    <SelectPrimitive.Icon asChild>
+      <ChevronDown className="h-4 w-4 opacity-50" />
+    </SelectPrimitive.Icon>
+  </SelectPrimitive.Trigger>
+))
+SelectTrigger.displayName = "SelectTrigger"
+
+const SelectContent = React.forwardRef(({ className, children, position = "popper", ...props }, ref) => (
+  <SelectPrimitive.Portal>
+    <SelectPrimitive.Content
+      ref={ref}
+      className="relative z-50 min-w-[8rem] overflow-hidden rounded-md border bg-popover text-popover-foreground shadow-md data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2"
+      position={position}
+      {...props}>
+      <SelectPrimitive.Viewport className="p-1">
+        {children}
+      </SelectPrimitive.Viewport>
+    </SelectPrimitive.Content>
+  </SelectPrimitive.Portal>
+))
+SelectContent.displayName = "SelectContent"
+
+const SelectItem = React.forwardRef(({ className, children, ...props }, ref) => (
+  <SelectPrimitive.Item
+    ref={ref}
+    className="relative flex w-full cursor-default select-none items-center rounded-sm py-1.5 pl-8 pr-2 text-sm outline-none focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50"
+    {...props}>
+    <span className="absolute left-2 flex h-3.5 w-3.5 items-center justify-center">
+      <SelectPrimitive.ItemIndicator>
+        <Check className="h-4 w-4" />
+      </SelectPrimitive.ItemIndicator>
+    </span>
+    <SelectPrimitive.ItemText>{children}</SelectPrimitive.ItemText>
+  </SelectPrimitive.Item>
+))
+SelectItem.displayName = "SelectItem"
+
+const SelectValue = SelectPrimitive.Value
+
+export { Select, SelectContent, SelectItem, SelectTrigger, SelectValue }
--- a/web/src/constants.jsx
+++ b/web/src/constants.jsx
@@ -1,36 +0,0 @@
-export const WebSocketState = {
-  CONNECTING: 'CONNECTING',
-  CONNECTED: 'CONNECTED',
-  DISCONNECTED: 'DISCONNECTED',
-};
-
-export const AgentState = {
-  UPDATE: 'UPDATE',
-  CONTEXT_APPROVAL: 'CONTEXT_APPROVAL',
-  INFERENCE: 'INFERENCE',
-  RESPONSE_APPROVAL: 'RESPONSE_APPROVAL',
-};
-
-export const Tabs = {
-  CONTEXT: 'context',
-  RESPONSE: 'response',
-  INPUT: 'input',
-  OUTPUT: 'output'
-};
-
-export const ClientMessageType = {
-  APPROVE_CONTEXT: 'APPROVE_CONTEXT',
-  MODIFY_RESPONSE: 'MODIFY_RESPONSE',
-  APPROVE_RESPONSE: 'APPROVE_RESPONSE',
-  SEND_INPUT: 'SEND_INPUT',
-  CLEAR_OUTPUT: 'CLEAR_OUTPUT',
-  AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
-};
-
-export const ServerMessageType = {
-  STATE_CHANGE: 'STATE_CHANGE',
-  CONTEXT_UPDATE: 'CONTEXT_UPDATE',
-  RESPONSE_UPDATE: 'RESPONSE_UPDATE',
-  OUTPUT_UPDATE: 'OUTPUT_UPDATE',
-  AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
-};
--- a/web/src/hooks/useWebSocket.jsx
+++ b/web/src/hooks/useWebSocket.jsx
@@ -1,5 +1,10 @@
 import { useState, useEffect, useRef, useCallback } from 'react';
-import { WebSocketState, MessageType } from '../constants';
+
+export const WebSocketState = {
+  CONNECTING: 'CONNECTING',
+  CONNECTED: 'CONNECTED',
+  DISCONNECTED: 'DISCONNECTED',
+};

 export const useWebSocket = (url) => {
  const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);