Multi-Agent/tools.py

"""工具注册表 — 每个工具是真实的可执行函数 + OpenAI function schema"""
import json, subprocess, sys, tempfile, os, re, time

try:
    from tavily import TavilyClient
except ImportError:
    TavilyClient = None

from config import TAVILY_API_KEY

# ── 工具函数 ──

def web_search(query: str, max_results: int = 2) -> str:
    """真实调用 Tavily Search API（带重试）"""
    if not TavilyClient:
        return "[错误] tavily-python 未安装"
    for attempt in range(3):
        try:
            client = TavilyClient(api_key=TAVILY_API_KEY)
            resp = client.search(query=query, max_results=max_results, search_depth="basic")
            results = resp.get("results", [])
            if not results:
                return "未找到相关结果"
            output = []
            for i, r in enumerate(results, 1):
                title = r.get('title', '')
                content = r.get('content', '')[:250]
                output.append(f"[{i}] {title}\n    {content}")
            return "\n\n".join(output)
        except Exception as e:
            if attempt < 2:
                time.sleep(1)
                continue
            return f"[搜索错误] {e}"
    return "[搜索错误] 重试3次均失败"


def calculator(expression: str) -> str:
    """安全数学计算"""
    allowed = set("0123456789+-*/.() %")
    cleaned = expression.replace("^", "**")
    if not all(c in allowed or c.isalpha() for c in cleaned):
        return "[错误] 不允许的字符"
    try:
        result = eval(cleaned, {"__builtins__": {}}, {"abs": abs, "round": round, "pow": pow})
        return f"{expression} = {result}"
    except Exception as e:
        return f"[计算错误] {e}"


def run_python(code: str) -> str:
    """在隔离子进程中执行 Python 代码（限时 10 秒）"""
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        f.write(code)
        f.flush()
        tmp_path = f.name
    try:
        result = subprocess.run(
            [sys.executable, tmp_path],
            capture_output=True, text=True, timeout=10,
            env={**os.environ, 'PYTHONIOENCODING': 'utf-8'}
        )
        output = result.stdout
        if result.returncode != 0:
            output += f"\n[STDERR] {result.stderr}"
        return output.strip() or "(无输出)"
    except subprocess.TimeoutExpired:
        return "[错误] 执行超时（10秒限制）"
    finally:
        os.unlink(tmp_path)


def text_analyze(text: str) -> str:
    """文本统计分析"""
    words = text.split()
    lines = text.strip().split('\n')
    chars = len(text)
    return (
        f"字符数: {chars}\n"
        f"词数: {len(words)}\n"
        f"行数: {len(lines)}\n"
        f"平均行长: {chars / max(len(lines), 1):.1f}"
    )

# ── OpenAI Function Schemas ──

TOOL_SCHEMAS = {
    "web_search": {
        "type": "function",
        "function": {
            "name": "web_search",
            "description": "搜索互联网获取实时信息。每次调用返回2条最相关的结果。建议一次搜索用精准关键词，避免多次搜索。",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "搜索关键词，建议用英文或中英混合"},
                    "max_results": {"type": "integer", "description": "返回结果数量，默认2", "default": 2}
                },
                "required": ["query"]
            }
        }
    },
    "calculator": {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "执行数学计算（支持加减乘除、幂运算、括号）",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "数学表达式，如 (3 + 5) * 7"}
                },
                "required": ["expression"]
            }
        }
    },
    "run_python": {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": "执行 Python 代码并返回输出。用于数据处理、分析、生成报告等",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "要执行的 Python 代码"}
                },
                "required": ["code"]
            }
        }
    },
    "text_analyze": {
        "type": "function",
        "function": {
            "name": "text_analyze",
            "description": "对文本进行统计分析（字符数、词数、行数等）",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {"type": "string", "description": "要分析的文本"}
                },
                "required": ["text"]
            }
        }
    },
}

# 函数映射
TOOL_FUNCTIONS = {
    "web_search": web_search,
    "calculator": calculator,
    "run_python": run_python,
    "text_analyze": text_analyze,
}

def execute_tool(name: str, args: dict) -> str:
    """统一工具执行入口"""
    fn = TOOL_FUNCTIONS.get(name)
    if not fn:
        return f"[错误] 未知工具: {name}"
    try:
        return fn(**args)
    except Exception as e:
        return f"[工具执行错误] {name}: {e}"