step B.1 complete: WS context + frame schemas
This commit is contained in:
68
app/core/ws_context.py
Normal file
68
app/core/ws_context.py
Normal file
@@ -0,0 +1,68 @@
|
||||
"""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
|
||||
|
||||
from contextvars import ContextVar
|
||||
from typing import Any, Callable, Coroutine
|
||||
from uuid import uuid4
|
||||
|
||||
# Holds the execute callback for the current WS session.
|
||||
# Set by the chat WS handler before the orchestrator runs; cleared after.
|
||||
_client_executor: ContextVar[Callable[[dict], Coroutine[Any, Any, dict]]] = ContextVar(
|
||||
"_client_executor"
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
|
||||
return await callback(payload)
|
||||
@@ -5,6 +5,7 @@ Mirrors the TypeScript types from the Electron app (src/shared/api-types.ts).
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
@@ -155,3 +156,54 @@ class PluginListResponse(BaseModel):
|
||||
|
||||
class PluginInstallRequest(BaseModel):
|
||||
plugin_id: str
|
||||
|
||||
|
||||
# ── WebSocket Frame Protocol ──────────────────────────────────────────
|
||||
|
||||
class WsFrameType(str, Enum):
|
||||
chat_request = "chat_request"
|
||||
text_chunk = "text_chunk"
|
||||
tool_call = "tool_call"
|
||||
tool_result = "tool_result"
|
||||
final = "final"
|
||||
ping = "ping"
|
||||
|
||||
|
||||
class WsToolCall(BaseModel):
|
||||
"""Server → Client: requests a CRUD/vector operation on the local DB."""
|
||||
|
||||
type: Literal[WsFrameType.tool_call] = WsFrameType.tool_call
|
||||
id: str
|
||||
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
|
||||
|
||||
|
||||
class WsToolResult(BaseModel):
|
||||
"""Client → Server: result of a CRUD/vector operation."""
|
||||
|
||||
type: Literal[WsFrameType.tool_result] = WsFrameType.tool_result
|
||||
id: str
|
||||
row: dict[str, Any] | None = None
|
||||
rows: list[dict[str, Any]] | None = None
|
||||
results: list[dict[str, Any]] | None = None
|
||||
deleted: bool | None = None
|
||||
ok: bool | None = None
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class WsTextChunk(BaseModel):
|
||||
"""Server → Client: incremental LLM response text."""
|
||||
|
||||
type: Literal[WsFrameType.text_chunk] = WsFrameType.text_chunk
|
||||
text: str
|
||||
|
||||
|
||||
class WsFinal(BaseModel):
|
||||
"""Server → Client: signals end of response with the complete text."""
|
||||
|
||||
type: Literal[WsFrameType.final] = WsFrameType.final
|
||||
response: str
|
||||
|
||||
Reference in New Issue
Block a user