theo-agent-dashboard

chat.py (6914B)

---

 """WebSocket chat endpoint — bridges client ↔ Hermes via SSE."""
 
 from __future__ import annotations
 
 import json
 import logging
 from typing import Any
 
 import aiohttp
 from fastapi import APIRouter, WebSocket, WebSocketDisconnect
 
 from app.config import get_config
 from app.hermes_client import _strip_untrusted_tags
 from app.routes.attachments import _fetch_metadata
 from app.sse_parser import parse_sse_events
 from app.ws_manager import ConnectionManager
 
 logger = logging.getLogger(__name__)
 
 chat_router = APIRouter()
 manager = ConnectionManager()
 
 
 async def _send_error(ws: WebSocket, message: str) -> None:  # pragma: no mutate: block
     """Send an error payload to the WebSocket client."""
     await ws.send_json({"type": "error", "message": message})
 
 
 def _map_sse_event_name(hermes_event: str) -> str:
     """Map Hermes SSE event names to frontend-expected names.
 
     The frontend's ChatView switch statement uses these exact type strings:
       run.started, message.started, assistant.delta, assistant.completed,
       tool.started, tool.completed, tool.failed, tool.progress, run.completed
 
     Hermes uses the same names, so this is a passthrough for known events.
     """
     return hermes_event
 
 
 def _map_sse_event_data(hermes_event: str, data: dict[str, Any]) -> dict[str, Any]:
     """Map Hermes SSE event data fields to what the frontend expects.
 
     Hermes field names → Frontend field names:
       assistant.delta:  message_id→id, delta→content
       message.started:  message{id,role}→{id,role} (unwrap nested)
       assistant.completed: message_id→id
       tool.started:     tool_name→name, args stays, preview→args
       tool.completed:   tool_name→name
       tool.failed:      tool_name→name
       tool.progress:    tool_name→name
     """
     if hermes_event == "assistant.delta":
         return {
             "id": data.get("message_id", ""),
             "content": data.get("delta", ""),
         }
     if hermes_event == "message.started":
         msg = data.get("message", data)
         return {
             "id": msg.get("id", ""),
             "role": msg.get("role", "assistant"),
             "model": msg.get("model"),
         }
     if hermes_event == "assistant.completed":
         return {
             "id": data.get("message_id", ""),
             "content": data.get("content", ""),
             "completed": data.get("completed", True),
         }
     if hermes_event in ("tool.started", "tool.completed", "tool.failed"):
         tool_id = data.get("tool_call_id", "") or data.get("tool_name", "")
         return {
             "messageId": data.get("message_id", ""),
             "toolCall": {
                 "id": tool_id,
                 "name": data.get("tool_name", ""),
                 "args": json.dumps(data.get("args", "")) if isinstance(data.get("args"), (dict, list)) else str(data.get("args", "")),
             },
             "toolCallId": tool_id,
             "result": _strip_untrusted_tags(str(data.get("preview", data.get("result", "")))),
             "error": data.get("error", ""),
         }
     if hermes_event == "tool.progress":
         return {
             "messageId": data.get("message_id", ""),
             "toolCallId": data.get("tool_call_id", data.get("tool_name", "")),
             "progress": data.get("delta", data.get("preview", "")),
         }
     # run.started, run.completed, error, done — pass through
     return data
 
 
 async def _stream_hermes_sse(  # pragma: no mutate: block
     session: aiohttp.ClientSession,
     session_id: str,
     content: str,
     ws: WebSocket,
     base_url: str,
     api_key: str,
 ) -> None:
     """POST to the SSE chat stream endpoint and forward events to the WS client.
 
     Hermes's session chat stream is a POST endpoint that accepts a message
     body and returns SSE events.  This function:
       1. POSTs the message to trigger the agent turn
       2. Reads the SSE response
       3. Forwards each event to the frontend with mapped field names
     """
     url = f"{base_url}/api/sessions/{session_id}/chat/stream"
     headers = {"Authorization": f"Bearer {api_key}"}
     payload = {"message": content}
     async with session.post(url, json=payload, headers=headers) as resp:
         resp.raise_for_status()
         text = await resp.text()
         for event in parse_sse_events(text):
             event_name = _map_sse_event_name(event["event"])
             event_data = _map_sse_event_data(event["event"], event["data"])
             await ws.send_json({"type": event_name, "data": event_data})
 
 
 async def _load_attachment_content(attachment_ids: list[str]) -> str:
     """Resolve attachment IDs to content for the Hermes message.
 
     Both images and non-image files have their local disk path appended.
     Returns a string to append to the message.
     """
     parts: list[str] = []
     for att_id in attachment_ids:
         meta = _fetch_metadata(att_id)
         if meta is None:
             continue
         disk_path = meta.get("disk_path", "")
         if disk_path:
             parts.append(disk_path)
     return "\n".join(parts)
 
 
 async def _handle_chat(  # pragma: no mutate: block
     ws: WebSocket,
     payload: dict,
     cfg: Any,
 ) -> None:
     """Process a single 'chat' message from the client."""
     # The WsClient wraps sends as {type, data}, so fields live in payload["data"]
     data = payload.get("data", payload)
     session_id: str = data["session_id"]
     content: str = data["content"]
     attachment_ids: list[str] | None = data.get("attachments")
 
     if attachment_ids:
         attachment_content = await _load_attachment_content(attachment_ids)
         if attachment_content:
             content = f"{content}\n\n{attachment_content}"
 
     async with aiohttp.ClientSession() as session:
         await _stream_hermes_sse(
             session, session_id, content, ws,
             cfg.hermes_base_url, cfg.hermes_api_key,
         )
 
 
 async def ws_chat_handler(ws: WebSocket) -> None:  # pragma: no mutate: block
     """Main WebSocket handler for the chat endpoint."""
     await ws.accept()
     await ws.send_json({"type": "status", "data": {"connected": True}})
     await manager.connect("default", ws)
     try:
         cfg = get_config()
         while True:
             payload = await ws.receive_json()
             msg_type = payload.get("type")
             if msg_type == "chat":
                 await _handle_chat(ws, payload, cfg)
             else:
                 await _send_error(ws, f"Unknown type: {msg_type}")
     except WebSocketDisconnect:
         pass
     except Exception:
         logger.exception("WebSocket chat error")
         await _send_error(ws, "Internal server error")
     finally:
         manager.disconnect("default")
 
 
 @chat_router.websocket("/ws/chat")
 async def ws_chat_route(websocket: WebSocket) -> None:  # pragma: no mutate: block
     """FastAPI WebSocket route at ``/ws/chat``."""
     await ws_chat_handler(websocket)