Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-agents/src/lexigram/ai/agents/tools/registry.py: 29%

95 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""Tool registry for managing agent tools with module visibility control. 

2 

3When module visibility is enabled, tool access is checked against 

4the compiled module graph before execution. An agent in 

5a specific module can only use tools exported by modules it imports. 

6 

7Boundary rule 

8------------- 

9**Tools are atomic, stateless, single-purpose operations** invoked directly by 

10the agent reasoning loop (ReAct, PlanExecute, etc.). They should complete in 

11milliseconds and have no cross-call state. 

12 

13**Skills** (``lexigram-ai-skills``) are composed multi-step workflows that may 

14be stateful, long-running, or permission-gated. Skills may invoke tools 

15internally. Tools must NOT invoke skills. 

16 

17The dependency direction is always: ``Skill → Tool``, never the reverse. 

18""" 

19 

20from __future__ import annotations 

21 

22from typing import TYPE_CHECKING, Any, cast 

23 

24from lexigram.ai.agents.exceptions import ( 

25 ToolAccessDeniedError, 

26 ToolExecutionError, 

27 ToolNotFoundError, 

28) 

29from lexigram.contracts.ai.agents import ToolError, ToolProtocol, ToolRegistryProtocol 

30from lexigram.logging import get_logger 

31from lexigram.primitives.registry import Registry 

32from lexigram.result import Err, Ok, Result 

33 

34logger = get_logger(__name__) 

35 

36if TYPE_CHECKING: 

37 from lexigram.contracts.core.module import CompiledModuleGraphProtocol 

38 

39 

40class ToolRegistryImpl(Registry[str, ToolProtocol], ToolRegistryProtocol): 

41 """Registry for atomic agent tools with optional module visibility control. 

42 

43 Tools are stateless, single-purpose functions invoked directly by the agent 

44 reasoning loop (ReAct, PlanExecute, etc.). They should be fast (<500ms) and 

45 have no side effects beyond their stated purpose. 

46 

47 **Boundary rule:** Tools do NOT invoke skills. Skills may invoke tools. 

48 

49 Extends :class:`Registry` for unified introspection, lifecycle hooks, 

50 and thread-safe storage while implementing :class:`ToolRegistryProtocol`. 

51 

52 When a ``CompiledModuleGraph`` is provided, tool access is checked against 

53 the module visibility map before execution. 

54 

55 Example:: 

56 

57 registry = ToolRegistry() 

58 registry.register(lookup_order, module_class=OrdersModule) 

59 registry.register(process_refund, module_class=PaymentsModule) 

60 

61 # Without visibility — all tools accessible 

62 result = await registry.execute("lookup_order", order_id="123") 

63 

64 # With visibility — only tools visible to the caller's module 

65 registry.set_module_graph(compiled_graph) 

66 registry.set_caller_module(SupportModule) 

67 result = await registry.execute("lookup_order", order_id="123") 

68 """ 

69 

70 def __init__(self) -> None: 

71 super().__init__(name="agent.tools", allow_overwrite=False) 

72 self._tool_modules: dict[str, type | None] = {} 

73 self._module_graph: CompiledModuleGraphProtocol | None = None 

74 self._caller_module: type | None = None 

75 

76 def register( # type: ignore[override] 

77 self, 

78 tool: ToolProtocol, 

79 module_class: type | None = None, 

80 ) -> None: 

81 """Register a tool. 

82 

83 Args: 

84 tool: Tool satisfying ``ToolProtocol``. 

85 module_class: Optional module that owns this tool 

86 (for visibility enforcement). 

87 

88 Raises: 

89 ValueError: If a tool with the same name is already registered. 

90 """ 

91 if self.has(tool.name): 

92 raise ValueError( 

93 f"Tool '{tool.name}' already registered. " 

94 f"Existing: {self.get(tool.name)!r}" 

95 ) 

96 super().register(tool.name, tool) 

97 self._tool_modules[tool.name] = module_class 

98 logger.debug( 

99 "tool_registered", 

100 tool=tool.name, 

101 module=module_class.__name__ if module_class else None, 

102 ) 

103 

104 def unregister(self, name: str) -> ToolProtocol | None: 

105 """Remove and return a tool by name.""" 

106 self._tool_modules.pop(name, None) 

107 return super().unregister(name) 

108 

109 def get(self, name: str) -> ToolProtocol | None: # type: ignore[override] 

110 """Get a tool by name, or None if not found.""" 

111 return super().get(name) 

112 

113 def list_tools(self) -> list[ToolProtocol]: 

114 """List all registered tools.""" 

115 return list(self.values()) 

116 

117 def list_tool_names(self) -> list[str]: 

118 """List all registered tool names.""" 

119 return list(self.keys()) 

120 

121 def list_visible_tools(self) -> list[ToolProtocol]: 

122 """List tools visible to the current caller module. 

123 

124 If no caller module is set, returns all tools (standalone mode). 

125 """ 

126 if not self._caller_module: 

127 return self.list_tools() 

128 

129 visible = [] 

130 for name, tool in self.items(): 

131 if self._is_tool_visible(name): 

132 visible.append(tool) 

133 return visible 

134 

135 def list_visible_tool_names(self) -> list[str]: 

136 """List names of tools visible to the current caller module.""" 

137 return [t.name for t in self.list_visible_tools()] 

138 

139 def list_tool_schemas(self) -> list[dict[str, Any]]: 

140 """Generate the tool schema list for LLM function calling. 

