105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
import asyncio
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
import json
|
|
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")
|
|
|
|
|
|
@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,
|
|
)
|
|
]
|
|
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()}
|