Coverage for src / dataknobs_llm / tools / registry.py: 24%
70 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-15 10:28 -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 ```python
26 # Create registry
27 registry = ToolRegistry()
29 # Register tools
30 registry.register_tool(CalculatorTool())
31 registry.register_tool(WebSearchTool())
33 # Check if tool exists
34 if registry.has_tool("calculator"):
35 tool = registry.get_tool("calculator")
36 result = await tool.execute(operation="add", a=5, b=3)
38 # Get all tools for LLM function calling
39 functions = registry.to_function_definitions()
41 # List available tools
42 tools = registry.list_tools()
43 for tool_info in tools:
44 print(f"{tool_info['name']}: {tool_info['description']}")
45 ```
46 """
48 def __init__(self):
49 """Initialize an empty tool registry."""
50 super().__init__(name="tools", enable_metrics=True)
52 def register_tool(self, tool: Tool) -> None:
53 """Register a tool.
55 Args:
56 tool: Tool instance to register
58 Raises:
59 OperationError: If a tool with the same name already exists
61 Example:
62 >>> calculator = CalculatorTool()
63 >>> registry.register_tool(calculator)
64 """
65 try:
66 self.register(
67 tool.name,
68 tool,
69 metadata={"description": tool.description, "schema": tool.schema},
70 )
71 except OperationError as e:
72 # Re-raise with more specific message for backward compatibility
73 raise OperationError(
74 f"Tool with name '{tool.name}' already registered. "
75 f"Use unregister() first or choose a different name.",
76 context=e.context,
77 ) from e
79 def register_many(self, tools: List[Tool]) -> None:
80 """Register multiple tools at once.
82 Args:
83 tools: List of Tool instances to register
85 Raises:
86 OperationError: If any tool name conflicts with existing tools
88 Example:
89 >>> tools = [CalculatorTool(), SearchTool(), WeatherTool()]
90 >>> registry.register_many(tools)
91 """
92 # Check for conflicts first
93 for tool in tools:
94 if self.has(tool.name):
95 raise OperationError(
96 f"Tool with name '{tool.name}' already registered",
97 context={"tool_name": tool.name, "registry": self.name},
98 )
100 # Register all tools
101 for tool in tools:
102 self.register(
103 tool.name,
104 tool,
105 metadata={"description": tool.description, "schema": tool.schema},
106 )
108 def get_tool(self, name: str) -> Tool:
109 """Get a tool by name.
111 Args:
112 name: Name of the tool to retrieve
114 Returns:
115 Tool instance
117 Raises:
118 NotFoundError: If no tool with the given name exists
120 Example:
121 >>> tool = registry.get_tool("calculator")
122 >>> result = await tool.execute(operation="add", a=5, b=3)
123 """
124 try:
125 return self.get(name)
126 except NotFoundError as e:
127 # Re-raise with more specific message for backward compatibility
128 raise NotFoundError(
129 f"Tool not found: {name}",
130 context=e.context,
131 ) from e
133 def has_tool(self, name: str) -> bool:
134 """Check if a tool with the given name exists.
136 Args:
137 name: Name of the tool to check
139 Returns:
140 True if tool exists, False otherwise
142 Example:
143 ```python
144 if registry.has_tool("calculator"):
145 print("Calculator available")
146 ```
147 """
148 return self.has(name)
150 def list_tools(self) -> List[Dict[str, Any]]:
151 """List all registered tools with their metadata.
153 Returns:
154 List of dictionaries containing tool information
156 Example:
157 ```python
158 tools = registry.list_tools()
159 for tool_info in tools:
160 print(f"{tool_info['name']}: {tool_info['description']}")
161 ```
163 Returns format:
164 [
165 {
166 "name": "calculator",
167 "description": "Performs arithmetic operations",
168 "schema": {...},
169 "metadata": {...}
170 },
171 ...
172 ]
173 """
174 return [
175 {
176 "name": tool.name,
177 "description": tool.description,
178 "schema": tool.schema,
179 "metadata": tool.metadata,
180 }
181 for tool in self.list_items()
182 ]
184 def get_tool_names(self) -> List[str]:
185 """Get list of all registered tool names.
187 Returns:
188 List of tool names
190 Example:
191 >>> names = registry.get_tool_names()
192 >>> print(names)
193 ['calculator', 'search', 'weather']
194 """
195 return self.list_keys()
197 def to_function_definitions(
198 self, include_only: Set[str] | None = None, exclude: Set[str] | None = None
199 ) -> List[Dict[str, Any]]:
200 """Convert tools to OpenAI function calling format.
202 Args:
203 include_only: If provided, only include tools with these names
204 exclude: If provided, exclude tools with these names
206 Returns:
207 List of function definition dictionaries
209 Example:
210 ```python
211 # Get all tools
212 functions = registry.to_function_definitions()
214 # Get only specific tools
215 functions = registry.to_function_definitions(
216 include_only={"calculator", "web_search"}
217 )
219 # Get all except specific tools
220 functions = registry.to_function_definitions(
221 exclude={"dangerous_tool"}
222 )
223 ```
224 """
225 tools_to_include = []
227 for name, tool in self.items():
228 # Apply filters
229 if include_only and name not in include_only:
230 continue
231 if exclude and name in exclude:
232 continue
234 tools_to_include.append(tool)
236 return [tool.to_function_definition() for tool in tools_to_include]
238 def to_anthropic_tool_definitions(
239 self, include_only: Set[str] | None = None, exclude: Set[str] | None = None
240 ) -> List[Dict[str, Any]]:
241 """Convert tools to Anthropic tool format.
243 Args:
244 include_only: If provided, only include tools with these names
245 exclude: If provided, exclude tools with these names
247 Returns:
248 List of tool definition dictionaries
250 Example:
251 ```python
252 tools = registry.to_anthropic_tool_definitions()
253 # Use with Anthropic API
254 response = client.messages.create(
255 model="claude-3-sonnet",
256 tools=tools,
257 messages=[...]
258 )
259 ```
260 """
261 tools_to_include = []
263 for name, tool in self.items():
264 # Apply filters
265 if include_only and name not in include_only:
266 continue
267 if exclude and name in exclude:
268 continue
270 tools_to_include.append(tool)
272 return [tool.to_anthropic_tool_definition() for tool in tools_to_include]
274 async def execute_tool(self, name: str, **kwargs: Any) -> Any:
275 """Execute a tool by name with given parameters.
277 This is a convenience method for getting and executing a tool
278 in a single call.
280 Args:
281 name: Name of the tool to execute
282 **kwargs: Parameters to pass to the tool
284 Returns:
285 Tool execution result
287 Raises:
288 NotFoundError: If tool not found
289 Exception: If tool execution fails
291 Example:
292 ```python
293 result = await registry.execute_tool(
294 "calculator",
295 operation="add",
296 a=5,
297 b=3
298 )
299 print(result)
300 # 8
301 ```
302 """
303 tool = self.get_tool(name)
304 return await tool.execute(**kwargs)
306 def filter_by_metadata(self, **filters: Any) -> List[Tool]:
307 """Filter tools by metadata attributes.
309 Args:
310 **filters: Key-value pairs to match in tool metadata
312 Returns:
313 List of tools matching all filters
315 Example:
316 ```python
317 # Get all tools with category="math"
318 math_tools = registry.filter_by_metadata(category="math")
320 # Get tools with multiple criteria
321 safe_tools = registry.filter_by_metadata(
322 category="utility",
323 safe=True
324 )
325 ```
326 """
327 matching_tools = []
329 for tool in self.list_items():
330 # Check if all filters match
331 matches = True
332 for key, value in filters.items():
333 if key not in tool.metadata or tool.metadata[key] != value:
334 matches = False
335 break
337 if matches:
338 matching_tools.append(tool)
340 return matching_tools
342 def clone(self) -> "ToolRegistry":
343 """Create a shallow copy of this registry.
345 Returns:
346 New ToolRegistry with same tools registered
348 Example:
349 >>> original = ToolRegistry()
350 >>> original.register_tool(CalculatorTool())
351 >>>
352 >>> copy = original.clone()
353 >>> copy.count()
354 1
355 """
356 new_registry = ToolRegistry()
357 for name, tool in self.items():
358 new_registry.register(name, tool, allow_overwrite=True)
359 return new_registry
361 def __repr__(self) -> str:
362 """String representation of registry."""
363 return f"ToolRegistry(tools={self.count()})"
365 def __str__(self) -> str:
366 """Human-readable string representation."""
367 if self.count() == 0:
368 return "ToolRegistry(empty)"
370 tool_names = ", ".join(self.list_keys())
371 return f"ToolRegistry({self.count()} tools: {tool_names})"
373 # Note: __len__, __contains__, and __iter__ are inherited from Registry base class