初始化提交
This commit is contained in:
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1 @@
|
||||
"""外部集成层。"""
|
||||
@@ -0,0 +1,187 @@
|
||||
"""ChromaDB 向量存储客户端。
|
||||
|
||||
实现 ``VectorStoreProtocol``。
|
||||
|
||||
内存模式(无外部依赖,本地开发默认):``chromadb.Client()`` 自动用 SQLite 持久化到
|
||||
本地文件,无需启动独立服务。生产环境配置 ``CHROMADB_HOST`` / ``CHROMADB_PORT`` 环境变量
|
||||
切换为 HTTP 客户端模式。
|
||||
|
||||
字符级隔离:每个 character_id 对应一个 ChromaDB collection,互不干扰。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import chromadb
|
||||
from chromadb.config import Settings as ChromaSettings
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.logging import logger
|
||||
|
||||
# 默认持久化路径(如果未通过 env 配置则 fallback)
|
||||
_DEFAULT_PERSIST_DIR = settings.chroma_persist_directory or str(
|
||||
settings.backend_dir / "../db/chroma"
|
||||
)
|
||||
|
||||
|
||||
class ChromaDBVectorStore:
|
||||
"""ChromaDB 向量存储后端。
|
||||
|
||||
用法::
|
||||
|
||||
store = ChromaDBVectorStore()
|
||||
store.add_memory(character_id="char_xxx", memory_id="mem_xxx",
|
||||
text="记忆内容", embedding=[0.1, 0.2, ...])
|
||||
results = store.search(character_id="char_xxx", query=[0.1, ...], limit=5)
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._client: chromadb.Client | None = None
|
||||
self._collections: dict[str, chromadb.Collection] = {}
|
||||
|
||||
# ── 生命周期 ──────────────────────────────────────────────────────────
|
||||
|
||||
@property
|
||||
def client(self) -> chromadb.Client:
|
||||
if self._client is None:
|
||||
self._client = self._create_client()
|
||||
return self._client
|
||||
|
||||
def _create_client(self) -> chromadb.Client:
|
||||
"""创建 ChromaDB 客户端(内存模式优先)。"""
|
||||
host = settings.chromadb_host or ""
|
||||
if host:
|
||||
return chromadb.HttpClient(
|
||||
host=host,
|
||||
port=settings.chromadb_port,
|
||||
)
|
||||
# 内存模式:SQLite 持久化到本地文件
|
||||
return chromadb.Client(
|
||||
ChromaSettings(
|
||||
persist_directory=_DEFAULT_PERSIST_DIR,
|
||||
anonymized_telemetry=False,
|
||||
),
|
||||
)
|
||||
|
||||
# ── Collection ────────────────────────────────────────────────────────
|
||||
|
||||
def _collection(self, character_id: str) -> chromadb.Collection:
|
||||
"""按 character_id 取 collection(延迟创建)。"""
|
||||
if character_id not in self._collections:
|
||||
name = f"char_{character_id}"
|
||||
try:
|
||||
self._collections[character_id] = self.client.get_or_create_collection(
|
||||
name=name,
|
||||
metadata={"character_id": character_id, "hnsw:space": "cosine"},
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(f"ChromaDB collection 创建失败: {exc}")
|
||||
raise
|
||||
return self._collections[character_id]
|
||||
|
||||
# ── 增 ────────────────────────────────────────────────────────────────
|
||||
|
||||
def add_memory(
|
||||
self,
|
||||
character_id: str,
|
||||
memory_id: str,
|
||||
text: str,
|
||||
embedding: list[float],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""添加一条记忆向量。"""
|
||||
try:
|
||||
coll = self._collection(character_id)
|
||||
coll.add(
|
||||
ids=[memory_id],
|
||||
embeddings=[embedding],
|
||||
metadatas=[{
|
||||
"text": text[:1000],
|
||||
**(metadata or {}),
|
||||
}],
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"vector_store_add_failed",
|
||||
memory_id=memory_id,
|
||||
error=str(exc),
|
||||
)
|
||||
|
||||
def update_memory(
|
||||
self,
|
||||
character_id: str,
|
||||
memory_id: str,
|
||||
text: str,
|
||||
embedding: list[float],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""更新记忆向量(覆盖原有 embedding)。"""
|
||||
try:
|
||||
coll = self._collection(character_id)
|
||||
coll.update(
|
||||
ids=[memory_id],
|
||||
embeddings=[embedding],
|
||||
metadatas=[{
|
||||
"text": text[:1000],
|
||||
**(metadata or {}),
|
||||
}],
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"vector_store_update_failed",
|
||||
memory_id=memory_id,
|
||||
error=str(exc),
|
||||
)
|
||||
|
||||
def delete_memory(self, character_id: str, memory_id: str) -> None:
|
||||
"""删除记忆向量。"""
|
||||
try:
|
||||
coll = self._collection(character_id)
|
||||
coll.delete(ids=[memory_id])
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"vector_store_delete_failed",
|
||||
memory_id=memory_id,
|
||||
error=str(exc),
|
||||
)
|
||||
|
||||
# ── 查 ────────────────────────────────────────────────────────────────
|
||||
|
||||
def search(
|
||||
self,
|
||||
character_id: str,
|
||||
query: list[float],
|
||||
*,
|
||||
limit: int = 5,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""向量相似度搜索,返回 ``[{id, score, metadata}]``。"""
|
||||
try:
|
||||
coll = self._collection(character_id)
|
||||
results = coll.query(
|
||||
query_embeddings=[query],
|
||||
n_results=limit,
|
||||
include=["distances", "metadatas"],
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("vector_store_search_failed", error=str(exc))
|
||||
return []
|
||||
|
||||
ids = results.get("ids", [[]])[0] or []
|
||||
distances = results.get("distances", [[]])[0] or []
|
||||
metadatas = results.get("metadatas", [[]])[0] or []
|
||||
|
||||
return [
|
||||
{"id": mid, "score": 1.0 - dist, "metadata": meta or {}}
|
||||
for mid, dist, meta in zip(ids, distances, metadatas, strict=False)
|
||||
]
|
||||
|
||||
def count(self, character_id: str) -> int:
|
||||
"""collection 中向量总数。"""
|
||||
try:
|
||||
coll = self._collection(character_id)
|
||||
return coll.count()
|
||||
except Exception:
|
||||
return 0
|
||||
@@ -0,0 +1,93 @@
|
||||
"""纯内存向量存储后端,供测试用。
|
||||
|
||||
不持久化、不依赖外部服务,每次进程重启清空。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class MemoryVectorStore:
|
||||
"""纯内存向量存储(测试用)。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._data: dict[str, dict[str, dict[str, Any]]] = {}
|
||||
"""{character_id: {memory_id: {"text": ..., "embedding": [...], "metadata": {...}}}}"""
|
||||
|
||||
# ── 增 ────────────────────────────────────────────────────────────────
|
||||
|
||||
def add_memory(
|
||||
self,
|
||||
character_id: str,
|
||||
memory_id: str,
|
||||
text: str,
|
||||
embedding: list[float],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
char_store = self._data.setdefault(character_id, {})
|
||||
if memory_id not in char_store:
|
||||
char_store[memory_id] = {}
|
||||
char_store[memory_id] = {
|
||||
"text": text,
|
||||
"embedding": embedding,
|
||||
"metadata": metadata or {},
|
||||
}
|
||||
|
||||
def update_memory(
|
||||
self,
|
||||
character_id: str,
|
||||
memory_id: str,
|
||||
text: str,
|
||||
embedding: list[float],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
self.add_memory(character_id, memory_id, text, embedding, metadata=metadata)
|
||||
|
||||
def delete_memory(self, character_id: str, memory_id: str) -> None:
|
||||
char_store = self._data.get(character_id)
|
||||
if char_store:
|
||||
char_store.pop(memory_id, None)
|
||||
|
||||
# ── 查 ────────────────────────────────────────────────────────────────
|
||||
|
||||
def search(
|
||||
self,
|
||||
character_id: str,
|
||||
query: list[float],
|
||||
*,
|
||||
limit: int = 5,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""余弦相似度暴力搜索(O(n),仅测试用)。"""
|
||||
char_store = self._data.get(character_id)
|
||||
if not char_store:
|
||||
return []
|
||||
|
||||
scored: list[tuple[float, str, dict[str, Any]]] = []
|
||||
for mid, entry in char_store.items():
|
||||
emb = entry.get("embedding", [])
|
||||
if not emb:
|
||||
continue
|
||||
score = _cosine_similarity(query, emb)
|
||||
scored.append((score, mid, entry.get("metadata", {})))
|
||||
|
||||
scored.sort(key=lambda x: x[0], reverse=True)
|
||||
return [
|
||||
{"id": mid, "score": score, "metadata": meta}
|
||||
for score, mid, meta in scored[:limit]
|
||||
]
|
||||
|
||||
def count(self, character_id: str) -> int:
|
||||
char_store = self._data.get(character_id)
|
||||
return len(char_store) if char_store else 0
|
||||
|
||||
|
||||
def _cosine_similarity(a: list[float], b: list[float]) -> float:
|
||||
dot = sum(x * y for x, y in zip(a, b, strict=False))
|
||||
norm_a = sum(x * x for x in a) ** 0.5
|
||||
norm_b = sum(x * x for x in b) ** 0.5
|
||||
if norm_a == 0 or norm_b == 0:
|
||||
return 0.0
|
||||
return dot / (norm_a * norm_b)
|
||||
@@ -0,0 +1 @@
|
||||
"""MCP server 接入(简化 HTTP client,非完整 JSON-RPC)。"""
|
||||
@@ -0,0 +1,106 @@
|
||||
"""MCP 客户端 — Streamable HTTP 与 stdio 均用的 JSON-RPC 解析工具。
|
||||
|
||||
两种传输层的具体实现在 ``transports/`` 子包中。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.core.config import settings
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
def _auth_headers(auth_token: str | None) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {auth_token}"} if auth_token else {}
|
||||
|
||||
|
||||
def _parse_call_result(data: dict) -> ToolResult:
|
||||
"""解析 JSON-RPC tools/call 响应为 ToolResult。
|
||||
|
||||
供 StreamableHttpTransport 和 StdioTransport 共用。
|
||||
"""
|
||||
if data.get("error"):
|
||||
return ToolResult(
|
||||
ok=False,
|
||||
stdout="",
|
||||
stderr=data["error"].get("message", "MCP 错误"),
|
||||
duration_ms=0,
|
||||
exit_code=1,
|
||||
timeout=False,
|
||||
truncated=False,
|
||||
trace_id="",
|
||||
)
|
||||
result = data.get("result") if isinstance(data, dict) else None
|
||||
if not isinstance(result, dict):
|
||||
return ToolResult(
|
||||
ok=False,
|
||||
stdout=str(result) if result else "",
|
||||
stderr="无效响应格式",
|
||||
duration_ms=0,
|
||||
exit_code=1,
|
||||
timeout=False,
|
||||
truncated=False,
|
||||
trace_id="",
|
||||
)
|
||||
content = result.get("content", [])
|
||||
stdout_parts = []
|
||||
for item in content:
|
||||
if item.get("type") in ("text", "json"):
|
||||
stdout_parts.append(item.get("text", ""))
|
||||
truncated = bool((result.get("meta") or {}).get("truncated", False))
|
||||
return ToolResult(
|
||||
ok=not result.get("isError", False),
|
||||
stdout="".join(stdout_parts),
|
||||
stderr="",
|
||||
duration_ms=0,
|
||||
exit_code=0 if not result.get("isError", False) else 1,
|
||||
timeout=False,
|
||||
truncated=truncated,
|
||||
trace_id="",
|
||||
)
|
||||
|
||||
|
||||
async def call_tool(
|
||||
http_client: httpx.AsyncClient,
|
||||
endpoint_url: str,
|
||||
tool_name: str,
|
||||
arguments: dict,
|
||||
auth_token: str | None = None,
|
||||
meta: dict | None = None,
|
||||
) -> ToolResult:
|
||||
"""通过 Streamable HTTP 调用 MCP 工具(供 _make_mcp_executor 使用)。"""
|
||||
body: dict[str, Any] = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": tool_name,
|
||||
"arguments": arguments,
|
||||
},
|
||||
}
|
||||
if meta:
|
||||
body["meta"] = meta
|
||||
try:
|
||||
resp = await http_client.post(
|
||||
endpoint_url,
|
||||
json=body,
|
||||
headers=_auth_headers(auth_token),
|
||||
timeout=settings.mcp_call_timeout_seconds,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
except (httpx.HTTPError, ValueError, TypeError) as exc:
|
||||
return ToolResult(
|
||||
ok=False,
|
||||
stdout="",
|
||||
stderr=str(exc),
|
||||
duration_ms=0,
|
||||
exit_code=None,
|
||||
timeout=False,
|
||||
truncated=False,
|
||||
trace_id="",
|
||||
)
|
||||
return _parse_call_result(data)
|
||||
@@ -0,0 +1,96 @@
|
||||
"""MCP stdio 子进程生命周期管理。
|
||||
|
||||
全局单例,维护 server_id → TransportProtocol 的映射,
|
||||
前端通过 API 控制 start/stop/restart/status。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from app.integrations.mcp.transports.base import TransportProtocol
|
||||
from app.integrations.mcp.transports.http import StreamableHttpTransport
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 模块级全局实例(避免依赖 request.app.state)
|
||||
_default_instance: ProcessManager | None = None
|
||||
|
||||
|
||||
def get_process_manager() -> ProcessManager | None:
|
||||
"""返回默认的全局 ProcessManager 实例。"""
|
||||
return _default_instance
|
||||
|
||||
|
||||
def set_process_manager(manager: ProcessManager) -> None:
|
||||
"""设置全局 ProcessManager 实例(由 lifespan 调用)。"""
|
||||
global _default_instance
|
||||
_default_instance = manager
|
||||
|
||||
|
||||
class ProcessManager:
|
||||
"""MCP 进程管理器(全局单例)。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._transports: dict[str, TransportProtocol] = {}
|
||||
|
||||
def register_transport(self, server_id: str, transport: TransportProtocol) -> None:
|
||||
"""注册传输层实例(不自动 start)。"""
|
||||
self._transports[server_id] = transport
|
||||
|
||||
def unregister(self, server_id: str) -> None:
|
||||
"""注销传输层实例(不自动 stop,调用方需先 stop)。"""
|
||||
self._transports.pop(server_id, None)
|
||||
|
||||
def get(self, server_id: str) -> TransportProtocol | None:
|
||||
return self._transports.get(server_id)
|
||||
|
||||
async def start(self, server_id: str) -> None:
|
||||
transport = self._transports.get(server_id)
|
||||
if transport is None:
|
||||
raise ValueError(f"server {server_id} 未注册")
|
||||
if isinstance(transport, StreamableHttpTransport):
|
||||
raise ValueError("streamable-http 无需启动")
|
||||
await transport.start()
|
||||
|
||||
async def stop(self, server_id: str) -> None:
|
||||
transport = self._transports.get(server_id)
|
||||
if transport is None:
|
||||
raise ValueError(f"server {server_id} 未注册")
|
||||
if isinstance(transport, StreamableHttpTransport):
|
||||
raise ValueError("streamable-http 无需停止")
|
||||
await transport.stop()
|
||||
|
||||
async def restart(self, server_id: str) -> None:
|
||||
transport = self._transports.get(server_id)
|
||||
if transport is None:
|
||||
raise ValueError(f"server {server_id} 未注册")
|
||||
if isinstance(transport, StreamableHttpTransport):
|
||||
raise ValueError("streamable-http 无需重启")
|
||||
await transport.restart()
|
||||
|
||||
async def status(self, server_id: str) -> dict[str, Any]:
|
||||
transport = self._transports.get(server_id)
|
||||
if transport is None:
|
||||
return {
|
||||
"server_id": server_id,
|
||||
"name": "unknown",
|
||||
"transport": "unknown",
|
||||
"status": "unregistered",
|
||||
"pid": None,
|
||||
"uptime_seconds": None,
|
||||
}
|
||||
return await transport.get_status()
|
||||
|
||||
async def shutdown_all(self) -> None:
|
||||
"""停止所有 stdio 子进程(lifespan 关闭时调用)。"""
|
||||
for sid, transport in list(self._transports.items()):
|
||||
try:
|
||||
await transport.stop()
|
||||
except Exception: # noqa: BLE001
|
||||
logger.warning(
|
||||
"停止 MCP 进程失败",
|
||||
extra={"server_id": sid},
|
||||
)
|
||||
self._transports.clear()
|
||||
@@ -0,0 +1,67 @@
|
||||
"""MCP 传输层工厂。
|
||||
|
||||
根据 server 的 transport 字段创建对应的 TransportProtocol 实例。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.integrations.mcp.transports.base import TransportProtocol
|
||||
from app.integrations.mcp.transports.http import StreamableHttpTransport
|
||||
from app.integrations.mcp.transports.stdio import StdioTransport
|
||||
|
||||
|
||||
def get_transport(
|
||||
server: Any,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
) -> TransportProtocol:
|
||||
"""根据 server 的 transport 字段创建传输层实例。"""
|
||||
transport = server.transport if hasattr(server, "transport") else "streamable-http"
|
||||
|
||||
if transport == "stdio":
|
||||
args: list[str] = []
|
||||
if hasattr(server, "args_json") and server.args_json:
|
||||
try:
|
||||
parsed = json.loads(server.args_json)
|
||||
if isinstance(parsed, list):
|
||||
args = [str(a) for a in parsed]
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
env: dict[str, str] | None = None
|
||||
if hasattr(server, "env_json") and server.env_json:
|
||||
try:
|
||||
parsed = json.loads(server.env_json)
|
||||
if isinstance(parsed, dict):
|
||||
env = {str(k): str(v) for k, v in parsed.items()}
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
cwd: str | None = None
|
||||
if hasattr(server, "cwd") and server.cwd:
|
||||
cwd = server.cwd
|
||||
|
||||
cmd = getattr(server, "command", None) or ""
|
||||
return StdioTransport(
|
||||
server_id=server.id,
|
||||
server_name=server.name,
|
||||
command=cmd,
|
||||
args=args,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
)
|
||||
|
||||
# 默认 streamable-http
|
||||
endpoint_url = getattr(server, "endpoint_url", None) or ""
|
||||
auth_token = getattr(server, "auth_token", None)
|
||||
return StreamableHttpTransport(
|
||||
server_id=server.id,
|
||||
server_name=server.name,
|
||||
endpoint_url=endpoint_url,
|
||||
auth_token=auth_token,
|
||||
http_client=http_client,
|
||||
)
|
||||
@@ -0,0 +1,57 @@
|
||||
"""TransportProtocol — MCP 传输层抽象接口。
|
||||
|
||||
定义所有传输层(streamable-http / stdio)共有的 list_tools / call_tool / 生命周期方法。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
class TransportProtocol(ABC):
|
||||
"""MCP 传输层协议接口。"""
|
||||
|
||||
server_id: str
|
||||
server_name: str
|
||||
|
||||
@abstractmethod
|
||||
async def list_tools(self) -> list[dict[str, Any]]:
|
||||
"""调用 tools/list,返回工具清单。"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
|
||||
"""调用 tools/call,返回 ToolResult。"""
|
||||
...
|
||||
|
||||
async def start(self) -> None:
|
||||
"""启动传输层(stdio 启子进程;HTTP no-op)。"""
|
||||
return
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""停止传输层(stdio 杀子进程;HTTP no-op)。"""
|
||||
return
|
||||
|
||||
async def restart(self) -> None:
|
||||
"""重启传输层。"""
|
||||
await self.stop()
|
||||
await self.start()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""传输层是否正在运行。"""
|
||||
return True
|
||||
|
||||
async def get_status(self) -> dict[str, Any]:
|
||||
"""返回当前状态信息。"""
|
||||
return {
|
||||
"server_id": self.server_id,
|
||||
"name": self.server_name,
|
||||
"transport": "unknown",
|
||||
"status": "unknown",
|
||||
"pid": None,
|
||||
"uptime_seconds": None,
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Streamable HTTP 传输层实现。
|
||||
|
||||
通过 HTTP POST /message 发送 JSON-RPC 2.0 请求的标准 MCP 传输协议。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.core.config import settings
|
||||
from app.integrations.mcp.client import _auth_headers, _parse_call_result
|
||||
from app.integrations.mcp.transports.base import TransportProtocol
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
|
||||
class StreamableHttpTransport(TransportProtocol):
|
||||
"""通过 HTTP POST /message 通信的 MCP 传输层。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_id: str,
|
||||
server_name: str,
|
||||
endpoint_url: str,
|
||||
auth_token: str | None = None,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
) -> None:
|
||||
self.server_id = server_id
|
||||
self.server_name = server_name
|
||||
self._endpoint_url = endpoint_url.rstrip("/")
|
||||
self._auth_token = auth_token
|
||||
self._http_client = http_client or httpx.AsyncClient(trust_env=False)
|
||||
self._own_client = http_client is None
|
||||
self._started_at: float | None = None
|
||||
|
||||
async def list_tools(self) -> list[dict[str, Any]]:
|
||||
body = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/list",
|
||||
"params": {},
|
||||
}
|
||||
try:
|
||||
resp = await self._http_client.post(
|
||||
self._endpoint_url,
|
||||
json=body,
|
||||
headers=_auth_headers(self._auth_token),
|
||||
timeout=settings.mcp_list_timeout_seconds,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
except (httpx.HTTPError, ValueError):
|
||||
return []
|
||||
result = data.get("result") if isinstance(data, dict) else None
|
||||
if not result:
|
||||
return []
|
||||
return result.get("tools") if isinstance(result.get("tools"), list) else []
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
|
||||
body = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": name,
|
||||
"arguments": arguments,
|
||||
},
|
||||
}
|
||||
try:
|
||||
resp = await self._http_client.post(
|
||||
self._endpoint_url,
|
||||
json=body,
|
||||
headers=_auth_headers(self._auth_token),
|
||||
timeout=settings.mcp_call_timeout_seconds,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
except (httpx.HTTPError, ValueError, TypeError) as exc:
|
||||
return ToolResult(
|
||||
ok=False,
|
||||
stdout="",
|
||||
stderr=str(exc),
|
||||
duration_ms=0,
|
||||
exit_code=None,
|
||||
timeout=False,
|
||||
truncated=False,
|
||||
trace_id="",
|
||||
)
|
||||
return _parse_call_result(data)
|
||||
|
||||
async def start(self) -> None:
|
||||
self._started_at = time.time()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return True
|
||||
|
||||
async def get_status(self) -> dict[str, Any]:
|
||||
uptime = None
|
||||
if self._started_at is not None:
|
||||
uptime = time.time() - self._started_at
|
||||
return {
|
||||
"server_id": self.server_id,
|
||||
"name": self.server_name,
|
||||
"transport": "streamable-http",
|
||||
"status": "running",
|
||||
"pid": None,
|
||||
"uptime_seconds": uptime,
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
"""stdio 传输层实现。
|
||||
|
||||
通过子进程 stdin/stdout 通信的标准 MCP 传输协议(本地 MCP 集成)。
|
||||
JSON-RPC 2.0 请求写入 stdin,从 stdout 读取响应,stderr 记日志。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from app.core.config import settings
|
||||
from app.integrations.mcp.client import _parse_call_result
|
||||
from app.integrations.mcp.transports.base import TransportProtocol
|
||||
from app.modules.tool_runtime.definition import ToolResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StdioTransport(TransportProtocol):
|
||||
"""通过子进程 stdin/stdout 通信的 MCP 传输层。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_id: str,
|
||||
server_name: str,
|
||||
command: str,
|
||||
args: list[str] | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
cwd: str | None = None,
|
||||
) -> None:
|
||||
self.server_id = server_id
|
||||
self.server_name = server_name
|
||||
self._command = command
|
||||
self._args = args or []
|
||||
self._env = env
|
||||
self._cwd = cwd
|
||||
self._process: asyncio.subprocess.Process | None = None
|
||||
self._request_id = 0
|
||||
self._started_at: float | None = None
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def start(self) -> None:
|
||||
if self._process is not None and self._process.returncode is None:
|
||||
return # already running
|
||||
|
||||
resolved_env = os.environ.copy()
|
||||
if self._env:
|
||||
resolved_env.update(self._env)
|
||||
|
||||
self._process = await asyncio.create_subprocess_exec(
|
||||
self._command,
|
||||
*self._args,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=resolved_env,
|
||||
cwd=self._cwd,
|
||||
)
|
||||
self._started_at = time.time()
|
||||
logger.info(
|
||||
"stdio MCP 子进程已启动",
|
||||
extra={
|
||||
"server_id": self.server_id,
|
||||
"server_name": self.server_name,
|
||||
"pid": self._process.pid,
|
||||
},
|
||||
)
|
||||
|
||||
# 后台读取 stderr,记日志
|
||||
asyncio.create_task(self._read_stderr())
|
||||
|
||||
async def stop(self) -> None:
|
||||
proc = self._process
|
||||
if proc is None or proc.returncode is not None:
|
||||
self._process = None
|
||||
self._started_at = None
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"停止 stdio MCP 子进程",
|
||||
extra={
|
||||
"server_id": self.server_id,
|
||||
"pid": proc.pid,
|
||||
},
|
||||
)
|
||||
try:
|
||||
proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
await asyncio.wait_for(proc.wait(), timeout=settings.mcp_stdio_stop_timeout_seconds)
|
||||
except TimeoutError:
|
||||
proc.send_signal(signal.SIGKILL)
|
||||
await proc.wait()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
self._process = None
|
||||
self._started_at = None
|
||||
|
||||
async def list_tools(self) -> list[dict[str, Any]]:
|
||||
result = await self._send_request("tools/list", {})
|
||||
if result is None:
|
||||
return []
|
||||
return result.get("tools") if isinstance(result.get("tools"), list) else []
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult:
|
||||
result = await self._send_request(
|
||||
"tools/call",
|
||||
{"name": name, "arguments": arguments},
|
||||
)
|
||||
if result is None:
|
||||
return ToolResult(
|
||||
ok=False,
|
||||
stdout="",
|
||||
stderr="stdio 子进程无响应",
|
||||
duration_ms=0,
|
||||
exit_code=1,
|
||||
timeout=False,
|
||||
truncated=False,
|
||||
trace_id="",
|
||||
)
|
||||
return _parse_call_result({"result": result})
|
||||
|
||||
async def restart(self) -> None:
|
||||
await self.stop()
|
||||
await self.start()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return (
|
||||
self._process is not None
|
||||
and self._process.returncode is None
|
||||
and self._process.stdin is not None
|
||||
)
|
||||
|
||||
async def get_status(self) -> dict[str, Any]:
|
||||
running = self.is_running
|
||||
uptime = None
|
||||
if running and self._started_at is not None:
|
||||
uptime = time.time() - self._started_at
|
||||
return {
|
||||
"server_id": self.server_id,
|
||||
"name": self.server_name,
|
||||
"transport": "stdio",
|
||||
"status": "running" if running else "stopped",
|
||||
"pid": self._process.pid if self._process else None,
|
||||
"uptime_seconds": uptime,
|
||||
}
|
||||
|
||||
async def _send_request(self, method: str, params: dict[str, Any]) -> dict[str, Any] | None:
|
||||
async with self._lock:
|
||||
self._request_id += 1
|
||||
req_id = self._request_id
|
||||
|
||||
proc = self._process
|
||||
if proc is None or proc.stdin is None or proc.stdout is None:
|
||||
return None
|
||||
|
||||
request = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": req_id,
|
||||
"method": method,
|
||||
"params": params,
|
||||
}
|
||||
line = (json.dumps(request) + "\n").encode("utf-8")
|
||||
|
||||
try:
|
||||
proc.stdin.write(line)
|
||||
await asyncio.wait_for(
|
||||
proc.stdin.drain(),
|
||||
timeout=settings.mcp_stdio_call_timeout_seconds,
|
||||
)
|
||||
|
||||
response_line = await asyncio.wait_for(
|
||||
proc.stdout.readline(),
|
||||
timeout=settings.mcp_stdio_call_timeout_seconds,
|
||||
)
|
||||
if not response_line:
|
||||
return None
|
||||
|
||||
data = json.loads(response_line.decode("utf-8"))
|
||||
except (TimeoutError, json.JSONDecodeError, OSError) as exc:
|
||||
logger.warning(
|
||||
"stdio MCP 请求失败",
|
||||
extra={
|
||||
"server_id": self.server_id,
|
||||
"method": method,
|
||||
"error": str(exc),
|
||||
},
|
||||
)
|
||||
return None
|
||||
|
||||
if isinstance(data, dict) and data.get("error"):
|
||||
logger.warning(
|
||||
"stdio MCP 返回 JSON-RPC 错误",
|
||||
extra={
|
||||
"server_id": self.server_id,
|
||||
"method": method,
|
||||
"error": data["error"],
|
||||
},
|
||||
)
|
||||
return None
|
||||
|
||||
return data.get("result") if isinstance(data, dict) else None
|
||||
|
||||
async def _read_stderr(self) -> None:
|
||||
proc = self._process
|
||||
if proc is None or proc.stderr is None:
|
||||
return
|
||||
try:
|
||||
while True:
|
||||
line = await proc.stderr.readline()
|
||||
if not line:
|
||||
break
|
||||
text = line.decode("utf-8", errors="replace").rstrip()
|
||||
if text:
|
||||
logger.info(
|
||||
"[stdio %s] %s",
|
||||
self.server_name,
|
||||
text,
|
||||
extra={
|
||||
"server_id": self.server_id,
|
||||
"server_name": self.server_name,
|
||||
},
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
@@ -0,0 +1,12 @@
|
||||
"""天目 OCR 集成:HTTP 客户端 + 响应模型。"""
|
||||
|
||||
from app.integrations.ocr.client import OCRBlock, OCRClient, OCRResult
|
||||
from app.integrations.ocr.schemas import OCRBlockSchema, OCRResultSchema
|
||||
|
||||
__all__ = [
|
||||
"OCRBlock",
|
||||
"OCRBlockSchema",
|
||||
"OCRClient",
|
||||
"OCRResult",
|
||||
"OCRResultSchema",
|
||||
]
|
||||
@@ -0,0 +1,269 @@
|
||||
"""
|
||||
天目 OCR HTTP 客户端。
|
||||
|
||||
封装对天目微服务的 HTTP 调用,提供同步/异步 OCR 接口。
|
||||
后端其他模块通过此客户端调用 OCR 能力,不直接依赖 HTTP 细节。
|
||||
|
||||
特性:
|
||||
- 自动重试(网络故障 / 5xx 错误,指数退避)
|
||||
- 连接池复用(httpx.AsyncClient 自动管理)
|
||||
- 端点级别超时配置(/healthz 短超时,/v1/ocr 适中,/v1/ocr/vl 长超时)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from app.core.logging import logger
|
||||
|
||||
TIANMU_DEFAULT_BASE = "http://localhost:8101"
|
||||
|
||||
# 各端点默认超时(秒)
|
||||
_TIMEOUTS = {
|
||||
"health": 5.0,
|
||||
"ocr": 60.0,
|
||||
"ocr_vl": 600.0,
|
||||
}
|
||||
|
||||
# 重试配置
|
||||
_MAX_RETRIES = 2
|
||||
_RETRYABLE_STATUSES = {502, 503, 504}
|
||||
_RETRY_BASE_DELAY = 0.5 # 秒
|
||||
|
||||
|
||||
@dataclass
|
||||
class OCRBlock:
|
||||
"""单行 OCR 识别结果。"""
|
||||
|
||||
text: str
|
||||
confidence: float
|
||||
bbox: list[list[float]] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class OCRResult:
|
||||
"""统一 OCR 响应。"""
|
||||
|
||||
engine: str
|
||||
text: str
|
||||
blocks: list[OCRBlock] = field(default_factory=list)
|
||||
duration_ms: float = 0.0
|
||||
error: str | None = None
|
||||
saved_md_path: str | None = None
|
||||
|
||||
@property
|
||||
def success(self) -> bool:
|
||||
return self.error is None
|
||||
|
||||
|
||||
class OCRClient:
|
||||
"""天目 OCR 服务 HTTP 客户端。
|
||||
|
||||
用法::
|
||||
|
||||
client = OCRClient()
|
||||
result = await client.ocr(image_bytes)
|
||||
await client.close()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str = "",
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
timeout: float = 600.0,
|
||||
max_retries: int = _MAX_RETRIES,
|
||||
) -> None:
|
||||
self.base_url = (base_url or TIANMU_DEFAULT_BASE).rstrip("/")
|
||||
self._max_retries = max_retries
|
||||
self._http = http_client or httpx.AsyncClient(timeout=timeout)
|
||||
self._own_http = http_client is None
|
||||
|
||||
async def close(self) -> None:
|
||||
"""释放 HTTP 客户端。"""
|
||||
if self._own_http:
|
||||
await self._http.aclose()
|
||||
|
||||
# ── 通用 OCR(PP-OCRv6 优先)─────────────────────────
|
||||
|
||||
async def ocr(
|
||||
self,
|
||||
image: bytes | str,
|
||||
*,
|
||||
engine: str = "auto",
|
||||
lang: list[str] | None = None,
|
||||
save_md: bool = False,
|
||||
) -> OCRResult:
|
||||
"""通用文字识别。
|
||||
|
||||
Args:
|
||||
image: 图片二进制数据或 HTTP URL。
|
||||
engine: 引擎选择 auto / ppocrv6 / paddleocr-vl。
|
||||
lang: 语言列表(默认中英文,仅 ppocrv6 生效)。
|
||||
save_md: 是否保存结果为 Markdown 文件。
|
||||
|
||||
Returns:
|
||||
OCRResult: 识别结果。
|
||||
"""
|
||||
body: dict[str, Any] = {"engine": engine}
|
||||
if lang:
|
||||
body["lang"] = lang
|
||||
if save_md:
|
||||
body["save_md"] = True
|
||||
|
||||
body["image"] = _encode_image(image)
|
||||
|
||||
data = await self._request_with_retry(
|
||||
"POST",
|
||||
f"{self.base_url}/v1/ocr",
|
||||
json=body,
|
||||
timeout=_TIMEOUTS["ocr"],
|
||||
)
|
||||
if data is None:
|
||||
return OCRResult(engine=engine, error="天目服务不可达(重试耗尽)")
|
||||
|
||||
blocks = [
|
||||
OCRBlock(
|
||||
text=b["text"],
|
||||
confidence=b["confidence"],
|
||||
bbox=b.get("bbox", {}).get("points", []),
|
||||
)
|
||||
for b in data.get("blocks", [])
|
||||
]
|
||||
|
||||
return OCRResult(
|
||||
engine=data.get("engine", ""),
|
||||
text=data.get("text", ""),
|
||||
blocks=blocks,
|
||||
duration_ms=data.get("duration_ms", 0),
|
||||
error=data.get("error"),
|
||||
saved_md_path=data.get("saved_md_path"),
|
||||
)
|
||||
|
||||
# ── 复杂文档解析(PaddleOCR-VL)──────────────────────
|
||||
|
||||
async def ocr_vl(
|
||||
self,
|
||||
image: bytes | str,
|
||||
*,
|
||||
output_format: str = "markdown",
|
||||
save_md: bool = False,
|
||||
) -> OCRResult:
|
||||
"""复杂文档解析(PaddleOCR-VL)。
|
||||
|
||||
Args:
|
||||
image: 图片二进制数据或 HTTP URL。
|
||||
output_format: 输出格式 markdown / json / text。
|
||||
save_md: 是否保存结果为 Markdown 文件。
|
||||
|
||||
Returns:
|
||||
OCRResult: 结构化解析结果。
|
||||
"""
|
||||
body: dict[str, Any] = {"output_format": output_format}
|
||||
if save_md:
|
||||
body["save_md"] = True
|
||||
|
||||
body["image"] = _encode_image(image)
|
||||
|
||||
data = await self._request_with_retry(
|
||||
"POST",
|
||||
f"{self.base_url}/v1/ocr/vl",
|
||||
json=body,
|
||||
timeout=_TIMEOUTS["ocr_vl"],
|
||||
)
|
||||
if data is None:
|
||||
return OCRResult(engine="paddleocr-vl", error="天目服务不可达(重试耗尽)")
|
||||
|
||||
return OCRResult(
|
||||
engine=data.get("engine", "paddleocr-vl"),
|
||||
text=data.get("text", ""),
|
||||
duration_ms=data.get("duration_ms", 0),
|
||||
error=data.get("error"),
|
||||
saved_md_path=data.get("saved_md_path"),
|
||||
)
|
||||
|
||||
# ── 健康检查 ──────────────────────────────────────────
|
||||
|
||||
async def health(self) -> dict:
|
||||
"""检查天目服务健康状态。"""
|
||||
resp = await self._http.get(
|
||||
f"{self.base_url}/healthz",
|
||||
timeout=_TIMEOUTS["health"],
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
# ── 内部:重试逻辑 ────────────────────────────────────
|
||||
|
||||
async def _request_with_retry(
|
||||
self,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
json: dict[str, Any] | None = None,
|
||||
timeout: float = 60.0,
|
||||
) -> dict[str, Any] | None:
|
||||
"""带指数退避重试的 HTTP 请求。
|
||||
|
||||
Args:
|
||||
method: HTTP 方法。
|
||||
url: 请求 URL。
|
||||
json: JSON 请求体。
|
||||
timeout: 请求超时(秒)。
|
||||
|
||||
Returns:
|
||||
响应 JSON 数据,重试耗尽返回 None。
|
||||
"""
|
||||
|
||||
for attempt in range(self._max_retries + 1):
|
||||
try:
|
||||
resp = await self._http.request(method, url, json=json, timeout=timeout)
|
||||
|
||||
# 可重试的服务器错误
|
||||
if resp.status_code in _RETRYABLE_STATUSES and attempt < self._max_retries:
|
||||
delay = _RETRY_BASE_DELAY * (2 ** attempt)
|
||||
logger.warning(
|
||||
"天目返回 {},{}s 后重试 ({}/{})",
|
||||
resp.status_code, delay, attempt + 1, self._max_retries,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
except (httpx.TimeoutException, httpx.ConnectError, httpx.RemoteProtocolError) as e:
|
||||
if attempt < self._max_retries:
|
||||
delay = _RETRY_BASE_DELAY * (2 ** attempt)
|
||||
logger.warning(
|
||||
"天目请求失败 ({},{}s 后重试): {}",
|
||||
type(e).__name__, delay, e,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
|
||||
logger.error("天目请求重试耗尽 ({}): {}", type(e).__name__, e)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _encode_image(image: bytes | str | Path) -> str:
|
||||
"""将图片编码为 base64 字符串。
|
||||
|
||||
Args:
|
||||
image: 图片二进制数据、文件路径或 HTTP URL。
|
||||
|
||||
Returns:
|
||||
base64 编码字符串(二进制输入)或原 URL。
|
||||
"""
|
||||
if isinstance(image, Path):
|
||||
image = image.read_bytes()
|
||||
if isinstance(image, str):
|
||||
return image
|
||||
|
||||
import base64
|
||||
return base64.b64encode(image).decode("ascii")
|
||||
@@ -0,0 +1,24 @@
|
||||
"""天目 OCR 响应模型(Pydantic 版,供 API 路由使用)。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class OCRBlockSchema(BaseModel):
|
||||
"""单行识别结果(Pydantic 版)。"""
|
||||
|
||||
text: str = Field(description="识别文字")
|
||||
confidence: float = Field(description="置信度 0~1", ge=0, le=1)
|
||||
bbox: list[list[float]] = Field(default_factory=list, description="文本框坐标")
|
||||
|
||||
|
||||
class OCRResultSchema(BaseModel):
|
||||
"""统一 OCR 响应模型。"""
|
||||
|
||||
engine: str = Field(description="实际使用的引擎")
|
||||
text: str = Field(default="", description="识别文本")
|
||||
blocks: list[OCRBlockSchema] = Field(default_factory=list, description="逐块结果")
|
||||
duration_ms: float = Field(default=0, description="推理耗时(毫秒)")
|
||||
error: str | None = Field(default=None, description="错误信息")
|
||||
saved_md_path: str | None = Field(default=None, description="保存的 Markdown 路径")
|
||||
@@ -0,0 +1,63 @@
|
||||
"""安全沙箱集成:封装对 Rust 沙箱服务的调用(司命授权 + 一炁事件回写)。"""
|
||||
|
||||
from app.core.config import SandboxSettings
|
||||
from app.integrations.sandbox.audit_sink import AuditEventSink
|
||||
from app.integrations.sandbox.auth import AuthorizationContext
|
||||
from app.integrations.sandbox.client import SandboxClient
|
||||
from app.integrations.sandbox.events import (
|
||||
CompositeEventSink,
|
||||
EventSink,
|
||||
LoggingEventSink,
|
||||
SandboxEvent,
|
||||
)
|
||||
from app.integrations.sandbox.exceptions import (
|
||||
SandboxBadRequest,
|
||||
SandboxError,
|
||||
SandboxExecFailed,
|
||||
SandboxNotFound,
|
||||
SandboxUnauthorized,
|
||||
SandboxUnavailable,
|
||||
)
|
||||
from app.integrations.sandbox.models import (
|
||||
CancelResponse,
|
||||
DirEntry,
|
||||
ExecRequest,
|
||||
ExecResponse,
|
||||
FileListRequest,
|
||||
FileListResponse,
|
||||
FileReadRequest,
|
||||
FileReadResponse,
|
||||
FileStatRequest,
|
||||
FileStatResponse,
|
||||
FileWriteRequest,
|
||||
FileWriteResponse,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AuthorizationContext",
|
||||
"AuditEventSink",
|
||||
"CancelResponse",
|
||||
"CompositeEventSink",
|
||||
"DirEntry",
|
||||
"EventSink",
|
||||
"ExecRequest",
|
||||
"ExecResponse",
|
||||
"FileListRequest",
|
||||
"FileListResponse",
|
||||
"FileReadRequest",
|
||||
"FileReadResponse",
|
||||
"FileStatRequest",
|
||||
"FileStatResponse",
|
||||
"FileWriteRequest",
|
||||
"FileWriteResponse",
|
||||
"LoggingEventSink",
|
||||
"SandboxBadRequest",
|
||||
"SandboxClient",
|
||||
"SandboxError",
|
||||
"SandboxEvent",
|
||||
"SandboxExecFailed",
|
||||
"SandboxNotFound",
|
||||
"SandboxSettings",
|
||||
"SandboxUnauthorized",
|
||||
"SandboxUnavailable",
|
||||
]
|
||||
@@ -0,0 +1,99 @@
|
||||
"""审计事件 sink:把沙箱 tool.completed / tool.failed 事件写入 AuditEvent 表(起居注)。
|
||||
|
||||
与 EventStoreSink 并列——EventStoreSink 落一炁流程事件(events 表,run 内 sequence),
|
||||
AuditEventSink 落治理审计(audit_events 表,含 decision / risk_level / argv 脱敏摘要)。
|
||||
两者经 CompositeEventSink 同时注入 SandboxClient,互不感知。
|
||||
|
||||
审计只记终态(completed/failed);started 不记(审计关注决策结果,非中间态)。
|
||||
argv 经 record_audit_event 的 ``_summary`` → ``redact_text`` + 截断 1000 脱敏,作为
|
||||
input_summary(审计核心价值:谁执行了什么命令;脱敏打码密钥/token 等敏感模式后保留可读性)。
|
||||
|
||||
per-call/per-run 绑定 run_id + character_id(与 EventStoreSink 构造一致);risk_level +
|
||||
argv 由 started 事件 payload 携带,completed/failed 时消费。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from app.integrations.sandbox.events import SandboxEvent
|
||||
from app.services.audit import record_audit_event
|
||||
|
||||
|
||||
class AuditEventSink:
|
||||
"""把沙箱事件转写为 AuditEvent(审计表),与 EventStoreSink 并列注入。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_factory: async_sessionmaker[AsyncSession],
|
||||
*,
|
||||
run_id: str | None = None,
|
||||
character_id: str | None = None,
|
||||
) -> None:
|
||||
self._session_factory = session_factory
|
||||
self._run_id = run_id
|
||||
self._character_id = character_id
|
||||
# 缓存最近一次 started 的上下文(risk_level + argv),completed/failed 写审计时消费。
|
||||
# sink 为 per-call 实例(随一次 exec 生命周期),状态隔离安全。
|
||||
self._pending_risk: str | None = None
|
||||
self._pending_argv: list[str] = []
|
||||
|
||||
async def emit(self, event: SandboxEvent) -> None:
|
||||
etype = event.type
|
||||
payload = event.payload
|
||||
|
||||
if etype == "tool.started":
|
||||
# 缓存 risk_level + argv(started 携带),completed/failed 写审计时消费
|
||||
self._pending_risk = str(payload.get("risk_level") or "L1")
|
||||
self._pending_argv = list(payload.get("argv") or [])
|
||||
return
|
||||
|
||||
if etype not in ("tool.completed", "tool.failed"):
|
||||
return
|
||||
|
||||
tool = str(payload.get("tool") or "")
|
||||
risk_level = self._pending_risk or "L1"
|
||||
argv = self._pending_argv
|
||||
|
||||
if etype == "tool.completed":
|
||||
action = "tool.executed"
|
||||
decision = "allow"
|
||||
reason: str | None = None
|
||||
output_summary = self._completed_summary(payload)
|
||||
else: # tool.failed
|
||||
action = "tool.failed"
|
||||
decision = "deny"
|
||||
reason = str(payload.get("error_code") or "") or None
|
||||
output_summary = str(payload.get("error_message") or "") or None
|
||||
|
||||
# input_summary = argv 空格拼接(record_audit_event 内部 _summary 会 redact + 截断 1000)
|
||||
input_summary = " ".join(argv) if argv else None
|
||||
|
||||
async with self._session_factory() as session:
|
||||
await record_audit_event(
|
||||
session,
|
||||
action=action,
|
||||
pivot="tool_runtime",
|
||||
character_id=self._character_id,
|
||||
run_id=self._run_id,
|
||||
risk_level=risk_level,
|
||||
decision=decision,
|
||||
reason=reason,
|
||||
input_summary=input_summary,
|
||||
output_summary=output_summary,
|
||||
resource_type="tool",
|
||||
resource_id=tool or None,
|
||||
trace_id=event.trace_id,
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
@staticmethod
|
||||
def _completed_summary(payload: dict[str, Any]) -> str:
|
||||
parts = [f"exit={payload.get('exit_code')}", f"dur={payload.get('duration_ms')}ms"]
|
||||
if payload.get("timeout"):
|
||||
parts.append("timeout")
|
||||
if payload.get("truncated"):
|
||||
parts.append("truncated")
|
||||
return " ".join(parts)
|
||||
@@ -0,0 +1,156 @@
|
||||
"""司命授权凭证:HMAC-SHA256 短期 token,与 Rust 沙箱端 ``auth.rs`` 共用算法与密钥。
|
||||
|
||||
凭证格式:``<payload_b64url>.<sig_b64url>``
|
||||
|
||||
- payload:JSON,字段顺序固定为
|
||||
``exp / trace_id / risk_level / tool / character_id / run_id / body_hash / rlimit_overrides / file_op_hash / isolation_mode``,
|
||||
必须与 Rust ``serde`` 字段顺序字节级一致(否则 HMAC 不匹配)。``isolation_mode`` 末尾追加,旧凭证向后兼容。
|
||||
- sig:``HMAC-SHA256(secret, payload_b64url)`` 的 base64url(无填充)。
|
||||
- body_hash:``compute_body_hash(argv, stdin)``,把凭证绑死到具体执行内容,防 TTL 内重放执行任意命令。
|
||||
|
||||
安全语义:客户端不判断风险,只把「经司命审核通过」的上下文连同请求内容摘要签成短期凭证;
|
||||
沙箱端验签 + body_hash 比对通过才执行。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import cast
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class AuthorizationContext:
|
||||
"""一次授权的上下文:由司命审核通过后随请求下发。"""
|
||||
|
||||
trace_id: str
|
||||
risk_level: str # L0–L4,见契约字典 risk_level
|
||||
tool: str
|
||||
character_id: str | None = None
|
||||
run_id: str | None = None
|
||||
# 司命授予的资源配额上限(key: fsize/as/cpu/nproc/nofile → u64)。
|
||||
# 随凭证下发,sandbox 验签后信任并 clamp 到全局天花板;None 表示不绑定分级配额。
|
||||
# 属授权属性,不进 body_hash(body_hash 只绑 argv + stdin)。
|
||||
rlimit_overrides: dict[str, int] | None = None
|
||||
# 隔离模式授权(L0-L2 → process,L3-L4 → container)。随凭证下发,沙箱据此路由;
|
||||
# None 表示司命未指定。属授权属性,不进 body_hash,与 rlimit_overrides 同链路。
|
||||
isolation_mode: str | None = None
|
||||
|
||||
|
||||
def _b64url_no_pad(raw: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
def _b64url_no_pad_decode(text: str) -> bytes:
|
||||
pad = "=" * (-len(text) % 4)
|
||||
return base64.urlsafe_b64decode(text + pad)
|
||||
|
||||
|
||||
def compute_body_hash(argv: list[str], stdin: str | None) -> str:
|
||||
"""计算请求内容摘要,与 Rust 沙箱端 ``compute_body_hash`` 字节级一致。
|
||||
|
||||
口径:``SHA256(canonical JSON {"argv": [...], "stdin": null|str})``——
|
||||
JSON 紧凑(``separators=(",",":")``)、不转义非 ASCII(``ensure_ascii=False``)、
|
||||
字段顺序固定 ``argv, stdin``;返回小写十六进制(64 字符)。
|
||||
"""
|
||||
payload = {"argv": argv, "stdin": stdin}
|
||||
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
return hashlib.sha256(encoded).hexdigest()
|
||||
|
||||
|
||||
def compute_file_op_hash(op: str, path: str) -> str:
|
||||
"""计算文件操作摘要,与 Rust 沙箱端 ``compute_file_op_hash`` 字节级一致。
|
||||
|
||||
口径:``SHA256(canonical JSON {"op": str, "path": str})``——把凭证绑死到具体文件操作
|
||||
(read/write/list/stat),防 TTL 内重放执行任意文件读写。字段顺序固定 ``op, path``。
|
||||
"""
|
||||
payload = {"op": op, "path": path}
|
||||
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
return hashlib.sha256(encoded).hexdigest()
|
||||
|
||||
|
||||
def compute_body_hash_code(language: str, code: str, stdin: str | None) -> str:
|
||||
"""计算 code 模式请求内容摘要,与 Rust 沙箱端 ``compute_body_hash_code`` 字节级一致。
|
||||
|
||||
口径:``SHA256(canonical JSON {"language": str, "code": str, "stdin": null|str})``——
|
||||
字段顺序固定 ``language, code, stdin``、紧凑 JSON、不转义非 ASCII;返回小写十六进制。
|
||||
code 模式凭证 body_hash 用此函数(argv 模式仍用 ``compute_body_hash``),把凭证绑死到
|
||||
具体代码内容,防 TTL 内重放执行任意代码。
|
||||
"""
|
||||
payload = {"language": language, "code": code, "stdin": stdin}
|
||||
encoded = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
||||
return hashlib.sha256(encoded).hexdigest()
|
||||
|
||||
|
||||
def sign_authorization(
|
||||
secret: str,
|
||||
ctx: AuthorizationContext,
|
||||
argv: list[str] | None = None,
|
||||
stdin: str | None = None,
|
||||
ttl_seconds: int = 30,
|
||||
*,
|
||||
code: str | None = None,
|
||||
language: str | None = None,
|
||||
file_op: tuple[str, str] | None = None,
|
||||
) -> str:
|
||||
"""签发授权凭证。
|
||||
|
||||
内容绑定(互斥,由调用方保证不同时传 argv 与 code):
|
||||
- exec 命令模式:传 ``argv``/``stdin``,按 ``compute_body_hash(argv, stdin)`` 绑定。
|
||||
- exec 代码模式:传 ``code``/``language``/``stdin``,按 ``compute_body_hash_code`` 绑定。
|
||||
- 低风险操作(cancel):都不传,不绑定 body_hash(file_op 单独走 file_op_hash)。
|
||||
|
||||
payload 字段顺序与 Rust ``AuthPayload`` 一致(isolation_mode 末尾追加,旧凭证向后兼容)。
|
||||
"""
|
||||
if not secret:
|
||||
raise ValueError("授权密钥不能为空")
|
||||
exp = int(time.time()) + ttl_seconds
|
||||
if code is not None:
|
||||
body_hash = compute_body_hash_code(language or "python", code, stdin)
|
||||
elif argv is not None:
|
||||
body_hash = compute_body_hash(argv, stdin)
|
||||
else:
|
||||
body_hash = None
|
||||
payload: dict[str, object] = {
|
||||
"exp": exp,
|
||||
"trace_id": ctx.trace_id,
|
||||
"risk_level": ctx.risk_level,
|
||||
"tool": ctx.tool,
|
||||
"character_id": ctx.character_id,
|
||||
"run_id": ctx.run_id,
|
||||
"body_hash": body_hash,
|
||||
# rlimit_overrides 内部按 key 排序,对齐 Rust serde BTreeMap 的字母序,
|
||||
# 保证多 key 时两端 payload JSON 字节一致 → HMAC 一致。
|
||||
"rlimit_overrides": (
|
||||
dict(sorted(ctx.rlimit_overrides.items())) if ctx.rlimit_overrides else None
|
||||
),
|
||||
# file_op_hash 绑定文件操作 {op, path}(exec 路径为 None);字段顺序对齐 Rust AuthPayload。
|
||||
"file_op_hash": compute_file_op_hash(file_op[0], file_op[1]) if file_op else None,
|
||||
# isolation_mode 末尾追加(字段顺序对齐 Rust AuthPayload;旧凭证无此字段仍可验签)。
|
||||
"isolation_mode": ctx.isolation_mode,
|
||||
}
|
||||
payload_json = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
|
||||
payload_b64 = _b64url_no_pad(payload_json.encode("utf-8"))
|
||||
digest = hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
|
||||
return f"{payload_b64}.{_b64url_no_pad(digest)}"
|
||||
|
||||
|
||||
def verify_authorization(secret: str, token: str) -> dict[str, object]:
|
||||
"""校验授权凭证(对称实现,供本地测试与跨端一致性校验)。"""
|
||||
payload_b64, sep, sig_b64 = token.rpartition(".")
|
||||
if not sep:
|
||||
raise ValueError("授权凭证格式不合法")
|
||||
expected = _b64url_no_pad(
|
||||
hmac.new(secret.encode("utf-8"), payload_b64.encode("ascii"), hashlib.sha256).digest()
|
||||
)
|
||||
if not hmac.compare_digest(expected, sig_b64):
|
||||
raise ValueError("授权凭证签名校验失败")
|
||||
payload = cast("dict[str, object]", json.loads(_b64url_no_pad_decode(payload_b64)))
|
||||
exp = payload.get("exp")
|
||||
if not isinstance(exp, int | float) or exp <= time.time():
|
||||
raise ValueError("授权凭证已过期")
|
||||
return payload
|
||||
@@ -0,0 +1,43 @@
|
||||
"""沙箱事件桥接:把 SandboxEvent 转写为一炁 Event 落入 EventStore。
|
||||
|
||||
实现 sandbox EventSink Protocol,使 SandboxClient 发出的 tool.started / tool.completed /
|
||||
tool.failed 事件自动进入一炁事件存储,无需改动客户端。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
from app.integrations.sandbox.events import SandboxEvent
|
||||
from app.modules.runtime.events import Event, EventStore
|
||||
|
||||
|
||||
class EventStoreSink:
|
||||
"""将沙箱事件适配为一炁事件并写入 EventStore。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: EventStore,
|
||||
*,
|
||||
run_id: str | None = None,
|
||||
character_id: str | None = None,
|
||||
) -> None:
|
||||
self._store = store
|
||||
self._run_id = run_id
|
||||
self._character_id = character_id
|
||||
|
||||
async def emit(self, event: SandboxEvent) -> None:
|
||||
await self._store.append(
|
||||
Event(
|
||||
id=f"event_{uuid.uuid4().hex}",
|
||||
type=event.type,
|
||||
source=event.source,
|
||||
payload=dict(event.payload),
|
||||
visibility="public",
|
||||
sensitivity=event.sensitivity, # type: ignore[arg-type]
|
||||
run_id=self._run_id,
|
||||
character_id=self._character_id,
|
||||
trace_id=event.trace_id,
|
||||
created_at=event.created_at,
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,475 @@
|
||||
"""沙箱客户端:封装对 Rust 沙箱服务的调用,含司命授权签发与一炁事件回写。
|
||||
|
||||
使用约束(见 backend/CLAUDE.md):禁止在调用处裸 ``async with httpx.AsyncClient()``,
|
||||
应复用注入的 ``AsyncClient``(长生命周期,由上层依赖注入管理)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, NoReturn, TypeVar
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from app.core.config import SandboxSettings
|
||||
from app.integrations.sandbox.auth import AuthorizationContext, sign_authorization
|
||||
from app.integrations.sandbox.events import (
|
||||
EventSink,
|
||||
LoggingEventSink,
|
||||
completed_event,
|
||||
failed_event,
|
||||
started_event,
|
||||
)
|
||||
from app.integrations.sandbox.exceptions import (
|
||||
SandboxBadRequest,
|
||||
SandboxError,
|
||||
SandboxExecFailed,
|
||||
SandboxNotFound,
|
||||
SandboxUnauthorized,
|
||||
SandboxUnavailable,
|
||||
)
|
||||
from app.integrations.sandbox.models import (
|
||||
CancelResponse,
|
||||
ExecRequest,
|
||||
ExecResponse,
|
||||
FileListRequest,
|
||||
FileListResponse,
|
||||
FileReadRequest,
|
||||
FileReadResponse,
|
||||
FileStatRequest,
|
||||
FileStatResponse,
|
||||
FileWriteRequest,
|
||||
FileWriteResponse,
|
||||
)
|
||||
|
||||
_T = TypeVar("_T", bound=BaseModel)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SandboxStreamEvent:
|
||||
"""沙箱流式事件:stdout/stderr 增量(base64 str,保真透传)或 done(结束汇总 dict)。"""
|
||||
|
||||
kind: str # "stdout" | "stderr" | "done"
|
||||
payload: str | dict[str, Any]
|
||||
|
||||
|
||||
class SandboxClient:
|
||||
"""沙箱服务异步客户端。
|
||||
|
||||
授权语义:调用方需先经司命审核(在本类之外完成),通过 ``AuthorizationContext`` 携带
|
||||
授权上下文;本类据此签发短期 HMAC 凭证(exec 含本次请求的 body_hash 绑定),放入
|
||||
``Authorization`` 头,由沙箱端校验。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
http_client: httpx.AsyncClient,
|
||||
settings: SandboxSettings,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> None:
|
||||
if settings.sandbox_auth_enabled and not settings.sandbox_auth_secret:
|
||||
raise ValueError("auth_enabled=True 时 auth_secret 不能为空(须与沙箱端共享)")
|
||||
self._http = http_client
|
||||
self._settings = settings
|
||||
self._sink: EventSink = event_sink or LoggingEventSink()
|
||||
|
||||
async def health(self) -> bool:
|
||||
"""健康探测,失败返回 False 而不抛异常。"""
|
||||
try:
|
||||
resp = await self._http.get(f"{self._settings.rust_sandbox_url}/healthz", timeout=5.0)
|
||||
except httpx.HTTPError:
|
||||
return False
|
||||
return resp.status_code == 200
|
||||
|
||||
async def exec(
|
||||
self,
|
||||
request: ExecRequest,
|
||||
auth_ctx: AuthorizationContext,
|
||||
*,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> ExecResponse:
|
||||
"""同步执行:收集全部输出后返回。授权失败、网络错误或沙箱错误分别抛出对应异常,
|
||||
并回写 tool.failed 事件。``event_sink`` 可按调用覆盖,便于编排层按 run 绑定上下文。"""
|
||||
sink = event_sink or self._sink
|
||||
trace_id = request.trace_id or auth_ctx.trace_id
|
||||
run_id = auth_ctx.run_id
|
||||
tool = auth_ctx.tool
|
||||
|
||||
# 请求体带上解析后的 trace_id,便于沙箱日志关联
|
||||
if request.trace_id is None and trace_id is not None:
|
||||
request.trace_id = trace_id
|
||||
|
||||
await sink.emit(
|
||||
started_event(
|
||||
tool=tool,
|
||||
argv=request.argv or [f"code:{request.language}", request.code],
|
||||
risk_level=auth_ctx.risk_level,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
resp = await self._http.post(
|
||||
f"{self._settings.rust_sandbox_url}/v1/exec",
|
||||
json=request.model_dump(exclude_none=True),
|
||||
headers=self._exec_auth_headers(auth_ctx, request),
|
||||
timeout=self._settings.http_timeout(request.timeout_ms),
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
await sink.emit(
|
||||
failed_event(
|
||||
tool=tool,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
error_code="sandbox_unavailable",
|
||||
error_message=str(exc),
|
||||
)
|
||||
)
|
||||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||||
|
||||
if resp.status_code >= 400:
|
||||
exc_type: type[SandboxError]
|
||||
if resp.status_code == 401:
|
||||
exc_type = SandboxUnauthorized
|
||||
elif resp.status_code == 400:
|
||||
exc_type = SandboxBadRequest
|
||||
else:
|
||||
exc_type = SandboxExecFailed
|
||||
await self._fail(tool, trace_id, run_id, resp, exc_type, sink)
|
||||
|
||||
result = ExecResponse.model_validate(resp.json())
|
||||
await sink.emit(
|
||||
completed_event(
|
||||
tool=tool,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
exit_code=result.exit_code,
|
||||
duration_ms=result.duration_ms,
|
||||
timeout=result.timeout,
|
||||
truncated=result.truncated,
|
||||
stdout_len=len(result.stdout),
|
||||
stderr_len=len(result.stderr),
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def exec_stream(
|
||||
self,
|
||||
request: ExecRequest,
|
||||
auth_ctx: AuthorizationContext,
|
||||
) -> AsyncIterator[SandboxStreamEvent]:
|
||||
"""流式执行:实时产出 stdout/stderr 增量(base64 str)与 done 汇总(dict)。
|
||||
|
||||
消费沙箱 ``POST /v1/exec/stream`` 的 SSE。不回写一炁事件——started/completed/failed
|
||||
由编排层 ``ToolRuntime.execute_stream`` 负责,本方法只透传执行级输出。
|
||||
"""
|
||||
if request.trace_id is None and auth_ctx.trace_id is not None:
|
||||
request.trace_id = auth_ctx.trace_id
|
||||
|
||||
try:
|
||||
async with self._http.stream(
|
||||
"POST",
|
||||
f"{self._settings.rust_sandbox_url}/v1/exec/stream",
|
||||
json=request.model_dump(exclude_none=True),
|
||||
headers=self._exec_auth_headers(auth_ctx, request),
|
||||
timeout=self._settings.http_timeout(request.timeout_ms),
|
||||
) as resp:
|
||||
if resp.status_code == 401:
|
||||
raise SandboxUnauthorized(await self._read_error(resp))
|
||||
if resp.status_code == 400:
|
||||
raise SandboxBadRequest(await self._read_error(resp))
|
||||
if resp.status_code >= 400:
|
||||
raise SandboxExecFailed(await self._read_error(resp))
|
||||
async for event_type, data in self._iter_sse(resp):
|
||||
if event_type in ("stdout", "stderr"):
|
||||
yield SandboxStreamEvent(event_type, data)
|
||||
elif event_type == "done":
|
||||
try:
|
||||
yield SandboxStreamEvent("done", json.loads(data))
|
||||
except ValueError:
|
||||
continue
|
||||
except httpx.HTTPError as exc:
|
||||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||||
|
||||
async def cancel(
|
||||
self,
|
||||
trace_id: str,
|
||||
auth_ctx: AuthorizationContext,
|
||||
) -> bool:
|
||||
"""取消运行中的任务(按 trace_id)。命中返回 True;任务不存在或已完成抛 SandboxNotFound。
|
||||
|
||||
cancel 不回写一炁事件(由编排层按需记录),凭证不绑定 body_hash(危害有限)。
|
||||
"""
|
||||
try:
|
||||
resp = await self._http.post(
|
||||
f"{self._settings.rust_sandbox_url}/v1/cancel",
|
||||
json={"trace_id": trace_id},
|
||||
headers=self._auth_headers(auth_ctx),
|
||||
timeout=10.0,
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||||
|
||||
if resp.status_code == 404:
|
||||
raise SandboxNotFound(f"任务不存在或已完成: {trace_id}")
|
||||
if resp.status_code >= 400:
|
||||
code, message = self._extract_error(resp)
|
||||
if resp.status_code == 401:
|
||||
raise SandboxUnauthorized(message)
|
||||
raise SandboxExecFailed(message, code=code)
|
||||
return CancelResponse.model_validate(resp.json()).cancelled
|
||||
|
||||
# ---- 文件 IO(不经 shell,二进制安全;与 Rust fileio.rs 契约对齐) ----
|
||||
|
||||
async def file_read(
|
||||
self,
|
||||
request: FileReadRequest,
|
||||
auth_ctx: AuthorizationContext,
|
||||
*,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> FileReadResponse:
|
||||
return await self._file_op(
|
||||
"read", "/v1/file/read", request, FileReadResponse, auth_ctx, event_sink=event_sink
|
||||
)
|
||||
|
||||
async def file_write(
|
||||
self,
|
||||
request: FileWriteRequest,
|
||||
auth_ctx: AuthorizationContext,
|
||||
*,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> FileWriteResponse:
|
||||
return await self._file_op(
|
||||
"write", "/v1/file/write", request, FileWriteResponse, auth_ctx, event_sink=event_sink
|
||||
)
|
||||
|
||||
async def file_list(
|
||||
self,
|
||||
request: FileListRequest,
|
||||
auth_ctx: AuthorizationContext,
|
||||
*,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> FileListResponse:
|
||||
return await self._file_op(
|
||||
"list", "/v1/file/list", request, FileListResponse, auth_ctx, event_sink=event_sink
|
||||
)
|
||||
|
||||
async def file_stat(
|
||||
self,
|
||||
request: FileStatRequest,
|
||||
auth_ctx: AuthorizationContext,
|
||||
*,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> FileStatResponse:
|
||||
return await self._file_op(
|
||||
"stat", "/v1/file/stat", request, FileStatResponse, auth_ctx, event_sink=event_sink
|
||||
)
|
||||
|
||||
async def _file_op(
|
||||
self,
|
||||
op: str,
|
||||
endpoint: str,
|
||||
request: FileReadRequest | FileWriteRequest | FileListRequest | FileStatRequest,
|
||||
response_cls: type[_T],
|
||||
auth_ctx: AuthorizationContext,
|
||||
*,
|
||||
event_sink: EventSink | None = None,
|
||||
) -> _T:
|
||||
"""文件操作通用流程:签发 file_op 凭证 → POST → 解析响应,并回写一炁事件。
|
||||
|
||||
与 exec 共用错误映射(``_fail``)与事件构造;argv 记 ``[file:<op>, <path>]``
|
||||
供审计 input_summary。文件操作在主进程运行(非 exec 子进程),无超时参数,
|
||||
用默认 HTTP 超时(``http_timeout(None)``)。
|
||||
"""
|
||||
sink = event_sink or self._sink
|
||||
trace_id = auth_ctx.trace_id
|
||||
run_id = auth_ctx.run_id
|
||||
tool = auth_ctx.tool
|
||||
path = request.path
|
||||
|
||||
await sink.emit(
|
||||
started_event(
|
||||
tool=tool,
|
||||
argv=[f"file:{op}", path],
|
||||
risk_level=auth_ctx.risk_level,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
)
|
||||
)
|
||||
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
resp = await self._http.post(
|
||||
f"{self._settings.rust_sandbox_url}{endpoint}",
|
||||
json=request.model_dump(exclude_none=True),
|
||||
headers=self._file_auth_headers(auth_ctx, op, path),
|
||||
timeout=self._settings.http_timeout(None),
|
||||
)
|
||||
except httpx.HTTPError as exc:
|
||||
await sink.emit(
|
||||
failed_event(
|
||||
tool=tool,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
error_code="sandbox_unavailable",
|
||||
error_message=str(exc),
|
||||
)
|
||||
)
|
||||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||||
|
||||
if resp.status_code >= 400:
|
||||
exc_type: type[SandboxError]
|
||||
if resp.status_code == 401:
|
||||
exc_type = SandboxUnauthorized
|
||||
elif resp.status_code == 400:
|
||||
exc_type = SandboxBadRequest
|
||||
elif resp.status_code == 404:
|
||||
exc_type = SandboxNotFound
|
||||
else:
|
||||
exc_type = SandboxExecFailed
|
||||
await self._fail(tool, trace_id, run_id, resp, exc_type, sink)
|
||||
|
||||
result = response_cls.model_validate(resp.json())
|
||||
# completed 事件:read 带 returned 字节数 + 截断标记,write 带 bytes_written,
|
||||
# list/stat 无字节数语义,以 0 占位(exit_code=0 表示操作成功)。
|
||||
stdout_len = getattr(result, "returned", None) or getattr(result, "bytes_written", 0)
|
||||
truncated = bool(getattr(result, "truncated", False))
|
||||
await sink.emit(
|
||||
completed_event(
|
||||
tool=tool,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
exit_code=0,
|
||||
duration_ms=int((time.monotonic() - t0) * 1000),
|
||||
timeout=False,
|
||||
truncated=truncated,
|
||||
stdout_len=stdout_len,
|
||||
stderr_len=0,
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
def _file_auth_headers(
|
||||
self,
|
||||
auth_ctx: AuthorizationContext,
|
||||
op: str,
|
||||
path: str,
|
||||
) -> dict[str, str]:
|
||||
"""签发文件操作凭证:file_op 绑定 {op, path}(exec 路径 body_hash 为 None,不绑定)。"""
|
||||
if not self._settings.sandbox_auth_enabled:
|
||||
return {}
|
||||
token = sign_authorization(
|
||||
self._settings.sandbox_auth_secret,
|
||||
auth_ctx,
|
||||
argv=None,
|
||||
stdin=None,
|
||||
ttl_seconds=self._settings.sandbox_auth_ttl_seconds,
|
||||
file_op=(op, path),
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
def _auth_headers(
|
||||
self,
|
||||
auth_ctx: AuthorizationContext,
|
||||
argv: list[str] | None = None,
|
||||
stdin: str | None = None,
|
||||
*,
|
||||
code: str | None = None,
|
||||
language: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""签发授权凭证。
|
||||
|
||||
- exec 命令模式:传 argv/stdin,按 compute_body_hash 绑定。
|
||||
- exec 代码模式:传 code/language/stdin,按 compute_body_hash_code 绑定。
|
||||
- cancel:都不传,不绑定 body_hash。
|
||||
"""
|
||||
if not self._settings.sandbox_auth_enabled:
|
||||
return {}
|
||||
token = sign_authorization(
|
||||
self._settings.sandbox_auth_secret,
|
||||
auth_ctx,
|
||||
argv=argv,
|
||||
stdin=stdin,
|
||||
code=code,
|
||||
language=language,
|
||||
ttl_seconds=self._settings.sandbox_auth_ttl_seconds,
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
def _exec_auth_headers(
|
||||
self,
|
||||
auth_ctx: AuthorizationContext,
|
||||
request: ExecRequest,
|
||||
) -> dict[str, str]:
|
||||
"""按 code/argv 模式分流签发 exec 授权头。"""
|
||||
if request.code is not None:
|
||||
return self._auth_headers(
|
||||
auth_ctx,
|
||||
stdin=request.stdin,
|
||||
code=request.code,
|
||||
language=request.language,
|
||||
)
|
||||
return self._auth_headers(auth_ctx, request.argv, request.stdin)
|
||||
|
||||
@staticmethod
|
||||
async def _read_error(resp: httpx.Response) -> str:
|
||||
"""流式响应的错误体:先 aread 再复用 _extract_error。"""
|
||||
await resp.aread()
|
||||
_, message = SandboxClient._extract_error(resp)
|
||||
return message
|
||||
|
||||
@staticmethod
|
||||
async def _iter_sse(resp: httpx.Response) -> AsyncIterator[tuple[str, str]]:
|
||||
"""极简 SSE 解析:按空行分界,聚合 ``event:`` 与 ``data:`` 行(忽略 id/retry/注释/心跳)。"""
|
||||
event_type = "message"
|
||||
data_lines: list[str] = []
|
||||
async for line in resp.aiter_lines():
|
||||
if line == "":
|
||||
if data_lines:
|
||||
yield event_type, "\n".join(data_lines)
|
||||
event_type = "message"
|
||||
data_lines = []
|
||||
elif line.startswith("event:"):
|
||||
event_type = line[len("event:") :].strip()
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line[len("data:") :].lstrip())
|
||||
# 其余行(id:/retry:/注释/心跳)忽略
|
||||
if data_lines:
|
||||
yield event_type, "\n".join(data_lines)
|
||||
|
||||
async def _fail(
|
||||
self,
|
||||
tool: str,
|
||||
trace_id: str | None,
|
||||
run_id: str | None,
|
||||
resp: httpx.Response,
|
||||
exc_type: type[SandboxError],
|
||||
sink: EventSink,
|
||||
) -> NoReturn:
|
||||
code, message = self._extract_error(resp)
|
||||
await sink.emit(
|
||||
failed_event(
|
||||
tool=tool,
|
||||
trace_id=trace_id,
|
||||
run_id=run_id,
|
||||
error_code=code,
|
||||
error_message=message,
|
||||
)
|
||||
)
|
||||
raise exc_type(message)
|
||||
|
||||
@staticmethod
|
||||
def _extract_error(resp: httpx.Response) -> tuple[str, str]:
|
||||
try:
|
||||
body = resp.json()
|
||||
except ValueError:
|
||||
return "exec_failed", resp.text
|
||||
err = body.get("error") if isinstance(body, dict) else None
|
||||
if isinstance(err, dict):
|
||||
return str(err.get("code", "exec_failed")), str(err.get("message", resp.text))
|
||||
return "exec_failed", resp.text
|
||||
@@ -0,0 +1,127 @@
|
||||
"""沙箱调用的一炁事件构造与回写接口。
|
||||
|
||||
事件结构对齐 docs/architecture/data-and-security-architecture.md:``type / source / payload / trace_id /
|
||||
sensitivity / created_at``。当前一炁事件存储尚未实现,``EventSink`` 用 Protocol 抽象,
|
||||
默认实现写 loguru 日志;后续接入真实事件存储时实现 ``EventSink`` 注入即可,无需改动客户端。
|
||||
|
||||
payload 仅记录摘要(argv、输出长度等),不含密钥、Authorization、签名 URL 或完整输出,
|
||||
遵循事件 payload 脱敏约定。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any, Protocol
|
||||
|
||||
from app.core.logging import logger
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SandboxEvent:
|
||||
"""沙箱领域事件(派生自一炁事件契约)。"""
|
||||
|
||||
type: str # tool.started / tool.completed / tool.failed
|
||||
source: str = "tool"
|
||||
trace_id: str | None = None
|
||||
sensitivity: str = "low" # none / low / medium / high
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
created_at: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
|
||||
|
||||
|
||||
class EventSink(Protocol):
|
||||
"""事件回写接口。"""
|
||||
|
||||
async def emit(self, event: SandboxEvent) -> None: ...
|
||||
|
||||
|
||||
class LoggingEventSink:
|
||||
"""默认事件回写:输出到 loguru(结构化字段,不含敏感原文)。"""
|
||||
|
||||
async def emit(self, event: SandboxEvent) -> None:
|
||||
logger.bind(
|
||||
event_type=event.type,
|
||||
source=event.source,
|
||||
trace_id=event.trace_id,
|
||||
sensitivity=event.sensitivity,
|
||||
payload=event.payload,
|
||||
).info("sandbox event: {}", event.type)
|
||||
|
||||
|
||||
class CompositeEventSink:
|
||||
"""扇出到多个 sink:一次 emit 顺序广播给所有子 sink。
|
||||
|
||||
供 ToolRuntime 同时落 EventStore(一炁流程事件)与 AuditEvent(治理审计),互不感知。
|
||||
任一子 sink 抛异常会中断后续 sink——审计写入失败应显式暴露(宁可事件漏写也不静默吞
|
||||
异常,便于及时发现 DB 问题)。
|
||||
"""
|
||||
|
||||
def __init__(self, sinks: list[EventSink]) -> None:
|
||||
self._sinks = sinks
|
||||
|
||||
async def emit(self, event: SandboxEvent) -> None:
|
||||
for sink in self._sinks:
|
||||
await sink.emit(event)
|
||||
|
||||
|
||||
def started_event(
|
||||
*,
|
||||
tool: str,
|
||||
argv: list[str],
|
||||
risk_level: str,
|
||||
trace_id: str | None,
|
||||
run_id: str | None,
|
||||
) -> SandboxEvent:
|
||||
return SandboxEvent(
|
||||
type="tool.started",
|
||||
trace_id=trace_id,
|
||||
payload={"tool": tool, "risk_level": risk_level, "run_id": run_id, "argv": argv},
|
||||
)
|
||||
|
||||
|
||||
def completed_event(
|
||||
*,
|
||||
tool: str,
|
||||
trace_id: str | None,
|
||||
run_id: str | None,
|
||||
exit_code: int | None,
|
||||
duration_ms: int,
|
||||
timeout: bool,
|
||||
truncated: bool,
|
||||
stdout_len: int,
|
||||
stderr_len: int,
|
||||
) -> SandboxEvent:
|
||||
return SandboxEvent(
|
||||
type="tool.completed",
|
||||
trace_id=trace_id,
|
||||
payload={
|
||||
"tool": tool,
|
||||
"run_id": run_id,
|
||||
"exit_code": exit_code,
|
||||
"duration_ms": duration_ms,
|
||||
"timeout": timeout,
|
||||
"truncated": truncated,
|
||||
"stdout_len": stdout_len,
|
||||
"stderr_len": stderr_len,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def failed_event(
|
||||
*,
|
||||
tool: str,
|
||||
trace_id: str | None,
|
||||
run_id: str | None,
|
||||
error_code: str,
|
||||
error_message: str,
|
||||
) -> SandboxEvent:
|
||||
return SandboxEvent(
|
||||
type="tool.failed",
|
||||
trace_id=trace_id,
|
||||
payload={
|
||||
"tool": tool,
|
||||
"run_id": run_id,
|
||||
"error_code": error_code,
|
||||
"error_message": error_message,
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,50 @@
|
||||
"""沙箱客户端异常。错误码遵循契约字典的 API 标准错误码口径。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class SandboxError(Exception):
|
||||
"""沙箱调用基础异常。"""
|
||||
|
||||
def __init__(self, message: str, *, code: str = "sandbox_error") -> None:
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.message = message
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.message
|
||||
|
||||
|
||||
class SandboxUnavailable(SandboxError):
|
||||
"""沙箱不可达(网络错误 / 连接超时)。"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, code="sandbox_unavailable")
|
||||
|
||||
|
||||
class SandboxUnauthorized(SandboxError):
|
||||
"""授权凭证被拒(HTTP 401)。"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, code="sandbox_unauthorized")
|
||||
|
||||
|
||||
class SandboxBadRequest(SandboxError):
|
||||
"""请求参数错误(HTTP 400)。"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, code="bad_request")
|
||||
|
||||
|
||||
class SandboxExecFailed(SandboxError):
|
||||
"""沙箱内部执行错误(HTTP 5xx,非 401)。"""
|
||||
|
||||
def __init__(self, message: str, *, code: str = "exec_failed") -> None:
|
||||
super().__init__(message, code=code)
|
||||
|
||||
|
||||
class SandboxNotFound(SandboxError):
|
||||
"""目标任务不存在或已完成(HTTP 404),例如取消一个未知 trace_id。"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, code="not_found")
|
||||
@@ -0,0 +1,143 @@
|
||||
"""沙箱执行请求 / 响应模型,与 Rust 端 JSON 契约对齐。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
|
||||
class ExecRequest(BaseModel):
|
||||
"""``POST /v1/exec`` 请求体。``argv``(命令式)与 ``code``+``language``(代码片段式)二选一。"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
argv: list[str] | None = Field(None, min_length=1)
|
||||
code: str | None = None
|
||||
language: str | None = None # code 模式语言,默认 python
|
||||
stdin: str | None = None
|
||||
timeout_ms: int | None = None
|
||||
workdir: str | None = None
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
max_output_bytes: int | None = None
|
||||
trace_id: str | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _enforce_argv_code_mutex(self) -> ExecRequest:
|
||||
has_argv = self.argv is not None
|
||||
has_code = self.code is not None
|
||||
if has_argv == has_code:
|
||||
raise ValueError("必须且只能指定 argv 或 code 之一")
|
||||
if has_code and self.language is None:
|
||||
self.language = "python"
|
||||
return self
|
||||
|
||||
|
||||
class ExecResponse(BaseModel):
|
||||
"""``POST /v1/exec`` 响应体。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
ok: bool
|
||||
exit_code: int | None = None
|
||||
stdout: str = ""
|
||||
stderr: str = ""
|
||||
duration_ms: int = 0
|
||||
timeout: bool = False
|
||||
cancelled: bool = False
|
||||
truncated: bool = False
|
||||
trace_id: str = ""
|
||||
isolation_mode: str = "process" # 实际生效的隔离模式(process/container)
|
||||
|
||||
|
||||
class CancelResponse(BaseModel):
|
||||
"""``POST /v1/cancel`` 响应体。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
cancelled: bool
|
||||
trace_id: str
|
||||
|
||||
|
||||
# ---- 文件 IO(不经 shell,二进制安全;与 Rust fileio.rs 契约对齐) ----
|
||||
|
||||
|
||||
class FileReadRequest(BaseModel):
|
||||
"""``POST /v1/file/read`` 请求。``path`` 为 base 内相对路径。"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
path: str
|
||||
offset: int | None = None
|
||||
limit: int | None = None
|
||||
|
||||
|
||||
class FileReadResponse(BaseModel):
|
||||
"""``POST /v1/file/read`` 响应。``content_b64`` 为 base64url-no-pad(二进制安全)。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
content_b64: str
|
||||
size: int
|
||||
offset: int
|
||||
returned: int
|
||||
truncated: bool
|
||||
|
||||
|
||||
class FileWriteRequest(BaseModel):
|
||||
"""``POST /v1/file/write`` 请求。``content_b64`` 为 base64url-no-pad。"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
path: str
|
||||
content_b64: str
|
||||
|
||||
|
||||
class FileWriteResponse(BaseModel):
|
||||
"""``POST /v1/file/write`` 响应。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
bytes_written: int
|
||||
|
||||
|
||||
class FileListRequest(BaseModel):
|
||||
"""``POST /v1/file/list`` 请求。``path`` 默认 base 根(``.``)。"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
path: str = "."
|
||||
|
||||
|
||||
class DirEntry(BaseModel):
|
||||
"""目录条目。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
name: str
|
||||
kind: str # "file" | "dir" | "other"
|
||||
size: int
|
||||
|
||||
|
||||
class FileListResponse(BaseModel):
|
||||
"""``POST /v1/file/list`` 响应(条目按名字排序)。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
entries: list[DirEntry]
|
||||
|
||||
|
||||
class FileStatRequest(BaseModel):
|
||||
"""``POST /v1/file/stat`` 请求。"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
path: str
|
||||
|
||||
|
||||
class FileStatResponse(BaseModel):
|
||||
"""``POST /v1/file/stat`` 响应。"""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
kind: str
|
||||
size: int
|
||||
modified_unix: int
|
||||
@@ -0,0 +1,84 @@
|
||||
"""向量存储接口与工厂。
|
||||
|
||||
通过 ``get_vector_store()`` 获取当前配置的后端实现:
|
||||
- ``chromadb``(默认):ChromaDB,内存/HTTP 双模式
|
||||
- ``memory``:纯内存实现(测试用)
|
||||
|
||||
通过环境变量 ``CHROMADB_HOST`` / ``CHROMADB_PORT`` 控制 ChromaDB 连接模式。
|
||||
通过 ``ENABLE_EMBEDDING=false`` 关闭 embedding(纯 SQL 检索)。
|
||||
通过 ``VECTOR_DB_TYPE=chroma`` 选择向量库后端。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VectorStoreProtocol(Protocol):
|
||||
"""向量存储接口。所有后端实现必须遵循此协议。"""
|
||||
|
||||
def add_memory(
|
||||
self,
|
||||
character_id: str,
|
||||
memory_id: str,
|
||||
text: str,
|
||||
embedding: list[float],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None: ...
|
||||
|
||||
def update_memory(
|
||||
self,
|
||||
character_id: str,
|
||||
memory_id: str,
|
||||
text: str,
|
||||
embedding: list[float],
|
||||
*,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None: ...
|
||||
|
||||
def delete_memory(self, character_id: str, memory_id: str) -> None: ...
|
||||
|
||||
def search(
|
||||
self,
|
||||
character_id: str,
|
||||
query: list[float],
|
||||
*,
|
||||
limit: int = 5,
|
||||
) -> list[dict[str, Any]]: ...
|
||||
|
||||
def count(self, character_id: str) -> int: ...
|
||||
|
||||
|
||||
_vector_store: VectorStoreProtocol | None = None
|
||||
|
||||
|
||||
def get_vector_store() -> VectorStoreProtocol | None:
|
||||
"""根据配置返回向量存储后端;embedding 关闭时返回 None。"""
|
||||
from app.core.config import settings
|
||||
|
||||
if not settings.vector_db_enabled:
|
||||
return None
|
||||
|
||||
global _vector_store # noqa: PLW0603
|
||||
if _vector_store is not None:
|
||||
return _vector_store
|
||||
|
||||
backend = settings.vector_db_type
|
||||
if backend == "memory":
|
||||
from app.integrations.chromadb.memory_store import MemoryVectorStore
|
||||
|
||||
_vector_store = MemoryVectorStore()
|
||||
else:
|
||||
# 默认 chromadb
|
||||
from app.integrations.chromadb.client import ChromaDBVectorStore
|
||||
|
||||
_vector_store = ChromaDBVectorStore()
|
||||
return _vector_store
|
||||
|
||||
|
||||
def reset_vector_store() -> None:
|
||||
"""重置单例(测试用)。"""
|
||||
global _vector_store # noqa: PLW0603
|
||||
_vector_store = None
|
||||
Reference in new issue
Block a user