初始化提交

This commit is contained in:
chenxin committed 2026-08-21 16:00:13 +08:00
commit 887c1f449e
963 files changed
+166869

No files matched your search

+1
View File
@@ -0,0 +1 @@
"""外部集成层。"""
+187
View File
@@ -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)
+1
View File
@@ -0,0 +1 @@
"""MCP server 接入(简化 HTTP client,非完整 JSON-RPC)。"""
+106
View File
@@ -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
+12
View File
@@ -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",
]
+269
View File
@@ -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")
+24
View File
@@ -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)
+156
View File
@@ -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,
)
)
+475
View File
@@ -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
+127
View File
@@ -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")
+143
View File
@@ -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
+84
View File
@@ -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