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

Fixed auto approver

Browse Source
This commit is contained in:
Niels Geens
2024-11-22 17:11:06 +01:00
parent 8766a945c0
commit 69f4f8a6fc
8 changed files with 343 additions and 244 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
sia/__main__.py
View File

@@ -1,7 +1,7 @@
from aiohttp import web
import asyncio
import mimetypes
from .auto_approver import AutoApprover
from .config import Config
from .hf_llm_engine import HfLlmEngine
from .iteration_logger import IterationLogger
@@ -10,11 +10,11 @@ from .mistral_llm_engine import MistralLlmEngine
from .openai_llm_engine import OpenAILlmEngine
from .response_parser import ResponseParser
from .system_metrics import SystemMetrics
from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer
from .web.api import Api
from .web.static import Static
from .web.websockts import Websockets
from .web_agent import WebAgent
from .web_io_buffer import WebIOBuffer
from .working_memory import WorkingMemory
from .xml_validator import XMLValidator
@@ -74,10 +74,11 @@ class Main:
parser=ResponseParser(self._io_buffer),
iteration_logger=IterationLogger(self._config.iterations_dir, self._system_prompt, self._action_schema),
)
self._auto_approver = AutoApprover(self._agent)
self._app = web.Application()
self._api = Api(self._app, self._agent, self._io_buffer)
self._websockets = Websockets(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._auto_approver)
self._static = Static(self._app, self._config)
return self

81
sia/auto_approver.py
View File

@@ -1,12 +1,13 @@
from threading import Thread, Event
from typing import Callable, TypeAlias, TypedDict
from .web_agent import WebAgent, WebAgentState
from .web_agent import WebAgent, LlmState
class AutoApproverConfig(TypedDict):
context_enabled: bool
response_enabled: bool
context_timeout: float
response_timeout: float
llm_name: str
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
@@ -20,6 +21,7 @@ class AutoApprover:
Initialize auto approver with a WebAgent instance.
"""
self.agent = agent
self._llm_name = next(iter(agent.llms.keys()))
self._context_timeout = 5.0
self._response_timeout = 10.0
self._context_enabled = False
@@ -29,7 +31,8 @@ class AutoApprover:
self._response_thread: Thread | None = None
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
def config(self) -> AutoApproverConfig:
@@ -37,25 +40,29 @@ class AutoApprover:
context_enabled=self._context_enabled,
response_enabled=self._response_enabled,
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:
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 config['llm_name'] not in self.agent.llms:
raise ValueError(f"Unknown LLM: {config['llm_name']}")
if self._context_enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
self._start_context_thread()
else:
self._stop_context_thread()
notify_config_change = self._notify_config_change
self._notify_config_change = lambda: None
try:
self.context_enabled = False
self.response_enabled = False
if self._response_enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
self._start_response_thread()
else:
self._stop_response_thread()
self.context_timeout = config['context_timeout']
self.response_timeout = config['response_timeout']
self.llm_name = config['llm_name']
self.context_enabled = config['context_enabled']
self.response_enabled = config['response_enabled']
finally:
self._notify_config_change = notify_config_change
self._notify_config_change()
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
@@ -98,7 +105,7 @@ class AutoApprover:
return
self._context_enabled = enabled
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._notify_config_change()
@@ -112,21 +119,36 @@ class AutoApprover:
return
self._response_enabled = enabled
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._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()
@property
def llm_name(self) -> str:
return self._llm_name
@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()
elif self._response_enabled:
else:
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:
if self._context_thread:
self._stop_event.set()
@@ -150,13 +172,12 @@ class AutoApprover:
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()
if self._context_enabled:
self.agent.run_inference(self._llm_name)
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()
self.agent.llms[self._llm_name] == LlmState.OUTPUT):
self.agent.approve_response(self._llm_name, self.agent.get_output(self._llm_name))

92
sia/web/api.py
View File

@@ -2,15 +2,16 @@ from aiohttp import web
import json
import asyncio
from ..auto_approver import AutoApprover
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):
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover):
self._app = app
self._agent = agent
self._io_buffer = io_buffer
self._auto_approver = auto_approver
self._init_routes()
@@ -23,10 +24,13 @@ class Api:
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
self._app.router.add_get("/api/auto_approver/config", self._get_auto_approver_config)
self._app.router.add_post("/api/auto_approver/config", self._set_auto_approver_config)
self._app.router.add_post("/api/auto_approver/context_enabled", self._set_context_enabled)
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:
"""Start inference on specified LLM."""
@@ -95,3 +99,79 @@ class Api:
}),
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
View 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

8
sia/web/websockts.py
View File

@@ -2,19 +2,23 @@ from aiohttp import web
from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer
from ..auto_approver import AutoApprover
from .llm_websocket import LlmWebSocket
from .context_websocket import ContextWebSocket
from .token_websocket import TokenWebSocket
from .stdout_websocket import StdoutWebSocket
from .auto_approver_websocket import AutoApproverWebSocket
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._context_ws = ContextWebSocket(agent)
self._token_ws = TokenWebSocket(agent)
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/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)
app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection)
app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)

173
sia/web_socket_manager.py
View File

@@ -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