vent_agent_2.0/tools/registry.py

"""工具的**唯一事实源**。

全仓只有这里能声明「有哪些工具」。其他地方（agent 组装、系统提示词、
工具名中文映射、SSE 事件渲染）都从这里派生，避免出现「提示词里说有、
实际没注册」这类漂移。

## 来源标注契约

工具返回「带来源」的结果时，统一用 :func:`sourced` 构造：

    return sourced(data=<业务数据>, sources=[source("标题", url, 摘要, 0.87)])

SSE 层（api/source_tracker.py）会从 ToolMessage 里抽取 ``sources`` 字段并
聚合去重，最终发 ``citations`` 事件。这是 2.0 相对 1.5 的关键改动：**来源
由工具侧结构化产出，而不是靠提示词要求模型自觉在正文里写引用**——模型可以
漏写、可以编造，工具不会有这两种失败模式。

## 扩展方式

新增业务工具：在 ``tools/`` 下新建模块，用 :func:`register` 装饰，
然后 import 该模块（见本文件底部的 ``_load_builtin_modules``）。
"""

from __future__ import annotations

import inspect
from dataclasses import dataclass, field
from typing import Any, Callable, Iterable

from obs import get_logger

logger = get_logger(__name__)

# 工具返回值里携带来源的字段名。改这里要同步改 api/source_tracker.py 的解析。
SOURCES_FIELD = "sources"
DATA_FIELD = "data"


@dataclass(frozen=True)
class ToolSpec:
    """一个工具的元数据。"""

    name: str
    func: Callable[..., Any]
    group: str
    label: str
    summary: str = ""

    @property
    def description(self) -> str:
        """给模型看的描述：优先用函数 docstring 首行。"""
        if self.summary:
            return self.summary
        doc = inspect.getdoc(self.func) or ""
        return doc.strip().split("\n", 1)[0]


@dataclass
class _Registry:
    """注册表本体。用类而不是全局 dict，方便测试里隔离。"""

    tools: dict[str, ToolSpec] = field(default_factory=dict)
    groups: dict[str, list[str]] = field(default_factory=dict)

    def add(self, spec: ToolSpec) -> None:
        if spec.name in self.tools:
            existing = self.tools[spec.name]
            raise ValueError(
                f"工具名冲突：{spec.name!r} 已由 {existing.func.__module__} 注册，"
                f"又来了 {spec.func.__module__}。工具名必须全局唯一。"
            )
        self.tools[spec.name] = spec
        self.groups.setdefault(spec.group, []).append(spec.name)


_registry = _Registry()


# ============================================================
# 注册
# ============================================================

