61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Callable, Dict, List, Optional
|
||
|
|
|
||
|
|
from .tool_trace import ToolTraceStore
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(slots=True)
|
||
|
|
class ToolDefinition:
|
||
|
|
name: str
|
||
|
|
description: str
|
||
|
|
parameters: Dict[str, Any]
|
||
|
|
handler: Callable[["ToolContext", Dict[str, Any]], Dict[str, Any]]
|
||
|
|
|
||
|
|
def to_openai_tool(self) -> Dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"type": "function",
|
||
|
|
"function": {
|
||
|
|
"name": self.name,
|
||
|
|
"description": self.description,
|
||
|
|
"parameters": self.parameters,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(slots=True)
|
||
|
|
class ToolContext:
|
||
|
|
workspace: Path
|
||
|
|
session: Dict[str, Any]
|
||
|
|
tool_trace_store: Optional[ToolTraceStore] = None
|
||
|
|
|
||
|
|
|
||
|
|
class ToolRegistry:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self._tools: Dict[str, ToolDefinition] = {}
|
||
|
|
|
||
|
|
def register(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
name: str,
|
||
|
|
description: str,
|
||
|
|
parameters: Dict[str, Any],
|
||
|
|
handler: Callable[[ToolContext, Dict[str, Any]], Dict[str, Any]],
|
||
|
|
) -> None:
|
||
|
|
self._tools[name] = ToolDefinition(name=name, description=description, parameters=parameters, handler=handler)
|
||
|
|
|
||
|
|
def definitions(self) -> List[Dict[str, Any]]:
|
||
|
|
return [tool.to_openai_tool() for tool in self._tools.values()]
|
||
|
|
|
||
|
|
def execute(self, name: str, args: Dict[str, Any], ctx: ToolContext) -> str:
|
||
|
|
if name not in self._tools:
|
||
|
|
raise KeyError(f"Unknown tool: {name}")
|
||
|
|
result = self._tools[name].handler(ctx, args)
|
||
|
|
return json.dumps(result, ensure_ascii=False, indent=2)
|
||
|
|
|
||
|
|
def names(self) -> List[str]:
|
||
|
|
return sorted(self._tools)
|