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

2
sia/__main__.py
View File

@@ -122,5 +122,5 @@ if __name__ == "__main__":
loop = asyncio.new_event_loop()
config = Config()
main = loop.run_until_complete(Main.create(config))
print(f"Web server started at http://{config.host}:{config.port}")
print(f"Web server started at http://localhost:{config.port}")
web.run_app(main.app, loop=loop, host=config.host, port=config.port)

162
sia/auto_approver.py Normal file
View File

@@ -0,0 +1,162 @@
from threading import Thread, Event
from typing import Callable, TypeAlias, TypedDict
from .web_agent import WebAgent, WebAgentState
class AutoApproverConfig(TypedDict):
context_enabled: bool
response_enabled: bool
context_timeout: float
response_timeout: float
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
class AutoApprover:
"""
Handles automatic approvals for WebAgent states after configurable timeouts.
"""
def __init__(self, agent: WebAgent):
"""
Initialize auto approver with a WebAgent instance.
"""
self.agent = agent
self._context_timeout = 5.0
self._response_timeout = 10.0
self._context_enabled = False
self._response_enabled = False
self._stop_event = Event()
self._context_thread: Thread | None = None
self._response_thread: Thread | None = None
self._config_change_handlers: list[ConfigChangeHandler] = []
self.agent.add_state_change_handler(self._handle_state_change)
@property
def config(self) -> AutoApproverConfig:
return AutoApproverConfig(
context_enabled=self._context_enabled,
response_enabled=self._response_enabled,
context_timeout=self._context_timeout,
response_timeout=self._response_timeout
)
def set_config(self, config: AutoApproverConfig) -> None:
self._context_enabled = config['context_enabled']
self._response_enabled = config['response_enabled']
self._context_timeout = config['context_timeout']
self._response_timeout = config['response_timeout']
if self._context_enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
self._start_context_thread()
else:
self._stop_context_thread()
if self._response_enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
self._start_response_thread()
else:
self._stop_response_thread()
self._notify_config_change()
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
self._config_change_handlers.append(handler)
def _notify_config_change(self) -> None:
current_config = self.config
for handler in self._config_change_handlers:
handler(current_config)
@property
def context_timeout(self) -> float:
return self._context_timeout
@context_timeout.setter
def context_timeout(self, timeout: float) -> None:
if self._context_enabled:
raise ValueError("Cannot change timeout while auto-approval is enabled")
self._context_timeout = timeout
self._notify_config_change()
@property
def response_timeout(self) -> float:
return self._response_timeout
@response_timeout.setter
def response_timeout(self, timeout: float) -> None:
if self._response_enabled:
raise ValueError("Cannot change timeout while auto-approval is enabled")
self._response_timeout = timeout
self._notify_config_change()
@property
def context_enabled(self) -> bool:
return self._context_enabled
@context_enabled.setter
def context_enabled(self, enabled: bool) -> None:
if enabled == self._context_enabled:
return
self._context_enabled = enabled
self._stop_context_thread()
if enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
self._start_context_thread()
self._notify_config_change()
@property
def response_enabled(self) -> bool:
return self._response_enabled
@response_enabled.setter
def response_enabled(self, enabled: bool) -> None:
if enabled == self._response_enabled:
return
self._response_enabled = enabled
self._stop_response_thread()
if enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
self._start_response_thread()
self._notify_config_change()
def _handle_state_change(self, new_state: WebAgentState) -> None:
if new_state == WebAgentState.CONTEXT_APPROVAL:
self._start_context_thread()
elif self._context_enabled:
self._stop_context_thread()
if new_state == WebAgentState.RESPONSE_APPROVAL:
self._start_response_thread()
elif self._response_enabled:
self._stop_response_thread()
def _stop_context_thread(self) -> None:
if self._context_thread:
self._stop_event.set()
self._context_thread = None
self._stop_event.clear()
def _stop_response_thread(self) -> None:
if self._response_thread:
self._stop_event.set()
self._response_thread = None
self._stop_event.clear()
def _start_context_thread(self) -> None:
self._context_thread = Thread(target=self._context_approval_thread)
self._context_thread.start()
def _start_response_thread(self) -> None:
self._response_thread = Thread(target=self._response_approval_thread)
self._response_thread.start()
def _context_approval_thread(self) -> None:
if self._stop_event.wait(self._context_timeout):
return
if (self._context_enabled and
self.agent.state == WebAgentState.CONTEXT_APPROVAL):
self.agent.approve_context()
def _response_approval_thread(self) -> None:
if self._stop_event.wait(self._response_timeout):
return
if (self._response_enabled and
self.agent.state == WebAgentState.RESPONSE_APPROVAL):
self.agent.approve_response()

