| 1 | """Tool registry for dynamic tool management.""" |
| 2 | |
| 3 | from contextvars import ContextVar |
| 4 | from typing import Any |
| 5 | |
| 6 | from nanobot.agent.tools.base import Tool |
| 7 | |
| 8 | _SESSION_METADATA_CTX: ContextVar[dict[str, Any] | None] = ContextVar( |
| 9 | "tool_registry_session_metadata", |
| 10 | default=None, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | def set_tool_session_metadata(metadata: dict[str, Any] | None) -> None: |
| 15 | """Bind per-turn session metadata for context-aware tools.""" |
| 16 | _SESSION_METADATA_CTX.set(metadata) |
| 17 | |
| 18 | |
| 19 | def get_tool_session_metadata() -> dict[str, Any] | None: |
| 20 | return _SESSION_METADATA_CTX.get() |
| 21 | |
| 22 | class ToolRegistry: |
| 23 | """ |
| 24 | Registry for agent tools. |
| 25 | |
| 26 | Allows dynamic registration and execution of tools. |
| 27 | """ |
| 28 | |
| 29 | def __init__(self): |
| 30 | self._tools: dict[str, Tool] = {} |
| 31 | self._cached_definitions: list[dict[str, Any]] | None = None |
| 32 | |
| 33 | def register(self, tool: Tool) -> None: |
| 34 | """Register a tool.""" |
| 35 | self._tools[tool.name] = tool |
| 36 | self._cached_definitions = None |
| 37 | |
| 38 | def unregister(self, name: str) -> None: |
| 39 | """Unregister a tool by name.""" |
| 40 | self._tools.pop(name, None) |
| 41 | self._cached_definitions = None |
| 42 | |
| 43 | def invalidate(self) -> None: |
| 44 | """Drop the cached tool definitions so live ``description``/``parameters`` are re-read. |
| 45 | |
| 46 | Needed when a tool's schema depends on external mutable state (e.g. the |
| 47 | active PE set), since :meth:`get_definitions` otherwise serves a snapshot. |
| 48 | """ |
| 49 | self._cached_definitions = None |
| 50 | |
| 51 | def get(self, name: str) -> Tool | None: |
| 52 | """Get a tool by name.""" |
| 53 | return self._tools.get(name) |
| 54 | |
| 55 | def has(self, name: str) -> bool: |
| 56 | """Check if a tool is registered.""" |
| 57 | return name in self._tools |
| 58 | |
| 59 | @staticmethod |
| 60 | def _schema_name(schema: dict[str, Any]) -> str: |
| 61 | """Extract a normalized tool name from either OpenAI or flat schemas.""" |
| 62 | fn = schema.get("function") |
| 63 | if isinstance(fn, dict): |
| 64 | name = fn.get("name") |
| 65 | if isinstance(name, str): |
| 66 | return name |
| 67 | name = schema.get("name") |
| 68 | return name if isinstance(name, str) else "" |
| 69 | |
| 70 | def get_definitions(self, session_metadata: dict[str, Any] | None = None) -> list[dict[str, Any]]: |
| 71 | """Get tool definitions with stable ordering for cache-friendly prompts. |
| 72 | |
| 73 | Built-in tools are sorted first as a stable prefix, then MCP tools are |
| 74 | sorted and appended. The result is cached until the next |
| 75 | register/unregister call. |
| 76 | """ |
| 77 | if self._cached_definitions is None: |
| 78 | definitions = [tool.to_schema() for tool in self._tools.values()] |
| 79 | builtins: list[dict[str, Any]] = [] |
| 80 | mcp_tools: list[dict[str, Any]] = [] |
| 81 | for schema in definitions: |
| 82 | name = self._schema_name(schema) |
| 83 | if name.startswith("mcp_"): |
| 84 | mcp_tools.append(schema) |
| 85 | else: |
| 86 | builtins.append(schema) |
| 87 | builtins.sort(key=self._schema_name) |
| 88 | mcp_tools.sort(key=self._schema_name) |
| 89 | self._cached_definitions = builtins + mcp_tools |
| 90 | |
| 91 | return list(self._cached_definitions) |
| 92 | |
| 93 | def prepare_call( |
| 94 | self, |
| 95 | name: str, |
| 96 | params: dict[str, Any], |
| 97 | session_metadata: dict[str, Any] | None = None, |
| 98 | ) -> tuple[Tool | None, dict[str, Any], str | None]: |
| 99 | """Resolve, cast, and validate one tool call.""" |
| 100 | # Guard against invalid parameter types (e.g., list instead of dict) |
| 101 | if not isinstance(params, dict) and name in ('write_file', 'read_file'): |
| 102 | return None, params, ( |
| 103 | f"Error: Tool '{name}' parameters must be a JSON object, got {type(params).__name__}. " |
| 104 | "Use named parameters: tool_name(param1=\"value1\", param2=\"value2\")" |
| 105 | ) |
| 106 | tool = self._tools.get(name) |
| 107 | if not tool: |
| 108 | return None, params, ( |
| 109 | f"Error: Tool '{name}' not found. Available: {', '.join(self.tool_names)}" |
| 110 | ) |
| 111 | |
| 112 | cast_params = tool.cast_params(params) |
| 113 | errors = tool.validate_params(cast_params) |
| 114 | if errors: |
| 115 | return tool, cast_params, ( |
| 116 | f"Error: Invalid parameters for tool '{name}': " + "; ".join(errors) |
| 117 | ) |
| 118 | return tool, cast_params, None |
| 119 | |
| 120 | async def execute(self, name: str, params: dict[str, Any]) -> Any: |
| 121 | """Execute a tool by name with given parameters.""" |
| 122 | retry_hint = "\n\n[Analyze the error above and try a different approach.]" |
| 123 | tool, params, error = self.prepare_call(name, params) |
| 124 | if error: |
| 125 | return error + retry_hint |
| 126 | |
| 127 | try: |
| 128 | assert tool is not None # guarded by prepare_call() |
| 129 | result = await tool.execute(**params) |
| 130 | if isinstance(result, str) and result.startswith("Error"): |
| 131 | return result + retry_hint |
| 132 | return result |
| 133 | except Exception as e: |
| 134 | return f"Error executing {name}: {str(e)}" + retry_hint |
| 135 | |
| 136 | @property |
| 137 | def tool_names(self) -> list[str]: |
| 138 | """Get list of registered tool names.""" |
| 139 | return list(self._tools.keys()) |
| 140 | |
| 141 | def __len__(self) -> int: |
| 142 | return len(self._tools) |
| 143 | |
| 144 | def __contains__(self, name: str) -> bool: |
| 145 | return name in self._tools |
| 146 |