Fixed auto approver
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
from aiohttp import web
|
from aiohttp import web
|
||||||
import asyncio
|
import asyncio
|
||||||
import mimetypes
|
|
||||||
|
|
||||||
|
from .auto_approver import AutoApprover
|
||||||
from .config import Config
|
from .config import Config
|
||||||
from .hf_llm_engine import HfLlmEngine
|
from .hf_llm_engine import HfLlmEngine
|
||||||
from .iteration_logger import IterationLogger
|
from .iteration_logger import IterationLogger
|
||||||
@@ -10,11 +10,11 @@ from .mistral_llm_engine import MistralLlmEngine
|
|||||||
from .openai_llm_engine import OpenAILlmEngine
|
from .openai_llm_engine import OpenAILlmEngine
|
||||||
from .response_parser import ResponseParser
|
from .response_parser import ResponseParser
|
||||||
from .system_metrics import SystemMetrics
|
from .system_metrics import SystemMetrics
|
||||||
from .web_agent import WebAgent
|
|
||||||
from .web_io_buffer import WebIOBuffer
|
|
||||||
from .web.api import Api
|
from .web.api import Api
|
||||||
from .web.static import Static
|
from .web.static import Static
|
||||||
from .web.websockts import Websockets
|
from .web.websockts import Websockets
|
||||||
|
from .web_agent import WebAgent
|
||||||
|
from .web_io_buffer import WebIOBuffer
|
||||||
from .working_memory import WorkingMemory
|
from .working_memory import WorkingMemory
|
||||||
from .xml_validator import XMLValidator
|
from .xml_validator import XMLValidator
|
||||||
|
|
||||||
@@ -74,10 +74,11 @@ class Main:
|
|||||||
parser=ResponseParser(self._io_buffer),
|
parser=ResponseParser(self._io_buffer),
|
||||||
iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
|
iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
|
||||||
)
|
)
|
||||||
|
self._auto_approver = AutoApprover(self._agent)
|
||||||
|
|
||||||
self._app = web.Application()
|
self._app = web.Application()
|
||||||
self._api = Api(self._app, self._agent, self._io_buffer)
|
self._api = Api(self._app, self._agent, self._io_buffer, self._auto_approver)
|
||||||
self._websockets = Websockets(self._app, self._agent, self._io_buffer)
|
self._websockets = Websockets(self._app, self._agent, self._io_buffer, self._auto_approver)
|
||||||
self._static = Static(self._app, self._config)
|
self._static = Static(self._app, self._config)
|
||||||
|
|
||||||
return self
|
return self
|
||||||
|
|||||||
@@ -1,12 +1,13 @@
|
|||||||
from threading import Thread, Event
|
from threading import Thread, Event
|
||||||
from typing import Callable, TypeAlias, TypedDict
|
from typing import Callable, TypeAlias, TypedDict
|
||||||
from .web_agent import WebAgent, WebAgentState
|
from .web_agent import WebAgent, LlmState
|
||||||
|
|
||||||
class AutoApproverConfig(TypedDict):
|
class AutoApproverConfig(TypedDict):
|
||||||
context_enabled: bool
|
context_enabled: bool
|
||||||
response_enabled: bool
|
response_enabled: bool
|
||||||
context_timeout: float
|
context_timeout: float
|
||||||
response_timeout: float
|
response_timeout: float
|
||||||
|
llm_name: str
|
||||||
|
|
||||||
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
|
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
|
||||||
|
|
||||||
@@ -20,6 +21,7 @@ class AutoApprover:
|
|||||||
Initialize auto approver with a WebAgent instance.
|
Initialize auto approver with a WebAgent instance.
|
||||||
"""
|
"""
|
||||||
self.agent = agent
|
self.agent = agent
|
||||||
|
self._llm_name = next(iter(agent.llms.keys()))
|
||||||
self._context_timeout = 5.0
|
self._context_timeout = 5.0
|
||||||
self._response_timeout = 10.0
|
self._response_timeout = 10.0
|
||||||
self._context_enabled = False
|
self._context_enabled = False
|
||||||
@@ -29,7 +31,8 @@ class AutoApprover:
|
|||||||
self._response_thread: Thread | None = None
|
self._response_thread: Thread | None = None
|
||||||
self._config_change_handlers: list[ConfigChangeHandler] = []
|
self._config_change_handlers: list[ConfigChangeHandler] = []
|
||||||
|
|
||||||
self.agent.add_llm_change_handler(self._handle_state_change)
|
self.agent.add_llm_change_handler(self._handle_llm_state_change)
|
||||||
|
self.agent.add_context_change_handler(self._handle_context_change)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def config(self) -> AutoApproverConfig:
|
def config(self) -> AutoApproverConfig:
|
||||||
@@ -37,25 +40,29 @@ class AutoApprover:
|
|||||||
context_enabled=self._context_enabled,
|
context_enabled=self._context_enabled,
|
||||||
response_enabled=self._response_enabled,
|
response_enabled=self._response_enabled,
|
||||||
context_timeout=self._context_timeout,
|
context_timeout=self._context_timeout,
|
||||||
response_timeout=self._response_timeout
|
response_timeout=self._response_timeout,
|
||||||
|
llm_name=self._llm_name
|
||||||
)
|
)
|
||||||
|
|
||||||
def set_config(self, config: AutoApproverConfig) -> None:
|
def set_config(self, config: AutoApproverConfig) -> None:
|
||||||
self._context_enabled = config['context_enabled']
|
if config['llm_name'] not in self.agent.llms:
|
||||||
self._response_enabled = config['response_enabled']
|
raise ValueError(f"Unknown LLM: {config['llm_name']}")
|
||||||
self._context_timeout = config['context_timeout']
|
|
||||||
self._response_timeout = config['response_timeout']
|
|
||||||
|
|
||||||
if self._context_enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
|
notify_config_change = self._notify_config_change
|
||||||
self._start_context_thread()
|
self._notify_config_change = lambda: None
|
||||||
else:
|
|
||||||
self._stop_context_thread()
|
try:
|
||||||
|
self.context_enabled = False
|
||||||
|
self.response_enabled = False
|
||||||
|
|
||||||
if self._response_enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
self.context_timeout = config['context_timeout']
|
||||||
self._start_response_thread()
|
self.response_timeout = config['response_timeout']
|
||||||
else:
|
self.llm_name = config['llm_name']
|
||||||
self._stop_response_thread()
|
|
||||||
|
self.context_enabled = config['context_enabled']
|
||||||
|
self.response_enabled = config['response_enabled']
|
||||||
|
finally:
|
||||||
|
self._notify_config_change = notify_config_change
|
||||||
self._notify_config_change()
|
self._notify_config_change()
|
||||||
|
|
||||||
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
|
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
|
||||||
@@ -98,7 +105,7 @@ class AutoApprover:
|
|||||||
return
|
return
|
||||||
self._context_enabled = enabled
|
self._context_enabled = enabled
|
||||||
self._stop_context_thread()
|
self._stop_context_thread()
|
||||||
if enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
|
if enabled and self.agent.llms[self._llm_name] == LlmState.NO_OUTPUT:
|
||||||
self._start_context_thread()
|
self._start_context_thread()
|
||||||
self._notify_config_change()
|
self._notify_config_change()
|
||||||
|
|
||||||
@@ -112,21 +119,36 @@ class AutoApprover:
|
|||||||
return
|
return
|
||||||
self._response_enabled = enabled
|
self._response_enabled = enabled
|
||||||
self._stop_response_thread()
|
self._stop_response_thread()
|
||||||
if enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
if enabled and self.agent.llms[self._llm_name] == LlmState.OUTPUT:
|
||||||
self._start_response_thread()
|
self._start_response_thread()
|
||||||
self._notify_config_change()
|
self._notify_config_change()
|
||||||
|
|
||||||
def _handle_state_change(self, new_state: WebAgentState) -> None:
|
@property
|
||||||
if new_state == WebAgentState.CONTEXT_APPROVAL:
|
def llm_name(self) -> str:
|
||||||
self._start_context_thread()
|
return self._llm_name
|
||||||
elif self._context_enabled:
|
|
||||||
self._stop_context_thread()
|
@llm_name.setter
|
||||||
|
def llm_name(self, name: str) -> None:
|
||||||
|
if name not in self.agent.llms:
|
||||||
|
raise ValueError(f"Unknown LLM: {name}")
|
||||||
|
self._llm_name = name
|
||||||
|
self._notify_config_change()
|
||||||
|
|
||||||
if new_state == WebAgentState.RESPONSE_APPROVAL:
|
def _handle_llm_state_change(self, llm_name: str, state: LlmState) -> None:
|
||||||
|
if llm_name != self._llm_name:
|
||||||
|
return
|
||||||
|
|
||||||
|
if state == LlmState.OUTPUT and self._response_enabled:
|
||||||
self._start_response_thread()
|
self._start_response_thread()
|
||||||
elif self._response_enabled:
|
else:
|
||||||
self._stop_response_thread()
|
self._stop_response_thread()
|
||||||
|
|
||||||
|
def _handle_context_change(self, context: str, generated: bool) -> None:
|
||||||
|
if generated and self._context_enabled:
|
||||||
|
self._start_context_thread()
|
||||||
|
else:
|
||||||
|
self._stop_context_thread()
|
||||||
|
|
||||||
def _stop_context_thread(self) -> None:
|
def _stop_context_thread(self) -> None:
|
||||||
if self._context_thread:
|
if self._context_thread:
|
||||||
self._stop_event.set()
|
self._stop_event.set()
|
||||||
@@ -150,13 +172,12 @@ class AutoApprover:
|
|||||||
def _context_approval_thread(self) -> None:
|
def _context_approval_thread(self) -> None:
|
||||||
if self._stop_event.wait(self._context_timeout):
|
if self._stop_event.wait(self._context_timeout):
|
||||||
return
|
return
|
||||||
if (self._context_enabled and
|
if self._context_enabled:
|
||||||
self.agent.state == WebAgentState.CONTEXT_APPROVAL):
|
self.agent.run_inference(self._llm_name)
|
||||||
self.agent.approve_context()
|
|
||||||
|
|
||||||
def _response_approval_thread(self) -> None:
|
def _response_approval_thread(self) -> None:
|
||||||
if self._stop_event.wait(self._response_timeout):
|
if self._stop_event.wait(self._response_timeout):
|
||||||
return
|
return
|
||||||
if (self._response_enabled and
|
if (self._response_enabled and
|
||||||
self.agent.state == WebAgentState.RESPONSE_APPROVAL):
|
self.agent.llms[self._llm_name] == LlmState.OUTPUT):
|
||||||
self.agent.approve_response()
|
self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name))
|
||||||
|
|||||||
@@ -2,15 +2,16 @@ from aiohttp import web
|
|||||||
import json
|
import json
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
|
from ..auto_approver import AutoApprover
|
||||||
from ..web_agent import WebAgent
|
from ..web_agent import WebAgent
|
||||||
from ..web_io_buffer import WebIOBuffer
|
from ..web_io_buffer import WebIOBuffer
|
||||||
from .util import wrap_async
|
|
||||||
|
|
||||||
class Api:
|
class Api:
|
||||||
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer):
|
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover):
|
||||||
self._app = app
|
self._app = app
|
||||||
self._agent = agent
|
self._agent = agent
|
||||||
self._io_buffer = io_buffer
|
self._io_buffer = io_buffer
|
||||||
|
self._auto_approver = auto_approver
|
||||||
|
|
||||||
self._init_routes()
|
self._init_routes()
|
||||||
|
|
||||||
@@ -23,10 +24,13 @@ class Api:
|
|||||||
self._app.router.add_post("/api/clear", self._clear_output)
|
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/output/{llm}", self._get_output)
|
||||||
self._app.router.add_get("/api/llms", self._get_llms)
|
self._app.router.add_get("/api/llms", self._get_llms)
|
||||||
|
self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config)
|
||||||
@property
|
self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config)
|
||||||
def app(self):
|
self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled)
|
||||||
return self._app
|
self._app.router.add_post("/api/auto_approver/response_enabled", self._set_response_enabled)
|
||||||
|
self._app.router.add_post("/api/auto_approver/context_timeout", self._set_context_timeout)
|
||||||
|
self._app.router.add_post("/api/auto_approver/response_timeout", self._set_response_timeout)
|
||||||
|
self._app.router.add_post("/api/auto_approver/llm", self._set_llm_name)
|
||||||
|
|
||||||
async def _run_inference(self, request: web.Request) -> web.Response:
|
async def _run_inference(self, request: web.Request) -> web.Response:
|
||||||
"""Start inference on specified LLM."""
|
"""Start inference on specified LLM."""
|
||||||
@@ -95,3 +99,79 @@ class Api:
|
|||||||
}),
|
}),
|
||||||
content_type="application/json"
|
content_type="application/json"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def _get_auto_approver_config(self, request: web.Request) -> web.Response:
|
||||||
|
"""Get current auto approver configuration."""
|
||||||
|
return web.Response(
|
||||||
|
text=json.dumps(self._auto_approver.config),
|
||||||
|
content_type="application/json"
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _set_auto_approver_config(self, request: web.Request) -> web.Response:
|
||||||
|
"""Update auto approver configuration."""
|
||||||
|
data = await request.json()
|
||||||
|
try:
|
||||||
|
self._auto_approver.set_config(data)
|
||||||
|
return web.Response(status=200)
|
||||||
|
except (ValueError, KeyError) as e:
|
||||||
|
return web.Response(status=400, text=str(e))
|
||||||
|
|
||||||
|
async def _set_context_enabled(self, request: web.Request) -> web.Response:
|
||||||
|
"""Set context auto-approval enabled state."""
|
||||||
|
data = await request.json()
|
||||||
|
enabled = data.get("enabled")
|
||||||
|
if enabled is None:
|
||||||
|
return web.Response(status=400, text="Missing enabled parameter")
|
||||||
|
try:
|
||||||
|
self._auto_approver.context_enabled = enabled
|
||||||
|
return web.Response(status=200)
|
||||||
|
except ValueError as e:
|
||||||
|
return web.Response(status=400, text=str(e))
|
||||||
|
|
||||||
|
async def _set_response_enabled(self, request: web.Request) -> web.Response:
|
||||||
|
"""Set response auto-approval enabled state."""
|
||||||
|
data = await request.json()
|
||||||
|
enabled = data.get("enabled")
|
||||||
|
if enabled is None:
|
||||||
|
return web.Response(status=400, text="Missing enabled parameter")
|
||||||
|
try:
|
||||||
|
self._auto_approver.response_enabled = enabled
|
||||||
|
return web.Response(status=200)
|
||||||
|
except ValueError as e:
|
||||||
|
return web.Response(status=400, text=str(e))
|
||||||
|
|
||||||
|
async def _set_context_timeout(self, request: web.Request) -> web.Response:
|
||||||
|
"""Set context auto-approval timeout."""
|
||||||
|
data = await request.json()
|
||||||
|
timeout = data.get("timeout")
|
||||||
|
if timeout is None:
|
||||||
|
return web.Response(status=400, text="Missing timeout parameter")
|
||||||
|
try:
|
||||||
|
self._auto_approver.context_timeout = float(timeout)
|
||||||
|
return web.Response(status=200)
|
||||||
|
except ValueError as e:
|
||||||
|
return web.Response(status=400, text=str(e))
|
||||||
|
|
||||||
|
async def _set_response_timeout(self, request: web.Request) -> web.Response:
|
||||||
|
"""Set response auto-approval timeout."""
|
||||||
|
data = await request.json()
|
||||||
|
timeout = data.get("timeout")
|
||||||
|
if timeout is None:
|
||||||
|
return web.Response(status=400, text="Missing timeout parameter")
|
||||||
|
try:
|
||||||
|
self._auto_approver.response_timeout = float(timeout)
|
||||||
|
return web.Response(status=200)
|
||||||
|
except ValueError as e:
|
||||||
|
return web.Response(status=400, text=str(e))
|
||||||
|
|
||||||
|
async def _set_llm_name(self, request: web.Request) -> web.Response:
|
||||||
|
"""Set LLM name for auto-approval."""
|
||||||
|
data = await request.json()
|
||||||
|
name = data.get("name")
|
||||||
|
if name is None:
|
||||||
|
return web.Response(status=400, text="Missing name parameter")
|
||||||
|
try:
|
||||||
|
self._auto_approver.llm_name = name
|
||||||
|
return web.Response(status=200)
|
||||||
|
except ValueError as e:
|
||||||
|
return web.Response(status=400, text=str(e))
|
||||||
53
sia/web/auto_approver_websocket.py
Normal file
53
sia/web/auto_approver_websocket.py
Normal file
@@ -0,0 +1,53 @@
|
|||||||
|
from aiohttp import web, WSMsgType
|
||||||
|
from typing import Dict, Set
|
||||||
|
|
||||||
|
from ..auto_approver import AutoApprover, AutoApproverConfig
|
||||||
|
from .util import wrap_async
|
||||||
|
|
||||||
|
class AutoApproverWebSocket:
|
||||||
|
"""
|
||||||
|
WebSocket handler for AutoApprover configuration.
|
||||||
|
Broadcasts config updates to all connected clients.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, auto_approver: AutoApprover):
|
||||||
|
self._auto_approver = auto_approver
|
||||||
|
self._clients: Set[web.WebSocketResponse] = set()
|
||||||
|
self._auto_approver.add_config_change_handler(wrap_async(self._handle_config_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_config_change(self, config: AutoApproverConfig):
|
||||||
|
"""Handle config changes from the AutoApprover."""
|
||||||
|
await self._broadcast_message({
|
||||||
|
"config": config
|
||||||
|
})
|
||||||
|
|
||||||
|
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 config
|
||||||
|
await ws.send_json({
|
||||||
|
"config": self._auto_approver.config
|
||||||
|
})
|
||||||
|
|
||||||
|
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
|
||||||
@@ -2,19 +2,23 @@ from aiohttp import web
|
|||||||
|
|
||||||
from ..web_agent import WebAgent
|
from ..web_agent import WebAgent
|
||||||
from ..web_io_buffer import WebIOBuffer
|
from ..web_io_buffer import WebIOBuffer
|
||||||
|
from ..auto_approver import AutoApprover
|
||||||
from .llm_websocket import LlmWebSocket
|
from .llm_websocket import LlmWebSocket
|
||||||
from .context_websocket import ContextWebSocket
|
from .context_websocket import ContextWebSocket
|
||||||
from .token_websocket import TokenWebSocket
|
from .token_websocket import TokenWebSocket
|
||||||
from .stdout_websocket import StdoutWebSocket
|
from .stdout_websocket import StdoutWebSocket
|
||||||
|
from .auto_approver_websocket import AutoApproverWebSocket
|
||||||
|
|
||||||
class Websockets:
|
class Websockets:
|
||||||
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer):
|
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover):
|
||||||
self._llm_ws = LlmWebSocket(agent)
|
self._llm_ws = LlmWebSocket(agent)
|
||||||
self._context_ws = ContextWebSocket(agent)
|
self._context_ws = ContextWebSocket(agent)
|
||||||
self._token_ws = TokenWebSocket(agent)
|
self._token_ws = TokenWebSocket(agent)
|
||||||
self._stdout_ws = StdoutWebSocket(io_buffer)
|
self._stdout_ws = StdoutWebSocket(io_buffer)
|
||||||
|
self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
|
||||||
|
|
||||||
app.router.add_get("/ws/llm", self._llm_ws.handle_connection)
|
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/context", self._context_ws.handle_connection)
|
||||||
app.router.add_get("/ws/token", self._token_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)
|
app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection)
|
||||||
|
app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)
|
||||||
|
|||||||
@@ -1,173 +0,0 @@
|
|||||||
from aiohttp import web, WSMsgType
|
|
||||||
from typing import Any, Dict, Set
|
|
||||||
import asyncio
|
|
||||||
import json
|
|
||||||
from dataclasses import dataclass
|
|
||||||
|
|
||||||
from .web_agent import WebAgent, WebAgentState
|
|
||||||
from .web_io_buffer import WebIOBuffer
|
|
||||||
from .auto_approver import AutoApprover, AutoApproverConfig
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class ClientMessage:
|
|
||||||
"""Messages that can be received from clients."""
|
|
||||||
APPROVE_CONTEXT = "APPROVE_CONTEXT"
|
|
||||||
APPROVE_RESPONSE = "APPROVE_RESPONSE"
|
|
||||||
MODIFY_RESPONSE = "MODIFY_RESPONSE"
|
|
||||||
SEND_INPUT = "SEND_INPUT"
|
|
||||||
CLEAR_OUTPUT = "CLEAR_OUTPUT"
|
|
||||||
AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG"
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class ServerMessage:
|
|
||||||
"""Messages that can be sent to clients."""
|
|
||||||
STATE_CHANGE = "STATE_CHANGE"
|
|
||||||
CONTEXT_UPDATE = "CONTEXT_UPDATE"
|
|
||||||
RESPONSE_UPDATE = "RESPONSE_UPDATE"
|
|
||||||
OUTPUT_UPDATE = "OUTPUT_UPDATE"
|
|
||||||
AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG"
|
|
||||||
|
|
||||||
class WebSocketManager:
|
|
||||||
"""
|
|
||||||
Manages WebSocket connections and integrates with the WebAgent
|
|
||||||
"""
|
|
||||||
@classmethod
|
|
||||||
async def create(
|
|
||||||
cls,
|
|
||||||
agent: WebAgent,
|
|
||||||
io_buffer: WebIOBuffer):
|
|
||||||
"""Initialize the web server."""
|
|
||||||
self = cls()
|
|
||||||
self.agent = agent
|
|
||||||
self.io_buffer = io_buffer
|
|
||||||
self._clients: Set[web.WebSocketResponse] = set()
|
|
||||||
self.auto_approver = AutoApprover(agent)
|
|
||||||
|
|
||||||
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))
|
|
||||||
self.auto_approver.add_config_change_handler(self._wrap_async(self._handle_auto_approver_config))
|
|
||||||
return self
|
|
||||||
|
|
||||||
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
|
|
||||||
|
|
||||||
def _wrap_async(self, coro_func):
|
|
||||||
"""
|
|
||||||
Wraps an async callback to be safely called from another thread.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
coro_func: The async function to wrap
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
A sync function that schedules the coroutine on the event loop
|
|
||||||
"""
|
|
||||||
loop = asyncio.get_event_loop()
|
|
||||||
def wrapper(*args, **kwargs):
|
|
||||||
loop.call_soon_threadsafe(lambda: asyncio.create_task(coro_func(*args, **kwargs)))
|
|
||||||
return wrapper
|
|
||||||
|
|
||||||
async def _handle_state_change(self, new_state: WebAgentState):
|
|
||||||
"""Handle state changes from the WebAgent."""
|
|
||||||
await self._broadcast_message({
|
|
||||||
"type": ServerMessage.STATE_CHANGE,
|
|
||||||
"state": new_state.name,
|
|
||||||
})
|
|
||||||
|
|
||||||
async def _handle_context_change(self, context: str):
|
|
||||||
"""Handle context changes from the WebAgent."""
|
|
||||||
await self._broadcast_message({
|
|
||||||
"type": ServerMessage.CONTEXT_UPDATE,
|
|
||||||
"context": context
|
|
||||||
})
|
|
||||||
|
|
||||||
async def _handle_response_change(self, response: str, validation_error: str):
|
|
||||||
"""Handle response changes from the WebAgent."""
|
|
||||||
await self._broadcast_message({
|
|
||||||
"type": ServerMessage.RESPONSE_UPDATE,
|
|
||||||
"response": response,
|
|
||||||
"validation_error": validation_error
|
|
||||||
})
|
|
||||||
|
|
||||||
async def _handle_stdout_change(self, output: str):
|
|
||||||
"""Handle output changes from the WebAgent."""
|
|
||||||
await self._broadcast_message({
|
|
||||||
"type": ServerMessage.OUTPUT_UPDATE,
|
|
||||||
"output": output
|
|
||||||
})
|
|
||||||
|
|
||||||
async def _handle_auto_approver_config(self, config: AutoApproverConfig):
|
|
||||||
"""Handle auto approver config changes."""
|
|
||||||
await self._broadcast_message({
|
|
||||||
"type": ServerMessage.AUTO_APPROVER_CONFIG,
|
|
||||||
"config": config
|
|
||||||
})
|
|
||||||
|
|
||||||
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
|
|
||||||
"""Handle incoming client message."""
|
|
||||||
message_type = data.get("type")
|
|
||||||
|
|
||||||
if message_type == ClientMessage.APPROVE_CONTEXT:
|
|
||||||
if self.agent.state == WebAgentState.CONTEXT_APPROVAL:
|
|
||||||
if data.get("context") is not None:
|
|
||||||
self.agent.modify_context(data.get("context"))
|
|
||||||
self.agent.approve_context()
|
|
||||||
elif message_type == ClientMessage.APPROVE_RESPONSE:
|
|
||||||
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
|
||||||
if data.get("response") is not None:
|
|
||||||
self.agent.modify_response(data.get("response"))
|
|
||||||
self.agent.approve_response()
|
|
||||||
elif message_type == ClientMessage.MODIFY_RESPONSE:
|
|
||||||
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
|
||||||
self.agent.modify_response(data.get("response"))
|
|
||||||
elif message_type == ClientMessage.SEND_INPUT:
|
|
||||||
self.io_buffer.append_stdin(data.get("input"))
|
|
||||||
elif message_type == ClientMessage.CLEAR_OUTPUT:
|
|
||||||
self.io_buffer.clear_stdout()
|
|
||||||
elif message_type == ClientMessage.AUTO_APPROVER_CONFIG:
|
|
||||||
config = data.get("config")
|
|
||||||
if isinstance(config, dict):
|
|
||||||
self.auto_approver.set_config(config)
|
|
||||||
|
|
||||||
async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
|
|
||||||
"""Handle new WebSocket connections and messages."""
|
|
||||||
ws = web.WebSocketResponse(heartbeat=30)
|
|
||||||
await ws.prepare(request)
|
|
||||||
|
|
||||||
self._clients.add(ws)
|
|
||||||
|
|
||||||
try:
|
|
||||||
await ws.send_json({
|
|
||||||
"type": ServerMessage.STATE_CHANGE,
|
|
||||||
"state": self.agent.state.name
|
|
||||||
})
|
|
||||||
await ws.send_json({
|
|
||||||
"type": ServerMessage.CONTEXT_UPDATE,
|
|
||||||
"context": self.agent.context
|
|
||||||
})
|
|
||||||
await ws.send_json({
|
|
||||||
"type": ServerMessage.RESPONSE_UPDATE,
|
|
||||||
"response": self.agent.response,
|
|
||||||
"validation_error": self.agent._validation_error
|
|
||||||
})
|
|
||||||
await ws.send_json({
|
|
||||||
"type": ServerMessage.AUTO_APPROVER_CONFIG,
|
|
||||||
"config": self.auto_approver.config
|
|
||||||
})
|
|
||||||
async for msg in ws:
|
|
||||||
if msg.type == WSMsgType.TEXT:
|
|
||||||
data = json.loads(msg.data)
|
|
||||||
await self._handle_client_message(request, data)
|
|
||||||
elif msg.type == WSMsgType.ERROR:
|
|
||||||
print(f"WebSocket connection closed with error: {ws.exception()}")
|
|
||||||
finally:
|
|
||||||
self._clients.remove(ws)
|
|
||||||
return ws
|
|
||||||
@@ -25,6 +25,15 @@ const App = () => {
|
|||||||
// LLM state
|
// LLM state
|
||||||
const [llms, setLlms] = useState({});
|
const [llms, setLlms] = useState({});
|
||||||
const [activeLlm, setActiveLlm] = useState(null);
|
const [activeLlm, setActiveLlm] = useState(null);
|
||||||
|
|
||||||
|
// Auto approver state
|
||||||
|
const [autoApproverConfig, setAutoApproverConfig] = useState({
|
||||||
|
context_enabled: false,
|
||||||
|
response_enabled: false,
|
||||||
|
context_timeout: 5.0,
|
||||||
|
response_timeout: 10.0,
|
||||||
|
llm_name: null
|
||||||
|
});
|
||||||
|
|
||||||
// UI state
|
// UI state
|
||||||
const [activeTab, setActiveTab] = useState(Tabs.CONTEXT);
|
const [activeTab, setActiveTab] = useState(Tabs.CONTEXT);
|
||||||
@@ -37,6 +46,7 @@ const App = () => {
|
|||||||
const contextWs = useWebSocket(`${wsRoot}/context`);
|
const contextWs = useWebSocket(`${wsRoot}/context`);
|
||||||
const tokenWs = useWebSocket(`${wsRoot}/token`);
|
const tokenWs = useWebSocket(`${wsRoot}/token`);
|
||||||
const stdoutWs = useWebSocket(`${wsRoot}/stdout`);
|
const stdoutWs = useWebSocket(`${wsRoot}/stdout`);
|
||||||
|
const autoApproverWs = useWebSocket(`${wsRoot}/auto_approver`);
|
||||||
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
|
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
|
||||||
|
|
||||||
// Handle llm state changes
|
// Handle llm state changes
|
||||||
@@ -58,9 +68,6 @@ const App = () => {
|
|||||||
return updated;
|
return updated;
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
if (activeLlm === null) {
|
|
||||||
setActiveLlm(data.llm);
|
|
||||||
}
|
|
||||||
});
|
});
|
||||||
}, []);
|
}, []);
|
||||||
|
|
||||||
@@ -96,6 +103,20 @@ const App = () => {
|
|||||||
});
|
});
|
||||||
}, []);
|
}, []);
|
||||||
|
|
||||||
|
// Handle auto approver config updates
|
||||||
|
useEffect(() => {
|
||||||
|
autoApproverWs.addMessageHandler((data) => {
|
||||||
|
setAutoApproverConfig(data.config);
|
||||||
|
});
|
||||||
|
}, []);
|
||||||
|
|
||||||
|
// Initialize active llm
|
||||||
|
useEffect(() => {
|
||||||
|
if (activeLlm === null && Object.keys(llms).length > 0) {
|
||||||
|
setActiveLlm(Object.keys(llms)[0]);
|
||||||
|
}
|
||||||
|
}, [llms, activeLlm]);
|
||||||
|
|
||||||
const handleNextLlm = () => {
|
const handleNextLlm = () => {
|
||||||
const llmNames = Object.keys(llms);
|
const llmNames = Object.keys(llms);
|
||||||
const currentIndex = llmNames.indexOf(activeLlm);
|
const currentIndex = llmNames.indexOf(activeLlm);
|
||||||
@@ -147,9 +168,25 @@ const App = () => {
|
|||||||
const resetLlms = {};
|
const resetLlms = {};
|
||||||
for (const llm of Object.keys(llms)) {
|
for (const llm of Object.keys(llms)) {
|
||||||
resetLlms[llm] = LlmState.NO_OUTPUT;
|
resetLlms[llm] = LlmState.NO_OUTPUT;
|
||||||
} setLlms(resetLlms);
|
}
|
||||||
|
setLlms(resetLlms);
|
||||||
setModifiedContext(context);
|
setModifiedContext(context);
|
||||||
}
|
};
|
||||||
|
|
||||||
|
const handleAutoApproverConfigChange = async (config) => {
|
||||||
|
try {
|
||||||
|
const response = await fetch('/api/auto_approver/config', {
|
||||||
|
method: 'POST',
|
||||||
|
headers: { 'Content-Type': 'application/json' },
|
||||||
|
body: JSON.stringify(config)
|
||||||
|
});
|
||||||
|
if (!response.ok) {
|
||||||
|
throw new Error('Failed to update auto approver config');
|
||||||
|
}
|
||||||
|
} catch (error) {
|
||||||
|
console.error('Error updating auto approver config:', error);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
const state = llms[activeLlm];
|
const state = llms[activeLlm];
|
||||||
@@ -168,14 +205,22 @@ const App = () => {
|
|||||||
}, [llms[activeLlm]]);
|
}, [llms[activeLlm]]);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
if (llmWs.wsState === WebSocketState.CONNECTED && contextWs.wsState === WebSocketState.CONNECTED && tokenWs.wsState === WebSocketState.CONNECTED && stdoutWs.wsState === WebSocketState.CONNECTED) {
|
if (llmWs.wsState === WebSocketState.CONNECTED &&
|
||||||
|
contextWs.wsState === WebSocketState.CONNECTED &&
|
||||||
|
tokenWs.wsState === WebSocketState.CONNECTED &&
|
||||||
|
stdoutWs.wsState === WebSocketState.CONNECTED &&
|
||||||
|
autoApproverWs.wsState === WebSocketState.CONNECTED) {
|
||||||
setWsState(WebSocketState.CONNECTED);
|
setWsState(WebSocketState.CONNECTED);
|
||||||
} else if (llmWs.wsState === WebSocketState.CONNECTING || contextWs.wsState === WebSocketState.CONNECTING || tokenWs.wsState === WebSocketState.CONNECTING || stdoutWs.wsState === WebSocketState.CONNECTING) {
|
} else if (llmWs.wsState === WebSocketState.CONNECTING ||
|
||||||
|
contextWs.wsState === WebSocketState.CONNECTING ||
|
||||||
|
tokenWs.wsState === WebSocketState.CONNECTING ||
|
||||||
|
stdoutWs.wsState === WebSocketState.CONNECTING ||
|
||||||
|
autoApproverWs.wsState === WebSocketState.CONNECTING) {
|
||||||
setWsState(WebSocketState.CONNECTING);
|
setWsState(WebSocketState.CONNECTING);
|
||||||
} else {
|
} else {
|
||||||
setWsState(WebSocketState.DISCONNECTED);
|
setWsState(WebSocketState.DISCONNECTED);
|
||||||
}
|
}
|
||||||
}, [llmWs.wsState, contextWs.wsState, tokenWs.wsState, stdoutWs.wsState]);
|
}, [llmWs.wsState, contextWs.wsState, tokenWs.wsState, stdoutWs.wsState, autoApproverWs.wsState]);
|
||||||
|
|
||||||
return (
|
return (
|
||||||
<div className="h-screen bg-gray-100 flex flex-col overflow-hidden">
|
<div className="h-screen bg-gray-100 flex flex-col overflow-hidden">
|
||||||
@@ -233,12 +278,14 @@ const App = () => {
|
|||||||
showDiff={showDiff}
|
showDiff={showDiff}
|
||||||
input={input}
|
input={input}
|
||||||
output={output}
|
output={output}
|
||||||
|
autoApproverConfig={autoApproverConfig}
|
||||||
onApprove={handleApprove}
|
onApprove={handleApprove}
|
||||||
onToggleDiff={() => setShowDiff(!showDiff)}
|
onToggleDiff={() => setShowDiff(!showDiff)}
|
||||||
onSendInput={handleSendInput}
|
onSendInput={handleSendInput}
|
||||||
onClearOutput={handleClearOutput}
|
onClearOutput={handleClearOutput}
|
||||||
onLlmChange={setActiveLlm}
|
onLlmChange={setActiveLlm}
|
||||||
onNextLlm={handleNextLlm}
|
onNextLlm={handleNextLlm}
|
||||||
|
onAutoApproverConfigChange={handleAutoApproverConfigChange}
|
||||||
/>
|
/>
|
||||||
|
|
||||||
<FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
|
<FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
import React from 'react';
|
import React from 'react';
|
||||||
import { Button } from '@/components/ui/button';
|
import { Button } from '@/components/ui/button';
|
||||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
|
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
|
||||||
|
import { Switch } from '@/components/ui/switch';
|
||||||
|
import { Input } from '@/components/ui/input';
|
||||||
import { LlmState } from '@/components/App';
|
import { LlmState } from '@/components/App';
|
||||||
|
|
||||||
export const Sidebar = ({
|
export const Sidebar = ({
|
||||||
@@ -11,12 +13,14 @@ export const Sidebar = ({
|
|||||||
showDiff,
|
showDiff,
|
||||||
input,
|
input,
|
||||||
output,
|
output,
|
||||||
|
autoApproverConfig,
|
||||||
onApprove,
|
onApprove,
|
||||||
onToggleDiff,
|
onToggleDiff,
|
||||||
onSendInput,
|
onSendInput,
|
||||||
onClearOutput,
|
onClearOutput,
|
||||||
onLlmChange,
|
onLlmChange,
|
||||||
onNextLlm,
|
onNextLlm,
|
||||||
|
onAutoApproverConfigChange,
|
||||||
}) => (
|
}) => (
|
||||||
<aside className={`
|
<aside className={`
|
||||||
fixed top-0 right-0 h-screen w-64 bg-white shadow-lg
|
fixed top-0 right-0 h-screen w-64 bg-white shadow-lg
|
||||||
@@ -25,6 +29,7 @@ export const Sidebar = ({
|
|||||||
${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
|
${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
|
||||||
`}>
|
`}>
|
||||||
<div className="p-4 space-y-4 overflow-y-auto h-full">
|
<div className="p-4 space-y-4 overflow-y-auto h-full">
|
||||||
|
{/* Existing controls */}
|
||||||
<Select value={activeLlm} onValueChange={onLlmChange}>
|
<Select value={activeLlm} onValueChange={onLlmChange}>
|
||||||
<SelectTrigger>
|
<SelectTrigger>
|
||||||
<SelectValue placeholder="Select LLM" />
|
<SelectValue placeholder="Select LLM" />
|
||||||
@@ -38,19 +43,11 @@ export const Sidebar = ({
|
|||||||
</SelectContent>
|
</SelectContent>
|
||||||
</Select>
|
</Select>
|
||||||
|
|
||||||
<Button
|
<Button variant="outline" onClick={onNextLlm} className="w-full">
|
||||||
variant="outline"
|
|
||||||
onClick={onNextLlm}
|
|
||||||
className="w-full"
|
|
||||||
>
|
|
||||||
Next LLM
|
Next LLM
|
||||||
</Button>
|
</Button>
|
||||||
|
|
||||||
<Button
|
<Button variant="outline" onClick={onToggleDiff} className="w-full">
|
||||||
variant="outline"
|
|
||||||
onClick={onToggleDiff}
|
|
||||||
className="w-full"
|
|
||||||
>
|
|
||||||
{showDiff ? 'Hide Diff' : 'Show Diff'}
|
{showDiff ? 'Hide Diff' : 'Show Diff'}
|
||||||
</Button>
|
</Button>
|
||||||
|
|
||||||
@@ -66,21 +63,90 @@ export const Sidebar = ({
|
|||||||
}
|
}
|
||||||
</Button>
|
</Button>
|
||||||
|
|
||||||
<Button
|
<Button onClick={onSendInput} disabled={!input.trim()} className="w-full">
|
||||||
onClick={onSendInput}
|
|
||||||
disabled={!input.trim()}
|
|
||||||
className="w-full"
|
|
||||||
>
|
|
||||||
Send Input
|
Send Input
|
||||||
</Button>
|
</Button>
|
||||||
|
|
||||||
<Button
|
<Button onClick={onClearOutput} disabled={!output.trim()} className="w-full">
|
||||||
onClick={onClearOutput}
|
|
||||||
disabled={!output.trim()}
|
|
||||||
className="w-full"
|
|
||||||
>
|
|
||||||
Clear Output
|
Clear Output
|
||||||
</Button>
|
</Button>
|
||||||
|
|
||||||
|
{/* Auto Approver Config */}
|
||||||
|
<div className="border-t pt-4">
|
||||||
|
<h3 className="font-medium mb-2">Auto Approver</h3>
|
||||||
|
|
||||||
|
<div className="space-y-2">
|
||||||
|
<div className="flex items-center justify-between">
|
||||||
|
<label>Context Auto-Approve</label>
|
||||||
|
<Switch
|
||||||
|
checked={autoApproverConfig.context_enabled}
|
||||||
|
onCheckedChange={(enabled) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
context_enabled: enabled
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="flex items-center justify-between">
|
||||||
|
<label>Response Auto-Approve</label>
|
||||||
|
<Switch
|
||||||
|
checked={autoApproverConfig.response_enabled}
|
||||||
|
onCheckedChange={(enabled) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
response_enabled: enabled
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="space-y-1">
|
||||||
|
<label className="text-sm">Context Timeout (s)</label>
|
||||||
|
<Input
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
step="0.1"
|
||||||
|
value={autoApproverConfig.context_timeout}
|
||||||
|
onChange={(e) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
context_timeout: parseFloat(e.target.value)
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="space-y-1">
|
||||||
|
<label className="text-sm">Response Timeout (s)</label>
|
||||||
|
<Input
|
||||||
|
type="number"
|
||||||
|
min="0"
|
||||||
|
step="0.1"
|
||||||
|
value={autoApproverConfig.response_timeout}
|
||||||
|
onChange={(e) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
response_timeout: parseFloat(e.target.value)
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="space-y-1">
|
||||||
|
<label className="text-sm">Auto-Approve LLM</label>
|
||||||
|
<Select
|
||||||
|
value={autoApproverConfig.llm_name}
|
||||||
|
onValueChange={(llm) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
llm_name: llm
|
||||||
|
})}
|
||||||
|
>
|
||||||
|
<SelectTrigger>
|
||||||
|
<SelectValue placeholder="Select LLM" />
|
||||||
|
</SelectTrigger>
|
||||||
|
<SelectContent>
|
||||||
|
{Object.keys(llms).map(llm => (
|
||||||
|
<SelectItem key={llm} value={llm}>{llm}</SelectItem>
|
||||||
|
))}
|
||||||
|
</SelectContent>
|
||||||
|
</Select>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
</div>
|
</div>
|
||||||
</aside>
|
</aside>
|
||||||
);
|
);
|
||||||
|
|||||||
Reference in New Issue
Block a user