22
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)

87
test/auto_approver_test.py Normal file
View File

@@ -0,0 +1,87 @@
import unittest
from unittest.mock import Mock, patch
from threading import Event
import time
from sia.auto_approver import AutoApprover
from sia.web_agent import WebAgent, WebAgentState
class AutoApproverTest(unittest.TestCase):
def setUp(self):
"""Create mock agent and approver for each test"""
self.mock_agent = Mock(spec=WebAgent)
self.mock_agent.state = WebAgentState.UPDATE
self.approver = AutoApprover(self.mock_agent)
def test_timeout_setters(self):
"""Test timeout setters block changes when enabled"""
self.approver.context_timeout = 10.0
self.approver.response_timeout = 15.0
self.approver.context_enabled = True
with self.assertRaises(ValueError):
self.approver.context_timeout = 20.0
self.approver.response_enabled = True
with self.assertRaises(ValueError):
self.approver.response_timeout = 25.0
def test_enable_starts_thread_in_correct_state(self):
"""Test enabling only starts thread in matching state"""
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
self.approver.context_enabled = True
time.sleep(0.1) # Allow thread to start
self.assertTrue(self.approver._context_thread.is_alive())
self.assertIsNone(self.approver._response_thread)
def test_state_change_stops_threads(self):
"""Test state changes stop irrelevant threads"""
# Start in context approval with thread
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
self.approver.context_enabled = True
time.sleep(0.1)
self.assertTrue(self.approver._context_thread.is_alive())
# Change state and verify thread stops
self.mock_agent.state = WebAgentState.INFERENCE
self.approver._handle_state_change(WebAgentState.INFERENCE)
time.sleep(0.1)
self.assertIsNone(self.approver._context_thread)
def test_quick_state_cycle(self):
"""Test thread stops when state changes before timeout"""
self.approver.context_timeout = 5.0
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
self.approver.context_enabled = True
# Change state quickly
time.sleep(0.1)
self.mock_agent.state = WebAgentState.INFERENCE
self.approver._handle_state_change(WebAgentState.INFERENCE)
# Verify no approval happened
self.mock_agent.approve_context.assert_not_called()
def test_disable_stops_thread(self):
"""Test disabling stops active thread"""
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
self.approver.context_enabled = True
time.sleep(0.1)
self.assertTrue(self.approver._context_thread.is_alive())
self.approver.context_enabled = False
self.assertIsNone(self.approver._context_thread)
def test_successful_auto_approval(self):
"""Test thread approves after timeout when conditions met"""
self.approver.context_timeout = 0.1
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
self.approver.context_enabled = True
time.sleep(0.2)
self.mock_agent.approve_context.assert_called_once()
def tearDown(self):
"""Clean up any running threads"""
self.approver.context_enabled = False
self.approver.response_enabled = False

108
test/web_socket_manager_test.py
View File

