feat(mcp): cache MCP tool vectors with lazy embedding in ToolManager

This commit is contained in:
fengyl07
2026-07-07 16:00:55 +08:00
parent 96b1fccf76
commit bf0831a664

View File

@@ -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.