diff --git a/run.sh b/run.sh index 1cd7320..bc52093 100755 --- a/run.sh +++ b/run.sh @@ -1,5 +1,7 @@ #!/bin/bash +set -e + docker build \ --tag sia \ . diff --git a/sia/__main__.py b/sia/__main__.py index 7742555..2fb7336 100644 --- 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() diff --git a/sia/auto_approver.py b/sia/auto_approver.py index dd3b4d8..898875a 100644 --- 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: diff --git a/sia/base_agent.py b/sia/base_agent.py index 00ea685..b115919 100644 --- 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 diff --git a/sia/config.py b/sia/config.py index 517967f..812cb0f 100644 --- 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 - - @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 \ No newline at end of file + 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 diff --git a/sia/mistral_llm_engine.py b/sia/mistral_llm_engine.py index d054f8d..9b4aeef 100644 --- 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() diff --git a/sia/web/api.py b/sia/web/api.py new file mode 100644 index 0000000..4b8619e --- /dev/null +++ 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" + ) diff --git a/sia/web/context_websocket.py b/sia/web/context_websocket.py new file mode 100644 index 0000000..f2ba962 --- /dev/null +++ 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 diff --git a/sia/web/llm_websocket.py b/sia/web/llm_websocket.py new file mode 100644 index 0000000..ccd9167 --- /dev/null +++ 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 \ No newline at end of file diff --git a/sia/web/static.py b/sia/web/static.py new file mode 100644 index 0000000..d6a150c --- /dev/null +++ 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" + ) \ No newline at end of file diff --git a/sia/web/stdout_websocket.py b/sia/web/stdout_websocket.py new file mode 100644 index 0000000..de69f33 --- /dev/null +++ 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 diff --git a/sia/web/token_websocket.py b/sia/web/token_websocket.py new file mode 100644 index 0000000..6dccd14 --- /dev/null +++ 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 diff --git a/sia/web/util.py b/sia/web/util.py new file mode 100644 index 0000000..1ddf4e2 --- /dev/null +++ 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 \ No newline at end of file diff --git a/sia/web/websockts.py b/sia/web/websockts.py new file mode 100644 index 0000000..f7bde14 --- /dev/null +++ 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) \ No newline at end of file diff --git a/sia/web_agent.py b/sia/web_agent.py index 2d1d50c..dfb80ff 100644 --- 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 - - @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() \ No newline at end of file + 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) diff --git a/sia/web_socket_manager.py b/sia/web_socket_manager.py index 3b2800c..c97d845 100644 --- 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)) diff --git a/test/web_agent_test.py b/test/web_agent_test.py index 8f4df57..2795dba 100644 --- 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("test reasoning") diff --git a/web/package.json b/web/package.json index 36f5bd2..6ec8d88 100644 --- 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", @@ -39,4 +40,4 @@ "tailwindcss": "^3.4.1", "vite": "^5.1.0" } -} \ No newline at end of file +} diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx index 9cf12af..66f0694 100644 --- 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); + } + }); + }, []); + + // Handle context changes + useEffect(() => { + contextWs.addMessageHandler((data) => { + setContextDirty(false); + setModifiedContext(data.context); + if (data.generated) { + setGeneratedContext(data.context); } }); - }, [addMessageHandler]); + }, []); - const handleAutoApproverConfigChange = (config) => { - sendMessage({ - type: ClientMessageType.AUTO_APPROVER_CONFIG, - config: config + // 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 (
@@ -131,19 +191,22 @@ const App = () => { {activeTab === Tabs.CONTEXT && ( setModifiedContext(value)} + onChange={handleContextEdit} validationError={null} /> )} {activeTab === Tabs.RESPONSE && ( 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 = () => { setShowDiff(!showDiff)} onSendInput={handleSendInput} onClearOutput={handleClearOutput} - autoApproverConfig={autoApproverConfig} - onAutoApproverConfigChange={handleAutoApproverConfigChange} + onLlmChange={setActiveLlm} + onNextLlm={handleNextLlm} /> setShowSidebar(!showSidebar)} /> @@ -183,4 +247,4 @@ const App = () => { ); }; -export default App; \ No newline at end of file +export default App; diff --git a/web/src/components/Header.jsx b/web/src/components/Header.jsx index 904217b..e4b0057 100644 --- 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, diff --git a/web/src/components/Sidebar.jsx b/web/src/components/Sidebar.jsx index c69d67d..f546004 100644 --- a/web/src/components/Sidebar.jsx +++ b/web/src/components/Sidebar.jsx @@ -1,72 +1,51 @@ 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, }) => (