roberto
24a9c1b752
- Replace langgraph create_react_agent with deepagents create_deep_agent - Sub-agents now configured as SubAgent dicts dispatched via built-in task tool - Stream filter updated: langgraph_node 'agent' → 'model' - Accept both AIMessage and AIMessageChunk in stream filter - Collector only captures write mutations (insert/update/delete) - Add deepagents>=0.4.10 to requirements.txt
101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
"""WebSocket client executor context.
|
|
|
|
Holds a per-request async callback that tools call to execute CRUD
|
|
operations on the Electron client's local SQLite / LanceDB databases.
|
|
The callback sends a `tool_call` WS frame and awaits the `tool_result`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from contextvars import ContextVar
|
|
from typing import Any, Callable, Coroutine
|
|
from uuid import uuid4
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Holds the execute callback for the current WS session.
|
|
# Set by the chat WS handler before the deep agent runs; cleared after.
|
|
_client_executor: ContextVar[Callable[[dict], Coroutine[Any, Any, dict]]] = ContextVar(
|
|
"_client_executor"
|
|
)
|
|
|
|
# Optional collector that captures raw execute_on_client results.
|
|
# Set by the deep agent tool loop to capture CRUD mutations.
|
|
_tool_result_collector: ContextVar[list[dict] | None] = ContextVar(
|
|
"_tool_result_collector", default=None
|
|
)
|
|
|
|
|
|
def set_tool_result_collector(lst: list[dict]) -> None:
|
|
"""Register *lst* as the collector for this async context."""
|
|
_tool_result_collector.set(lst)
|
|
|
|
|
|
def clear_tool_result_collector() -> None:
|
|
"""Clear the collector (best-effort)."""
|
|
_tool_result_collector.set(None)
|
|
|
|
|
|
def set_client_executor(fn: Callable[[dict], Coroutine[Any, Any, dict]]) -> None:
|
|
"""Bind *fn* as the executor for the current async context (task/coroutine)."""
|
|
_client_executor.set(fn)
|
|
|
|
|
|
def clear_client_executor() -> None:
|
|
"""Remove the executor binding (best-effort; ContextVar resets on task exit)."""
|
|
try:
|
|
_client_executor.set(None) # type: ignore[arg-type]
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
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 CRUD/vector operation to the Electron client and return the result.
|
|
|
|
Builds a ``tool_call`` payload, invokes the per-session WS callback,
|
|
and returns the ``tool_result`` dict from Electron.
|
|
|
|
Raises ``RuntimeError`` if no executor is set (i.e. called outside a WS session).
|
|
"""
|
|
callback = _client_executor.get(None)
|
|
if callback is None:
|
|
raise RuntimeError(
|
|
"execute_on_client() called outside a WebSocket session — "
|
|
"no client executor is set."
|
|
)
|
|
|
|
payload: dict[str, Any] = {"id": str(uuid4()), "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
|
|
|
|
logger.info("execute_on_client: sending payload action=%s table=%s id=%s", action, table, payload["id"])
|
|
result = await callback(payload)
|
|
if result is None:
|
|
logger.error("execute_on_client: callback returned None for action=%s table=%s id=%s", action, table, payload["id"])
|
|
else:
|
|
logger.info("execute_on_client: got result type=%s keys=%s", type(result).__name__, list(result.keys()) if isinstance(result, dict) else "N/A")
|
|
collector = _tool_result_collector.get(None)
|
|
if collector is not None and action in ("insert", "update", "delete"):
|
|
collector.append({
|
|
"action": action,
|
|
"table": table,
|
|
"data": data or {},
|
|
})
|
|
return result
|