diff --git a/agent/tools/tool_manager.py b/agent/tools/tool_manager.py index 823e8b8e..9f9db462 100644 --- a/agent/tools/tool_manager.py +++ b/agent/tools/tool_manager.py @@ -71,6 +71,22 @@ class ToolManager: if not hasattr(self, '_mcp_active_configs'): # server_name -> normalized config dict, for diff-based reload. self._mcp_active_configs: dict = {} + if not hasattr(self, '_mcp_tool_vectors'): + # mcp_tool_name -> embedding vector, used by on-demand tool + # retrieval. Populated lazily on first retrieval so users who + # never enable the feature pay zero embedding cost. + self._mcp_tool_vectors: dict = {} + if not hasattr(self, '_mcp_vector_lock'): + # Guards incremental index builds so concurrent turns don't + # double-embed the same newly-loaded MCP tools. + self._mcp_vector_lock = threading.Lock() + if not hasattr(self, '_embedding_provider_initialized'): + # The embedding provider is created once, lazily, and reused for + # both tool-index and per-query embeddings. None means keyword-only + # mode (no provider configured) — retrieval then falls back to full + # injection at the caller. + self._embedding_provider_initialized = False + self._embedding_provider = None def load_tools(self, tools_dir: str = "", config_dict=None): """ @@ -574,6 +590,91 @@ class ToolManager: return (sorted(added), sorted(removed)) + # ------------------------------------------------------------------ + # On-demand MCP tool retrieval support + # + # The vector index and the embedding provider are owned here (singleton, + # process-wide, aligned with the MCP tool lifecycle). The context-aware + # selection itself lives in agent.tools.mcp.tool_retrieval, driven by the + # executor which is the only place that knows the conversation context. + # ------------------------------------------------------------------ + + def count_mcp_tools(self) -> int: + """Return the number of currently loaded MCP tools.""" + return len(self._mcp_tool_instances) + + def get_mcp_tool_vectors(self) -> dict: + """Return ``{mcp_tool_name: vector}`` for currently loaded MCP tools. + + Lazily embeds any MCP tools not yet in the cache (MCP servers load + asynchronously, so tools may appear over time). Returns an empty dict + when no embedding provider is available or embedding fails — the caller + then falls back to full injection. Never raises. + """ + try: + self._ensure_mcp_tool_vectors() + except Exception as e: + logger.debug(f"[ToolManager] MCP tool vector build skipped: {e}") + return dict(self._mcp_tool_vectors) + + def embed_query(self, text: str): + """Embed a retrieval query with the shared provider. + + Returns the embedding vector, or None if no provider is available or + the call fails (caller falls back to full injection). Never raises. + """ + if not text: + return None + provider = self._get_embedding_provider() + if provider is None: + return None + try: + return provider.embed_query(text) + except Exception as e: + logger.debug(f"[ToolManager] query embedding failed: {e}") + return None + + def _ensure_mcp_tool_vectors(self) -> None: + """Incrementally embed MCP tools that are not yet cached.""" + # Snapshot to avoid concurrent-mutation while the async loader runs. + current = dict(self._mcp_tool_instances) + missing = [name for name in current if name not in self._mcp_tool_vectors] + if not missing: + return + + provider = self._get_embedding_provider() + if provider is None: + return + + with self._mcp_vector_lock: + # Re-check under lock: another thread may have filled these in. + missing = [name for name in current if name not in self._mcp_tool_vectors] + if not missing: + return + texts = [self._mcp_tool_embed_text(current[name]) for name in missing] + vectors = provider.embed_batch(texts) + for name, vec in zip(missing, vectors): + self._mcp_tool_vectors[name] = vec + + @staticmethod + def _mcp_tool_embed_text(tool) -> str: + """Build the text that represents an MCP tool for embedding.""" + name = getattr(tool, "name", "") or "" + description = getattr(tool, "description", "") or "" + return f"{name}: {description}".strip() + + def _get_embedding_provider(self): + """Lazily create and cache the shared embedding provider (or None).""" + if not self._embedding_provider_initialized: + try: + from agent.memory.embedding import create_default_embedding_provider + self._embedding_provider = create_default_embedding_provider() + except Exception as e: + logger.warning(f"[ToolManager] embedding provider init failed: {e}") + self._embedding_provider = None + self._embedding_provider_initialized = True + return self._embedding_provider + def create_tool(self, name: str) -> BaseTool: """ Get a new instance of a tool by name.