vent_agent_2.0/tools/registry.py

"""工具元数据、运行目录与框架工具转换。

业务扩展通过 setup(registry) 统一声明，校验后发布到这里。
Agent 组装、工具名中文映射和 SSE 事件渲染从同一目录读取。

## 来源标注契约

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

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

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

## 扩展方式

新增业务工具：使用 define_tool 描述元数据，并在 extensions/<业务>/extension.py
的 setup(registry) 中声明工具与子 Agent。register 保留为兼容入口。
"""

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__)

# 工具返回值里携带来源的字段名。改这里要同步改 runtime/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 = ""
    requires_approval: bool = False

    @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 tool_spec(func: Callable[..., Any], *, group: str | None = None, name: str | None = None,
              label: str | None = None, summary: str | None = None,
              requires_approval: bool | None = None) -> ToolSpec:
    """读取不带注册副作用的工具元数据；显式参数优先。"""
    meta = getattr(func, "__vent_tool__", {})
    selected = name if name is not None else meta.get("name") or getattr(func, "name", None) or getattr(func, "__name__", "")
    approval = requires_approval if requires_approval is not None else meta.get("requires_approval", False)
    if not isinstance(approval, bool):
        raise ValueError(f"工具 {selected} 的 requires_approval 必须是布尔值")
    return ToolSpec(name=selected, func=func,
                    group=group if group is not None else meta.get("group", "业务工具"),
                    label=label if label is not None else meta.get("label") or selected,
                    summary=summary if summary is not None else meta.get("summary", ""),
                    requires_approval=approval)


def define_tool(*, group: str = "业务工具", label: str = "", name: str | None = None,
                summary: str = "", requires_approval: bool = False):
    """声明工具元数据，不在 import 时注册；函数仍可以直接调用和独立测试。"""
    def decorator(func):
        func.__vent_tool__ = dict(group=group, label=label, name=name, summary=summary,
                                  requires_approval=requires_approval)
        return func
    return decorator


def spec_as_tool(spec: ToolSpec) -> Any:
    """保留注册别名、完整 docstring 和异步执行方式，转换成框架工具。"""
    from langchain_core.tools import BaseTool, StructuredTool
    if isinstance(spec.func, BaseTool):
        return spec.func.model_copy(update={"name": spec.name})
    description = inspect.getdoc(spec.func) or spec.summary
    if not description:
        raise ValueError(f"工具 {spec.name} 缺少 docstring 或 summary")
    is_async = inspect.iscoroutinefunction(spec.func)
    return StructuredTool.from_function(
        func=None if is_async else spec.func, coroutine=spec.func if is_async else None,
        name=spec.name, description=description,
    )


def install_specs(specs: Iterable[ToolSpec]) -> None:
    """校验完成后一次性替换全局目录，加载失败不发布半个注册表。"""
    global _registry
    replacement = _Registry()
    for spec in specs:
        replacement.add(spec)
    _registry = replacement


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

def register(
    group: str,
    label: str = "",
    name: str | None = None,
    summary: str = "",
    requires_approval: bool = False,
) -> 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,
                requires_approval=requires_approval,
            )
        )
        return func

    return decorator


def register_external(
    func: Callable[..., Any],
    group: str,
    label: str = "",
    name: str | None = None,
    summary: str = "",
    requires_approval: bool = False,
) -> 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,
            requires_approval=requires_approval,
        )
    )
    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 [spec_as_tool(spec) 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(spec_as_tool(spec))
    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,
                    "requires_approval": _registry.tools[name].requires_approval,
                }
                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


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

def load_builtin_tools() -> int:
    """兼容旧调用方；工具和子 Agent 统一由业务扩展声明。"""
    from extensions.registry import load_extensions
    load_extensions()
    return len(_registry.tools)


def reset_for_tests() -> None:
    """仅供测试：清空发布目录；下次重新执行 setup，不依赖 import 重放。"""
    from extensions.registry import reset_for_tests as reset_extensions
    _registry.tools.clear()
    _registry.groups.clear()
    reset_extensions()