my_agent/tools/registry.py

70 lines
1.9 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 core_agent.dispatch import TaskDispatcher
from core_agent.memory import SimpleMemoryStore
from core_agent.skills import SkillStore
@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
skill_store: SkillStore
dispatcher: TaskDispatcher
memory_store: Optional[SimpleMemoryStore]
session: Dict[str, Any]
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)