Server side implementation for interactive entry edits

sia/web/api.py

@@ -5,11 +5,21 @@ import asyncio
 from ..auto_approver import AutoApprover
 from ..web_agent import WebAgent
 from ..web_io_buffer import WebIOBuffer
+from ..working_memory import WorkingMemory
+from ..entry.entry_factory import EntryFactory
 
 class Api:
-    def __init__(self, app: web.Application, agent: WebAgent, io_buffer: WebIOBuffer, auto_approver: AutoApprover):
+    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
         
@@ -31,6 +41,11 @@ class Api:
         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._update_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)
 
     async def _run_inference(self, request: web.Request) -> web.Response:
         """Start inference on specified LLM."""
@@ -174,4 +189,53 @@ class Api:
             self._auto_approver.llm_name = name
             return web.Response(status=200)
         except ValueError as e:
-            return web.Response(status=400, text=str(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 _update_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)
+            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)

sia/web/memory_websocket.py

@@ -0,0 +1,59 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+
+from ..entry import Entry
+from ..web_agent import WebAgent
+from ..working_memory import WorkingMemory
+from .util import wrap_async
+
+class MemoryWebSocket:
+    """
+    WebSocket handler for working memory changes.
+    Broadcasts memory updates to all connected clients.
+    """
+    
+    def __init__(self, working_memory: WorkingMemory):
+        self._working_memory = working_memory
+        self._clients: Set[web.WebSocketResponse] = set()
+        self._working_memory.add_change_handler(wrap_async(self._handle_memory_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_memory_change(self):
+        """Handle working memory changes."""
+        entries = self._working_memory.get_entries()
+        await self._broadcast_message({
+            "type": "memory_state",
+            "entries": [e.serialize() for e in entries]
+        })
+
+    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+        
+        try:
+            # Send initial memory state
+            entries = self._working_memory.get_entries()
+            await ws.send_json({
+                "type": "memory_state",
+                "entries": [e.serialize() for e in entries]
+            })
+        
+            self._clients.add(ws)
+                
+            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

sia/web/websockts.py

@@ -1,24 +1,28 @@
 from aiohttp import web
 
+from ..auto_approver import AutoApprover
 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 ..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):
+    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)