@@ -68,6 +68,11 @@ class WebSocketManagerTest(AioHTTPTestCase):
app.router.add_get('/ws', self.ws_manager.handle_websocket)
return app
async def _skip_initial_messages(self, ws):
"""Helper to skip past the initial connection messages"""
for _ in range(4): # Now 4 messages including auto approver config
await ws.receive_json()
@unittest_run_loop
async def test_initial_connection(self):
"""Test initial WebSocket connection and state messages"""
@@ -88,25 +93,34 @@ class WebSocketManagerTest(AioHTTPTestCase):
self.assertEqual(msg['response'], '')
self.assertIsNone(msg['validation_error'])
# Verify initial auto approver config
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
self.assertIn('config', msg)
@unittest_run_loop
async def test_approve_context(self):
"""Test context approval flow"""
async with self.client.ws_connect('/ws') as ws:
for _ in range(3):
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
await self._skip_initial_messages(ws)
await ws.send_json({
'type': ClientMessage.APPROVE_CONTEXT
})
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
self.assertEqual(msg['response'], "")
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)
@@ -116,53 +130,51 @@ class WebSocketManagerTest(AioHTTPTestCase):
test_input = "test input"
async with self.client.ws_connect('/ws') as ws:
for _ in range(3):
await ws.receive_json()
await self._skip_initial_messages(ws)
# Send input
await ws.send_json({
'type': ClientMessage.SEND_INPUT,
'input': test_input
})
# Delay for processing
await asyncio.sleep(0.5)
# Verify input was added to buffer
self.assertEqual(self.io_buffer.read(), test_input)
@unittest_run_loop
async def test_modify_response(self):
"""Test modifying response in RESPONSE_APPROVAL state"""
# Set agent to RESPONSE_APPROVAL state
async with self.client.ws_connect('/ws') as ws:
for _ in range(3):
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
await self._skip_initial_messages(ws)
await ws.send_json({
'type': ClientMessage.APPROVE_CONTEXT
})
# Skip to RESPONSE_APPROVAL state
for _ in range(4):
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
await ws.receive_json()
modified_response = "<reasoning>modified test</reasoning>"
await ws.send_json({
'type': ClientMessage.MODIFY_RESPONSE,
'response': modified_response
})
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
self.assertEqual(msg['response'], modified_response)
@unittest_run_loop
async def test_multiple_clients(self):
"""Test broadcasting to multiple WebSocket clients"""
# Connect first client
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
for _ in range(3):
await asyncio.wait_for(ws1.receive_json(), timeout=0.5)
await asyncio.wait_for(ws2.receive_json(), timeout=0.5)
await self._skip_initial_messages(ws1)
await self._skip_initial_messages(ws2)
await ws1.send_json({
'type': ClientMessage.APPROVE_CONTEXT
})
msg1 = await ws1.receive_json()
msg2 = await ws2.receive_json()
self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE)
@@ -170,6 +182,60 @@ class WebSocketManagerTest(AioHTTPTestCase):
self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)
@unittest_run_loop
async def test_auto_approver_config(self):
"""Test setting and receiving auto approver config"""
async with self.client.ws_connect('/ws') as ws:
await self._skip_initial_messages(ws)
test_config = {
'context_enabled': True,
'response_enabled': True,
'context_timeout': 10.0,
'response_timeout': 20.0
}
await ws.send_json({
'type': ClientMessage.AUTO_APPROVER_CONFIG,
'config': test_config
})
msg = await ws.receive_json()
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
self.assertEqual(msg['config'], test_config)
self.ws_manager.auto_approver.context_enabled = False
self.ws_manager.auto_approver.response_enabled = False
@unittest_run_loop
async def test_auto_approver_broadcast(self):
"""Test broadcasting auto approver config changes to multiple clients"""
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
await self._skip_initial_messages(ws1)
await self._skip_initial_messages(ws2)
test_config = {
'context_enabled': True,
'response_enabled': False,
'context_timeout': 15.0,
'response_timeout': 25.0
}
await ws1.send_json({
'type': ClientMessage.AUTO_APPROVER_CONFIG,
'config': test_config
})
msg1 = await ws1.receive_json()
msg2 = await ws2.receive_json()
self.assertEqual(msg1['type'], ServerMessage.AUTO_APPROVER_CONFIG)
self.assertEqual(msg2['type'], ServerMessage.AUTO_APPROVER_CONFIG)
self.assertEqual(msg1['config'], test_config)
self.assertEqual(msg2['config'], test_config)
self.ws_manager.auto_approver.context_enabled = False
self.ws_manager.auto_approver.response_enabled = False
def tearDown(self):
"""Clean up resources"""
super().tearDown()

18
web/src/components/App.jsx
View File

