app/agents/tools.py

"""Pydantic AI 工具集 (供 LangGraph tool_executor 节点调用).

工具列表:
- summarize_document: 总结指定文档
- compare_documents: 对比 2 个文档的某方面
- calculate: 安全数学计算 (用 AST 而不是 eval)
- list_documents: 列出已上传的文档
- get_current_time: 当前时间 (调试用)
"""
from __future__ import annotations

import ast
import datetime as _dt
import logging
import operator
from typing import Any

from pydantic import BaseModel, Field

from app.models import db
from app.services.embedding import get_embedder
from app.services.vector_store import hybrid_query

logger = logging.getLogger(__name__)


# ========== Tool schemas (OpenAI function calling 格式) ==========
TOOL_SCHEMAS: list[dict[str, Any]] = [
    {
        "type": "function",
        "function": {
            "name": "summarize_document",
            "description": "根据 doc_id 总结指定文档的主要内容 (取前若干 chunks 拼接 + LLM 摘要).",
            "parameters": {
                "type": "object",
                "properties": {
                    "doc_id": {
                        "type": "string",
                        "description": "要总结的文档 ID",
                    },
                    "max_chunks": {
                        "type": "integer",
                        "description": "最多取前多少个 chunk (默认 30)",
                        "default": 30,
                    },
                },
                "required": ["doc_id"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "compare_documents",
            "description": "对比两个文档在某方面的异同. 返回结构化对比结果.",
            "parameters": {
                "type": "object",
                "properties": {
                    "doc_id_a": {"type": "string"},
                    "doc_id_b": {"type": "string"},
                    "aspect": {
                        "type": "string",
                        "description": "对比的方面, 如 '技术选型', '营收', '风险' 等",
                    },
                },
                "required": ["doc_id_a", "doc_id_b", "aspect"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "calculate",
            "description": "安全地计算数学表达式 (加减乘除/括号/比较), 用 AST 解析避免注入.",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "数学表达式, 如 '(2+3)*4'",
                    },
                },
                "required": ["expression"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "list_documents",
            "description": "列出已上传的所有文档 (ID + 文件名 + 状态).",
            "parameters": {
                "type": "object",
                "properties": {},
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "获取服务器当前时间 (UTC / 本地时区).",
            "parameters": {
                "type": "object",
                "properties": {
                    "tz": {
                        "type": "string",
                        "description": "时区, 如 'UTC', 'Asia/Shanghai'. 默认 UTC",
                        "default": "UTC",
                    },
                },
            },
        },
    },
]


# ========== 工具实现 ==========
async def tool_summarize_document(args: dict[str, Any]) -> str:
    doc_id = args.get("doc_id")
    if not doc_id:
        return "❌ 缺少 doc_id"
    doc = db.doc_get(doc_id)
    if doc is None:
        return f"❌ 文档 {doc_id} 不存在"

    chunks = db.chunk_get_by_doc(doc_id, limit=int(args.get("max_chunks", 30)))
    if not chunks:
        return f"❌ 文档 {doc_id} 暂无内容"

    text = "\n\n".join(c["text"] for c in chunks)
    # 简单截断前 6000 字给 LLM 摘要
    snippet = text[:6000]

    from app.llm.factory import get_llm
    llm = get_llm()
    resp = await llm.chat(
        messages=[
            {"role": "system", "content": "你是文档摘要助手. 用中文输出简洁摘要, 列出主要章节和关键事实."},
            {"role": "user", "content": f"请摘要以下内容 (来自 {doc.get('filename', doc_id)}):\n\n{snippet}"},
        ],
        temperature=0.3,
        max_tokens=600,
    )
    return resp.content or "(空)"


async def tool_compare_documents(args: dict[str, Any]) -> str:
    doc_a_id = args.get("doc_id_a")
    doc_b_id = args.get("doc_id_b")
    aspect = args.get("aspect", "整体")
    if not (doc_a_id and doc_b_id):
        return "❌ 缺少 doc_id_a 或 doc_id_b"

    chunks_a = db.chunk_get_by_doc(doc_a_id, limit=20)
    chunks_b = db.chunk_get_by_doc(doc_b_id, limit=20)
    if not chunks_a or not chunks_b:
        return f"❌ 文档 {doc_a_id if not chunks_a else doc_b_id} 无内容"

    text_a = "\n".join(c["text"] for c in chunks_a)[:3000]
    text_b = "\n".join(c["text"] for c in chunks_b)[:3000]

    from app.llm.factory import get_llm
    llm = get_llm()
    resp = await llm.chat(
        messages=[
            {"role": "system", "content": "你是文档对比助手. 输出 Markdown 表格, 列出差異与相同点."},
            {"role": "user", "content": (
                f"对比以下两份文档的「{aspect}」方面.\n\n"
                f"## 文档 A\n{text_a}\n\n## 文档 B\n{text_b}\n\n"
                f"输出: 1) 关键差异 (表格) 2) 相同点 3) 建议"
            )},
        ],
        temperature=0.3,
        max_tokens=800,
    )
    return resp.content or "(空)"


_BIN_OPS = {
    ast.Add: operator.add, ast.Sub: operator.sub,
    ast.Mult: operator.mul, ast.Div: operator.truediv,
    ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod,
    ast.Pow: operator.pow,
}
_CMP_OPS = {
    ast.Eq: operator.eq, ast.NotEq: operator.ne,
    ast.Lt: operator.lt, ast.LtE: operator.le,
    ast.Gt: operator.gt, ast.GtE: operator.ge,
}
_UNARY_OPS = {ast.UAdd: operator.pos, ast.USub: operator.neg}


def _safe_eval(node: ast.AST) -> float | bool:
    if isinstance(node, ast.Expression):
        return _safe_eval(node.body)
    if isinstance(node, ast.Constant):
        if isinstance(node.value, (int, float)):
            return node.value
        raise ValueError(f"不支持的字面量: {node.value!r}")
    if isinstance(node, ast.BinOp):
        op = _BIN_OPS.get(type(node.op))
        if op is None:
            raise ValueError(f"不支持的二元运算: {type(node.op).__name__}")
        return op(_safe_eval(node.left), _safe_eval(node.right))
    if isinstance(node, ast.UnaryOp):
        op = _UNARY_OPS.get(type(node.op))
        if op is None:
            raise ValueError(f"不支持的一元运算: {type(node.op).__name__}")
        return op(_safe_eval(node.operand))
    if isinstance(node, ast.Compare):
        left = _safe_eval(node.left)
        for cmp_op, right_node in zip(node.ops, node.comparators):
            op = _CMP_OPS.get(type(cmp_op))
            if op is None:
                raise ValueError(f"不支持的比较: {type(cmp_op).__name__}")
            left = op(left, _safe_eval(right_node))
            if not left:
                return False
        return True
    raise ValueError(f"不支持的语法: {ast.dump(node)}")


async def tool_calculate(args: dict[str, Any]) -> str:
    expr = args.get("expression", "")
    if not expr:
        return "❌ 缺少 expression"
    try:
        tree = ast.parse(expr, mode="eval")
        result = _safe_eval(tree)
        return f"{expr} = {result}"
    except Exception as e:  # noqa: BLE001
        return f"❌ 计算失败: {e}"


async def tool_list_documents(args: dict[str, Any]) -> str:
    docs = db.doc_list(limit=200)
    if not docs:
        return "暂无已上传文档"
    lines = [f"- {d['id'][:8]}…  {d['filename']}  (chunks={d.get('chunk_count', 0)}, status={d.get('status', '?')})" for d in docs]
    return "已上传文档:\n" + "\n".join(lines)


async def tool_get_current_time(args: dict[str, Any]) -> str:
    tz = args.get("tz", "UTC")
    try:
        from zoneinfo import ZoneInfo
        now = _dt.datetime.now(ZoneInfo(tz))
    except Exception:  # noqa: BLE001
        now = _dt.datetime.utcnow()
    return now.strftime("%Y-%m-%d %H:%M:%S %Z")


# ========== 工具路由表 ==========
TOOL_FUNCS = {
    "summarize_document": tool_summarize_document,
    "compare_documents": tool_compare_documents,
    "calculate": tool_calculate,
    "list_documents": tool_list_documents,
    "get_current_time": tool_get_current_time,
}


async def execute_tool(name: str, args: dict[str, Any]) -> str:
    """统一工具执行入口."""
    fn = TOOL_FUNCS.get(name)
    if fn is None:
        return f"❌ 未知工具: {name}"
    try:
        return await fn(args)
    except Exception as e:  # noqa: BLE001
        logger.exception("Tool %s failed", name)
        return f"❌ 工具执行失败: {e}"


__all__ = ["TOOL_SCHEMAS", "TOOL_FUNCS", "execute_tool"]