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

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

45

46 def __init__(self):

47 """Initialize an empty tool registry."""

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

49

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

51 """Register a tool.

52

53 Args:

54 tool: Tool instance to register

55

56 Raises:

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

58

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

76

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

78 """Register multiple tools at once.

79

80 Args:

81 tools: List of Tool instances to register

82

83 Raises:

84 OperationError: If any tool name conflicts with existing tools

85

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 )

97

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 )

105

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

107 """Get a tool by name.

108

109 Args:

110 name: Name of the tool to retrieve

111

112 Returns:

113 Tool instance

114

115 Raises:

116 NotFoundError: If no tool with the given name exists

117

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

130

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

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

133

134 Args:

135 name: Name of the tool to check

136

137 Returns:

138 True if tool exists, False otherwise

139

140 Example:

141 >>> if registry.has_tool("calculator"):

142 ... print("Calculator available")

143 """

144 return self.has(name)

145

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

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

148

149 Returns:

150 List of dictionaries containing tool information

151

152 Example:

153 >>> tools = registry.list_tools()

154 >>> for tool_info in tools:

155 ... print(f"{tool_info['name']}: {tool_info['description']}")

156

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 ]

177

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

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

180

181 Returns:

182 List of tool names

183

184 Example:

185 >>> names = registry.get_tool_names()

186 >>> print(names)

187 ['calculator', 'search', 'weather']

188 """

189 return self.list_keys()

190

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.

195

196 Args:

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

198 exclude: If provided, exclude tools with these names

199

200 Returns:

201 List of function definition dictionaries

202

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

218

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

225

226 tools_to_include.append(tool)

227

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

229

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.

234

235 Args:

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

237 exclude: If provided, exclude tools with these names

238

239 Returns:

240 List of tool definition dictionaries

241

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

252

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

259

260 tools_to_include.append(tool)

261

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

263

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

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

266

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

268 in a single call.

269

270 Args:

271 name: Name of the tool to execute

272 **kwargs: Parameters to pass to the tool

273

274 Returns:

275 Tool execution result

276

277 Raises:

278 NotFoundError: If tool not found

279 Exception: If tool execution fails

280

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)

293

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

295 """Filter tools by metadata attributes.

296

297 Args:

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

299

300 Returns:

301 List of tools matching all filters

302

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

314

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

322

323 if matches:

324 matching_tools.append(tool)

325

326 return matching_tools

327

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

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

330

331 Returns:

332 New ToolRegistry with same tools registered

333

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

346

347 def __repr__(self) -> str:

348 """String representation of registry."""

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

350

351 def __str__(self) -> str:

352 """Human-readable string representation."""

353 if self.count() == 0:

354 return "ToolRegistry(empty)"

355

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

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

358

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