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

Approve context

Browse Source
This commit is contained in:
geens
2024-11-02 19:17:27 +01:00
parent 96d14fc46d
commit 2a15708145
9 changed files with 296 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
test/mock_web_agent.py
View File

@@ -1,4 +1,3 @@
import asyncio
from typing import Optional, List, Callable
from sia.web_agent import WebAgentState
@@ -21,35 +20,34 @@ class MockWebAgent:
def add_response_change_handler(self, handler):
self._response_handlers.append(handler)
def set_input(self, input_text):
pass
async def _notify_state_handlers(self, new_state: WebAgentState):
def _notify_state_handlers(self, new_state: WebAgentState):
"""Notify all state handlers of a state change."""
for handler in self._state_handlers:
await handler(new_state)
handler(new_state)
async def _notify_response_handlers(self, response: str):
def _notify_response_handlers(self, response: str):
"""Notify all response handlers of a response change."""
for handler in self._response_handlers:
await handler(response)
handler(response)
async def _set_response(self, response: str):
def _set_response(self, response: str):
"""Set response and notify handlers."""
self.response = response
await self._notify_response_handlers(response)
self._notify_response_handlers(response)
async def approve_context(self):
def approve_context(self):
"""Handle context approval with immediate state updates."""
# Change to INFERENCE state and notify
self._current_state = WebAgentState.INFERENCE
for handler in self._state_handlers:
await handler(WebAgentState.INFERENCE)
self._notify_state_handlers(WebAgentState.INFERENCE)
# Set response and notify
self.response = "<reasoning>test</reasoning>"
for handler in self._response_handlers:
await handler(self.response)
self._set_response("<reasoning>test</reasoning>")
# Change to RESPONSE_APPROVAL state and notify
self._current_state = WebAgentState.RESPONSE_APPROVAL
for handler in self._state_handlers:
await handler(WebAgentState.RESPONSE_APPROVAL)
self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL)

5
test/standard_io_buffer_test.py
View File

@@ -159,11 +159,6 @@ if __name__ == '__main__':
stdout, stderr, code = self.run_script('read', special_chars)
self.assertEqual(code, 0, f"Script failed with stderr: {stderr}")
expected = f'Read: {special_chars}'
# Print actual bytes for debugging
if stdout != expected:
print("\nDebug output:")
print("Expected bytes:", expected.encode('utf-8'))
print("Actual bytes:", stdout.encode('utf-8'))
self.assertEqual(stdout, expected)
def test_empty_input(self):

26
test/web_agent_test.py
View File

@@ -1,3 +1,4 @@
import asyncio
import unittest
from datetime import datetime
from unittest.mock import Mock, MagicMock, patch, call
@@ -123,7 +124,7 @@ class WebAgentTest(unittest.TestCase):
self.agent._current_state = WebAgentState.UPDATE
with self.assertRaises(Exception) as context:
self.agent.approve_context()
asyncio.run(self.agent.approve_context())
self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
def test_approve_response_state_error(self):
@@ -132,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
self.agent.approve_response()
self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
def test_context_approval_flow(self):
async def test_context_approval_flow(self):
"""Test complete context approval flow with state transitions."""
# Register handlers
self.agent.add_state_change_handler(self.state_change_handler)
@@ -140,6 +141,7 @@ class WebAgentTest(unittest.TestCase):
# Setup LLM response
def mock_infer(prompt: str, context: str) -> Iterator[str]:
# Return complete response at once for testing
yield "<reasoning>test reasoning</reasoning>"
self.mock_llm.infer.side_effect = mock_infer
@@ -149,6 +151,9 @@ class WebAgentTest(unittest.TestCase):
# Approve context
self.agent.approve_context()
# Wait for response tasks to complete
await asyncio.gather(*self.web_server._response_tasks)
# Verify state transitions
expected_states = [
WebAgentState.INFERENCE,
@@ -156,10 +161,11 @@ class WebAgentTest(unittest.TestCase):
]
self.assertEqual(self.state_changes, expected_states)
# Verify response updates
self.assertEqual(len(self.response_changes), 2) # Empty + final response
self.assertEqual(self.response_changes[0], "")
self.assertEqual(self.response_changes[1], "<reasoning>test reasoning</reasoning>")
# Updated test: Should now expect streaming updates
self.assertEqual(len(self.response_changes), len(self.received_messages))
for i, msg in enumerate(self.received_messages):
self.assertEqual(msg["type"], "RESPONSE_UPDATE")
self.assertEqual(msg["data"], self.response_changes[i])
def test_response_approval_flow_command(self):
"""Test response approval flow with command execution."""
@@ -200,7 +206,7 @@ class WebAgentTest(unittest.TestCase):
# Set initial state and response
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
self.agent._response = "<reasoning>test reasoning</reasoning>"
self.agent.response = "<reasoning>test reasoning</reasoning>"
# Approve response
self.agent.approve_response()
@@ -218,7 +224,7 @@ class WebAgentTest(unittest.TestCase):
]
self.assertEqual(self.state_changes, expected_states)
def test_response_validation(self):
async def test_response_validation(self):
"""Test response validation during response setting."""
# Register handlers
self.agent.add_response_change_handler(self.response_change_handler)
@@ -227,8 +233,8 @@ class WebAgentTest(unittest.TestCase):
error_message = "Invalid XML"
self.mock_validator.validate.return_value = error_message
# Set response
self.agent._set_response("<invalid>")
# Set response using async method
await self.agent._set_response("<invalid>")
# Verify validation error stored
self.assertEqual(self.agent._validation_error, error_message)

