vent_agent_2.0/tools/demo_tools.py

"""示例工具：演示三种注册形态。

骨架刻意只放三个工具，每个演示一件不同的事，业务工具照这个模子加：

1. :func:`get_current_time` —— 最简形态：无参、返回字符串
2. :func:`search_demo_knowledge` —— 带来源标注：返回 :func:`tools.registry.sourced`
   包裹的结构，SSE 层会自动把它变成 ``citations`` 事件
3. :func:`calculate` —— 确定性计算：把「不该交给模型算」的事放到工具侧，
   顺带演示参数校验与错误返回

写工具的三条纪律：
- **docstring 是给模型看的**——第一行说清「做什么」，正文说清「什么时候用、
  参数含义、返回什么」。模型选错工具多半是因为 docstring 写得含糊。
- **参数要写类型注解**——工具 schema 由签名推导，缺注解会导致参数变成
  「任意字符串」。
- **失败要返回可读的原因**，不要抛裸异常：抛异常会中断这一轮工具调用，
  返回错误说明能让模型自己纠正（比如改个参数重试）。
"""

from __future__ import annotations

import ast
import datetime
import operator
import re
from typing import Any

from tools.registry import register, source, sourced

# ============================================================
# 1. 最简形态
# ============================================================


@register(group="系统", label="获取当前时间", summary="返回服务器当前本地时间")
def get_current_time() -> str:
    """返回服务器当前的本地日期时间。

    当用户问「现在几点」「今天几号」，或任务需要按当前日期做推算时使用。
    不要凭记忆回答时间——你的训练数据里没有「现在」。
    """
    now = datetime.datetime.now()
    weekday = "一二三四五六日"[now.weekday()]
    return f"{now.strftime('%Y-%m-%d %H:%M:%S')}（周{weekday}）"


# ============================================================
# 2. 带来源标注
# ============================================================

# 演示用的假知识库。真实项目里这里会替换成 MCP / HTTP 调用。
_DEMO_KB: list[dict[str, Any]] = [
    {
        "title": "矿井通风基础知识",
        "url": "https://example.invalid/kb/ventilation-basics",
        "text": "矿井通风的任务是向井下各用风地点连续输送足够的新鲜空气，"
        "稀释并排出瓦斯、粉尘和有害气体，同时为作业人员提供适宜的气候条件。",
        "keywords": ["通风", "矿井", "新鲜空气", "瓦斯"],
    },
    {
        "title": "需风量计算方法",
        "url": "https://example.invalid/kb/air-volume",
        "text": "采煤工作面需风量按瓦斯涌出量、同时作业人数、气象条件、风速上限"
        "四项分别计算，取其中最大值。",
        "keywords": ["需风量", "采煤", "瓦斯", "风速"],
    },
    {
        "title": "deepagents 框架简介",
        "url": "https://example.invalid/kb/deepagents",
        "text": "deepagents 是构建深度智能体的框架，提供待办清单、文件系统、"
        "子智能体、技能与上下文压缩等开箱能力。",
        "keywords": ["deepagents", "智能体", "子智能体", "技能"],
    },
]


@register(
    group="示例检索",
    label="检索演示知识库",
    summary="在内置演示知识库里按关键词检索，返回命中片段与来源",
)
def search_demo_knowledge(query: str, limit: int = 3) -> dict[str, Any]:
    """在内置演示知识库中检索。

    Args:
        query: 检索关键词，支持空格分隔多个词（命中任意一个即算命中）。
        limit: 最多返回几条，默认 3。

    Returns:
        含 ``data``（命中片段列表）与 ``sources``（来源列表）的结构。
        ``sources`` 会被上层自动渲染成引用标注，因此必须如实填写，
        不要编造 url。
    """
    terms = [t for t in re.split(r"\s+", (query or "").strip()) if t]
    if not terms:
        # 返回可读原因而非抛异常：模型看到原因后可以补上关键词重试
        return sourced(data=[], sources=[], error="query 为空，请提供至少一个关键词")

    hits: list[dict[str, Any]] = []
    for item in _DEMO_KB:
        haystack = f"{item['title']} {' '.join(item['keywords'])} {item['text']}"
        matched = [t for t in terms if t.lower() in haystack.lower()]
        if matched:
            # 命中词数 / 总词数，作为粗糙的相关度打分
            hits.append(
                {
                    "title": item["title"],
                    "text": item["text"],
                    "matched": matched,
                    "score": round(len(matched) / len(terms), 4),
                }
            )

    hits.sort(key=lambda h: h["score"], reverse=True)
    top = hits[: max(1, int(limit or 3))]

    return sourced(
        data=top,
        sources=[
            source(
                title=hit["title"],
                url=next(i["url"] for i in _DEMO_KB if i["title"] == hit["title"]),
                snippet=hit["text"][:80],
                score=hit["score"],
            )
            for hit in top
        ],
        matched_count=len(hits),
    )