@@ -23,6 +23,12 @@ const App = () => {
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('ws://localhost:8080/ws');
@@ -44,10 +50,20 @@ const App = () => {
case ServerMessageType.OUTPUT_UPDATE:
setOutput(message.output);
break;
case ServerMessageType.AUTO_APPROVER_CONFIG:
setAutoApproverConfig(message.config);
break;
}
});
}, [addMessageHandler]);
const handleAutoApproverConfigChange = (config) => {
sendMessage({
type: ClientMessageType.AUTO_APPROVER_CONFIG,
config: config
});
};
const handleApprove = () => {
if (agentState === AgentState.CONTEXT_APPROVAL) {
sendMessage({
@@ -158,6 +174,8 @@ const App = () => {
onToggleDiff={() => setShowDiff(!showDiff)}
onSendInput={handleSendInput}
onClearOutput={handleClearOutput}
autoApproverConfig={autoApproverConfig}
onAutoApproverConfigChange={handleAutoApproverConfigChange}
/>
<FloatingButton onClick={() => setShowSidebar(!showSidebar)} />

48
web/src/components/Sidebar.jsx
View File

@@ -1,5 +1,7 @@
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';
export const Sidebar = ({
@@ -7,10 +9,12 @@ export const Sidebar = ({
agentState,
showDiff,
input,
autoApproverConfig,
onApprove,
onToggleDiff,
onSendInput,
onClearOutput,
onAutoApproverConfigChange,
}) => (
<aside className={`
fixed top-0 right-0 h-screen w-64 bg-white shadow-lg
@@ -19,6 +23,50 @@ export const Sidebar = ({
${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
`}>
<div className="p-4 space-y-4 md:pt-0">
<div className="space-y-2">
<div className="flex items-center justify-between">
<label>Auto-approve Context</label>
<Switch
checked={autoApproverConfig.context_enabled}
onCheckedChange={(checked) => onAutoApproverConfigChange({
...autoApproverConfig,
context_enabled: checked
})}
/>
</div>
<Input
type="number"
value={autoApproverConfig.context_timeout}
onChange={(e) => onAutoApproverConfigChange({
...autoApproverConfig,
context_timeout: parseFloat(e.target.value)
})}
disabled={autoApproverConfig.context_enabled}
/>
</div>
<div className="space-y-2">
<div className="flex items-center justify-between">
<label>Auto-approve Response</label>
<Switch
checked={autoApproverConfig.response_enabled}
onCheckedChange={(checked) => onAutoApproverConfigChange({
...autoApproverConfig,
response_enabled: checked
})}
/>
</div>
<Input
type="number"
value={autoApproverConfig.response_timeout}
onChange={(e) => onAutoApproverConfigChange({
...autoApproverConfig,
response_timeout: parseFloat(e.target.value)
})}
disabled={autoApproverConfig.response_enabled}
/>
</div>
<Button
variant="outline"
onClick={onToggleDiff}

21
web/src/components/ui/switch.jsx Normal file
View File

@@ -0,0 +1,21 @@
import React from 'react';
export const Switch = ({ checked, onCheckedChange, disabled }) => (
<label className="relative inline-flex items-center cursor-pointer">
<input
type="checkbox"
className="sr-only peer"
checked={checked}
onChange={(e) => onCheckedChange(e.target.checked)}
disabled={disabled}
/>
<div className={`
w-11 h-6 bg-gray-200 rounded-full peer
peer-checked:after:translate-x-full peer-checked:bg-blue-600
after:content-[''] after:absolute after:top-[2px] after:left-[2px]
after:bg-white after:rounded-full after:h-5 after:w-5
after:transition-all
${disabled ? 'opacity-50 cursor-not-allowed' : ''}
`}></div>
</label>
);

2
web/src/constants.jsx
View File

@@ -24,6 +24,7 @@ export const ClientMessageType = {
APPROVE_RESPONSE: 'APPROVE_RESPONSE',
SEND_INPUT: 'SEND_INPUT',
CLEAR_OUTPUT: 'CLEAR_OUTPUT',
AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
};
export const ServerMessageType = {
@@ -31,4 +32,5 @@ export const ServerMessageType = {
CONTEXT_UPDATE: 'CONTEXT_UPDATE',
RESPONSE_UPDATE: 'RESPONSE_UPDATE',
OUTPUT_UPDATE: 'OUTPUT_UPDATE',
AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
};