36
test/web_server_test.py
View File

@@ -6,6 +6,7 @@ from sia.web_agent import WebAgentState
from unittest.mock import patch
import asyncio
import asyncio
import time
import json
import json
import shutil
@@ -29,6 +30,7 @@ class WebServerTest(AioHTTPTestCase):
self.static_dir.mkdir(exist_ok=True)
# Create web server with temp static directory
self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
self.web_server = WebServer(
self.agent,
self.io_buffer,
@@ -54,7 +56,6 @@ class WebServerTest(AioHTTPTestCase):
elif msg.type == WSMsgType.BINARY:
messages.append(json.loads(msg.data.decode()))
elif msg.type == WSMsgType.CLOSED:
print("WebSocket closed normally")
break
elif msg.type == WSMsgType.ERROR:
self.fail(f"WebSocket error while waiting for messages: {msg.data}")
@@ -108,6 +109,7 @@ class WebServerTest(AioHTTPTestCase):
await ws.send_json({
"type": ClientMessage.APPROVE_CONTEXT,
})
await asyncio.sleep(0.1) # Let state changes propagate
# Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state
messages = await self.wait_for_messages(ws, count=3)
@@ -124,7 +126,7 @@ class WebServerTest(AioHTTPTestCase):
self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)
@unittest_run_loop
@unittest_run_loop
async def test_validation_error(self):
"""Test handling of validation errors."""
async with self.client.ws_connect("/ws") as ws:
@@ -134,7 +136,7 @@ class WebServerTest(AioHTTPTestCase):
# Set validation error and response
self.agent._validation_error = "Invalid XML"
await self.agent._set_response("<invalid>")
self.agent._set_response("<invalid>") # Now synchronous
await asyncio.sleep(0.1) # Let changes propagate
# Should receive response update with error
@@ -204,24 +206,34 @@ class WebServerTest(AioHTTPTestCase):
await ws1.close()
await ws2.close()
@unittest_run_loop
async def wait_for_buffer_content(self, timeout=1.0):
start = time.time()
while time.time() - start < timeout:
if self.io_buffer.buffer_length() > 0:
return True
await asyncio.sleep(0.01)
return False
async def test_send_input(self):
"""Test sending input to agent."""
async with self.client.ws_connect("/ws") as ws:
# Drain initial messages
await self.drain_messages(ws)
await asyncio.sleep(0.1) # Let connection settle
# Send input
await asyncio.sleep(0.1)
test_input = "test input"
await ws.send_json({
"type": ClientMessage.SEND_INPUT,
"data": test_input
})
await asyncio.sleep(0.1) # Let input process
# Wait for buffer to have content
self.assertTrue(await self.wait_for_buffer_content())
# Verify input added to buffer with newline
self.assertEqual(self.io_buffer.read(), test_input + "\n")
content = self.io_buffer.read()
self.assertEqual(
content,
test_input + "\n",
f"Expected '{test_input}\\n' but got '{content}'"
)
@unittest_run_loop
async def test_invalid_json(self):