141 

142 Returns only tools visible to the current caller module. 

143 """ 

144 tools = self.list_visible_tools() 

145 return [ 

146 { 

147 "type": "function", 

148 "function": { 

149 "name": t.name, 

150 "description": t.description, 

151 "parameters": t.parameters_schema, 

152 }, 

153 } 

154 for t in tools 

155 ] 

156 

157 def set_module_graph(self, graph: CompiledModuleGraphProtocol | None) -> None: 

158 """Set the compiled module graph for visibility enforcement. 

159 

160 Args: 

161 graph: A ``CompiledModuleGraph`` from the module compiler. 

162 """ 

163 self._module_graph = graph 

164 logger.debug("tool_registry_module_graph_set") 

165 

166 def set_caller_module(self, module_class: type | None) -> None: 

167 """Set the calling module for visibility checks. 

168 

169 Args: 

170 module_class: The module class of the agent using this registry. 

171 ``None`` means standalone (no visibility restrictions). 

172 """ 

173 self._caller_module = module_class 

174 

175 async def execute( 

176 self, 

177 name: str, 

178 **kwargs: Any, 

179 ) -> Result[Any, ToolError]: 

180 """Execute a tool by name with visibility and error handling. 

181 

182 Checks module visibility before execution when a module graph 

183 is configured. 

184 

185 Args: 

186 name: Tool name to execute. 

187 **kwargs: Arguments to pass to the tool. 

188 

189 Returns: 

190 ``Ok(result)`` on success, ``Err(ToolError)`` on failure. 

191 """ 

192 tool = self.get(name) 

193 if tool is None: 

194 return Err( 

195 ToolNotFoundError( 

196 message=f"Tool '{name}' not found", 

197 tool_name=name, 

198 available_tools=self.list_visible_tool_names(), 

199 ) 

200 ) 

201 

202 # Module visibility check 

203 if not self._is_tool_visible(name): 

204 tool_module = self._tool_modules.get(name) 

205 tool_module_name = tool_module.__name__ if tool_module else "unknown" 

206 caller_name = ( 

207 self._caller_module.__name__ if self._caller_module else "unknown" 

208 ) 

209 logger.warning( 

210 "tool_access_denied", 

211 tool=name, 

212 tool_module=tool_module_name, 

213 caller_module=caller_name, 

214 ) 

215 return Err( 

216 ToolAccessDeniedError( 

217 message=f"Tool '{name}' is not accessible from module " 

218 f"'{caller_name}'. It belongs to module " 

219 f"'{tool_module_name}' which is not imported.", 

220 tool_name=name, 

221 agent_module=caller_name, 

222 tool_module=tool_module_name, 

223 ) 

224 ) 

225 

226 try: 

227 result = await tool.execute(**kwargs) 

228 logger.debug("tool_executed", tool=name) 

229 return Ok(result) 

230 except (RuntimeError, TypeError, ValueError, AttributeError, OSError) as e: 

231 logger.warning( 

232 "tool_execution_failed", 

233 tool=name, 

234 error=str(e), 

235 error_type=type(e).__name__, 

236 ) 

237 return Err( 

238 ToolExecutionError( 

239 message=f"Tool '{name}' failed: {e}", 

240 tool_name=name, 

241 arguments=kwargs, 

242 cause=e, 

243 ) 

244 ) 

245 

246 def _is_tool_visible(self, tool_name: str) -> bool: 

247 """Check if a tool is visible to the current caller module. 

248 

249 Visibility rules: 

250 1. No caller module → all tools visible (standalone agent) 

251 2. Tool has no module → visible to everyone 

252 3. Tool module == caller module → visible (same module) 

253 4. Tool module exports are visible to caller → visible 

254 5. Tool module is global → visible 

255 6. Otherwise → not visible 

256 """ 

257 if not self._caller_module: 

258 return True 

259 

260 tool_module = self._tool_modules.get(tool_name) 

261 if tool_module is None: 

262 return True # unowned tools are always visible 

263 

264 if tool_module == self._caller_module: 

265 return True # same module 

266 

267 # Check module graph visibility 

268 try: 

269 graph_any = cast("Any", self._module_graph) 

270 caller_node = graph_any.get_module(self._caller_module) 

271 tool_node = graph_any.get_module(tool_module) 

272 

273 if tool_node and getattr(tool_node, "is_global", False): 

274 return True 

275 

276 if caller_node: 

277 imports = getattr(caller_node, "imports", []) 

278 if tool_module in imports: 

279 return True 

280 except (RuntimeError, TypeError, ValueError, AttributeError, OSError): 

281 # If module graph lookup fails, deny access (fail closed) 

282 logger.warning( 

283 "module_graph_visibility_check_failed", 

284 tool=tool_name, 

285 ) 

286 return False 

287 

288 return False 

289 

290 def clear(self) -> None: 

291 """Remove all registered tools.""" 

292 super().clear() 

293 self._tool_modules.clear() 

294 

295 def __repr__(self) -> str: 

296 visible = len(self.list_visible_tools()) 

297 total = len(self) 

298 if visible == total: 

299 return f"ToolRegistry(tools={total})" 

300 return f"ToolRegistry(tools={total}, visible={visible})"