long_articles_agentic_supply/supply/core/agent/tool.py

"""Tool 抽象基类 + ToolRegistry —— Agent 的工具系统

每个 Tool 是 Agent 可调用的原子能力单元。
ToolRegistry 管理工具注册与发现（实例级，非全局）。
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


# ==================== ToolResult ====================


@dataclass
class ToolResult:
    """Tool 执行结果"""

    success: bool = True
    data: Any = None
    error: Optional[str] = None
    # 调用耗时（毫秒）
    duration_ms: int = 0
    # 额外元数据
    metadata: Dict[str, Any] = field(default_factory=dict)

    @classmethod
    def ok(cls, data: Any, **kwargs) -> "ToolResult":
        return cls(success=True, data=data, **kwargs)

    @classmethod
    def fail(cls, error: str, **kwargs) -> "ToolResult":
        return cls(success=False, error=error, **kwargs)


# ==================== Tool 抽象基类 ====================


class Tool(ABC):
    """Tool 抽象基类

    每个 Tool 是一个原子能力单元，例如：
      - WebSearchTool: 网络搜索
      - DatabaseQueryTool: 数据库查询
      - HttpRequestTool: HTTP 调用
      - LLMCallTool: LLM 推理

    Tool 的 parameters_schema 描述输入参数的 JSON Schema，
    供 Agent 框架生成 function calling 的 tool definition。

    使用方式：

        class WebSearchTool(Tool):
            name = "web_search"
            description = "搜索互联网获取最新信息"

            @property
            def parameters_schema(self) -> dict:
                return {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "搜索关键词"},
                        "max_results": {"type": "integer", "default": 5},
                    },
                    "required": ["query"],
                }

            async def execute(self, **kwargs) -> ToolResult:
                results = await search_api(kwargs["query"])
                return ToolResult.ok(results)
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Tool 唯一标识，用于 Agent 引用"""
        ...

    @property
    @abstractmethod
    def description(self) -> str:
        """Tool 功能描述，写入 LLM prompt"""
        ...

    @property
    @abstractmethod
    def parameters_schema(self) -> dict:
        """输入参数 JSON Schema

        Returns:
            {
                "type": "object",
                "properties": {...},
                "required": [...],
            }
        """
        ...

    @abstractmethod
    async def execute(self, **kwargs) -> ToolResult:
        """执行 Tool，返回结果"""
        ...

    def to_openai_tool(self) -> dict:
        """导出为 OpenAI function calling 格式"""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters_schema,
            },
        }

    def __repr__(self) -> str:
        return f"<Tool name={self.name}>"


# ==================== ToolRegistry ====================


class ToolRegistry:
    """Tool 注册表 —— 实例级，可注入

    替代全局 dict，支持测试隔离和多 Agent 场景。

    使用方式：
        registry = ToolRegistry()
        registry.register(WebSearchTool())
        registry.register(FactCheckTool())

        tool = registry.get("web_search")
        result = await tool.execute(query="AI 安全")
    """

    def __init__(self):
        self._tools: Dict[str, Tool] = {}

    def register(self, tool: Tool) -> "ToolRegistry":
        """注册 Tool，链式调用"""
        if tool.name in self._tools:
            raise ValueError(f"Tool '{tool.name}' 已注册")
        self._tools[tool.name] = tool
        return self

    def get(self, name: str) -> Tool:
        """按名获取 Tool"""
        tool = self._tools.get(name)
        if tool is None:
            raise KeyError(f"Tool '{name}' 未注册，可用: {self.list_names()}")
        return tool

    def list_names(self) -> List[str]:
        return list(self._tools.keys())

    def list_all(self) -> List[Tool]:
        return list(self._tools.values())

    def to_openai_tools(self) -> List[dict]:
        """导出全部 Tool 为 OpenAI function calling 格式"""
        return [t.to_openai_tool() for t in self._tools.values()]

    def __len__(self) -> int:
        return len(self._tools)

    def __contains__(self, name: str) -> bool:
        return name in self._tools