more_dots/services/tool_router.py

"""
工具路由器模块

支持动态注册和管理工具
"""

import json
import logging
import time
from typing import Any, Callable, Dict, List, Optional, Type

from langchain_core.tools import BaseTool

from tools.calculator import CalculatorTool
from tools.web_search import WebSearchTool
from tools.rest_api_tool import RestApiTool
from tools.sr_api_tool import SrApiQueryTool
from core.registry import ToolRegistry, ToolMetadata

logger = logging.getLogger(__name__)


class ToolRouter:
    """
    工具路由器：统一调用入口
    
    支持特性：
    - 动态注册工具
    - 工具元数据管理
    - 执行监控
    """
    
    def __init__(self, tools: Optional[List[BaseTool]] = None):
        self._tools: Dict[str, BaseTool] = {}
        self._tool_metadata: Dict[str, ToolMetadata] = {}
        self._execution_stats: Dict[str, Dict[str, Any]] = {}
        
        if tools is not None:
            for tool in tools:
                self.register_tool(tool)
        else:
            self._register_default_tools()
    
    def _register_default_tools(self) -> None:
        """注册默认工具"""
        default_tools = [
            CalculatorTool(),
            WebSearchTool(),
            RestApiTool(),
            SrApiQueryTool(),
        ]
        for tool in default_tools:
            self.register_tool(tool)
    
    def register_tool(
        self,
        tool: BaseTool,
        description: str = "",
        version: str = "1.0.0",
        timeout: int = 30,
        retry: int = 0,
        tags: Optional[List[str]] = None,
    ) -> None:
        """
        注册工具
        
        Args:
            tool: 工具实例
            description: 描述（默认使用 tool.description）
            version: 版本
            timeout: 超时时间
            retry: 重试次数
            tags: 标签
        """
        name = tool.name
        metadata = ToolMetadata(
            name=name,
            description=description or tool.description,
            version=version,
            timeout=timeout,
            retry=retry,
            tags=tags or [],
        )
        
        self._tools[name] = tool
        self._tool_metadata[name] = metadata
        self._execution_stats[name] = {
            "total_calls": 0,
            "success_calls": 0,
            "failed_calls": 0,
            "total_time_ms": 0,
        }
        
        ToolRegistry._entries[name] = type(
            "RegistryEntry",
            (),
            {"instance": tool, "metadata": {"tool_metadata": metadata}}
        )()
        
        logger.info(f"Registered tool: {name} (v{version})")
    
    def unregister_tool(self, name: str) -> bool:
        """
        注销工具
        
        Args:
            name: 工具名称
            
        Returns:
            是否成功注销
        """
        if name in self._tools:
            del self._tools[name]
            del self._tool_metadata[name]
            del self._execution_stats[name]
            ToolRegistry.unregister(name)
            logger.info(f"Unregistered tool: {name}")
            return True
        return False
    
    def get_tool(self, name: str) -> Optional[BaseTool]:
        """获取工具实例"""
        return self._tools.get(name)
    
    def get_tool_metadata(self, name: str) -> Optional[ToolMetadata]:
        """获取工具元数据"""
        return self._tool_metadata.get(name)
    
    def list_tools(self) -> List[str]:
        """列出可用工具名称"""
        return list(self._tools.keys())
    
    def get_tool_info(self, name: str) -> Optional[Dict[str, Any]]:
        """获取工具详细信息"""
        if name not in self._tools:
            return None
        
        tool = self._tools[name]
        metadata = self._tool_metadata.get(name)
        stats = self._execution_stats.get(name, {})
        
        return {
            "name": name,
            "description": metadata.description if metadata else tool.description,
            "version": metadata.version if metadata else "unknown",
            "timeout": metadata.timeout if metadata else 30,
            "tags": metadata.tags if metadata else [],
            "stats": {
                "total_calls": stats.get("total_calls", 0),
                "success_rate": self._calculate_success_rate(name),
            },
        }
    
    def call(self, tool_name: str, payload: Any) -> Dict[str, Any]:
        """
        调用工具并返回标准化结果
        
        Args:
            tool_name: 工具名称
            payload: 输入参数
            
        Returns:
            标准化结果 {ok, data, error}
        """
        tool = self._tools.get(tool_name)
        if not tool:
            return {"ok": False, "data": None, "error": f"工具不存在: {tool_name}"}
        
        start_time = time.time()
        
        try:
            if isinstance(payload, (dict, list)):
                input_value = json.dumps(payload, ensure_ascii=False)
            elif payload is None:
                input_value = ""
            else:
                input_value = str(payload)
            
            result = tool.run(input_value)
            
            self._record_success(tool_name, time.time() - start_time)
            
            return {"ok": True, "data": result, "error": None}
        
        except Exception as e:
            self._record_failure(tool_name, time.time() - start_time)
            return {"ok": False, "data": None, "error": str(e)}
    
    def call_with_metadata(
        self,
        tool_name: str,
        payload: Any,
    ) -> Dict[str, Any]:
        """
        调用工具并返回包含元数据的结果
        
        Args:
            tool_name: 工具名称
            payload: 输入参数
            
        Returns:
            包含元数据的结果
        """
        result = self.call(tool_name, payload)
        metadata = self.get_tool_metadata(tool_name)
        
        return {
            **result,
            "tool_name": tool_name,
            "tool_version": metadata.version if metadata else "unknown",
            "execution_time_ms": self._execution_stats.get(tool_name, {}).get("last_time_ms", 0),
        }
    
    def _record_success(self, tool_name: str, elapsed: float) -> None:
        """记录成功执行"""
        if tool_name in self._execution_stats:
            stats = self._execution_stats[tool_name]
            stats["total_calls"] += 1
            stats["success_calls"] += 1
            stats["total_time_ms"] += elapsed * 1000
            stats["last_time_ms"] = elapsed * 1000
    
    def _record_failure(self, tool_name: str, elapsed: float) -> None:
        """记录失败执行"""
        if tool_name in self._execution_stats:
            stats = self._execution_stats[tool_name]
            stats["total_calls"] += 1
            stats["failed_calls"] += 1
            stats["total_time_ms"] += elapsed * 1000
            stats["last_time_ms"] = elapsed * 1000
    
    def _calculate_success_rate(self, tool_name: str) -> float:
        """计算成功率"""
        stats = self._execution_stats.get(tool_name)
        if not stats or stats["total_calls"] == 0:
            return 0.0
        return stats["success_calls"] / stats["total_calls"]
    
    def get_all_stats(self) -> Dict[str, Dict[str, Any]]:
        """获取所有工具的执行统计"""
        result = {}
        for name in self._tools:
            result[name] = {
                **self._execution_stats.get(name, {}),
                "success_rate": self._calculate_success_rate(name),
            }
        return result
    
    def register_function(
        self,
        name: str,
        func: Callable,
        description: str = "",
        timeout: int = 30,
    ) -> None:
        """
        将普通函数注册为工具
        
        Args:
            name: 工具名称
            func: 函数
            description: 描述
            timeout: 超时时间
        """
        from langchain_core.tools import Tool
        
        tool = Tool(
            name=name,
            description=description,
            func=func,
        )
        self.register_tool(tool, description=description, timeout=timeout)