90018af311
WS Gateway:
- WebSocket lifecycle handler with RS256 JWT auth
- Redis bridge: device registry, frame publishing, tool_result routing
- Inbound routing: tool_result→LPUSH, home/floating→chat pub/sub
- Outbound: subscribes to ws:out:{user_id}, forwards to Electron
- Single-worker Dockerfile (long-lived WS connections)
Chat Service:
- Redis consumer: subscribes to chat:request:* pattern
- Redis-based ws_context: tool_call→publish, BRPOP tool_result (30s timeout)
- deep_agent: single-agent runner with home/floating/stream variants
- memory_middleware: core/associative/episodic/proactive memory with Fernet
- Domain agents: task (8 tools), note (5), project (6), timeline (4)
- LLM factory via LiteLLM (100+ providers)
- Output formatter (StreamFormatter)
- POST /chat REST fallback with Traefik header auth
- Multi-worker Dockerfile with 120s timeout for LLM calls
116 lines
3.3 KiB
Python
116 lines
3.3 KiB
Python
"""WebSocket context for Chat Service — Redis-based tool call round-trip.
|
|
|
|
Replaces the monolith's ws_context.py. Instead of calling Electron directly
|
|
via WebSocket, this publishes tool_call frames to Redis (ws:out:{user_id})
|
|
and awaits the result via BRPOP on tool:result:{call_id}.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from contextvars import ContextVar
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from shared.redis import redis_client, tool_result_key, ws_out_channel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_TOOL_CALL_TIMEOUT = 30 # seconds — BRPOP timeout
|
|
|
|
# Per-request user_id context var (set before agent runs)
|
|
_current_user_id: ContextVar[str | None] = ContextVar("_current_user_id", default=None)
|
|
|
|
# Optional collector for debug
|
|
_tool_result_collector: ContextVar[list[dict] | None] = ContextVar(
|
|
"_tool_result_collector", default=None
|
|
)
|
|
|
|
|
|
def set_current_user(user_id: str) -> None:
|
|
_current_user_id.set(user_id)
|
|
|
|
|
|
def clear_current_user() -> None:
|
|
_current_user_id.set(None)
|
|
|
|
|
|
def set_tool_result_collector(lst: list[dict]) -> None:
|
|
_tool_result_collector.set(lst)
|
|
|
|
|
|
def clear_tool_result_collector() -> None:
|
|
_tool_result_collector.set(None)
|
|
|
|
|
|
async def execute_on_client(
|
|
action: str,
|
|
table: str | None = None,
|
|
data: dict[str, Any] | None = None,
|
|
filters: dict[str, Any] | None = None,
|
|
vector: list[float] | None = None,
|
|
limit: int | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Send a tool_call to Electron via Redis and await the result.
|
|
|
|
1. Build tool_call payload
|
|
2. Publish to ws:out:{user_id} (WS Gateway forwards to Electron)
|
|
3. BRPOP on tool:result:{call_id} (WS Gateway pushes when Electron replies)
|
|
4. Return result dict
|
|
|
|
Raises RuntimeError if no user_id is set or if the call times out.
|
|
"""
|
|
user_id = _current_user_id.get()
|
|
if not user_id:
|
|
raise RuntimeError(
|
|
"execute_on_client() called without a user_id — "
|
|
"set_current_user() must be called first."
|
|
)
|
|
|
|
call_id = str(uuid4())
|
|
payload: dict[str, Any] = {
|
|
"type": "tool_call",
|
|
"id": call_id,
|
|
"action": action,
|
|
}
|
|
if table is not None:
|
|
payload["table"] = table
|
|
if data is not None:
|
|
payload["data"] = data
|
|
if filters is not None:
|
|
payload["filters"] = {k: v for k, v in filters.items() if v is not None}
|
|
if vector is not None:
|
|
payload["vector"] = vector
|
|
if limit is not None:
|
|
payload["limit"] = limit
|
|
|
|
# Publish tool_call to WS Gateway → Electron
|
|
channel = ws_out_channel(user_id)
|
|
await redis_client.publish(channel, json.dumps(payload))
|
|
|
|
# Wait for Electron's tool_result
|
|
result_key = tool_result_key(call_id)
|
|
response = await redis_client.brpop(result_key, timeout=_TOOL_CALL_TIMEOUT)
|
|
|
|
if response is None:
|
|
raise RuntimeError(
|
|
f"Tool call {call_id} timed out after {_TOOL_CALL_TIMEOUT}s — "
|
|
f"device may be offline or unresponsive."
|
|
)
|
|
|
|
# response is (key, value) tuple
|
|
_, raw = response
|
|
result = json.loads(raw)
|
|
|
|
# Collect for debug if requested
|
|
collector = _tool_result_collector.get(None)
|
|
if collector is not None:
|
|
collector.append({
|
|
"action": action,
|
|
"table": table,
|
|
"data": result,
|
|
})
|
|
|
|
return result
|