# ============================================================
# 3. 确定性计算
# ============================================================

# 白名单运算符：只允许算术与幂，杜绝 eval 带来的代码执行面
_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,
}
_UNARY_OPS = {ast.UAdd: operator.pos, ast.USub: operator.neg}
_MAX_POW_EXPONENT = 1000


def _eval_node(node: ast.AST) -> float:
    """递归求值 AST 节点。只放行算术表达式，其余一律拒绝。"""
    if isinstance(node, ast.Expression):
        return _eval_node(node.body)
    if isinstance(node, ast.Constant):
        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
            raise ValueError(f"不支持的字面量：{node.value!r}")
        return float(node.value)
    if isinstance(node, ast.BinOp):
        op_type = type(node.op)
        if op_type not in _BIN_OPS:
            raise ValueError(f"不支持的运算符：{op_type.__name__}")
        left, right = _eval_node(node.left), _eval_node(node.right)
        if op_type is ast.Pow and abs(right) > _MAX_POW_EXPONENT:
            # 挡住 9**9**9 这类会把进程算死的表达式
            raise ValueError(f"幂指数过大（上限 {_MAX_POW_EXPONENT}）")
        return _BIN_OPS[op_type](left, right)
    if isinstance(node, ast.UnaryOp):
        op_type = type(node.op)
        if op_type not in _UNARY_OPS:
            raise ValueError(f"不支持的一元运算符：{op_type.__name__}")
        return _UNARY_OPS[op_type](_eval_node(node.operand))
    raise ValueError(f"不支持的表达式元素：{type(node).__name__}")


@register(
    group="示例计算",
    label="安全算术计算",
    summary="对纯算术表达式做精确求值，避免心算出错",
)
def calculate(expression: str) -> dict[str, Any]:
    """对纯算术表达式求值，支持 + - * / // % ** 与括号。

    Args:
        expression: 算术表达式，例如 ``"(1200 + 300) * 0.85"``。

    Returns:
        ``{"expression": ..., "result": ...}``；表达式非法时返回 ``error`` 字段。

    需要精确数值（配风量核算、成本换算、单位折算）时使用本工具，
    不要自己心算——多位数乘除很容易出错，而且错了没人能发现。
    """
    expr = (expression or "").strip()
    if not expr:
        return {"error": "expression 为空"}
    if len(expr) > 200:
        return {"error": "表达式过长（上限 200 字符）"}
    try:
        tree = ast.parse(expr, mode="eval")
        result = _eval_node(tree)
    except SyntaxError:
        return {"error": f"表达式语法错误：{expr!r}"}
    except ZeroDivisionError:
        return {"error": "除数为零"}
    except (ValueError, OverflowError) as exc:
        return {"error": str(exc)}
    except Exception as exc:  # pragma: no cover - 兜底，避免工具异常打断整轮
        return {"error": f"计算失败：{type(exc).__name__}: {exc}"}
    # 整数结果去掉多余小数位，便于阅读与后续引用
    if result == int(result) and abs(result) < 1e15:
        result = int(result)
    return {"expression": expr, "result": result}


# ============================================================
# 4. 人工审批（HITL 管道）
# ============================================================

# 这个工具的名字会被 agents/agent_factory 放进 interrupt_on：
# 调用它时官方 HumanInTheLoopMiddleware 会在**执行前**拦下，向用户抛出
# 审批请求；用户批准后才真正执行函数体。
PLAN_APPROVAL_TOOL = "request_plan_approval"


@register(
    group="人机协同",
    label="请求计划审批",
    summary="把待执行的计划提交给用户确认，用户批准后才继续",
)
def request_plan_approval(plan: str) -> str:
    """把计划提交给用户审批，等待其批准后再继续执行。

    Args:
        plan: 待审批的计划全文。要写清步骤、涉及的操作、预期产出。
            不要只写一句「我将开始处理」——用户没有信息就无法判断该不该批准。

    当执行模式为 plan、或任务包含不可逆操作（删除、覆盖、对外发送、
    消耗额度）需要用户点头时调用。普通问答与只读查询不要调用，
    无谓的确认会打断用户。

    Returns:
        用户批准后的确认文本；用户拒绝时本函数不会被执行。
    """
    # 正常情况下走不到这里：调用被 interrupt_on 拦下，用户批准后才会执行函数体。
    # 这条 return 是「批准」路径的返回值。
    return f"用户已批准以下计划，请按计划执行：\n{plan}"