Files
SIA/sia/web/api.py
2024-11-24 20:35:28 +01:00

274 lines
12 KiB
Python

from aiohttp import web
import json
import asyncio
from ..auto_approver import AutoApprover
from ..entry.entry_factory import EntryFactory
from ..iteration_parser import IterationParser
from ..web_agent import WebAgent
from ..web_io_buffer import WebIOBuffer
from ..working_memory import WorkingMemory
class Api:
def __init__(
self,
app: web.Application,
agent: WebAgent,
io_buffer: WebIOBuffer,
working_memory: WorkingMemory,
auto_approver: AutoApprover
):
self._app = app
self._agent = agent
self._working_memory = working_memory
self._io_buffer = io_buffer
self._auto_approver = auto_approver
self._init_routes()
def _init_routes(self):
"""Initialize REST API and WebSocket routes."""
self._app.router.add_post("/api/inference/{llm}", self._run_inference)
self._app.router.add_post("/api/approve/{llm}", self._approve_response)
self._app.router.add_post("/api/context", self._modify_context)
self._app.router.add_post("/api/input", self._send_input)
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)
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)
self._app.router.add_get("/api/memory", self._get_memory)
self._app.router.add_post("/api/memory/entry", self._create_entry)
self._app.router.add_put("/api/memory/entry/{id}", self._save_entry)
self._app.router.add_delete("/api/memory/entry/{id}", self._delete_entry)
self._app.router.add_post("/api/memory/entry/{id}/reset", self._reset_entry)
self._app.router.add_post("/api/memory/entry/{id}/update", self._update_entry)
self._app.router.add_post("/api/memory/load_iteration", self._load_iteration)
async def _run_inference(self, request: web.Request) -> web.Response:
"""Start inference on specified LLM."""
llm_name = request.match_info["llm"]
try:
await asyncio.get_event_loop().run_in_executor(None, self._agent.run_inference, llm_name)
return web.Response(status=200)
except (ValueError, RuntimeError) as e:
return web.Response(status=400, text=str(e))
async def _approve_response(self, request: web.Request) -> web.Response:
"""Approve response from specified LLM."""
llm_name = request.match_info["llm"]
data = await request.json()
response = data.get("response")
if not response:
return web.Response(status=400, text="Missing response in request body")
try:
self._agent.approve_response(llm_name, response)
return web.Response(status=200)
except ValueError as e:
return web.Response(status=400, text=str(e))
async def _modify_context(self, request: web.Request) -> web.Response:
"""Modify the current context."""
data = await request.json()
context = data.get("context")
if not context:
return web.Response(status=400, text="Missing context in request body")
self._agent.modify_context(context)
return web.Response(status=200)
async def _send_input(self, request: web.Request) -> web.Response:
"""Send input to the IO buffer."""
data = await request.json()
input_text = data.get("input")
if not input_text:
return web.Response(status=400, text="Missing input in request body")
self._io_buffer.append_stdin(input_text)
return web.Response(status=200)
async def _clear_output(self, request: web.Request) -> web.Response:
"""Clear the stdout buffer."""
self._io_buffer.clear_stdout()
return web.Response(status=200)
async def _get_output(self, request: web.Request) -> web.Response:
"""Get complete output for specified LLM."""
llm_name = request.match_info["llm"]
try:
output = self._agent.get_output(llm_name)
return web.Response(
text=json.dumps({"output": output}),
content_type="application/json"
)
except ValueError as e:
return web.Response(status=400, text=str(e))
async def _get_llms(self, request: web.Request) -> web.Response:
"""Get all LLMs and their current states."""
states = self._agent.llms
return web.Response(
text=json.dumps({
name: state.name
for name, state in states.items()
}),
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))
async def _get_memory(self, request: web.Request) -> web.Response:
"""Get complete working memory state."""
entries = self._working_memory.get_entries()
return web.Response(
text=json.dumps([e.serialize() for e in entries]),
content_type="application/json"
)
async def _create_entry(self, request: web.Request) -> web.Response:
"""Create a new entry in working memory."""
data = await request.json()
try:
entry = EntryFactory.create_entry(data, self._io_buffer)
self._working_memory.add_entry(entry)
return web.Response(
text=json.dumps({"id": entry.id}),
content_type="application/json"
)
except (ValueError, TypeError) as e:
return web.Response(status=400, text=str(e))
async def _save_entry(self, request: web.Request) -> web.Response:
"""Update properties of an existing entry."""
entry_id = request.match_info["id"]
data = await request.json()
entry = self._working_memory.get_entry(entry_id)
if not entry:
return web.Response(status=404, text="Entry not found")
try:
EntryFactory.update_entry(entry, data)
entry.notify_change()
return web.Response(status=200)
except ValueError as e:
return web.Response(status=400, text=str(e))
async def _delete_entry(self, request: web.Request) -> web.Response:
"""Delete an entry from working memory."""
entry_id = request.match_info["id"]
self._working_memory.remove_entry(entry_id)
return web.Response(status=200)
async def _reset_entry(self, request: web.Request) -> web.Response:
"""Reset an entry's state."""
entry_id = request.match_info["id"]
entry = self._working_memory.get_entry(entry_id)
if not entry:
return web.Response(status=404, text="Entry not found")
entry.reset()
return web.Response(status=200)
async def _update_entry(self, request: web.Request) -> web.Response:
"""Update an entry's state."""
entry_id = request.match_info["id"]
entry = self._working_memory.get_entry(entry_id)
if not entry:
return web.Response(status=404, text="Entry not found")
try:
entry.update()
entry.notify_change()
return web.Response(status=200)
except ValueError as e:
return web.Response(status=400, text=str(e))
async def _load_iteration(self, request: web.Request) -> web.Response:
"""Load entries from iteration XML content into working memory"""
data = await request.json()
content = data.get("content")
if not content:
return web.Response(status=400, text="Missing content in request body")
entries = IterationParser.parse_iteration(content, self._io_buffer)
for entry in entries:
self._working_memory.add_entry(entry)
return web.Response(status=200)