def register(
    group: str,
    label: str = "",
    name: str | None = None,
    summary: str = "",
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """把函数注册为工具。

    ``group`` 是给 UI 分组展示用的中文名；``label`` 是工具的中文显示名
    （SSE 的 tool_call 事件会带出去）；``name`` 不填则用函数名。
    """

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        tool_name = name or func.__name__
        _registry.add(
            ToolSpec(
                name=tool_name,
                func=func,
                group=group,
                label=label or tool_name,
                summary=summary,
            )
        )
        return func

    return decorator


def register_external(
    func: Callable[..., Any],
    group: str,
    label: str = "",
    name: str | None = None,
    summary: str = "",
) -> Callable[..., Any]:
    """以函数式（非装饰器）方式注册，供不便加装饰器的场景使用。"""
    tool_name = name or getattr(func, "__name__", "unnamed_tool")
    _registry.add(
        ToolSpec(
            name=tool_name,
            func=func,
            group=group,
            label=label or tool_name,
            summary=summary,
        )
    )
    return func


# ============================================================
# 查询
# ============================================================

def get_all_tool_functions() -> list[Callable[..., Any]]:
    """全部工具函数，交给 ``create_deep_agent(tools=...)``。"""
    return [spec.func for spec in _registry.tools.values()]


def as_tool(func: Callable[..., Any]) -> Any:
    """把裸函数包装成 ``BaseTool``。

    为什么两种形态都要有：主 agent 的 ``tools=`` 接受裸可调用对象（框架
    内部会包），但**子智能体的 ``tools`` 会被直接塞进 ``ToolNode``**，
    那里要求对象带 ``.name``——传裸函数会炸在 ``tool_.name`` 上，而且报错
    发生在图编译期、离真正的原因很远。所以对外统一走 ``BaseTool``。
    """
    from langchain_core.tools import StructuredTool

    if hasattr(func, "name"):  # 已经是 BaseTool
        return func
    return StructuredTool.from_function(func)


def get_all_tools() -> list[Any]:
    """全部工具，以 ``BaseTool`` 形态返回。主 agent 用这个。"""
    return [as_tool(spec.func) for spec in _registry.tools.values()]


def get_tools_by_name(names: Iterable[str]) -> list[Any]:
    """按名取工具（``BaseTool`` 形态）。子智能体用这个。

    名字未注册时记警告并跳过——子智能体少一个工具是能力降级，不该让整个
    agent 编译失败。
    """
    tools: list[Any] = []
    for name in names:
        spec = _registry.tools.get(name)
        if spec is None:
            logger.warning("子智能体引用了未注册的工具 %r，已跳过", name)
            continue
        tools.append(as_tool(spec.func))
    return tools


def get_tool_names() -> list[str]:
    """全部工具名。"""
    return list(_registry.tools.keys())


def get_spec(name: str) -> ToolSpec | None:
    """按名取工具元数据。"""
    return _registry.tools.get(name)


def get_group_map() -> dict[str, list[str]]:
    """分组 → 工具名列表。"""
    return {group: list(names) for group, names in _registry.groups.items()}


def get_label(name: str) -> str:
    """工具的中文显示名；未注册时回落工具名本身。"""
    spec = _registry.tools.get(name)
    return spec.label if spec else name


def describe() -> dict[str, Any]:
    """整个注册表的可序列化描述，供 ``GET /api/model`` 或调试接口使用。"""
    return {
        "count": len(_registry.tools),
        "groups": {
            group: [
                {
                    "name": name,
                    "label": _registry.tools[name].label,
                    "description": _registry.tools[name].description,
                }
                for name in names
                if name in _registry.tools
            ]
            for group, names in _registry.groups.items()
        },
    }


# ============================================================
# 来源标注
# ============================================================

def source(
    title: str,
    url: str = "",
    snippet: str = "",
    score: float | None = None,
    **extra: Any,
) -> dict[str, Any]:
    """构造一条来源记录。"""
    record: dict[str, Any] = {"title": title, "url": url, "snippet": snippet}
    if score is not None:
        record["score"] = score
    record.update(extra)
    return record


def sourced(
    data: Any,
    sources: Iterable[dict[str, Any]] | None = None,
    **extra: Any,
) -> dict[str, Any]:
    """构造「带来源」的工具返回值。

    用这个而不是自己拼 dict，是为了让来源标注契约只有一个出口——
    将来改字段名只改这里，而不用满仓搜 ``"sources"``。
    """
    payload: dict[str, Any] = {DATA_FIELD: data, SOURCES_FIELD: list(sources or [])}
    payload.update(extra)
    return payload


# ============================================================
# 内置工具模块装载
# ============================================================

_BUILTIN_MODULES = ("tools.demo_tools",)


def load_builtin_tools() -> int:
    """import 内置工具模块以触发注册。返回注册到的工具数。

    这里要处理一个 Python 导入语义的坑：模块的注册是靠**模块体里的装饰器**
    完成的，而 ``import_module`` 对已导入的模块是空操作——模块一旦在
    ``sys.modules`` 里，再 import 不会重放装饰器。所以「注册表被清空后重新
    装载」这个场景（测试隔离、或将来做热重载）必须显式 ``reload``，
    否则注册表会一直是空的，表现为「工具莫名其妙全没了」。
    """
    import importlib
    import sys

    for module_name in _BUILTIN_MODULES:
        if module_name not in sys.modules:
            importlib.import_module(module_name)
        elif not _registry.tools:
            # 模块已在缓存里但注册表是空的 → 重放模块体以重新注册。
            # 生产路径不会走到这里（首次 import 就已填充注册表），
            # 只在测试 reset 之后触发。
            importlib.reload(sys.modules[module_name])
    return len(_registry.tools)


def reset_for_tests() -> None:
    """仅供测试：清空注册表，使下一次 ``load_builtin_tools`` 重新注册。"""
    _registry.tools.clear()
    _registry.groups.clear()