ml-knowledge-platform/knowledge_platform/tools/registry.py
2026-02-16 04:50:51 -08:00

112 lines
3.2 KiB
Python

"""Tool registry for Crystal's agentic framework.
Provides discovery and registration of tools by name. Supports both
explicit registration and a decorator-based API for automatic registration.
"""
from __future__ import annotations
from typing import Any
from .base import Tool
class ToolRegistry:
"""Central registry for tool discovery and retrieval.
Tools are registered by their ``name`` class attribute. Duplicate
registrations for the same name raise ``ValueError``.
Example::
registry = ToolRegistry()
registry.register(ReadFileTool)
tool = registry.get("read_file")
schema = registry.get_all_schemas()
"""
def __init__(self) -> None:
self._tools: dict[str, Tool] = {}
def register(self, tool_class: type[Tool]) -> None:
"""Register a tool class (instantiates it).
Args:
tool_class: A Tool subclass to register.
Raises:
TypeError: If tool_class is not a Tool subclass.
ValueError: If a tool with the same name is already registered.
"""
if not isinstance(tool_class, type) or not issubclass(tool_class, Tool):
raise TypeError(f"Expected a Tool subclass, got {tool_class!r}")
instance = tool_class()
name = instance.name
if name in self._tools:
raise ValueError(
f"Tool '{name}' is already registered "
f"(existing: {self._tools[name].__class__.__name__}, "
f"new: {tool_class.__name__})"
)
self._tools[name] = instance
def get(self, name: str) -> Tool | None:
"""Retrieve a registered tool by name.
Returns:
The Tool instance, or None if not found.
"""
return self._tools.get(name)
def list_tools(self) -> list[Tool]:
"""Return all registered tools, sorted by name."""
return sorted(self._tools.values(), key=lambda t: t.name)
def list_names(self) -> list[str]:
"""Return sorted list of registered tool names."""
return sorted(self._tools.keys())
def get_all_schemas(self) -> list[dict[str, Any]]:
"""Return Anthropic-compatible schemas for all registered tools.
Suitable for passing directly to the ``tools`` parameter of the
Anthropic Messages API.
"""
return [tool.to_anthropic_schema() for tool in self.list_tools()]
def __len__(self) -> int:
return len(self._tools)
def __contains__(self, name: str) -> bool:
return name in self._tools
def __repr__(self) -> str:
names = ", ".join(self.list_names())
return f"<ToolRegistry tools=[{names}]>"
# Module-level default registry for decorator usage.
_default_registry = ToolRegistry()
def register_tool(cls: type[Tool]) -> type[Tool]:
"""Class decorator that registers a Tool with the default registry.
Example::
@register_tool
class MyTool(Tool):
name = "my_tool"
...
"""
_default_registry.register(cls)
return cls
def get_default_registry() -> ToolRegistry:
"""Return the module-level default registry."""
return _default_registry