simple-chat-api/tools.py

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")