forked from admin/PlantAI-rag-new
feat(capabilities): S1+S2 机制地基+多轮记忆(能力框架/热加载/MCP Provider/history 线程化)
- capabilities/{manifests,registry,providers,install_manager}.py + tool_runner.py:
三层机制(注册层/调用层/安装管理层),四 Provider,热加载差分刷新
- configs/capabilities.yaml: 唯一注册入口(明文热加载面)
- api.py: QARequest+history/capabilities, /capabilities/install loopback 路由
- llm_client.py: complete() tools 透传 + complete_raw_messages() 工具调用循环支持
- pipeline.py/answer_composer.py: history 线程化(只进 Dialogue Context,检索零回归)
- sdk: services/history.py(3轮/16k预算) + routers/qa.py session_id 剥离注入
- compile_qa.sh: rsync 允许子目录 + capabilities.yaml 入 install 清单
- tooling/cap/: fake_mcp_server.py + cap_selfcheck.py 机制自检
- requirements.lock: mcp 2.1.1 + 依赖 12 包
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,39 @@
|
||||
# capabilities.yaml —— Capability Gateway 唯一注册入口(明文、热加载面)
|
||||
#
|
||||
# 每个能力 = 一段 manifest(provider/permission/schema/默认开关)。
|
||||
# 坏条目跳过并记入 load_warnings;一条坏配置不影响其余能力,更不影响 QA 主体启动。
|
||||
#
|
||||
# provider: http | mcp | python | prompt_pack
|
||||
# permission: read_only | write(write 本批拒绝注册)
|
||||
# allowed_llm_phases: [answer](本批仅答案阶段可调用工具)
|
||||
#
|
||||
# 聊天内自助安装:install_manager 校验后追加写本文件 → reload() 热加载(不重启)。
|
||||
# 手工编辑同样生效:下次 reload(或下次安装操作)后差分刷新。
|
||||
#
|
||||
# 示例(http provider):
|
||||
# - name: web_search
|
||||
# provider: http
|
||||
# permission: read_only
|
||||
# description: 检索 PubMed 获取最新文献
|
||||
# session_default: false
|
||||
# schema:
|
||||
# query: {type: string, required: true, description: 检索关键词}
|
||||
# http:
|
||||
# endpoint: https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi
|
||||
# method: GET
|
||||
# timeout: 8
|
||||
#
|
||||
# 示例(mcp provider):
|
||||
# - name: my_mcp_tool
|
||||
# provider: mcp
|
||||
# permission: read_only
|
||||
# description: 由外部 MCP server 提供的工具
|
||||
# session_default: false
|
||||
# mcp:
|
||||
# command: npx
|
||||
# args: [-y, @scope/server-name]
|
||||
# # 或 streamable-http:
|
||||
# # mcp:
|
||||
# # url: https://example.com/mcp
|
||||
#
|
||||
capabilities: []
|
||||
@@ -87,6 +87,7 @@ Do not use unsupported general knowledge. By default, prefer not to use general
|
||||
correction_note: str = "",
|
||||
agent_handoff: Dict[str, Any] | None = None,
|
||||
subquestion_evidence_groups: List[Dict[str, Any]] | None = None,
|
||||
history: List[Dict[str, str]] | None = None,
|
||||
) -> str:
|
||||
context_data = self._build_context_data(evidence, graph_chain_context)
|
||||
allowed_evidence_markers = {
|
||||
@@ -112,6 +113,7 @@ Do not use unsupported general knowledge. By default, prefer not to use general
|
||||
if correction_note
|
||||
else ""
|
||||
)
|
||||
dialogue_block = self._build_dialogue_context(history or [])
|
||||
handoff_block = (
|
||||
"--- Agent-approved final handoff ---\n\n"
|
||||
+ json.dumps(agent_handoff, ensure_ascii=False, indent=2)
|
||||
@@ -174,8 +176,7 @@ If the provided data are insufficient, say so directly.
|
||||
|
||||
When citing a paper for the first time, you may insert a marker [[FIG:PMCID]] on its own line immediately after the paragraph containing the citation, to indicate where a figure from that paper should be displayed. Use at most one marker per paper, at most 6 markers total. Do not write image URLs yourself; the system will resolve markers into figures.
|
||||
|
||||
{correction_block}
|
||||
Now answer the following question:
|
||||
{correction_block}{dialogue_block}Now answer the following question:
|
||||
|
||||
{question}
|
||||
""",
|
||||
@@ -263,6 +264,22 @@ Now answer the following question:
|
||||
sections.append("Accepted chain references: none")
|
||||
return "\n".join(sections).strip() + "\n\n"
|
||||
|
||||
@staticmethod
|
||||
def _build_dialogue_context(history: List[Dict[str, str]]) -> str:
|
||||
"""S2:Dialogue Context 段——仅供指代消解,不补充新事实。"""
|
||||
if not history:
|
||||
return ""
|
||||
lines = [
|
||||
"--- Dialogue Context ---",
|
||||
"",
|
||||
"以下是本会话此前的问答记录,仅供消解指代(如「它」「该基因」)使用。",
|
||||
"除指代消解外,不得据此补充新事实;一切事实必须来自 Data Report / GraphChains。",
|
||||
]
|
||||
for msg in history:
|
||||
role_label = "用户" if msg.get("role") == "user" else "助手"
|
||||
lines.append(f"{role_label}: {msg.get('content', '')}")
|
||||
return "\n".join(lines) + "\n\n"
|
||||
|
||||
@staticmethod
|
||||
def _build_context_data(evidence: List[EvidenceItem], graph_chain_context: str) -> str:
|
||||
sections: List[str] = []
|
||||
|
||||
@@ -6,7 +6,7 @@ import queue
|
||||
import threading
|
||||
from functools import lru_cache
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
@@ -44,6 +44,9 @@ class QARequest(BaseModel):
|
||||
endpoint_path_max_expansions: int = Field(default=24, ge=10, le=1000)
|
||||
endpoint_path_max_frontier_states: int = Field(default=120, ge=20, le=2000)
|
||||
endpoint_path_time_budget_seconds: float = Field(default=5.0, ge=1.0, le=30.0)
|
||||
# ---- Capability Gateway(S1):纯数据字段,无 validator(Cython 编译安全)----
|
||||
history: List[Dict[str, str]] = Field(default_factory=list)
|
||||
capabilities: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
@@ -176,6 +179,8 @@ async def chat(request: QARequest) -> Dict[str, Any]:
|
||||
endpoint_path_max_expansions=request.endpoint_path_max_expansions,
|
||||
endpoint_path_max_frontier_states=request.endpoint_path_max_frontier_states,
|
||||
endpoint_path_time_budget_seconds=request.endpoint_path_time_budget_seconds,
|
||||
history=request.history,
|
||||
capabilities=request.capabilities,
|
||||
)
|
||||
except QAAnswerError as exc:
|
||||
raise HTTPException(
|
||||
@@ -239,6 +244,8 @@ async def chat_stream(request: QARequest) -> StreamingResponse:
|
||||
endpoint_path_max_expansions=request.endpoint_path_max_expansions,
|
||||
endpoint_path_max_frontier_states=request.endpoint_path_max_frontier_states,
|
||||
endpoint_path_time_budget_seconds=request.endpoint_path_time_budget_seconds,
|
||||
history=request.history,
|
||||
capabilities=request.capabilities,
|
||||
progress_callback=emit,
|
||||
)
|
||||
events.put(("complete", result))
|
||||
@@ -319,3 +326,55 @@ async def search(request: QARequest) -> Dict[str, Any]:
|
||||
if isinstance(response, dict):
|
||||
response["answer"] = ""
|
||||
return response
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Capability Gateway 安装管理(S1):loopback,热加载生效不重启
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class CapabilityInstallRequest(BaseModel):
|
||||
action: str = "register" # register | unregister
|
||||
manifest: str = "" # register:一段 manifest YAML
|
||||
name: str = "" # unregister:能力名
|
||||
apply: bool = True # False = 只校验不落盘(预检)
|
||||
|
||||
|
||||
@router.post("/capabilities/install")
|
||||
async def capabilities_install(request: CapabilityInstallRequest) -> Dict[str, Any]:
|
||||
"""Capability Gateway 安装管理(loopback)。热加载生效,不重启进程。"""
|
||||
from .capabilities.install_manager import InstallError, get_install_manager
|
||||
|
||||
manager = get_install_manager()
|
||||
if request.action == "unregister":
|
||||
if not request.name:
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail={"error_code": "bad_request", "answer": "", "message": "unregister 需 name"},
|
||||
)
|
||||
try:
|
||||
removed = manager.unregister(request.name)
|
||||
warnings, diff = manager.reload()
|
||||
return {"ok": True, "removed": removed, "warnings": warnings, "diff": diff}
|
||||
except Exception as exc:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error_code": "capability_unregister_failed", "answer": "", "message": str(exc)},
|
||||
) from exc
|
||||
try:
|
||||
manifests, warnings = manager.register_manifest(request.manifest, apply=request.apply)
|
||||
reload_warnings: List[str] = []
|
||||
diff: Dict[str, Any] = {}
|
||||
if request.apply:
|
||||
reload_warnings, diff = manager.reload()
|
||||
return {
|
||||
"ok": True,
|
||||
"installed": [m.name for m in manifests],
|
||||
"warnings": warnings + reload_warnings,
|
||||
"diff": diff,
|
||||
"hot_reload": request.apply,
|
||||
}
|
||||
except InstallError as exc:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error_code": "capability_install_failed", "answer": "", "message": str(exc)},
|
||||
) from exc
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""Capability Gateway 注册/调用/安装管理三层(机制核心)。
|
||||
|
||||
- manifests.py : capabilities.yaml → CapabilityManifest 列表(坏条目跳过 + load_warnings)
|
||||
- registry.py : 运行时注册表(manifest → OpenAI function schema → LLM 工具表;差分刷新)
|
||||
- providers.py : 四 Provider(http / mcp / python / prompt_pack)
|
||||
- install_manager.py: 聊天内自助安装(白名单校验、单写锁+原子替换、reload() 热加载)
|
||||
- tool_runner.py : LLM 工具调用循环(原生 tool_calls + JSON 契约 fallback)
|
||||
|
||||
热加载硬要求:QA 冷启动 ~167s,新增/停用能力全程不重启进程。
|
||||
"""
|
||||
@@ -0,0 +1,206 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""install_manager.py —— 聊天内自助安装管理(机制第三层)。
|
||||
|
||||
- register_manifest(yaml 文本):白名单校验 → 单写锁 → 原子替换写 capabilities.yaml → 审计
|
||||
- unregister(name):移除指定能力段(文件保留)
|
||||
- reload():热加载(重解析 → registry 差分刷新 → MCP 子进程按 diff 启停,已加载进程不重启)
|
||||
- 审计:每次操作记 JSONL(install_log.jsonl,只追加)
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import yaml
|
||||
|
||||
from .manifests import MAX_YAML_BYTES, parse_capabilities_text
|
||||
from .registry import get_registry
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
AUDIT_FILE = "capabilities_install_log.jsonl"
|
||||
AUDIT_MAX_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class InstallError(RuntimeError):
|
||||
"""安装失败(原因回传给调用方)。"""
|
||||
|
||||
|
||||
class InstallManager:
|
||||
"""capabilities.yaml 安装管理(单写锁 + 原子替换)。"""
|
||||
|
||||
def __init__(self, config_path: Path | str) -> None:
|
||||
self._config_path = Path(config_path)
|
||||
self._lock = threading.Lock()
|
||||
self._audit_path = self._config_path.parent / AUDIT_FILE
|
||||
|
||||
# ---- 内部 ----
|
||||
|
||||
def _read_entries(self) -> List[Dict[str, Any]]:
|
||||
try:
|
||||
text = self._config_path.read_text(encoding="utf-8")
|
||||
data = yaml.safe_load(text) or {}
|
||||
except Exception:
|
||||
data = {}
|
||||
entries = data.get("capabilities") if isinstance(data, dict) else []
|
||||
return entries if isinstance(entries, list) else []
|
||||
|
||||
def _write_entries(self, entries: List[Dict[str, Any]]) -> None:
|
||||
payload = {"capabilities": entries}
|
||||
text = yaml.safe_dump(payload, allow_unicode=True, sort_keys=False)
|
||||
size = len(text.encode("utf-8"))
|
||||
if size > MAX_YAML_BYTES:
|
||||
raise InstallError(f"capabilities.yaml 超出大小上限({size}B > {MAX_YAML_BYTES}B)")
|
||||
# 原子替换:临时文件同目录 + os.replace
|
||||
fd, tmp_name = tempfile.mkstemp(
|
||||
dir=str(self._config_path.parent), prefix=".capabilities.", suffix=".tmp"
|
||||
)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as handle:
|
||||
handle.write(text)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(tmp_name, self._config_path)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_name)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
def _audit(self, action: str, name: str, detail: str, ok: bool) -> None:
|
||||
try:
|
||||
if self._audit_path.is_file() and self._audit_path.stat().st_size > AUDIT_MAX_BYTES:
|
||||
self._audit_path.unlink(missing_ok=True)
|
||||
with self._audit_path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(
|
||||
json.dumps(
|
||||
{
|
||||
"ts": time.strftime("%Y-%m-%dT%H:%M:%S"),
|
||||
"action": action,
|
||||
"name": name,
|
||||
"detail": detail[:2000],
|
||||
"ok": ok,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
except Exception as exc:
|
||||
LOGGER.warning("能力安装审计写入失败: %s", exc)
|
||||
|
||||
# ---- 公开操作 ----
|
||||
|
||||
def register_manifest(
|
||||
self, yaml_text: str, *, source: str = "api", apply: bool = True
|
||||
) -> Tuple[List[Any], List[str]]:
|
||||
"""校验并(可选)落盘一段 manifest YAML。返回 (manifests, warnings)。
|
||||
|
||||
校验先行(parse_capabilities_text 白名单),通过后与现有条目合并;
|
||||
同名能力 = 更新(覆盖旧段)。写盘用单写锁 + 原子替换。
|
||||
"""
|
||||
manifests, warnings = parse_capabilities_text(yaml_text)
|
||||
bad_names = {w for w in warnings if "跳过" in w}
|
||||
if warnings and not manifests:
|
||||
self._audit("register", yaml_text[:80], "; ".join(warnings[:5]), ok=False)
|
||||
raise InstallError("manifest 校验失败: " + "; ".join(warnings[:8]))
|
||||
if not manifests:
|
||||
self._audit("register", yaml_text[:80], "无有效条目", ok=False)
|
||||
raise InstallError("manifest 为空或全部条目被拒")
|
||||
|
||||
with self._lock:
|
||||
existing = self._read_entries()
|
||||
by_name: Dict[str, Dict[str, Any]] = {}
|
||||
for entry in existing:
|
||||
if isinstance(entry, dict) and entry.get("name"):
|
||||
by_name[str(entry["name"])] = entry
|
||||
for manifest in manifests:
|
||||
raw = {
|
||||
"name": manifest.name,
|
||||
"provider": manifest.provider,
|
||||
"permission": manifest.permission,
|
||||
"description": manifest.description,
|
||||
"allowed_llm_phases": manifest.allowed_llm_phases,
|
||||
"session_default": manifest.session_default,
|
||||
}
|
||||
if manifest.schema:
|
||||
raw["schema"] = manifest.schema
|
||||
for section in ("http", "mcp", "python", "prompt_pack"):
|
||||
if section in manifest.config:
|
||||
raw[section] = manifest.config[section]
|
||||
by_name[manifest.name] = raw
|
||||
merged = list(by_name.values())
|
||||
if not apply:
|
||||
return manifests, warnings
|
||||
try:
|
||||
self._write_entries(merged)
|
||||
except Exception as exc:
|
||||
self._audit("register", yaml_text[:80], f"写盘失败: {exc}", ok=False)
|
||||
raise InstallError(f"写盘失败: {exc}") from exc
|
||||
self._audit("register", yaml_text[:80], f"安装/更新 {len(manifests)} 个能力", ok=True)
|
||||
return manifests, warnings
|
||||
|
||||
def unregister(self, name: str) -> bool:
|
||||
"""移除指定能力段;文件保留其余条目。"""
|
||||
with self._lock:
|
||||
existing = self._read_entries()
|
||||
kept = [
|
||||
entry
|
||||
for entry in existing
|
||||
if not (isinstance(entry, dict) and str(entry.get("name")) == name)
|
||||
]
|
||||
if len(kept) == len(existing):
|
||||
self._audit("unregister", name, "不存在", ok=False)
|
||||
return False
|
||||
self._write_entries(kept)
|
||||
self._audit("unregister", name, "已移除", ok=True)
|
||||
return True
|
||||
|
||||
def reload(self, registry=None) -> Tuple[List[str], Dict[str, Any]]:
|
||||
"""热加载:registry 差分刷新;MCP 连接按 diff 启停(已加载进程不重启)。
|
||||
|
||||
registry 可注入(测试用);默认用进程内单例 get_registry()。
|
||||
"""
|
||||
registry = registry if registry is not None else get_registry()
|
||||
warnings, diff = registry.reload()
|
||||
if diff.get("removed") or diff.get("changed"):
|
||||
# 停用/变更的能力:关闭其 MCP 连接(新连接下次调用时按需拉起)
|
||||
from .providers import get_dispatcher
|
||||
|
||||
names = list(set(diff["removed"]) | set(diff["changed"]))
|
||||
try:
|
||||
get_dispatcher().shutdown_connections(names)
|
||||
except Exception as exc:
|
||||
LOGGER.warning("MCP 连接关闭失败: %s", exc)
|
||||
return warnings, diff
|
||||
|
||||
# ---- 查询 ----
|
||||
|
||||
@property
|
||||
def config_path(self) -> Path:
|
||||
return self._config_path
|
||||
|
||||
|
||||
_install_manager: Optional[InstallManager] = None
|
||||
_install_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_install_manager() -> InstallManager:
|
||||
global _install_manager
|
||||
with _install_lock:
|
||||
if _install_manager is None:
|
||||
root = Path(os.environ.get("GENE_FAMILY_ROOT", "/data/gene-family-mas-v2"))
|
||||
_install_manager = InstallManager(root / "configs" / "capabilities.yaml")
|
||||
return _install_manager
|
||||
|
||||
|
||||
def set_install_manager(manager: InstallManager) -> None:
|
||||
"""测试注入用。"""
|
||||
global _install_manager
|
||||
with _install_lock:
|
||||
_install_manager = manager
|
||||
@@ -0,0 +1,190 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""manifests.py —— capabilities.yaml 解析与校验。
|
||||
|
||||
- 每个能力 = 一段 manifest(provider/permission/schema/默认开关 + provider 专属配置)
|
||||
- 坏条目跳过并记入 load_warnings;一条坏配置不影响其余能力,更不影响 QA 主体启动
|
||||
- 大小上限:单条目 ≤8KB(序列化后)、yaml 总大小 ≤1MB
|
||||
- provider 白名单:http | mcp | python | prompt_pack
|
||||
- phase 白名单:answer(本批只允许答案阶段调用)
|
||||
- permission 白名单:read_only | write(write 本批拒绝注册)
|
||||
"""
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import yaml
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
ALLOWED_PROVIDERS = ("http", "mcp", "python", "prompt_pack")
|
||||
ALLOWED_PHASES = ("answer",)
|
||||
ALLOWED_PERMISSIONS = ("read_only", "write")
|
||||
MAX_ENTRY_BYTES = 8 * 1024
|
||||
MAX_YAML_BYTES = 1024 * 1024
|
||||
# 注册 schema 的已知键;越界键丢弃(防注入)
|
||||
ALLOWED_SCHEMA_TYPES = ("string", "number", "boolean", "integer")
|
||||
|
||||
|
||||
@dataclass
|
||||
class CapabilityManifest:
|
||||
name: str
|
||||
provider: str
|
||||
permission: str = "read_only"
|
||||
description: str = ""
|
||||
allowed_llm_phases: List[str] = field(default_factory=lambda: ["answer"])
|
||||
session_default: bool = False
|
||||
schema: Dict[str, Any] = field(default_factory=dict)
|
||||
config: Dict[str, Any] = field(default_factory=dict) # provider 专属段(http:/mcp:/python:/prompt_pack: 归并)
|
||||
|
||||
def to_tool_spec(self) -> Dict[str, Any]:
|
||||
"""转 OpenAI function schema(tools 参数形态)。"""
|
||||
properties: Dict[str, Any] = {}
|
||||
required: List[str] = []
|
||||
for key, spec in self.schema.items():
|
||||
if not isinstance(spec, dict):
|
||||
continue
|
||||
prop: Dict[str, Any] = {}
|
||||
ptype = spec.get("type")
|
||||
if ptype in ALLOWED_SCHEMA_TYPES:
|
||||
prop["type"] = ptype
|
||||
else:
|
||||
prop["type"] = "string"
|
||||
if spec.get("description"):
|
||||
prop["description"] = str(spec["description"])[:500]
|
||||
if spec.get("enum") and isinstance(spec["enum"], list):
|
||||
prop["enum"] = [str(v)[:100] for v in spec["enum"][:20]]
|
||||
properties[key] = prop
|
||||
if spec.get("required") is True:
|
||||
required.append(key)
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": self.name[:64],
|
||||
"description": (self.description or self.name)[:1000],
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": properties,
|
||||
"required": required,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _clean_text(value: Any, limit: int) -> str:
|
||||
return str(value or "")[:limit].strip()
|
||||
|
||||
|
||||
def _parse_entry(raw: Any, index: int, load_warnings: List[str]) -> CapabilityManifest | None:
|
||||
if not isinstance(raw, dict):
|
||||
load_warnings.append(f"entry[{index}] 非对象,跳过")
|
||||
return None
|
||||
name = _clean_text(raw.get("name"), 64)
|
||||
if not name:
|
||||
load_warnings.append(f"entry[{index}] 缺少 name,跳过")
|
||||
return None
|
||||
provider = _clean_text(raw.get("provider"), 32)
|
||||
if provider not in ALLOWED_PROVIDERS:
|
||||
load_warnings.append(f"能力 {name} provider={provider!r} 不在白名单,跳过")
|
||||
return None
|
||||
permission = _clean_text(raw.get("permission", "read_only"), 32)
|
||||
if permission not in ALLOWED_PERMISSIONS:
|
||||
load_warnings.append(f"能力 {name} permission={permission!r} 不在白名单,跳过")
|
||||
return None
|
||||
if permission == "write":
|
||||
# 本批安全决策:write 权限一律拒绝(聊天内自助安装不允许改文件/起进程以外的写操作)
|
||||
load_warnings.append(f"能力 {name} permission=write 本批拒绝注册,跳过")
|
||||
return None
|
||||
phases = raw.get("allowed_llm_phases") or ["answer"]
|
||||
phases = [
|
||||
_clean_text(p, 32)
|
||||
for p in (phases if isinstance(phases, list) else [phases])
|
||||
if _clean_text(p, 32) in ALLOWED_PHASES
|
||||
]
|
||||
if not phases:
|
||||
phases = ["answer"]
|
||||
schema = raw.get("schema") if isinstance(raw.get("schema"), dict) else {}
|
||||
config: Dict[str, Any] = {}
|
||||
for section in ("http", "mcp", "python", "prompt_pack"):
|
||||
if isinstance(raw.get(section), dict):
|
||||
config[section] = raw[section]
|
||||
config["session_default"] = bool(raw.get("session_default", False))
|
||||
manifest = CapabilityManifest(
|
||||
name=name,
|
||||
provider=provider,
|
||||
permission=permission,
|
||||
description=_clean_text(raw.get("description"), 1000),
|
||||
allowed_llm_phases=phases,
|
||||
session_default=bool(raw.get("session_default", False)),
|
||||
schema=schema,
|
||||
config=config,
|
||||
)
|
||||
# 大小上限:单条目序列化后 ≤8KB
|
||||
import json
|
||||
|
||||
size = len(json.dumps(manifest.to_tool_spec(), ensure_ascii=False).encode("utf-8"))
|
||||
if size > MAX_ENTRY_BYTES:
|
||||
load_warnings.append(f"能力 {name} 条目过大({size}B > {MAX_ENTRY_BYTES}B),跳过")
|
||||
return None
|
||||
return manifest
|
||||
|
||||
|
||||
def parse_capabilities_file(path: Path) -> tuple[List[CapabilityManifest], List[str]]:
|
||||
"""解析 capabilities.yaml → (清单, load_warnings)。坏条目跳过,主体不受影响。"""
|
||||
warnings: List[str] = []
|
||||
if not path.is_file():
|
||||
return [], [f"capabilities.yaml 不存在:{path}"]
|
||||
try:
|
||||
size = path.stat().st_size
|
||||
if size > MAX_YAML_BYTES:
|
||||
return [], [f"capabilities.yaml 过大({size}B > {MAX_YAML_BYTES}B),整文件跳过"]
|
||||
text = path.read_text(encoding="utf-8")
|
||||
data = yaml.safe_load(text) or {}
|
||||
except Exception as exc:
|
||||
return [], [f"capabilities.yaml 解析失败:{type(exc).__name__}: {exc}"]
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return [], ["capabilities.yaml 顶层非映射,跳过"]
|
||||
entries = data.get("capabilities")
|
||||
if not isinstance(entries, list):
|
||||
return [], ["capabilities.yaml 缺 capabilities 列表"]
|
||||
|
||||
manifests: List[CapabilityManifest] = []
|
||||
seen: set[str] = set()
|
||||
for index, raw in enumerate(entries):
|
||||
manifest = _parse_entry(raw, index, warnings)
|
||||
if manifest is None:
|
||||
continue
|
||||
if manifest.name in seen:
|
||||
warnings.append(f"能力 {manifest.name} 重名,仅保留第一个")
|
||||
continue
|
||||
seen.add(manifest.name)
|
||||
manifests.append(manifest)
|
||||
return manifests, warnings
|
||||
|
||||
|
||||
def parse_capabilities_text(text: str) -> tuple[List[CapabilityManifest], List[str]]:
|
||||
"""解析 capabilities.yaml 文本(install_manager 写入前预校验用)。"""
|
||||
warnings: List[str] = []
|
||||
try:
|
||||
data = yaml.safe_load(text) or {}
|
||||
except Exception as exc:
|
||||
return [], [f"YAML 解析失败:{type(exc).__name__}: {exc}"]
|
||||
if not isinstance(data, dict):
|
||||
return [], ["顶层非映射"]
|
||||
entries = data.get("capabilities")
|
||||
if not isinstance(entries, list):
|
||||
return [], ["缺 capabilities 列表"]
|
||||
manifests: List[CapabilityManifest] = []
|
||||
seen: set[str] = set()
|
||||
for index, raw in enumerate(entries):
|
||||
manifest = _parse_entry(raw, index, warnings)
|
||||
if manifest is None:
|
||||
continue
|
||||
if manifest.name in seen:
|
||||
warnings.append(f"能力 {manifest.name} 重名")
|
||||
continue
|
||||
seen.add(manifest.name)
|
||||
manifests.append(manifest)
|
||||
return manifests, warnings
|
||||
@@ -0,0 +1,468 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""providers.py —— 四个 Provider(机制调用层的执行端)。
|
||||
|
||||
- HttpToolProvider : HTTP JSON 工具(PubMed/Crossref 即此形态);仅 https(或本机 http)
|
||||
- McpToolProvider : 官方 mcp SDK(lazy import),stdio 子进程 + streamable-http 两种连接;
|
||||
tools/list 握手 → tools/call;空闲回收、崩溃限次重启(3 次熔断)
|
||||
- PythonToolProvider: module:function 内置 handler(本批预留形态)
|
||||
- PromptPackProvider: 读 SKILL.md 风格 md,注入指定提示词尾部(Superpowers 即此形态)
|
||||
|
||||
所有 provider 调用均为同步形态(QA pipeline 跑在线程池里);MCP 的异步 SDK
|
||||
通过 anyio portal 在专用线程内运行。任何 provider 异常一律转 ToolExecutionError,
|
||||
由 tool_runner 捕获后作为 tool result 回传给 LLM(不炸主流程)。
|
||||
"""
|
||||
import asyncio
|
||||
import importlib
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import requests
|
||||
|
||||
from .manifests import CapabilityManifest
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
MCP_IDLE_SECONDS = 300.0
|
||||
MCP_MAX_RESTARTS = 3
|
||||
MCP_COOLDOWN_SECONDS = 60.0
|
||||
TOOL_RESULT_MAX_CHARS = 4000
|
||||
|
||||
|
||||
class ToolExecutionError(RuntimeError):
|
||||
"""Provider 执行失败(tool_runner 捕获后作为 tool result 回传 LLM)。"""
|
||||
|
||||
|
||||
def _truncate(value: Any, limit: int = TOOL_RESULT_MAX_CHARS) -> str:
|
||||
text = json.dumps(value, ensure_ascii=False, default=str) if not isinstance(value, str) else value
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return text[:limit] + f"...(截断,原长 {len(text)} 字符)"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HttpToolProvider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class HttpToolProvider:
|
||||
"""HTTP JSON 工具:GET/POST 一个 endpoint,参数从 LLM tool 入参映射到 query/json。"""
|
||||
|
||||
@staticmethod
|
||||
def execute(manifest: CapabilityManifest, arguments: Dict[str, Any]) -> str:
|
||||
cfg = manifest.config.get("http") or {}
|
||||
endpoint = str(cfg.get("endpoint") or "").strip()
|
||||
if not endpoint:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} 未配置 http.endpoint")
|
||||
method = str(cfg.get("method") or "GET").upper()
|
||||
if method not in ("GET", "POST"):
|
||||
raise ToolExecutionError(f"工具 {manifest.name} 不支持的 HTTP 方法 {method}")
|
||||
timeout = max(1.0, min(float(cfg.get("timeout") or 8), 30.0))
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "PlantAI-rag-new-capability/1.0",
|
||||
}
|
||||
for key, value in (cfg.get("headers") or {}).items():
|
||||
headers[str(key)] = str(value)
|
||||
query: Dict[str, str] = {}
|
||||
body: Dict[str, Any] = {}
|
||||
for key, value in arguments.items():
|
||||
if method == "GET":
|
||||
query[key] = str(value)
|
||||
else:
|
||||
body[key] = value
|
||||
try:
|
||||
if method == "GET":
|
||||
response = requests.get(
|
||||
endpoint, params=query, headers=headers, timeout=timeout
|
||||
)
|
||||
else:
|
||||
headers["Content-Type"] = "application/json"
|
||||
response = requests.post(
|
||||
endpoint, json=body, headers=headers, timeout=timeout
|
||||
)
|
||||
except requests.Timeout as exc:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} 请求超时({timeout}s)") from exc
|
||||
except requests.RequestException as exc:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} 网络失败: {exc}") from exc
|
||||
if response.status_code >= 400:
|
||||
raise ToolExecutionError(
|
||||
f"工具 {manifest.name} 上游返回 HTTP {response.status_code}"
|
||||
)
|
||||
content_type = response.headers.get("Content-Type", "")
|
||||
if "json" in content_type:
|
||||
try:
|
||||
payload = response.json()
|
||||
except Exception as exc:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} 响应非 JSON: {exc}") from exc
|
||||
else:
|
||||
payload = {"text": response.text[:TOOL_RESULT_MAX_CHARS]}
|
||||
return _truncate(payload)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# McpToolProvider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _McpConnection:
|
||||
"""一个 MCP server 连接(stdio 子进程或 streamable-http URL)。
|
||||
|
||||
生命周期:首次调用时拉起 → 空闲 MCP_IDLE_SECONDS 后回收 → 崩溃限次重启
|
||||
(MCP_MAX_RESTARTS 次后熔断 MCP_COOLDOWN_SECONDS)。
|
||||
|
||||
实现:专用 asyncio 循环线程。mcp 的 stdio/streamable_http client 是 async
|
||||
context manager,其内部 cancel scope/task group 必须在同一事件循环的同一
|
||||
任务内 enter/exit——所以连接建立、工具调用、teardown 全部交给该线程里的
|
||||
一个长驻 _loop_main 任务串行执行;同步侧用 run_coroutine_threadsafe 投递。
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, config: Dict[str, Any]) -> None:
|
||||
self.name = name
|
||||
self.config = config
|
||||
self._lock = threading.Lock()
|
||||
self._loop: Any = None # 专用 asyncio loop
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._session: Any = None
|
||||
self._server_tools: List[str] = []
|
||||
self._started_at = 0.0
|
||||
self._last_used = 0.0
|
||||
self._restarts = 0
|
||||
self._open_until = 0.0
|
||||
|
||||
# ---- loop 线程管理 ----
|
||||
|
||||
def _ensure_loop(self) -> None:
|
||||
if self._loop is not None and not self._loop.is_closed():
|
||||
return
|
||||
import asyncio
|
||||
|
||||
ready = threading.Event()
|
||||
holder: dict = {}
|
||||
|
||||
def _run() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
holder["loop"] = loop
|
||||
ready.set()
|
||||
loop.run_forever()
|
||||
|
||||
thread = threading.Thread(target=_run, name=f"mcp-{self.name}", daemon=True)
|
||||
thread.start()
|
||||
ready.wait(timeout=10)
|
||||
self._loop = holder.get("loop")
|
||||
self._thread = thread
|
||||
if self._loop is None:
|
||||
raise ToolExecutionError(f"工具 {self.name} MCP 事件循环线程启动失败")
|
||||
|
||||
def _submit(self, coro: Any, timeout: float = 120.0) -> Any:
|
||||
self._ensure_loop()
|
||||
future = asyncio.run_coroutine_threadsafe(coro, self._loop)
|
||||
return future.result(timeout=timeout)
|
||||
|
||||
# ---- 连接(在专用 loop 里执行)----
|
||||
|
||||
async def _connect_async(self) -> None:
|
||||
from mcp import ClientSession
|
||||
|
||||
if "url" in self.config:
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
|
||||
ctx = streamable_http_client(str(self.config["url"]))
|
||||
else:
|
||||
from mcp.client.stdio import stdio_client
|
||||
from mcp import StdioServerParameters
|
||||
|
||||
command = str(self.config.get("command") or "").strip()
|
||||
if not command:
|
||||
raise ToolExecutionError(f"工具 {self.name} 缺 mcp.command/mcp.url")
|
||||
params = StdioServerParameters(
|
||||
command=command,
|
||||
args=[str(a) for a in (self.config.get("args") or [])],
|
||||
env=None,
|
||||
)
|
||||
ctx = stdio_client(params)
|
||||
# ctx 与 session 的 enter/exit 都在本任务内完成(同一 loop、同一 task)
|
||||
read_stream, write_stream = await ctx.__aenter__()
|
||||
session = ClientSession(read_stream, write_stream)
|
||||
await session.__aenter__()
|
||||
try:
|
||||
await session.initialize()
|
||||
tools_result = await session.list_tools()
|
||||
except Exception:
|
||||
await session.__aexit__(None, None, None)
|
||||
await ctx.__aexit__(None, None, None)
|
||||
raise
|
||||
self._ctx = ctx
|
||||
self._server_tools = [t.name for t in tools_result.tools]
|
||||
self._session = session
|
||||
self._started_at = time.time()
|
||||
|
||||
async def _call_async(self, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||||
session = self._session
|
||||
if session is None:
|
||||
raise ToolExecutionError(f"工具 {self.name} MCP 会话未建立")
|
||||
result = await session.call_tool(tool_name, arguments)
|
||||
if getattr(result, "is_error", False):
|
||||
raise ToolExecutionError(
|
||||
f"MCP server {self.name} 返回 isError: {getattr(result, 'content', result)}"
|
||||
)
|
||||
parts = []
|
||||
for item in result.content:
|
||||
if hasattr(item, "text"):
|
||||
parts.append(item.text)
|
||||
else:
|
||||
parts.append(str(item))
|
||||
return _truncate("\n".join(parts) or json.dumps(result, ensure_ascii=False, default=str))
|
||||
|
||||
async def _teardown_async(self) -> None:
|
||||
session, self._session = self._session, None
|
||||
ctx = getattr(self, "_ctx", None)
|
||||
self._ctx = None
|
||||
self._server_tools = []
|
||||
if session is not None:
|
||||
try:
|
||||
await session.__aexit__(None, None, None)
|
||||
except Exception:
|
||||
pass
|
||||
if ctx is not None:
|
||||
try:
|
||||
await ctx.__aexit__(None, None, None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 同步面 ----
|
||||
|
||||
def _ensure_session(self) -> None:
|
||||
now = time.time()
|
||||
if self._session is not None:
|
||||
self._last_used = now
|
||||
return
|
||||
if now < self._open_until:
|
||||
raise ToolExecutionError(
|
||||
f"工具 {self.name} 已熔断(崩溃 {MCP_MAX_RESTARTS} 次),"
|
||||
f"{int(self._open_until - now)}s 后可重试"
|
||||
)
|
||||
try:
|
||||
self._submit(self._connect_async())
|
||||
except Exception as exc:
|
||||
self._restarts += 1
|
||||
if self._restarts >= MCP_MAX_RESTARTS:
|
||||
self._open_until = now + MCP_COOLDOWN_SECONDS
|
||||
self._restarts = 0
|
||||
raise ToolExecutionError(
|
||||
f"工具 {self.name} MCP 连接失败(第 {self._restarts} 次): {type(exc).__name__}: {exc}"
|
||||
) from exc
|
||||
self._last_used = time.time()
|
||||
|
||||
def _call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||||
try:
|
||||
return self._submit(self._call_async(tool_name, arguments))
|
||||
except Exception as exc:
|
||||
self._submit(self._teardown_async())
|
||||
raise ToolExecutionError(f"工具 {self.name} MCP 调用失败: {exc}") from exc
|
||||
|
||||
def available_tools(self) -> List[str]:
|
||||
try:
|
||||
self._ensure_session()
|
||||
except ToolExecutionError:
|
||||
return []
|
||||
return list(self._server_tools)
|
||||
|
||||
def call(self, tool_name: str, arguments: Dict[str, Any]) -> str:
|
||||
self._ensure_session()
|
||||
try:
|
||||
return self._call_tool(tool_name, arguments)
|
||||
finally:
|
||||
self._last_used = time.time()
|
||||
|
||||
def reap_if_idle(self, now: Optional[float] = None) -> bool:
|
||||
now = now if now is not None else time.time()
|
||||
if self._session is not None and (now - self._last_used) >= MCP_IDLE_SECONDS:
|
||||
try:
|
||||
self._submit(self._teardown_async())
|
||||
except Exception:
|
||||
pass
|
||||
return True
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
loop, self._loop = self._loop, None
|
||||
self._thread = None
|
||||
if loop is not None and not loop.is_closed():
|
||||
try:
|
||||
future = asyncio.run_coroutine_threadsafe(self._teardown_async(), loop)
|
||||
future.result(timeout=10)
|
||||
except Exception:
|
||||
pass
|
||||
loop.call_soon_threadsafe(loop.stop)
|
||||
|
||||
|
||||
class McpToolProvider:
|
||||
"""MCP 工具管理器:连接按能力名缓存、空闲回收、崩溃熔断。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._connections: Dict[str, _McpConnection] = {}
|
||||
|
||||
def _connection(self, manifest: CapabilityManifest) -> _McpConnection:
|
||||
with self._lock:
|
||||
conn = self._connections.get(manifest.name)
|
||||
if conn is None:
|
||||
conn = _McpConnection(manifest.name, manifest.config.get("mcp") or {})
|
||||
self._connections[manifest.name] = conn
|
||||
return conn
|
||||
|
||||
def execute(self, manifest: CapabilityManifest, arguments: Dict[str, Any]) -> str:
|
||||
cfg = manifest.config.get("mcp") or {}
|
||||
tool_name = str(cfg.get("tool") or manifest.name)
|
||||
conn = self._connection(manifest)
|
||||
return conn.call(tool_name, arguments)
|
||||
|
||||
def available_tools(self, manifest: CapabilityManifest) -> List[str]:
|
||||
cfg = manifest.config.get("mcp") or {}
|
||||
tool_name = str(cfg.get("tool") or manifest.name)
|
||||
conn = self._connection(manifest)
|
||||
tools = conn.available_tools()
|
||||
if not tools:
|
||||
return []
|
||||
# 与声明 tool 对比:若声明 tool 不在 server 列表,也返回(调用时会报清晰错误)
|
||||
if tool_name not in tools:
|
||||
return tools + [f"<缺失声明工具 {tool_name}>"]
|
||||
return tools
|
||||
|
||||
def reap_idle(self) -> int:
|
||||
reaped = 0
|
||||
with self._lock:
|
||||
for name, conn in list(self._connections.items()):
|
||||
if conn.reap_if_idle():
|
||||
reaped += 1
|
||||
LOGGER.info("MCP 连接空闲回收: %s", name)
|
||||
return reaped
|
||||
|
||||
def shutdown_connections(self, names: Optional[List[str]] = None) -> None:
|
||||
"""热加载差分:停用/删除的能力对应连接显式关闭(未变的进程不重启)。"""
|
||||
with self._lock:
|
||||
targets = names if names is not None else list(self._connections)
|
||||
for name in targets:
|
||||
conn = self._connections.pop(name, None)
|
||||
if conn is not None:
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PythonToolProvider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PythonToolProvider:
|
||||
"""module:function 内置 handler(本批无实例,形态留位)。"""
|
||||
|
||||
@staticmethod
|
||||
def execute(manifest: CapabilityManifest, arguments: Dict[str, Any]) -> str:
|
||||
cfg = manifest.config.get("python") or {}
|
||||
target = str(cfg.get("function") or "").strip()
|
||||
if not target:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} 未配置 python.function")
|
||||
if "." not in target:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} function 需 module:function 形态")
|
||||
module_name, _, func_name = target.rpartition(".")
|
||||
try:
|
||||
module = importlib.import_module(module_name)
|
||||
func = getattr(module, func_name)
|
||||
except Exception as exc:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} handler 加载失败: {exc}") from exc
|
||||
try:
|
||||
result = func(arguments)
|
||||
except Exception as exc:
|
||||
raise ToolExecutionError(f"工具 {manifest.name} handler 执行失败: {exc}") from exc
|
||||
return _truncate(result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PromptPackProvider
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class PromptPackProvider:
|
||||
"""读 SKILL.md 风格 md 文件(Superpowers 即此形态)。
|
||||
|
||||
config.prompt_pack.file: md 路径(相对 QA 根或绝对路径);白名单:仅允许
|
||||
QA_ROOT/prompts/ 目录下文件(防任意文件读取)。
|
||||
"""
|
||||
|
||||
_QA_ROOT_HINT = Path("/data/gene-family-mas-v2")
|
||||
|
||||
@classmethod
|
||||
def load(cls, manifest: CapabilityManifest, qa_root: Optional[Path] = None) -> str:
|
||||
cfg = manifest.config.get("prompt_pack") or {}
|
||||
file_name = str(cfg.get("file") or "").strip()
|
||||
if not file_name:
|
||||
return ""
|
||||
root = qa_root or cls._QA_ROOT_HINT
|
||||
path = Path(file_name)
|
||||
if not path.is_absolute():
|
||||
path = root / "prompts" / path.name
|
||||
# 白名单:必须位于 root/prompts/ 之下
|
||||
try:
|
||||
path.resolve().relative_to((root / "prompts").resolve())
|
||||
except ValueError as exc:
|
||||
raise ToolExecutionError(
|
||||
f"prompt_pack 文件越界(只允许 QA_ROOT/prompts/ 下): {path}"
|
||||
) from exc
|
||||
try:
|
||||
return path.read_text(encoding="utf-8").strip()
|
||||
except FileNotFoundError:
|
||||
LOGGER.warning("prompt_pack 文件缺失,静默跳过: %s", path)
|
||||
return ""
|
||||
except Exception as exc:
|
||||
raise ToolExecutionError(f"prompt_pack 读取失败 {path}: {exc}") from exc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 分发
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ToolDispatcher:
|
||||
"""按 manifest.provider 分发执行;prompt_pack 不参与工具调用(注入提示词用)。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._mcp = McpToolProvider()
|
||||
|
||||
def execute(self, manifest: CapabilityManifest, arguments: Dict[str, Any]) -> str:
|
||||
provider = manifest.provider
|
||||
if provider == "http":
|
||||
return HttpToolProvider.execute(manifest, arguments)
|
||||
if provider == "mcp":
|
||||
return self._mcp.execute(manifest, arguments)
|
||||
if provider == "python":
|
||||
return PythonToolProvider.execute(manifest, arguments)
|
||||
raise ToolExecutionError(f"工具 {manifest.name} provider={provider} 不支持调用")
|
||||
|
||||
def reap_idle(self) -> int:
|
||||
return self._mcp.reap_idle()
|
||||
|
||||
def shutdown_connections(self, names: Optional[List[str]] = None) -> None:
|
||||
self._mcp.shutdown_connections(names)
|
||||
|
||||
|
||||
# 进程内单例(registry 热加载后按 diff 启停用)
|
||||
_dispatcher: Optional[ToolDispatcher] = None
|
||||
_dispatcher_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_dispatcher() -> ToolDispatcher:
|
||||
global _dispatcher
|
||||
with _dispatcher_lock:
|
||||
if _dispatcher is None:
|
||||
_dispatcher = ToolDispatcher()
|
||||
return _dispatcher
|
||||
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""registry.py —— 运行时能力注册表。
|
||||
|
||||
- 加载 capabilities.yaml → CapabilityManifest → OpenAI function schema
|
||||
- 热加载:reload() 差分刷新(新增/更新/删除按 diff 处理;未变的条目不动)
|
||||
- get/enabled/list 供 QA 主流程与 /api/capabilities 使用
|
||||
"""
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .manifests import CapabilityManifest, parse_capabilities_file
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
# 每会话能力开关:capabilities_dict name → bool(请求里由 SDK 注入)
|
||||
# 注册表级默认值只做新会话回退;env 总闸与 session 开关在 pipeline 里共同决策
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolSpec:
|
||||
name: str
|
||||
description: str
|
||||
function_schema: Dict[str, Any]
|
||||
manifest: CapabilityManifest
|
||||
|
||||
|
||||
class CapabilityRegistry:
|
||||
"""能力注册表(进程内单例,热加载差分刷新)。"""
|
||||
|
||||
def __init__(self, config_path: Path | str | None = None) -> None:
|
||||
self._lock = threading.RLock()
|
||||
self._tools: Dict[str, ToolSpec] = {}
|
||||
self._warnings: List[str] = []
|
||||
self._config_path = Path(config_path) if config_path else None
|
||||
|
||||
# ---- 加载 / 热加载 ----
|
||||
|
||||
def load(self, path: Path | str | None = None) -> List[str]:
|
||||
"""初载或重载 capabilities.yaml;返回 load_warnings。"""
|
||||
if path is not None:
|
||||
self._config_path = Path(path)
|
||||
if self._config_path is None:
|
||||
return ["registry 未配置 capabilities.yaml 路径"]
|
||||
manifests, warnings = parse_capabilities_file(self._config_path)
|
||||
with self._lock:
|
||||
old = dict(self._tools)
|
||||
new_tools: Dict[str, ToolSpec] = {}
|
||||
for manifest in manifests:
|
||||
new_tools[manifest.name] = ToolSpec(
|
||||
name=manifest.name,
|
||||
description=manifest.description,
|
||||
function_schema=manifest.to_tool_spec(),
|
||||
manifest=manifest,
|
||||
)
|
||||
added = sorted(set(new_tools) - set(old))
|
||||
removed = sorted(set(old) - set(new_tools))
|
||||
changed = sorted(
|
||||
name
|
||||
for name in set(new_tools) & set(old)
|
||||
if new_tools[name].function_schema != old[name].function_schema
|
||||
)
|
||||
self._tools = new_tools
|
||||
self._warnings = list(warnings)
|
||||
if added or removed or changed:
|
||||
LOGGER.info(
|
||||
"capabilities reload: added=%s removed=%s changed=%s",
|
||||
added,
|
||||
removed,
|
||||
changed,
|
||||
)
|
||||
return list(warnings)
|
||||
|
||||
def reload(self) -> tuple[List[str], Dict[str, Any]]:
|
||||
"""热加载:重解析 yaml → 差分刷新。返回 (warnings, diff)。"""
|
||||
manifests, warnings = parse_capabilities_file(self._config_path or Path("."))
|
||||
with self._lock:
|
||||
old = dict(self._tools)
|
||||
new_tools: Dict[str, ToolSpec] = {}
|
||||
for manifest in manifests:
|
||||
new_tools[manifest.name] = ToolSpec(
|
||||
name=manifest.name,
|
||||
description=manifest.description,
|
||||
function_schema=manifest.to_tool_spec(),
|
||||
manifest=manifest,
|
||||
)
|
||||
diff = {
|
||||
"added": sorted(set(new_tools) - set(old)),
|
||||
"removed": sorted(set(old) - set(new_tools)),
|
||||
"changed": sorted(
|
||||
name
|
||||
for name in set(new_tools) & set(old)
|
||||
if new_tools[name].function_schema != old[name].function_schema
|
||||
),
|
||||
}
|
||||
self._tools = new_tools
|
||||
self._warnings = list(warnings)
|
||||
return list(warnings), diff
|
||||
|
||||
# ---- 查询 ----
|
||||
|
||||
def get(self, name: str) -> Optional[ToolSpec]:
|
||||
with self._lock:
|
||||
return self._tools.get(name)
|
||||
|
||||
def list_all(self) -> List[ToolSpec]:
|
||||
with self._lock:
|
||||
return sorted(self._tools.values(), key=lambda t: t.name)
|
||||
|
||||
@property
|
||||
def load_warnings(self) -> List[str]:
|
||||
with self._lock:
|
||||
return list(self._warnings)
|
||||
|
||||
def enabled_tools(
|
||||
self,
|
||||
session_capabilities: Optional[Dict[str, Any]] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""返回当前会话可用的 function schema 列表(session 开关优先,回退注册表默认)。"""
|
||||
session_capabilities = session_capabilities or {}
|
||||
tools: List[Dict[str, Any]] = []
|
||||
with self._lock:
|
||||
for name, spec in self._tools.items():
|
||||
enabled = bool(session_capabilities.get(name, spec.manifest.session_default))
|
||||
if enabled:
|
||||
tools.append(spec.function_schema)
|
||||
return tools
|
||||
|
||||
|
||||
# 进程内单例(api.py / pipeline.py 引用)
|
||||
_registry: Optional[CapabilityRegistry] = None
|
||||
_registry_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_registry() -> CapabilityRegistry:
|
||||
global _registry
|
||||
with _registry_lock:
|
||||
if _registry is None:
|
||||
import os
|
||||
|
||||
root = Path(os.environ.get("GENE_FAMILY_ROOT", "/data/gene-family-mas-v2"))
|
||||
_registry = CapabilityRegistry(root / "configs" / "capabilities.yaml")
|
||||
_registry.load()
|
||||
return _registry
|
||||
|
||||
|
||||
def set_registry(registry: CapabilityRegistry) -> None:
|
||||
"""测试注入用。"""
|
||||
global _registry
|
||||
with _registry_lock:
|
||||
_registry = registry
|
||||
@@ -68,6 +68,7 @@ class LLMClient:
|
||||
max_tokens: int = 2000,
|
||||
json_mode: bool = False,
|
||||
request_attempts: int | None = None,
|
||||
tools: List[Dict[str, Any]] | None = None,
|
||||
) -> str:
|
||||
if not self.api_key:
|
||||
raise RuntimeError("QA V1 LLM API key is not configured")
|
||||
@@ -83,6 +84,8 @@ class LLMClient:
|
||||
}
|
||||
if json_mode:
|
||||
request_body["response_format"] = {"type": "json_object"}
|
||||
if tools:
|
||||
request_body["tools"] = tools
|
||||
attempt_limit = self.max_attempts
|
||||
if request_attempts is not None:
|
||||
attempt_limit = max(1, min(int(request_attempts), self.max_attempts))
|
||||
@@ -143,6 +146,7 @@ class LLMClient:
|
||||
choice = choices[0] if isinstance(choices[0], dict) else {}
|
||||
message = choice.get("message") if isinstance(choice, dict) else {}
|
||||
content = self._message_text(message)
|
||||
has_tool_calls = bool((message or {}).get("tool_calls"))
|
||||
finish_reason = str(choice.get("finish_reason") or "")
|
||||
if finish_reason.strip().lower() in {"length", "max_tokens"}:
|
||||
last_error = (
|
||||
@@ -150,8 +154,8 @@ class LLMClient:
|
||||
f"content_chars={len(content)}"
|
||||
)
|
||||
raise LLMResponseError(last_error)
|
||||
if content:
|
||||
if attempt > 1:
|
||||
if content or has_tool_calls:
|
||||
if content and attempt > 1:
|
||||
LOGGER.warning(
|
||||
"LLM call recovered on attempt %s/%s after %.2fs; "
|
||||
"finish_reason=%s content_chars=%s",
|
||||
@@ -192,6 +196,106 @@ class LLMClient:
|
||||
f"LLM failed after {attempt_limit} attempts: {last_error}"
|
||||
)
|
||||
|
||||
def complete_raw_messages(
|
||||
self,
|
||||
system: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
*,
|
||||
tools: List[Dict[str, Any]] | None = None,
|
||||
temperature: float = 0.1,
|
||||
max_tokens: int = 2000,
|
||||
json_mode: bool = False,
|
||||
request_attempts: int | None = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Generic multi-message completion for the tool-calling loop.
|
||||
|
||||
Returns {"content", "tool_calls", "finish_reason"}; tool_calls is a list of
|
||||
{"id", "name", "arguments"}. Retry policy mirrors complete().
|
||||
"""
|
||||
if not self.api_key:
|
||||
raise RuntimeError("QA V1 LLM API key is not configured")
|
||||
request_body = {
|
||||
"model": self.model,
|
||||
"messages": [{"role": "system", "content": system}] + list(messages),
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": False,
|
||||
}
|
||||
if json_mode:
|
||||
request_body["response_format"] = {"type": "json_object"}
|
||||
if tools:
|
||||
request_body["tools"] = tools
|
||||
attempt_limit = self.max_attempts
|
||||
if request_attempts is not None:
|
||||
attempt_limit = max(1, min(int(request_attempts), self.max_attempts))
|
||||
last_error = "unknown upstream failure"
|
||||
for attempt in range(1, attempt_limit + 1):
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{self.base_url}/chat/completions",
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=request_body,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
payload = response.json()
|
||||
if response.status_code >= 400:
|
||||
last_error = f"LLM HTTP {response.status_code}: {payload}"
|
||||
if response.status_code not in {
|
||||
408, 409, 425, 429, 500, 502, 503, 504,
|
||||
}:
|
||||
raise LLMNonRetryableError(last_error)
|
||||
raise requests.HTTPError(last_error, response=response)
|
||||
except LLMNonRetryableError as exc:
|
||||
last_error = f"{type(exc).__name__}: {exc}"
|
||||
LOGGER.warning(
|
||||
"LLM raw call failed without retry after %.2fs: %s",
|
||||
time.perf_counter() - started,
|
||||
last_error[:800],
|
||||
)
|
||||
raise
|
||||
except Exception as exc:
|
||||
last_error = f"{type(exc).__name__}: {exc}"
|
||||
LOGGER.warning(
|
||||
"LLM raw call attempt %s/%s failed after %.2fs: %s",
|
||||
attempt, attempt_limit, time.perf_counter() - started, last_error[:800],
|
||||
)
|
||||
if attempt >= attempt_limit:
|
||||
break
|
||||
self._sleep_before_retry(attempt)
|
||||
continue
|
||||
if isinstance(payload, dict) and payload.get("choices"):
|
||||
choice = payload["choices"][0]
|
||||
message = choice.get("message") or {}
|
||||
content = self._message_text(message)
|
||||
raw_tool_calls = message.get("tool_calls") or choice.get("tool_calls")
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
if isinstance(raw_tool_calls, list):
|
||||
for call in raw_tool_calls:
|
||||
if not isinstance(call, dict):
|
||||
continue
|
||||
fn = call.get("function") or {}
|
||||
tool_calls.append({
|
||||
"id": str(call.get("id") or ""),
|
||||
"name": str(fn.get("name") or ""),
|
||||
"arguments": str(fn.get("arguments") or "{}"),
|
||||
})
|
||||
return {
|
||||
"content": content,
|
||||
"tool_calls": tool_calls,
|
||||
"finish_reason": str(choice.get("finish_reason") or ""),
|
||||
}
|
||||
last_error = "LLM response contains no choices"
|
||||
if attempt >= attempt_limit:
|
||||
break
|
||||
self._sleep_before_retry(attempt)
|
||||
raise LLMResponseError(
|
||||
f"LLM failed after {attempt_limit} attempts: {last_error}"
|
||||
)
|
||||
|
||||
def complete_json(
|
||||
self,
|
||||
system: str,
|
||||
|
||||
@@ -1228,6 +1228,8 @@ class QAV1Pipeline:
|
||||
endpoint_path_max_expansions: int = 24,
|
||||
endpoint_path_max_frontier_states: int = 120,
|
||||
endpoint_path_time_budget_seconds: float = 5.0,
|
||||
history: List[Dict[str, str]] | None = None,
|
||||
capabilities: Dict[str, Any] | None = None,
|
||||
retrieval_only: bool = False,
|
||||
progress_callback: Callable[[Dict[str, Any]], None] | None = None,
|
||||
) -> Dict[str, Any]:
|
||||
@@ -2239,6 +2241,7 @@ class QAV1Pipeline:
|
||||
composer_graph_chain_text,
|
||||
agentic_review.get("final_handoff") if agentic_review else None,
|
||||
subquestion_evidence_groups,
|
||||
history=history,
|
||||
)
|
||||
else:
|
||||
answer = ""
|
||||
@@ -2338,6 +2341,10 @@ class QAV1Pipeline:
|
||||
"cluster_audit": cluster_audit[:200],
|
||||
"rerank_audit": rerank_audit[:80],
|
||||
"timings": timings,
|
||||
"capabilities_audit": {
|
||||
"requested": dict(capabilities or {}),
|
||||
"registry_enabled": _capability_registry_names(),
|
||||
},
|
||||
"flow": [
|
||||
"question",
|
||||
"provisional_answer",
|
||||
@@ -2371,6 +2378,7 @@ class QAV1Pipeline:
|
||||
graph_chain_context: str,
|
||||
agent_handoff: Dict[str, Any] | None = None,
|
||||
subquestion_evidence_groups: List[Dict[str, Any]] | None = None,
|
||||
history: List[Dict[str, str]] | None = None,
|
||||
) -> tuple[str, Dict[str, object]]:
|
||||
public_evidence = [
|
||||
item
|
||||
@@ -2430,9 +2438,10 @@ class QAV1Pipeline:
|
||||
answer = self.answer_composer.compose(
|
||||
*compose_args,
|
||||
subquestion_evidence_groups=subquestion_evidence_groups,
|
||||
history=history,
|
||||
)
|
||||
else:
|
||||
answer = self.answer_composer.compose(*compose_args)
|
||||
answer = self.answer_composer.compose(*compose_args, history=history)
|
||||
except Exception as exc:
|
||||
generation_errors.append(f"{type(exc).__name__}: {exc}")
|
||||
previous_validation = {
|
||||
@@ -3413,3 +3422,13 @@ class QAV1Pipeline:
|
||||
]
|
||||
)
|
||||
return "\n".join(lines).strip()
|
||||
|
||||
|
||||
def _capability_registry_names() -> list:
|
||||
"""Capability Gateway(S1):注册表工具名列表(registry 未安装/失败时空列表)。"""
|
||||
try:
|
||||
from .capabilities.registry import get_registry
|
||||
|
||||
return [t.name for t in get_registry().list_all()]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
"""test_capabilities.py —— manifests 解析/registry 热加载/install_manager 单测。
|
||||
|
||||
注意:QA 树编译时 test_* 自动排除(compile_qa.sh 只收 *.py 但由 Cython 白名单
|
||||
排除 test 文件——见编译清单)。测试在 venv 里直接 pytest 跑(源码树,不依赖 .so)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73.capabilities.manifests import ( # noqa: E402
|
||||
CapabilityManifest,
|
||||
parse_capabilities_text,
|
||||
parse_capabilities_file,
|
||||
)
|
||||
from rag.qa_v1_1_73.capabilities.registry import CapabilityRegistry # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.install_manager import InstallError, InstallManager # noqa: E402
|
||||
|
||||
VALID_YAML = """
|
||||
capabilities:
|
||||
- name: web_search
|
||||
provider: http
|
||||
permission: read_only
|
||||
description: 检索 PubMed
|
||||
session_default: false
|
||||
schema:
|
||||
query: {type: string, required: true}
|
||||
http:
|
||||
endpoint: https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi
|
||||
method: GET
|
||||
timeout: 8
|
||||
"""
|
||||
|
||||
|
||||
class TestManifests:
|
||||
def test_parse_valid(self):
|
||||
manifests, warnings = parse_capabilities_text(VALID_YAML)
|
||||
assert warnings == []
|
||||
assert len(manifests) == 1
|
||||
m = manifests[0]
|
||||
assert m.name == "web_search"
|
||||
assert m.provider == "http"
|
||||
spec = m.to_tool_spec()
|
||||
assert spec["type"] == "function"
|
||||
assert spec["function"]["name"] == "web_search"
|
||||
assert spec["function"]["parameters"]["required"] == ["query"]
|
||||
|
||||
def test_bad_entry_skipped(self):
|
||||
text = """
|
||||
capabilities:
|
||||
- name: good_one
|
||||
provider: http
|
||||
- name: bad_provider
|
||||
provider: sneaky
|
||||
- name: write_denied
|
||||
provider: http
|
||||
permission: write
|
||||
- provider: http
|
||||
"""
|
||||
manifests, warnings = parse_capabilities_text(text)
|
||||
assert [m.name for m in manifests] == ["good_one"]
|
||||
assert len(warnings) == 3 # sneaky 拒绝 + write 拒绝 + 缺 name 跳过
|
||||
|
||||
def test_unknown_schema_keys_dropped(self):
|
||||
text = """
|
||||
capabilities:
|
||||
- name: t
|
||||
provider: http
|
||||
schema:
|
||||
q: {type: string, required: true, evil: "x"}
|
||||
"""
|
||||
manifests, _ = parse_capabilities_text(text)
|
||||
spec = manifests[0].to_tool_spec()
|
||||
assert "evil" not in spec["function"]["parameters"]["properties"]["q"]
|
||||
|
||||
def test_oversize_entry_skipped(self):
|
||||
# schema 无字段级截断 → 巨大 schema 触发单条目 ≤8KB 上限
|
||||
huge_props = ",".join(
|
||||
f'p{i}: {{type: string, description: "{"x" * 300}"}}' for i in range(40)
|
||||
)
|
||||
text = f"""
|
||||
capabilities:
|
||||
- name: big
|
||||
provider: http
|
||||
schema: {{{huge_props}}}
|
||||
"""
|
||||
manifests, warnings = parse_capabilities_text(text)
|
||||
assert manifests == []
|
||||
assert any("过大" in w for w in warnings)
|
||||
|
||||
def test_duplicate_name_kept_once(self):
|
||||
text = """
|
||||
capabilities:
|
||||
- name: dup
|
||||
provider: http
|
||||
- name: dup
|
||||
provider: http
|
||||
"""
|
||||
manifests, warnings = parse_capabilities_text(text)
|
||||
assert len(manifests) == 1
|
||||
assert any("重名" in w for w in warnings)
|
||||
|
||||
def test_parse_missing_file(self, tmp_path):
|
||||
manifests, warnings = parse_capabilities_file(tmp_path / "nope.yaml")
|
||||
assert manifests == []
|
||||
assert warnings
|
||||
|
||||
|
||||
class TestRegistry:
|
||||
def test_load_and_diff(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
registry = CapabilityRegistry(path)
|
||||
assert registry.load() == []
|
||||
assert registry.list_all() == []
|
||||
|
||||
path.write_text(VALID_YAML, encoding="utf-8")
|
||||
warnings, diff = registry.reload()
|
||||
assert warnings == []
|
||||
assert diff["added"] == ["web_search"]
|
||||
assert len(registry.list_all()) == 1
|
||||
|
||||
# 空能力表热加载:删除
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
_, diff = registry.reload()
|
||||
assert diff["removed"] == ["web_search"]
|
||||
|
||||
def test_enabled_tools_session_switch(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text(VALID_YAML, encoding="utf-8")
|
||||
registry = CapabilityRegistry(path)
|
||||
registry.load()
|
||||
assert registry.enabled_tools({}) == [] # session_default false
|
||||
assert len(registry.enabled_tools({"web_search": True})) == 1
|
||||
assert registry.enabled_tools({"web_search": False}) == []
|
||||
|
||||
|
||||
class TestInstallManager:
|
||||
def test_register_unregister_reload(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
manager = InstallManager(path)
|
||||
registry = CapabilityRegistry(path)
|
||||
registry.load()
|
||||
|
||||
manifests, warnings = manager.register_manifest(VALID_YAML)
|
||||
assert len(manifests) == 1 and warnings == []
|
||||
assert "web_search" in path.read_text(encoding="utf-8")
|
||||
|
||||
manager.reload(registry=registry) # 热加载注入测试 registry
|
||||
assert registry.get("web_search") is not None
|
||||
|
||||
assert manager.unregister("web_search") is True
|
||||
assert manager.unregister("web_search") is False
|
||||
assert "web_search" not in path.read_text(encoding="utf-8")
|
||||
|
||||
def test_register_invalid_rejected(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
manager = InstallManager(path)
|
||||
with pytest.raises(InstallError):
|
||||
manager.register_manifest("capabilities:\n - name: x\n provider: evil")
|
||||
assert "evil" not in path.read_text(encoding="utf-8")
|
||||
|
||||
def test_register_same_name_updates(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
manager = InstallManager(path)
|
||||
manager.register_manifest(VALID_YAML)
|
||||
manager.register_manifest(VALID_YAML.replace("检索 PubMed", "检索 PubMed v2"))
|
||||
text = path.read_text(encoding="utf-8")
|
||||
assert text.count("name: web_search") == 1
|
||||
assert "v2" in text
|
||||
|
||||
def test_atomic_write_no_temp_left(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
manager = InstallManager(path)
|
||||
manager.register_manifest(VALID_YAML)
|
||||
leftovers = list(tmp_path.glob(".capabilities.*.tmp"))
|
||||
assert leftovers == []
|
||||
|
||||
def test_audit_written(self, tmp_path):
|
||||
path = tmp_path / "capabilities.yaml"
|
||||
path.write_text("capabilities: []", encoding="utf-8")
|
||||
manager = InstallManager(path)
|
||||
manager.register_manifest(VALID_YAML)
|
||||
audit = tmp_path / "capabilities_install_log.jsonl"
|
||||
assert audit.is_file()
|
||||
assert "register" in audit.read_text(encoding="utf-8")
|
||||
@@ -0,0 +1,132 @@
|
||||
"""test_tool_runner.py —— 工具调用循环单测(原生 tool_calls / JSON fallback / 轮次封顶)。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent.parent / "src"))
|
||||
|
||||
from rag.qa_v1_1_73.tool_runner import ToolRunner # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.manifests import CapabilityManifest # noqa: E402
|
||||
from rag.qa_v1_1_73.capabilities.providers import ToolExecutionError, get_dispatcher # noqa: E402
|
||||
|
||||
|
||||
def make_manifest(name: str = "echo") -> CapabilityManifest:
|
||||
return CapabilityManifest(name=name, provider="python", description="test")
|
||||
|
||||
|
||||
class FakeDispatcher:
|
||||
def __init__(self, results=None):
|
||||
self.results = results or {}
|
||||
self.calls = []
|
||||
|
||||
def execute(self, manifest, arguments):
|
||||
self.calls.append((manifest.name, arguments))
|
||||
if manifest.name in self.results and isinstance(self.results[manifest.name], Exception):
|
||||
raise self.results[manifest.name]
|
||||
return self.results.get(manifest.name, f"result:{manifest.name}")
|
||||
|
||||
|
||||
class FakeLLM:
|
||||
"""按脚本返回:[(content, tool_calls, error), ...],最后一轮自动无工具。"""
|
||||
|
||||
def __init__(self, script):
|
||||
self.script = list(script)
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, system, messages, tools):
|
||||
self.calls.append({"system": system, "messages": messages, "tools": tools})
|
||||
if self.script:
|
||||
return self.script.pop(0)
|
||||
return "final-answer", [], None
|
||||
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "echo", "description": "echo", "parameters": {"type": "object", "properties": {}}},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class TestToolRunner:
|
||||
def test_native_tool_call_flow(self):
|
||||
llm = FakeLLM(
|
||||
[
|
||||
("", [{"id": "c1", "name": "echo", "arguments": {"x": 1}}], None),
|
||||
("answered-with-tool", [], None),
|
||||
]
|
||||
)
|
||||
runner = ToolRunner(llm, dispatcher=FakeDispatcher({"echo": "E"}), max_rounds=3)
|
||||
content, audit = runner.run("sys", "user", TOOLS, {"echo": make_manifest()})
|
||||
assert audit["tool_calls"] == 1
|
||||
assert audit["used_json_fallback"] is False
|
||||
assert audit["tool_names"] == ["echo"]
|
||||
# 第一轮原生 tools 传入
|
||||
assert llm.calls[0]["tools"] == TOOLS
|
||||
|
||||
def test_json_fallback_on_4xx(self):
|
||||
llm = FakeLLM(
|
||||
[
|
||||
("", [], "HTTP 400 tools not supported"),
|
||||
("", [{"id": "c1", "name": "echo", "arguments": "{}"}], None),
|
||||
("done", [], None),
|
||||
]
|
||||
)
|
||||
runner = ToolRunner(llm, dispatcher=FakeDispatcher(), max_rounds=3)
|
||||
content, audit = runner.run("sys", "user", TOOLS, {"echo": make_manifest()})
|
||||
assert audit["used_json_fallback"] is True
|
||||
# fallback 时 tools 清空(system 追加契约)
|
||||
assert llm.calls[1]["tools"] == []
|
||||
assert "JSON" in llm.calls[1]["system"]
|
||||
|
||||
def test_rounds_capped(self):
|
||||
script = []
|
||||
for i in range(5):
|
||||
script.append(("", [{"id": f"c{i}", "name": "echo", "arguments": "{}"}], None))
|
||||
llm = FakeLLM(script)
|
||||
runner = ToolRunner(llm, dispatcher=FakeDispatcher(), max_rounds=3)
|
||||
content, audit = runner.run("sys", "user", TOOLS, {"echo": make_manifest()})
|
||||
assert audit["rounds"] == 3
|
||||
assert audit["tool_calls"] == 3
|
||||
|
||||
def test_tool_error_becomes_result(self):
|
||||
llm = FakeLLM(
|
||||
[
|
||||
("", [{"id": "c1", "name": "echo", "arguments": "{}"}], None),
|
||||
("final", [], None),
|
||||
]
|
||||
)
|
||||
runner = ToolRunner(llm, dispatcher=FakeDispatcher({"echo": ToolExecutionError("boom")}), max_rounds=3)
|
||||
content, audit = runner.run("sys", "user", TOOLS, {"echo": make_manifest()})
|
||||
assert audit["errors"] == []
|
||||
# 工具结果以文本回传,主流程不炸
|
||||
tool_message = [m for m in llm.calls[0]["messages"] if m["role"] == "tool"]
|
||||
assert tool_message and "boom" in tool_message[-1]["content"]
|
||||
|
||||
def test_unknown_tool_not_registered(self):
|
||||
llm = FakeLLM(
|
||||
[
|
||||
("", [{"id": "c1", "name": "ghost", "arguments": "{}"}], None),
|
||||
("final", [], None),
|
||||
]
|
||||
)
|
||||
runner = ToolRunner(llm, dispatcher=FakeDispatcher(), max_rounds=3)
|
||||
content, audit = runner.run("sys", "user", TOOLS, {"echo": make_manifest()})
|
||||
tool_message = [m for m in llm.calls[0]["messages"] if m["role"] == "tool"]
|
||||
assert "未注册" in tool_message[-1]["content"]
|
||||
|
||||
def test_string_arguments_parsed(self):
|
||||
llm = FakeLLM(
|
||||
[
|
||||
("", [{"id": "c1", "name": "echo", "arguments": '{"x": 42}'}], None),
|
||||
("final", [], None),
|
||||
]
|
||||
)
|
||||
dispatcher = FakeDispatcher()
|
||||
runner = ToolRunner(llm, dispatcher=dispatcher, max_rounds=3)
|
||||
runner.run("sys", "user", TOOLS, {"echo": make_manifest()})
|
||||
assert dispatcher.calls[0][1] == {"x": 42}
|
||||
@@ -0,0 +1,159 @@
|
||||
from __future__ import annotations
|
||||
|
||||
"""tool_runner.py —— LLM 工具调用循环(机制调用层统一入口)。
|
||||
|
||||
complete_with_tools(system, messages, tools, max_rounds=3):
|
||||
1. 原生 tool_calls:LLM 返回 message.tool_calls → 逐个执行 → 把结果以
|
||||
role=tool 消息追加 → 循环(max_rounds 封顶)
|
||||
2. JSON 契约 fallback:LLM 侧不支持 tools(4xx)或三轮内无 tool_calls 时,
|
||||
system 追加 JSON 契约说明再试一次;audit 记 used_json_fallback
|
||||
|
||||
结果始终是纯文本 content;工具异常转 tool result 文本回传 LLM(不炸主流程)。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .capabilities.manifests import CapabilityManifest
|
||||
from .capabilities.providers import ToolDispatcher, ToolExecutionError, get_dispatcher
|
||||
|
||||
LOGGER = logging.getLogger(__name__)
|
||||
|
||||
JSON_FALLBACK_CONTRACT = (
|
||||
"\n\nYou may use tools by replying with a JSON object:\n"
|
||||
'{"tool_calls": [{"name": "<tool_name>", "arguments": {<params>}}]}'
|
||||
"\nThe tool result will be appended to the conversation. "
|
||||
"If you do not need a tool, reply normally."
|
||||
)
|
||||
|
||||
|
||||
class ToolRunner:
|
||||
"""通用工具调用循环。
|
||||
|
||||
llm_complete 是注入的调用函数:
|
||||
(system, messages, tools, json_mode) -> (content, tool_calls, error)
|
||||
其中 tool_calls 是原生 tool_calls 列表;error 非空表示上游 4xx(触发 fallback)。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
llm_complete,
|
||||
dispatcher: Optional[ToolDispatcher] = None,
|
||||
max_rounds: int = 3,
|
||||
) -> None:
|
||||
self._complete = llm_complete
|
||||
self._dispatcher = dispatcher or get_dispatcher()
|
||||
self.max_rounds = max_rounds
|
||||
|
||||
def run(
|
||||
self,
|
||||
system: str,
|
||||
user: str,
|
||||
tools: List[Dict[str, Any]],
|
||||
manifests: Dict[str, CapabilityManifest],
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""执行工具调用循环。返回 (content, audit)。
|
||||
|
||||
audit: {"rounds": n, "tool_calls": k, "used_json_fallback": bool,
|
||||
"tool_names": [...], "errors": [...]}
|
||||
"""
|
||||
audit: Dict[str, Any] = {
|
||||
"rounds": 0,
|
||||
"tool_calls": 0,
|
||||
"used_json_fallback": False,
|
||||
"tool_names": [],
|
||||
"errors": [],
|
||||
}
|
||||
messages: List[Dict[str, Any]] = [{"role": "user", "content": user}]
|
||||
system_text = system
|
||||
for round_index in range(1, self.max_rounds + 1):
|
||||
audit["rounds"] = round_index
|
||||
content, tool_calls, error = self._complete(system_text, messages, tools)
|
||||
if error:
|
||||
# 上游 4xx:切 JSON 契约 fallback 重试一次
|
||||
if round_index == 1 and not audit["used_json_fallback"]:
|
||||
audit["used_json_fallback"] = True
|
||||
system_text = system + JSON_FALLBACK_CONTRACT
|
||||
tools = []
|
||||
continue
|
||||
audit["errors"].append(str(error)[:400])
|
||||
break
|
||||
if not tool_calls:
|
||||
break
|
||||
assistant_message: Dict[str, Any] = {"role": "assistant", "content": content or ""}
|
||||
tool_call_blocks: List[Dict[str, Any]] = []
|
||||
for call in tool_calls:
|
||||
name = str(call.get("name") or call.get("function", {}).get("name") or "")
|
||||
raw_args = call.get("arguments") or call.get("function", {}).get("arguments") or {}
|
||||
if isinstance(raw_args, str):
|
||||
try:
|
||||
arguments = json.loads(raw_args) if raw_args.strip() else {}
|
||||
except Exception:
|
||||
arguments = {"_raw": raw_args[:2000]}
|
||||
else:
|
||||
arguments = dict(raw_args or {})
|
||||
audit["tool_calls"] += 1
|
||||
audit["tool_names"].append(name)
|
||||
manifest = manifests.get(name)
|
||||
if manifest is None:
|
||||
result_text = f"工具 {name} 未注册"
|
||||
else:
|
||||
try:
|
||||
result_text = self._dispatcher.execute(manifest, arguments)
|
||||
except ToolExecutionError as exc:
|
||||
result_text = f"工具执行失败: {exc}"
|
||||
except Exception as exc:
|
||||
result_text = f"工具执行异常: {type(exc).__name__}: {exc}"
|
||||
call_id = str(call.get("id") or f"call_{round_index}_{len(tool_call_blocks)}")
|
||||
tool_call_blocks.append(
|
||||
{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {"name": name or "unknown", "arguments": json.dumps(arguments, ensure_ascii=False)},
|
||||
}
|
||||
)
|
||||
messages.append({"role": "tool", "tool_call_id": call_id, "content": result_text})
|
||||
assistant_message["tool_calls"] = tool_call_blocks
|
||||
messages.append(assistant_message)
|
||||
continue
|
||||
# 最后一轮:拿最终文本(无工具或轮次用尽后由调用方用 messages 再 complete 一次)
|
||||
return self._final_text(system_text, messages, tools), audit
|
||||
|
||||
def _final_text(self, system: str, messages: List[Dict[str, Any]], tools: List[Dict[str, Any]]) -> str:
|
||||
content, _tool_calls, error = self._complete(system, messages, tools if not isinstance(tools, list) else [])
|
||||
if error:
|
||||
return content or ""
|
||||
return content or ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# llm_client 适配(QA 树内使用)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_llm_complete(llm_client: Any):
|
||||
"""把 QA 树 LLMClient 包装成 ToolRunner 需要的调用函数。
|
||||
|
||||
返回 (content, tool_calls, error):
|
||||
- 原生路径:临时扩展 client.complete 的 tools 透传(S1 llm_client 已支持)
|
||||
- 4xx → error 非空 → ToolRunner 自动切 JSON fallback
|
||||
"""
|
||||
|
||||
def _complete(
|
||||
system: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
tools: List[Dict[str, Any]],
|
||||
) -> Tuple[str, List[Dict[str, Any]], Optional[str]]:
|
||||
try:
|
||||
raw = llm_client.complete_raw_messages(
|
||||
system=system,
|
||||
messages=messages,
|
||||
tools=tools or None,
|
||||
)
|
||||
except Exception as exc:
|
||||
return "", [], f"{type(exc).__name__}: {exc}"
|
||||
content = raw.get("content") or ""
|
||||
tool_calls = raw.get("tool_calls") or []
|
||||
return content, tool_calls, None
|
||||
|
||||
return _complete
|
||||
@@ -54,6 +54,11 @@ class Settings(BaseSettings):
|
||||
db_path: str = "data/plantai.db" # 相对 sdk/(BFF cwd),bff_restart.sh 保证
|
||||
health_cache_ttl: float = 5.0
|
||||
|
||||
# ---- 多轮记忆(S2)----
|
||||
bff_history_rounds: int = 3 # 携带最近 N 轮问答
|
||||
bff_history_max_question_chars: int = 4000 # 单条 user 文本截断
|
||||
bff_history_max_answer_chars: int = 12000 # 单条 assistant 答案截断
|
||||
|
||||
# ---- 生产形态:BFF 静态托管 ui/dist(ARCHITECTURE.md §3,M6 挂载)----
|
||||
frontend_dist_dir: str = "ui/dist" # 生产形态:静态托管 ui/dist
|
||||
|
||||
|
||||
+45
-1
@@ -20,6 +20,7 @@ from fastapi import APIRouter, Request
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from ..services import sse_relay
|
||||
from ..services.history import assemble_history
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -105,7 +106,7 @@ async def _proxy(
|
||||
normalizer = normalizer or sse_relay.get_qa_normalizer()
|
||||
url = f"{sse_relay.qa_base_url()}{_upstream_path(endpoint)}"
|
||||
timeout = _read_timeout(read_timeout)
|
||||
json_body = body if body is not None else {}
|
||||
json_body = await _inject_history(body)
|
||||
|
||||
for attempt in (1, 2):
|
||||
try:
|
||||
@@ -208,6 +209,7 @@ async def qa_chat_stream(request: Request):
|
||||
except Exception:
|
||||
return _envelope("VALIDATION", 422)
|
||||
evidence_first = bool(body.pop("evidence_first", False))
|
||||
body = await _inject_history(body)
|
||||
return StreamingResponse(
|
||||
sse_relay.relay_stream(body, evidence_first=evidence_first),
|
||||
media_type="text/event-stream",
|
||||
@@ -223,3 +225,45 @@ async def qa_search(request: Request):
|
||||
return await _proxy(
|
||||
"search", body, read_timeout=SEARCH_READ_TIMEOUT, retry_invalid_citation=False, mode="search"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 多轮记忆(S2):剥离 session_id → 组装 history 注入上游
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def _inject_history(body: Optional[dict]) -> dict:
|
||||
"""剥离 BFF 契约字段 session_id,读会话历史注入 history。
|
||||
|
||||
- 不传 session_id / 会话不存在 / 无历史消息 → 行为与现役一致(history 缺省)
|
||||
- history 只在 memory 开启时注入(会话 meta.capabilities.memory,缺省 True,
|
||||
与 S4 设置面板的默认一致;meta 缺失或非 dict 时按默认开启)
|
||||
"""
|
||||
from ..db import db as bff_db
|
||||
from ..db import loads_json
|
||||
|
||||
body = dict(body or {})
|
||||
session_id = str(body.pop("session_id", None) or "").strip()
|
||||
if not session_id:
|
||||
return body
|
||||
if body.get("history"):
|
||||
return body # 显式传入的历史优先(上游直传场景)
|
||||
try:
|
||||
row = await bff_db.fetchone(
|
||||
"SELECT meta FROM sessions WHERE id = ?", (session_id,)
|
||||
)
|
||||
if row is None:
|
||||
return body
|
||||
meta = loads_json(row["meta"], default={})
|
||||
if not isinstance(meta, dict):
|
||||
meta = {}
|
||||
caps = meta.get("capabilities")
|
||||
if not isinstance(caps, dict):
|
||||
caps = {}
|
||||
if bool(caps.get("memory", True)) is False:
|
||||
return body
|
||||
history = await assemble_history(bff_db, session_id)
|
||||
if history:
|
||||
body["history"] = history
|
||||
except Exception as exc:
|
||||
logger.warning("多轮记忆注入失败(降级为无历史): %s", exc)
|
||||
return body
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
"""services/history.py —— 多轮记忆组装(S2,SDK 层)。
|
||||
|
||||
assemble_history(db, session_id, rounds=3, budget):
|
||||
读 messages.text(用户问题 + 助手答案全文),组装 (user, assistant) 轮次对,
|
||||
只保留最近 N 轮,16k 字符硬预算(超预算丢最旧轮,永远保留最近一轮完整 user);
|
||||
status=streaming 排除;查询失败 → 空 history(记忆是增强,不阻塞问答)。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, List
|
||||
|
||||
from ..config import get_settings
|
||||
|
||||
|
||||
async def assemble_history(
|
||||
db,
|
||||
session_id: str,
|
||||
rounds: int | None = None,
|
||||
max_question_chars: int | None = None,
|
||||
max_answer_chars: int | None = None,
|
||||
) -> List[dict]:
|
||||
"""组装 history:List[{role, content}],旧→新,不含当前这轮提问。
|
||||
|
||||
预算硬约束:单条 question ≤ max_question_chars、answer ≤ max_answer_chars,
|
||||
总字符约 ≤16k(超预算丢最旧轮,最近轮永远保留)。
|
||||
"""
|
||||
settings = get_settings()
|
||||
rounds = rounds if rounds is not None else settings.bff_history_rounds
|
||||
max_question_chars = (
|
||||
max_question_chars
|
||||
if max_question_chars is not None
|
||||
else settings.bff_history_max_question_chars
|
||||
)
|
||||
max_answer_chars = (
|
||||
max_answer_chars
|
||||
if max_answer_chars is not None
|
||||
else settings.bff_history_max_answer_chars
|
||||
)
|
||||
if not session_id:
|
||||
return []
|
||||
try:
|
||||
rows = await db.fetchall(
|
||||
"SELECT role, status, text FROM messages "
|
||||
"WHERE session_id = ? AND role IN ('user','assistant') "
|
||||
"AND status != 'streaming' "
|
||||
"ORDER BY created_at ASC, rowid ASC",
|
||||
(session_id,),
|
||||
)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
# 1. 倒序组装 (user, assistant) 轮次对(旧→新输出)
|
||||
pairs: List[tuple] = [] # 新→旧
|
||||
answer_parts: List[str] = []
|
||||
for row in reversed(rows):
|
||||
role = row["role"]
|
||||
text = (row["text"] or "").strip()
|
||||
if not text:
|
||||
continue
|
||||
if role == "assistant":
|
||||
answer_parts.append(text)
|
||||
else:
|
||||
answer = "\n".join(reversed(answer_parts)) if answer_parts else ""
|
||||
answer_parts = []
|
||||
pairs.append((text, answer))
|
||||
pairs.reverse() # 旧→新
|
||||
|
||||
# 2. 只保留最近 N 轮
|
||||
if len(pairs) > rounds:
|
||||
pairs = pairs[-rounds:]
|
||||
|
||||
# 3. 组装 + 16k 字符预算(超预算丢最旧轮,最近轮永远保留)
|
||||
budget = 16 * 1024
|
||||
history: List[dict] = []
|
||||
for user_text, answer_text in pairs:
|
||||
u = user_text[:max_question_chars]
|
||||
a = answer_text[:max_answer_chars]
|
||||
history.append({"role": "user", "content": u})
|
||||
if a:
|
||||
history.append({"role": "assistant", "content": a})
|
||||
used = sum(len(m["content"]) for m in history)
|
||||
while used > budget and history:
|
||||
# 丢最旧一轮(user + 其后的 assistant),保最近轮
|
||||
dropped = history.pop(0)
|
||||
if dropped["role"] == "user" and history and history[0]["role"] == "assistant":
|
||||
history.pop(0)
|
||||
used = sum(len(m["content"]) for m in history)
|
||||
if not history:
|
||||
break
|
||||
return history
|
||||
@@ -0,0 +1,105 @@
|
||||
"""test_history.py —— 多轮记忆组装单测(S2,SDK 层)。
|
||||
|
||||
覆盖:空 messages→[] / N 轮截断 / 预算截断保留最近一轮 user / streaming 排除 /
|
||||
查询异常→空 history。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from app.services.history import assemble_history # noqa: E402
|
||||
|
||||
|
||||
class FakeRow(dict):
|
||||
"""aiosqlite.Row 的最小替身(支持键访问)。"""
|
||||
|
||||
def __getattr__(self, key):
|
||||
return self[key]
|
||||
|
||||
|
||||
class FakeDB:
|
||||
def __init__(self, rows=None, error=None):
|
||||
self.rows = rows or []
|
||||
self.error = error
|
||||
|
||||
async def fetchall(self, sql, params):
|
||||
if self.error:
|
||||
raise self.error
|
||||
# 模拟 SQL 的 status != 'streaming' 过滤(与真实查询一致)
|
||||
return [r for r in self.rows if r.get("status") != "streaming"]
|
||||
|
||||
|
||||
def row(role, text, status="done"):
|
||||
return FakeRow(role=role, text=text, status=status)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_messages():
|
||||
db = FakeDB([])
|
||||
assert await assemble_history(db, "s1") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_db_error_returns_empty():
|
||||
db = FakeDB(error=RuntimeError("boom"))
|
||||
assert await assemble_history(db, "s1") == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_round():
|
||||
db = FakeDB([row("user", "水稻SPL是什么?"), row("assistant", "答案是……")])
|
||||
history = await assemble_history(db, "s1", rounds=3)
|
||||
assert [m["role"] for m in history] == ["user", "assistant"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rounds_truncation():
|
||||
rows = []
|
||||
for i in range(5):
|
||||
rows.append(row("user", f"问题{i}"))
|
||||
rows.append(row("assistant", f"答案{i}"))
|
||||
db = FakeDB(rows)
|
||||
history = await assemble_history(db, "s1", rounds=3)
|
||||
# 3 轮 = 3 user + 3 assistant(最近 3 轮)
|
||||
assert [m["content"] for m in history if m["role"] == "user"] == ["问题2", "问题3", "问题4"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_truncation_keeps_last_user():
|
||||
rows = [
|
||||
row("user", "Q1"),
|
||||
row("assistant", "A" * 20000), # 超大答案被截断
|
||||
row("user", "Q2"),
|
||||
]
|
||||
db = FakeDB(rows)
|
||||
history = await assemble_history(db, "s1", rounds=3, max_answer_chars=12000)
|
||||
# 最近一轮 user(Q2)必须保留
|
||||
assert history[-1]["role"] == "user"
|
||||
assert history[-1]["content"] == "Q2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_excluded():
|
||||
rows = [
|
||||
row("user", "Q1"),
|
||||
row("assistant", "流式中…", status="streaming"),
|
||||
row("user", "Q2"),
|
||||
row("assistant", "完成答案"),
|
||||
]
|
||||
db = FakeDB(rows)
|
||||
history = await assemble_history(db, "s1", rounds=3)
|
||||
contents = [m["content"] for m in history]
|
||||
assert "流式中…" not in contents
|
||||
assert "完成答案" in contents
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_session_id():
|
||||
db = FakeDB([])
|
||||
assert await assemble_history(db, "") == []
|
||||
assert await assemble_history(db, None) == []
|
||||
@@ -0,0 +1,153 @@
|
||||
#!/usr/bin/env python3
|
||||
"""tooling/cap/cap_selfcheck.py —— 机制端到端自检(假 MCP server + 热加载证明)。
|
||||
|
||||
流程(全程同一 QA 进程,不重启):
|
||||
1. 加载现有 capabilities.yaml → registry 工具数
|
||||
2. 临时注册假 MCP echo 工具(manifest 落盘)→ reload → 工具表出现 echo
|
||||
3. 实际调用 echo(stdio 子进程)→ 返回正确
|
||||
4. 停用(移除 manifest)→ reload → 工具表消失、MCP 连接关闭
|
||||
5. 恢复原 capabilities.yaml
|
||||
|
||||
用法:
|
||||
QA_ENV_FILE=<prod.env> python3 cap_selfcheck.py <qa-config-dir>
|
||||
|
||||
<qa-config-dir> = deploy 树根(含 configs/capabilities.yaml),
|
||||
即 $GENE_FAMILY_ROOT 的值。测试中的 yaml 落盘用其真实路径,因此会短暂修改
|
||||
配置文件;本脚本保证恢复。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
TOOL_DIR = Path(__file__).parent
|
||||
FAKE_MCP = TOOL_DIR / "fake_mcp_server.py"
|
||||
|
||||
FAKE_MANIFEST = """
|
||||
capabilities:
|
||||
- name: fake_echo
|
||||
provider: mcp
|
||||
permission: read_only
|
||||
description: 自检用 echo 工具
|
||||
session_default: true
|
||||
schema:
|
||||
text: {type: string, required: true, description: 要回显的文本}
|
||||
mcp:
|
||||
command: {python}
|
||||
args: ["{fake_mcp}"]
|
||||
tool: echo
|
||||
"""
|
||||
|
||||
|
||||
def _python() -> str:
|
||||
return os.environ.get(
|
||||
"QA_PYTHON", "/home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python"
|
||||
)
|
||||
|
||||
|
||||
def _qa_import(root: Path):
|
||||
"""在目标 QA 树环境里跑一段 python(用 deploy 树的 venv + PYTHONPATH)。"""
|
||||
env = dict(os.environ)
|
||||
env["PYTHONPATH"] = f"{root}:{root}/src" + (
|
||||
f":{env['PYTHONPATH']}" if env.get("PYTHONPATH") else ""
|
||||
)
|
||||
return env
|
||||
|
||||
|
||||
def run_step(root: Path, code: str) -> str:
|
||||
env = _qa_import(root)
|
||||
proc = subprocess.run(
|
||||
[_python(), "-c", code],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
env=env,
|
||||
cwd=str(root),
|
||||
)
|
||||
if proc.returncode != 0:
|
||||
raise RuntimeError(f"step failed: {proc.stderr[-2000:]}")
|
||||
return proc.stdout.strip()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
root = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else None
|
||||
if root is None or not (root / "configs").is_dir():
|
||||
print(__doc__)
|
||||
return 2
|
||||
caps_yaml = root / "configs" / "capabilities.yaml"
|
||||
original = caps_yaml.read_text(encoding="utf-8") if caps_yaml.is_file() else "capabilities: []\n"
|
||||
|
||||
manifest = FAKE_MANIFEST.format(python=_python(), fake_mcp=FAKE_MCP)
|
||||
results = {}
|
||||
try:
|
||||
# 1. 初始工具表
|
||||
out = run_step(root, (
|
||||
"from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
|
||||
"r=get_registry();"
|
||||
"print(json.dumps({'tools':[t.name for t in r.list_all()],"
|
||||
"'warnings':r.load_warnings}, ensure_ascii=False))"
|
||||
))
|
||||
initial = json.loads(out)
|
||||
results["initial_tools"] = initial["tools"]
|
||||
assert "fake_echo" not in initial["tools"], "预置状态异常:fake_echo 已在工具表"
|
||||
|
||||
# 2. 注册 + reload(热加载)
|
||||
caps_yaml.write_text(manifest, encoding="utf-8")
|
||||
out = run_step(root, (
|
||||
"from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
|
||||
"w,d=get_registry().reload();"
|
||||
"print(json.dumps({'warnings':w,'diff':d}, ensure_ascii=False))"
|
||||
))
|
||||
reloaded = json.loads(out)
|
||||
assert "fake_echo" in reloaded["diff"]["added"], "热加载未新增 fake_echo"
|
||||
results["hot_reload_added"] = reloaded["diff"]["added"]
|
||||
|
||||
# 3. 实际调用(MCP stdio 子进程)
|
||||
out = run_step(root, (
|
||||
"from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
|
||||
"from src.rag.qa_v1_1_73.capabilities.providers import get_dispatcher;"
|
||||
"m=get_registry().get('fake_echo').manifest;"
|
||||
"r=get_dispatcher().execute(m, {'text':'你好机制'});"
|
||||
"print(r)"
|
||||
))
|
||||
assert "你好机制" in out, f"echo 返回异常: {out}"
|
||||
results["mcp_call_ok"] = True
|
||||
|
||||
# 4. 停用 + reload(热卸载)
|
||||
caps_yaml.write_text("capabilities: []\n", encoding="utf-8")
|
||||
out = run_step(root, (
|
||||
"from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
|
||||
"from src.rag.qa_v1_1_73.capabilities.providers import get_dispatcher;"
|
||||
"w,d=get_registry().reload();"
|
||||
"get_dispatcher().shutdown_connections(d['removed']);"
|
||||
"print(json.dumps({'warnings':w,'diff':d}, ensure_ascii=False))"
|
||||
))
|
||||
unloaded = json.loads(out)
|
||||
assert "fake_echo" in unloaded["diff"]["removed"], "热卸载未移除 fake_echo"
|
||||
results["hot_reload_removed"] = unloaded["diff"]["removed"]
|
||||
|
||||
# 5. 恢复原文件
|
||||
caps_yaml.write_text(original, encoding="utf-8")
|
||||
run_step(root, (
|
||||
"from src.rag.qa_v1_1_73.capabilities.registry import get_registry;"
|
||||
"get_registry().reload();print('restored')"
|
||||
))
|
||||
results["restored"] = True
|
||||
|
||||
print("SELFCHECK PASS:", json.dumps(results, ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
except Exception as exc:
|
||||
print(f"SELFCHECK FAIL: {exc}")
|
||||
return 1
|
||||
finally:
|
||||
caps_yaml.write_text(original, encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,131 @@
|
||||
#!/usr/bin/env python3
|
||||
"""fake_mcp_server.py —— 自检用假 MCP server(stdio,一个 echo 工具)。
|
||||
|
||||
用官方 mcp SDK 2.x 的 lowlevel Server + add_request_handler 实现;
|
||||
若 mcp SDK 未安装则退化为纯 stdio JSON-RPC 手工实现(保底)。
|
||||
S1 端到端自检用它验证「热加载 + 工具调用」全链。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
|
||||
|
||||
def _serve_manual() -> int:
|
||||
"""无 mcp SDK 保底:JSON-RPC over stdio(initialize/tools-list/tools-call)。"""
|
||||
for line in sys.stdin:
|
||||
try:
|
||||
message = json.loads(line)
|
||||
except Exception:
|
||||
continue
|
||||
method = message.get("method")
|
||||
msg_id = message.get("id")
|
||||
if method == "initialize":
|
||||
reply = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": msg_id,
|
||||
"result": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "fake-echo", "version": "0.1.0"},
|
||||
},
|
||||
}
|
||||
elif method == "notifications/initialized":
|
||||
continue
|
||||
elif method == "tools/list":
|
||||
reply = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": msg_id,
|
||||
"result": {
|
||||
"tools": [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "回显文本",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {"type": "string", "description": "要回显的文本"}
|
||||
},
|
||||
"required": ["text"],
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
elif method == "tools/call":
|
||||
params = message.get("params") or {}
|
||||
arguments = params.get("arguments") or {}
|
||||
text = str(arguments.get("text", ""))
|
||||
reply = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": msg_id,
|
||||
"result": {
|
||||
"content": [{"type": "text", "text": f"echo: {text}"}],
|
||||
"isError": False,
|
||||
},
|
||||
}
|
||||
else:
|
||||
reply = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": msg_id,
|
||||
"error": {"code": -32601, "message": f"unknown method {method}"},
|
||||
}
|
||||
sys.stdout.write(json.dumps(reply, ensure_ascii=False) + "\n")
|
||||
sys.stdout.flush()
|
||||
return 0
|
||||
|
||||
|
||||
def _serve_mcp() -> int:
|
||||
from mcp.server.lowlevel import Server
|
||||
from mcp.server.stdio import stdio_server
|
||||
import mcp.types as types
|
||||
import anyio
|
||||
|
||||
server = Server("fake-echo")
|
||||
|
||||
async def handle_list_tools(ctx, params: types.PaginatedRequestParams):
|
||||
return types.ListToolsResult(
|
||||
tools=[
|
||||
types.Tool(
|
||||
name="echo",
|
||||
description="回显文本",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"text": {"type": "string", "description": "要回显的文本"}
|
||||
},
|
||||
"required": ["text"],
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
async def handle_call_tool(ctx, params: types.CallToolRequestParams):
|
||||
arguments = params.arguments or {}
|
||||
text = str(arguments.get("text", ""))
|
||||
return types.CallToolResult(
|
||||
content=[types.TextContent(type="text", text=f"echo: {text}")],
|
||||
is_error=False,
|
||||
)
|
||||
|
||||
server.add_request_handler("tools/list", types.PaginatedRequestParams, handle_list_tools)
|
||||
server.add_request_handler("tools/call", types.CallToolRequestParams, handle_call_tool)
|
||||
|
||||
async def run() -> None:
|
||||
async with stdio_server() as (read_stream, write_stream):
|
||||
await server.run(read_stream, write_stream, server.create_initialization_options())
|
||||
|
||||
anyio.run(run)
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
try:
|
||||
import mcp # noqa: F401
|
||||
except ImportError:
|
||||
return _serve_manual()
|
||||
return _serve_mcp()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -38,9 +38,9 @@ mkdir -p \
|
||||
|
||||
echo "[QA 1/6] Copying the audited runtime surface"
|
||||
install -m 0750 "$SOURCE_ROOT/$ENTRY_SCRIPT" "$OUTPUT_DIR/$ENTRY_SCRIPT"
|
||||
rsync -a --include='*.py' --exclude='*' \
|
||||
rsync -a --include='*/' --include='*.py' --exclude='*' \
|
||||
"$SOURCE_ROOT/src/rag/qa_v1_1_73/" "$OUTPUT_DIR/src/rag/qa_v1_1_73/"
|
||||
rsync -a --include='*.py' --exclude='*' \
|
||||
rsync -a --include='*/' --include='*.py' --exclude='*' \
|
||||
"$SOURCE_ROOT/src/rag/retrieval/" "$OUTPUT_DIR/src/rag/retrieval/"
|
||||
install -m 0640 \
|
||||
"$SOURCE_ROOT/src/rag/dual_kb_taxonomy_rag/entity_alias.py" \
|
||||
@@ -52,6 +52,7 @@ install -m 0640 \
|
||||
for relative in \
|
||||
configs/rag/entity_alias.yaml \
|
||||
configs/rag/relation_taxonomy.yaml \
|
||||
configs/capabilities.yaml \
|
||||
pyproject.toml \
|
||||
requirements.txt \
|
||||
requirements-rag.txt
|
||||
|
||||
@@ -9,9 +9,11 @@ anyio==4.14.2
|
||||
attrs==26.1.0
|
||||
biopython==1.87
|
||||
certifi==2026.7.22
|
||||
cffi==2.1.1
|
||||
charset-normalizer==3.5.1
|
||||
click==8.4.2
|
||||
coverage==7.15.4
|
||||
cryptography==50.0.1
|
||||
Cython==3.2.8
|
||||
distro==1.9.0
|
||||
et_xmlfile==2.0.0
|
||||
@@ -41,6 +43,8 @@ jiter==0.16.0
|
||||
joblib==1.5.3
|
||||
jsonpatch==1.33
|
||||
jsonpointer==3.1.1
|
||||
jsonschema==4.26.0
|
||||
jsonschema-specifications==2025.9.1
|
||||
langchain==1.3.2
|
||||
langchain-classic==1.0.7
|
||||
langchain-community==0.4.2
|
||||
@@ -57,6 +61,8 @@ loguru==0.7.3
|
||||
lxml==6.1.2
|
||||
markdown-it-py==4.2.0
|
||||
MarkupSafe==3.0.3
|
||||
mcp==2.1.1
|
||||
mcp-types==2.1.1
|
||||
mdurl==0.1.2
|
||||
mpmath==1.3.0
|
||||
multidict==6.7.1
|
||||
@@ -66,6 +72,7 @@ networkx==3.6.1
|
||||
numpy==2.4.6
|
||||
openai==3.3.1
|
||||
openpyxl==3.1.5
|
||||
opentelemetry-api==1.44.0
|
||||
orjson==3.11.9
|
||||
ormsgpack==1.12.2
|
||||
packaging==26.3
|
||||
@@ -74,10 +81,12 @@ pluggy==1.6.0
|
||||
portalocker==3.2.0
|
||||
propcache==0.5.2
|
||||
protobuf==7.35.0
|
||||
pycparser==3.0
|
||||
pydantic==2.13.4
|
||||
pydantic-settings==2.15.0
|
||||
pydantic_core==2.46.4
|
||||
Pygments==2.21.0
|
||||
PyJWT==2.13.0
|
||||
pymupdf==1.28.2
|
||||
pytest==9.1.1
|
||||
pytest-asyncio==1.4.0
|
||||
@@ -91,11 +100,13 @@ PyYAML==6.0.3
|
||||
qdrant-client==1.18.0
|
||||
rank-bm25==0.2.2
|
||||
redis==8.0.0
|
||||
referencing==0.37.0
|
||||
regex==2026.5.9
|
||||
requests==2.34.2
|
||||
requests-toolbelt==1.0.0
|
||||
respx==0.23.1
|
||||
rich==15.0.0
|
||||
rpds-py==2026.6.3
|
||||
safetensors==0.8.0
|
||||
scikit-learn==1.9.0
|
||||
scipy==1.17.1
|
||||
@@ -105,6 +116,7 @@ shellingham==1.5.4
|
||||
six==1.17.0
|
||||
sniffio==1.3.1
|
||||
SQLAlchemy==2.0.50
|
||||
sse-starlette==3.4.8
|
||||
starlette==1.6.0
|
||||
sympy==1.14.0
|
||||
tenacity==9.1.4
|
||||
|
||||
Reference in New Issue
Block a user