Coverage for src/dataknobs_llm/tools/registry.py: 87%
70 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-08 13:50 -0700
1"""Tool registry for managing available LLM tools.
3This module provides the ToolRegistry class for registering, discovering,
4and managing tools that can be used with LLMs. Now built on the common
5Registry pattern from dataknobs_common.
6"""
8from typing import Any, Dict, List, Set
10from dataknobs_common import NotFoundError, OperationError, Registry
12from dataknobs_llm.tools.base import Tool
15class ToolRegistry(Registry[Tool]):
16 """Registry for managing available tools/functions.
18 The ToolRegistry provides a central place to register and discover
19 tools that can be called by LLMs. It supports tool registration,
20 retrieval, listing, and conversion to function calling formats.
22 Built on dataknobs_common.Registry for consistency across the ecosystem.
24 Example:
25 >>> # Create registry
26 >>> registry = ToolRegistry()
27 >>>
28 >>> # Register tools
29 >>> registry.register_tool(CalculatorTool())
30 >>> registry.register_tool(WebSearchTool())
31 >>>
32 >>> # Check if tool exists
33 >>> if registry.has_tool("calculator"):
34 ... tool = registry.get_tool("calculator")
35 ... result = await tool.execute(operation="add", a=5, b=3)
36 >>>
37 >>> # Get all tools for LLM function calling
38 >>> functions = registry.to_function_definitions()
39 >>>
40 >>> # List available tools
41 >>> tools = registry.list_tools()
42 >>> for tool_info in tools:
43 ... print(f"{tool_info['name']}: {tool_info['description']}")
44 """
46 def __init__(self):
47 """Initialize an empty tool registry."""
48 super().__init__(name="tools", enable_metrics=True)
50 def register_tool(self, tool: Tool) -> None:
51 """Register a tool.
53 Args:
54 tool: Tool instance to register
56 Raises:
57 OperationError: If a tool with the same name already exists
59 Example:
60 >>> calculator = CalculatorTool()
61 >>> registry.register_tool(calculator)
62 """
63 try:
64 self.register(
65 tool.name,
66 tool,
67 metadata={"description": tool.description, "schema": tool.schema},
68 )
69 except OperationError as e:
70 # Re-raise with more specific message for backward compatibility
71 raise OperationError(
72 f"Tool with name '{tool.name}' already registered. "
73 f"Use unregister() first or choose a different name.",
74 context=e.context,
75 ) from e
77 def register_many(self, tools: List[Tool]) -> None:
78 """Register multiple tools at once.
80 Args:
81 tools: List of Tool instances to register
83 Raises:
84 OperationError: If any tool name conflicts with existing tools
86 Example:
87 >>> tools = [CalculatorTool(), SearchTool(), WeatherTool()]
88 >>> registry.register_many(tools)
89 """
90 # Check for conflicts first
91 for tool in tools:
92 if self.has(tool.name):
93 raise OperationError(
94 f"Tool with name '{tool.name}' already registered",
95 context={"tool_name": tool.name, "registry": self.name},
96 )
98 # Register all tools
99 for tool in tools:
100 self.register(
101 tool.name,
102 tool,
103 metadata={"description": tool.description, "schema": tool.schema},
104 )
106 def get_tool(self, name: str) -> Tool:
107 """Get a tool by name.
109 Args:
110 name: Name of the tool to retrieve
112 Returns:
113 Tool instance
115 Raises:
116 NotFoundError: If no tool with the given name exists
118 Example:
119 >>> tool = registry.get_tool("calculator")
120 >>> result = await tool.execute(operation="add", a=5, b=3)
121 """
122 try:
123 return self.get(name)
124 except NotFoundError as e:
125 # Re-raise with more specific message for backward compatibility
126 raise NotFoundError(
127 f"Tool not found: {name}",
128 context=e.context,
129 ) from e
131 def has_tool(self, name: str) -> bool:
132 """Check if a tool with the given name exists.
134 Args:
135 name: Name of the tool to check
137 Returns:
138 True if tool exists, False otherwise
140 Example:
141 >>> if registry.has_tool("calculator"):
142 ... print("Calculator available")
143 """
144 return self.has(name)
146 def list_tools(self) -> List[Dict[str, Any]]:
147 """List all registered tools with their metadata.
149 Returns:
150 List of dictionaries containing tool information
152 Example:
153 >>> tools = registry.list_tools()
154 >>> for tool_info in tools:
155 ... print(f"{tool_info['name']}: {tool_info['description']}")
157 Returns format:
158 [
159 {
160 "name": "calculator",
161 "description": "Performs arithmetic operations",
162 "schema": {...},
163 "metadata": {...}
164 },
165 ...
166 ]
167 """
168 return [
169 {
170 "name": tool.name,
171 "description": tool.description,
172 "schema": tool.schema,
173 "metadata": tool.metadata,
174 }
175 for tool in self.list_items()
176 ]
178 def get_tool_names(self) -> List[str]:
179 """Get list of all registered tool names.
181 Returns:
182 List of tool names
184 Example:
185 >>> names = registry.get_tool_names()
186 >>> print(names)
187 ['calculator', 'search', 'weather']
188 """
189 return self.list_keys()
191 def to_function_definitions(
192 self, include_only: Set[str] | None = None, exclude: Set[str] | None = None
193 ) -> List[Dict[str, Any]]:
194 """Convert tools to OpenAI function calling format.
196 Args:
197 include_only: If provided, only include tools with these names
198 exclude: If provided, exclude tools with these names
200 Returns:
201 List of function definition dictionaries
203 Example:
204 >>> # Get all tools
205 >>> functions = registry.to_function_definitions()
206 >>>
207 >>> # Get only specific tools
208 >>> functions = registry.to_function_definitions(
209 ... include_only={"calculator", "web_search"}
210 ... )
211 >>>
212 >>> # Get all except specific tools
213 >>> functions = registry.to_function_definitions(
214 ... exclude={"dangerous_tool"}
215 ... )
216 """
217 tools_to_include = []
219 for name, tool in self.items():
220 # Apply filters
221 if include_only and name not in include_only:
222 continue
223 if exclude and name in exclude:
224 continue
226 tools_to_include.append(tool)
228 return [tool.to_function_definition() for tool in tools_to_include]
230 def to_anthropic_tool_definitions(
231 self, include_only: Set[str] | None = None, exclude: Set[str] | None = None
232 ) -> List[Dict[str, Any]]:
233 """Convert tools to Anthropic tool format.
235 Args:
236 include_only: If provided, only include tools with these names
237 exclude: If provided, exclude tools with these names
239 Returns:
240 List of tool definition dictionaries
242 Example:
243 >>> tools = registry.to_anthropic_tool_definitions()
244 >>> # Use with Anthropic API
245 >>> response = client.messages.create(
246 ... model="claude-3-sonnet",
247 ... tools=tools,
248 ... messages=[...]
249 ... )
250 """
251 tools_to_include = []
253 for name, tool in self.items():
254 # Apply filters
255 if include_only and name not in include_only:
256 continue
257 if exclude and name in exclude:
258 continue
260 tools_to_include.append(tool)
262 return [tool.to_anthropic_tool_definition() for tool in tools_to_include]
264 async def execute_tool(self, name: str, **kwargs: Any) -> Any:
265 """Execute a tool by name with given parameters.
267 This is a convenience method for getting and executing a tool
268 in a single call.
270 Args:
271 name: Name of the tool to execute
272 **kwargs: Parameters to pass to the tool
274 Returns:
275 Tool execution result
277 Raises:
278 NotFoundError: If tool not found
279 Exception: If tool execution fails
281 Example:
282 >>> result = await registry.execute_tool(
283 ... "calculator",
284 ... operation="add",
285 ... a=5,
286 ... b=3
287 ... )
288 >>> print(result)
289 8
290 """
291 tool = self.get_tool(name)
292 return await tool.execute(**kwargs)
294 def filter_by_metadata(self, **filters: Any) -> List[Tool]:
295 """Filter tools by metadata attributes.
297 Args:
298 **filters: Key-value pairs to match in tool metadata
300 Returns:
301 List of tools matching all filters
303 Example:
304 >>> # Get all tools with category="math"
305 >>> math_tools = registry.filter_by_metadata(category="math")
306 >>>
307 >>> # Get tools with multiple criteria
308 >>> safe_tools = registry.filter_by_metadata(
309 ... category="utility",
310 ... safe=True
311 ... )
312 """
313 matching_tools = []
315 for tool in self.list_items():
316 # Check if all filters match
317 matches = True
318 for key, value in filters.items():
319 if key not in tool.metadata or tool.metadata[key] != value:
320 matches = False
321 break
323 if matches:
324 matching_tools.append(tool)
326 return matching_tools
328 def clone(self) -> "ToolRegistry":
329 """Create a shallow copy of this registry.
331 Returns:
332 New ToolRegistry with same tools registered
334 Example:
335 >>> original = ToolRegistry()
336 >>> original.register_tool(CalculatorTool())
337 >>>
338 >>> copy = original.clone()
339 >>> copy.count()
340 1
341 """
342 new_registry = ToolRegistry()
343 for name, tool in self.items():
344 new_registry.register(name, tool, allow_overwrite=True)
345 return new_registry
347 def __repr__(self) -> str:
348 """String representation of registry."""
349 return f"ToolRegistry(tools={self.count()})"
351 def __str__(self) -> str:
352 """Human-readable string representation."""
353 if self.count() == 0:
354 return "ToolRegistry(empty)"
356 tool_names = ", ".join(self.list_keys())
357 return f"ToolRegistry({self.count()} tools: {tool_names})"
359 # Note: __len__, __contains__, and __iter__ are inherited from Registry base class