增加工具使用功能

2026-07-03 20:33:22 +08:00
parent 8c05311628
commit 04fefcf488
10 changed files with 662 additions and 71 deletions
@@ -0,0 +1,171 @@
+import ast
+import asyncio
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime
+import json
+import math
+import operator
+from typing import Any
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+
+
+class CurrentTimeArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    timezone: str = Field(default="UTC", description="IANA timezone, e.g. Asia/Shanghai")
+
+
+class CalculatorArguments(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    expression: str = Field(description="Arithmetic expression to evaluate")
+
+
+@dataclass(frozen=True, slots=True)
+class ToolSpec:
+    name: str
+    description: str
+    arguments_model: type[BaseModel]
+    handler: Callable[[BaseModel], dict[str, object]]
+
+    def api_definition(self) -> dict[str, object]:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.arguments_model.model_json_schema(),
+            },
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class ToolExecutionResult:
+    content: str
+    is_error: bool
+
+
+class ToolRegistry:
+    def __init__(self, timeout_seconds: float = 5.0) -> None:
+        if timeout_seconds <= 0:
+            raise ValueError("tool timeout must be greater than zero")
+        self.timeout_seconds = timeout_seconds
+        specs = [
+            ToolSpec(
+                name="get_current_time",
+                description="Get the current date and time in an IANA timezone.",
+                arguments_model=CurrentTimeArguments,
+                handler=_get_current_time,
+            ),
+            ToolSpec(
+                name="calculate",
+                description="Safely evaluate a basic arithmetic expression.",
+                arguments_model=CalculatorArguments,
+                handler=_calculate,
+            ),
+        ]
+        self._specs = {spec.name: spec for spec in specs}
+
+    def definitions(self) -> list[dict[str, object]]:
+        return [spec.api_definition() for spec in self._specs.values()]
+
+    async def execute(self, name: str, arguments: str) -> ToolExecutionResult:
+        spec = self._specs.get(name)
+        if spec is None:
+            return _tool_error("unknown_tool", f"Unknown tool: {name}")
+
+        try:
+            raw_arguments = json.loads(arguments)
+            if not isinstance(raw_arguments, dict):
+                raise ValueError("arguments must be a JSON object")
+            parsed_arguments = spec.arguments_model.model_validate(raw_arguments)
+        except (json.JSONDecodeError, ValidationError, ValueError) as exc:
+            return _tool_error("invalid_arguments", str(exc))
+
+        try:
+            result = await asyncio.wait_for(
+                asyncio.to_thread(spec.handler, parsed_arguments),
+                timeout=self.timeout_seconds,
+            )
+            return ToolExecutionResult(
+                content=json.dumps(result, ensure_ascii=False),
+                is_error=False,
+            )
+        except TimeoutError:
+            return _tool_error("timeout", f"Tool {name} timed out")
+        except Exception:
+            return _tool_error("execution_failed", f"Tool {name} failed")
+
+
+def _tool_error(code: str, message: str) -> ToolExecutionResult:
+    return ToolExecutionResult(
+        content=json.dumps(
+            {"error": {"code": code, "message": message}},
+            ensure_ascii=False,
+        ),
+        is_error=True,
+    )
+
+
+def _get_current_time(arguments: BaseModel) -> dict[str, object]:
+    assert isinstance(arguments, CurrentTimeArguments)
+    try:
+        timezone = ZoneInfo(arguments.timezone)
+    except ZoneInfoNotFoundError as exc:
+        raise ValueError("unknown timezone") from exc
+    now = datetime.now(timezone)
+    return {"timezone": arguments.timezone, "datetime": now.isoformat()}
+
+
+_BINARY_OPERATORS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
+    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_OPERATORS: dict[type[ast.unaryop], Callable[[Any], Any]] = {
+    ast.UAdd: operator.pos,
+    ast.USub: operator.neg,
+}
+_MAX_EXPRESSION_LENGTH = 200
+_MAX_ABSOLUTE_RESULT = 1e100
+_MAX_EXPONENT = 100
+
+
+def _calculate(arguments: BaseModel) -> dict[str, object]:
+    assert isinstance(arguments, CalculatorArguments)
+    expression = arguments.expression.strip()
+    if not expression or len(expression) > _MAX_EXPRESSION_LENGTH:
+        raise ValueError("expression is empty or too long")
+    tree = ast.parse(expression, mode="eval")
+    result = _evaluate_node(tree.body)
+    if isinstance(result, float) and not math.isfinite(result):
+        raise ValueError("result is not finite")
+    if abs(result) > _MAX_ABSOLUTE_RESULT:
+        raise ValueError("result is too large")
+    return {"expression": expression, "result": result}
+
+
+def _evaluate_node(node: ast.AST) -> int | float:
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
+            raise ValueError("only numeric constants are allowed")
+        return node.value
+    if isinstance(node, ast.UnaryOp) and type(node.op) in _UNARY_OPERATORS:
+        return _UNARY_OPERATORS[type(node.op)](_evaluate_node(node.operand))
+    if isinstance(node, ast.BinOp) and type(node.op) in _BINARY_OPERATORS:
+        left = _evaluate_node(node.left)
+        right = _evaluate_node(node.right)
+        if isinstance(node.op, ast.Pow) and abs(right) > _MAX_EXPONENT:
+            raise ValueError("exponent is too large")
+        result = _BINARY_OPERATORS[type(node.op)](left, right)
+        if abs(result) > _MAX_ABSOLUTE_RESULT:
+            raise ValueError("intermediate result is too large")
+        return result
+    raise ValueError("unsupported expression")