Shared response buffer backend
This commit is contained in:
Niels Geens
594
sia/web/api.py
594
sia/web/api.py
@@ -1,286 +1,308 @@
|
||||
from pathlib import Path
|
||||
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,
|
||||
work_dir: Path,
|
||||
app: web.Application,
|
||||
agent: WebAgent,
|
||||
io_buffer: WebIOBuffer,
|
||||
working_memory: WorkingMemory,
|
||||
auto_approver: AutoApprover
|
||||
):
|
||||
self._work_dir = work_dir
|
||||
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/inference/{llm}/stop", self._stop_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 _stop_inference(self, request: web.Request) -> web.Response:
|
||||
"""Stop inference on specified LLM."""
|
||||
llm_name = request.match_info["llm"]
|
||||
try:
|
||||
self._agent.stop_inference(llm_name)
|
||||
return web.Response(status=200)
|
||||
except ValueError 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._work_dir, 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._work_dir, self._io_buffer)
|
||||
|
||||
for entry in entries:
|
||||
self._working_memory.add_entry(entry)
|
||||
|
||||
return web.Response(status=200)
|
||||
from pathlib import Path
|
||||
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,
|
||||
work_dir: Path,
|
||||
app: web.Application,
|
||||
agent: WebAgent,
|
||||
io_buffer: WebIOBuffer,
|
||||
working_memory: WorkingMemory,
|
||||
auto_approver: AutoApprover
|
||||
):
|
||||
self._work_dir = work_dir
|
||||
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/inference/{llm}/stop", self._stop_inference)
|
||||
self._app.router.add_post("/api/response", self._set_response)
|
||||
self._app.router.add_post("/api/response/approve", self._approve_response)
|
||||
self._app.router.add_get("/api/response", self._get_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 _stop_inference(self, request: web.Request) -> web.Response:
|
||||
"""Stop inference on specified LLM."""
|
||||
llm_name = request.match_info["llm"]
|
||||
try:
|
||||
self._agent.stop_inference(llm_name)
|
||||
return web.Response(status=200)
|
||||
except ValueError as e:
|
||||
return web.Response(status=400, text=str(e))
|
||||
|
||||
async def _set_response(self, request: web.Request) -> web.Response:
|
||||
"""Edit the shared response buffer"""
|
||||
try:
|
||||
data = await request.json()
|
||||
text = data.get("response")
|
||||
if text is None:
|
||||
return web.Response(status=400, text="Missing response text in request body")
|
||||
try:
|
||||
self._agent._response_buffer.set_text(text)
|
||||
return web.Response(status=200)
|
||||
except ValueError as e:
|
||||
return web.Response(status=400, text=str(e))
|
||||
except json.JSONDecodeError:
|
||||
return web.Response(status=400, text="Invalid JSON in request body")
|
||||
|
||||
async def _approve_response(self, request: web.Request) -> web.Response:
|
||||
"""Approve current buffer content"""
|
||||
data = await request.json()
|
||||
try:
|
||||
self._agent.approve_response()
|
||||
return web.Response(status=200)
|
||||
except ValueError as e:
|
||||
return web.Response(status=400, text=str(e))
|
||||
|
||||
async def _get_response(self, request: web.Request) -> web.Response:
|
||||
"""Get current buffer content"""
|
||||
return web.Response(
|
||||
text=json.dumps({
|
||||
"response": self._agent.output_buffer.get_text(),
|
||||
}),
|
||||
content_type="application/json"
|
||||
)
|
||||
|
||||
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._work_dir, 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._work_dir, self._io_buffer)
|
||||
|
||||
for entry in entries:
|
||||
self._working_memory.add_entry(entry)
|
||||
|
||||
return web.Response(status=200)
|
||||
|
||||
@@ -1,24 +1,21 @@
|
||||
from aiohttp import web, WSMsgType
|
||||
from typing import Dict, Set
|
||||
import asyncio
|
||||
import json
|
||||
from aiohttp import web
|
||||
|
||||
from .util import wrap_async
|
||||
from ..web_agent import WebAgent
|
||||
from .util import wrap_async
|
||||
|
||||
class TokenWebSocket:
|
||||
"""
|
||||
WebSocket handler for LLM token streaming.
|
||||
Broadcasts new tokens to all connected clients.
|
||||
"""
|
||||
|
||||
|
||||
class ResponseWebSocket:
|
||||
def __init__(self, agent: WebAgent):
|
||||
self._agent = agent
|
||||
self._clients: Set[web.WebSocketResponse] = set()
|
||||
self._agent.add_token_handler(wrap_async(self._handle_new_token))
|
||||
|
||||
async def _broadcast_message(self, message: Dict):
|
||||
"""Broadcast message to all connected clients."""
|
||||
self._clients = set()
|
||||
self._agent._response_buffer.add_observer(wrap_async(self._handle_buffer_change))
|
||||
|
||||
async def _handle_buffer_change(self, buffer):
|
||||
await self._broadcast_message({
|
||||
"buffer": buffer,
|
||||
})
|
||||
|
||||
async def _broadcast_message(self, message):
|
||||
disconnected = set()
|
||||
for ws in self._clients:
|
||||
try:
|
||||
@@ -26,26 +23,22 @@ class TokenWebSocket:
|
||||
except ConnectionResetError:
|
||||
disconnected.add(ws)
|
||||
self._clients -= disconnected
|
||||
|
||||
async def _handle_new_token(self, llm_name: str, token: str):
|
||||
"""Handle new tokens from the WebAgent."""
|
||||
await self._broadcast_message({
|
||||
"llm": llm_name,
|
||||
"token": token
|
||||
})
|
||||
|
||||
|
||||
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:
|
||||
await ws.send_json({
|
||||
"buffer": self._agent.response_buffer.get_text(),
|
||||
})
|
||||
|
||||
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
|
||||
return ws
|
||||
@@ -1,28 +1,28 @@
|
||||
from aiohttp import web
|
||||
|
||||
from ..auto_approver import AutoApprover
|
||||
from ..web_agent import WebAgent
|
||||
from ..web_io_buffer import WebIOBuffer
|
||||
from ..working_memory import WorkingMemory
|
||||
from .auto_approver_websocket import AutoApproverWebSocket
|
||||
from .context_websocket import ContextWebSocket
|
||||
from .llm_websocket import LlmWebSocket
|
||||
from .memory_websocket import MemoryWebSocket
|
||||
from .stdout_websocket import StdoutWebSocket
|
||||
from .token_websocket import TokenWebSocket
|
||||
|
||||
class Websockets:
|
||||
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
|
||||
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)
|
||||
self._memory_ws = MemoryWebSocket(working_memory)
|
||||
|
||||
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/auto_approver", self._auto_approver_ws.handle_connection)
|
||||
app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
|
||||
from aiohttp import web
|
||||
|
||||
from ..auto_approver import AutoApprover
|
||||
from ..web_agent import WebAgent
|
||||
from ..web_io_buffer import WebIOBuffer
|
||||
from ..working_memory import WorkingMemory
|
||||
from .auto_approver_websocket import AutoApproverWebSocket
|
||||
from .context_websocket import ContextWebSocket
|
||||
from .llm_websocket import LlmWebSocket
|
||||
from .memory_websocket import MemoryWebSocket
|
||||
from .response_websocket import ResponseWebSocket
|
||||
from .stdout_websocket import StdoutWebSocket
|
||||
|
||||
class Websockets:
|
||||
def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
|
||||
self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
|
||||
self._context_ws = ContextWebSocket(agent)
|
||||
self._llm_ws = LlmWebSocket(agent)
|
||||
self._memory_ws = MemoryWebSocket(working_memory)
|
||||
self._response_ws = ResponseWebSocket(agent)
|
||||
self._stdout_ws = StdoutWebSocket(io_buffer)
|
||||
|
||||
app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)
|
||||
app.router.add_get("/ws/context", self._context_ws.handle_connection)
|
||||
app.router.add_get("/ws/llm", self._llm_ws.handle_connection)
|
||||
app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
|
||||
app.router.add_get("/ws/response", self._response_ws.handle_connection)
|
||||
app.router.add_get("/ws/stdout", self._stdout_ws.handle_connection)
|
||||
Reference in New Issue
Block a user