Coverage for src/dataknobs_llm/tools/registry.py: 24%

1"""Tool registry for managing available LLM tools.

2

3 This module provides the ToolRegistry class for registering, discovering,

4 and managing tools that can be used with LLMs. Now built on the common

5 Registry pattern from dataknobs_common.

6"""

7

8 from typing import Any, Dict, List, Set

9

10 from dataknobs_common import NotFoundError, OperationError, Registry

11

12 from dataknobs_llm.tools.base import Tool

13

14

15 class ToolRegistry(Registry[Tool]):

16 """Registry for managing available tools/functions.

17

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.

21

22 Built on dataknobs_common.Registry for consistency across the ecosystem.

23

24 Example:

25 ```python

26 # Create registry

27 registry = ToolRegistry()

28

29 # Register tools

30 registry.register_tool(CalculatorTool())

31 registry.register_tool(WebSearchTool())

32

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)

37

38 # Get all tools for LLM function calling

39 functions = registry.to_function_definitions()

40

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 """

47

48 def __init__(self):

49 """Initialize an empty tool registry."""

50 super().__init__(name="tools", enable_metrics=True)

51

52 def register_tool(self, tool: Tool) -> None:

53 """Register a tool.

54

55 Args:

56 tool: Tool instance to register

57

58 Raises:

59 OperationError: If a tool with the same name already exists

60

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

78

79 def register_many(self, tools: List[Tool]) -> None:

80 """Register multiple tools at once.

81

82 Args:

83 tools: List of Tool instances to register

84

85 Raises:

86 OperationError: If any tool name conflicts with existing tools

87

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 )

99

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 )

107

108 def get_tool(self, name: str) -> Tool:

109 """Get a tool by name.

110

111 Args:

112 name: Name of the tool to retrieve

113

114 Returns:

115 Tool instance

116

117 Raises:

118 NotFoundError: If no tool with the given name exists

119

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

132

133 def has_tool(self, name: str) -> bool:

134 """Check if a tool with the given name exists.

135

136 Args:

137 name: Name of the tool to check

138

139 Returns:

140 True if tool exists, False otherwise

141

142 Example:

143 ```python

144 if registry.has_tool("calculator"):

145 print("Calculator available")

146 ```

147 """

148 return self.has(name)

149

150 def list_tools(self) -> List[Dict[str, Any]]:

151 """List all registered tools with their metadata.

152

153 Returns:

154 List of dictionaries containing tool information

155

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 ```

162

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 ]

183

184 def get_tool_names(self) -> List[str]:

185 """Get list of all registered tool names.

186

187 Returns:

188 List of tool names

189

190 Example:

191 >>> names = registry.get_tool_names()

192 >>> print(names)

193 ['calculator', 'search', 'weather']

194 """

195 return self.list_keys()

196

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.

201

202 Args:

203 include_only: If provided, only include tools with these names

204 exclude: If provided, exclude tools with these names

205

206 Returns:

207 List of function definition dictionaries

208

209 Example:

210 ```python

211 # Get all tools

212 functions = registry.to_function_definitions()

213

214 # Get only specific tools

215 functions = registry.to_function_definitions(

216 include_only={"calculator", "web_search"}

217 )

218

219 # Get all except specific tools

220 functions = registry.to_function_definitions(

221 exclude={"dangerous_tool"}

222 )

223 ```

224 """

225 tools_to_include = []

226

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

233

234 tools_to_include.append(tool)

235

236 return [tool.to_function_definition() for tool in tools_to_include]

237

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.

242

243 Args:

244 include_only: If provided, only include tools with these names

245 exclude: If provided, exclude tools with these names

246

247 Returns:

248 List of tool definition dictionaries

249

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 = []

262

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

269

270 tools_to_include.append(tool)

271

272 return [tool.to_anthropic_tool_definition() for tool in tools_to_include]

273

274 async def execute_tool(self, name: str, **kwargs: Any) -> Any:

275 """Execute a tool by name with given parameters.

276

277 This is a convenience method for getting and executing a tool

278 in a single call.

279

280 Args:

281 name: Name of the tool to execute

282 **kwargs: Parameters to pass to the tool

283

284 Returns:

285 Tool execution result

286

287 Raises:

288 NotFoundError: If tool not found

289 Exception: If tool execution fails

290

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)

305

306 def filter_by_metadata(self, **filters: Any) -> List[Tool]:

307 """Filter tools by metadata attributes.

308

309 Args:

310 **filters: Key-value pairs to match in tool metadata

311

312 Returns:

313 List of tools matching all filters

314

315 Example:

316 ```python

317 # Get all tools with category="math"

318 math_tools = registry.filter_by_metadata(category="math")

319

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 = []

328

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

336

337 if matches:

338 matching_tools.append(tool)

339

340 return matching_tools

341

342 def clone(self) -> "ToolRegistry":

343 """Create a shallow copy of this registry.

344

345 Returns:

346 New ToolRegistry with same tools registered

347

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

360

361 def __repr__(self) -> str:

362 """String representation of registry."""

363 return f"ToolRegistry(tools={self.count()})"

364

365 def __str__(self) -> str:

366 """Human-readable string representation."""

367 if self.count() == 0:

368 return "ToolRegistry(empty)"

369

370 tool_names = ", ".join(self.list_keys())

371 return f"ToolRegistry({self.count()} tools: {tool_names})"

372

373 # Note: __len__, __contains__, and __iter__ are inherited from Registry base class