llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add auto_approver

Browse Source
This commit is contained in:
Niels Geens
2024-11-15 17:09:07 +01:00
parent 4ce421bbce
commit 9dfd5ddc39
9 changed files with 449 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
sia/web_socket_manager.py
View File

@@ -6,6 +6,7 @@ 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:
@@ -15,6 +16,7 @@ class ClientMessage:
MODIFY_RESPONSE = "MODIFY_RESPONSE"
SEND_INPUT = "SEND_INPUT"
CLEAR_OUTPUT = "CLEAR_OUTPUT"
AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG"
@dataclass
class ServerMessage:
@@ -23,6 +25,7 @@ class ServerMessage:
CONTEXT_UPDATE = "CONTEXT_UPDATE"
RESPONSE_UPDATE = "RESPONSE_UPDATE"
OUTPUT_UPDATE = "OUTPUT_UPDATE"
AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG"
class WebSocketManager:
"""
@@ -38,11 +41,13 @@ class WebSocketManager:
self.agent = agent
self.io_buffer = io_buffer
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_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):
@@ -99,6 +104,13 @@ class WebSocketManager:
"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")
@@ -117,9 +129,13 @@ class WebSocketManager:
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"))
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."""
@@ -142,6 +158,10 @@ class WebSocketManager:
"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)
@@ -150,4 +170,4 @@ class WebSocketManager:
print(f"WebSocket connection closed with error: {ws.exception()}")
finally:
self._clients.remove(ws)
return ws
return ws