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
« 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.
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.
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.
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.
17The dependency direction is always: ``Skill → Tool``, never the reverse.
18"""
20from __future__ import annotations
22from typing import TYPE_CHECKING, Any, cast
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
34logger = get_logger(__name__)
36if TYPE_CHECKING:
37 from lexigram.contracts.core.module import CompiledModuleGraphProtocol
40class ToolRegistryImpl(Registry[str, ToolProtocol], ToolRegistryProtocol):
41 """Registry for atomic agent tools with optional module visibility control.
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.
47 **Boundary rule:** Tools do NOT invoke skills. Skills may invoke tools.
49 Extends :class:`Registry` for unified introspection, lifecycle hooks,
50 and thread-safe storage while implementing :class:`ToolRegistryProtocol`.
52 When a ``CompiledModuleGraph`` is provided, tool access is checked against
53 the module visibility map before execution.
55 Example::
57 registry = ToolRegistry()
58 registry.register(lookup_order, module_class=OrdersModule)
59 registry.register(process_refund, module_class=PaymentsModule)
61 # Without visibility — all tools accessible
62 result = await registry.execute("lookup_order", order_id="123")
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 """
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
76 def register( # type: ignore[override]
77 self,
78 tool: ToolProtocol,
79 module_class: type | None = None,
80 ) -> None:
81 """Register a tool.
83 Args:
84 tool: Tool satisfying ``ToolProtocol``.
85 module_class: Optional module that owns this tool
86 (for visibility enforcement).
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 )
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)
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)
113 def list_tools(self) -> list[ToolProtocol]:
114 """List all registered tools."""
115 return list(self.values())
117 def list_tool_names(self) -> list[str]:
118 """List all registered tool names."""
119 return list(self.keys())
121 def list_visible_tools(self) -> list[ToolProtocol]:
122 """List tools visible to the current caller module.
124 If no caller module is set, returns all tools (standalone mode).
125 """
126 if not self._caller_module:
127 return self.list_tools()
129 visible = []
130 for name, tool in self.items():
131 if self._is_tool_visible(name):
132 visible.append(tool)
133 return visible
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()]
139 def list_tool_schemas(self) -> list[dict[str, Any]]:
140 """Generate the tool schema list for LLM function calling.
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 ]
157 def set_module_graph(self, graph: CompiledModuleGraphProtocol | None) -> None:
158 """Set the compiled module graph for visibility enforcement.
160 Args:
161 graph: A ``CompiledModuleGraph`` from the module compiler.
162 """
163 self._module_graph = graph
164 logger.debug("tool_registry_module_graph_set")
166 def set_caller_module(self, module_class: type | None) -> None:
167 """Set the calling module for visibility checks.
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
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.
182 Checks module visibility before execution when a module graph
183 is configured.
185 Args:
186 name: Tool name to execute.
187 **kwargs: Arguments to pass to the tool.
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 )
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 )
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 )
246 def _is_tool_visible(self, tool_name: str) -> bool:
247 """Check if a tool is visible to the current caller module.
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
260 tool_module = self._tool_modules.get(tool_name)
261 if tool_module is None:
262 return True # unowned tools are always visible
264 if tool_module == self._caller_module:
265 return True # same module
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)
273 if tool_node and getattr(tool_node, "is_global", False):
274 return True
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
288 return False
290 def clear(self) -> None:
291 """Remove all registered tools."""
292 super().clear()
293 self._tool_modules.clear()
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})"