mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-17 11:07:11 +08:00
feat(mcp): cache MCP tool vectors with lazy embedding in ToolManager
This commit is contained in:
@@ -71,6 +71,22 @@ class ToolManager:
|
|||||||
if not hasattr(self, '_mcp_active_configs'):
|
if not hasattr(self, '_mcp_active_configs'):
|
||||||
# server_name -> normalized config dict, for diff-based reload.
|
# server_name -> normalized config dict, for diff-based reload.
|
||||||
self._mcp_active_configs: dict = {}
|
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):
|
def load_tools(self, tools_dir: str = "", config_dict=None):
|
||||||
"""
|
"""
|
||||||
@@ -574,6 +590,91 @@ class ToolManager:
|
|||||||
|
|
||||||
return (sorted(added), sorted(removed))
|
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:
|
def create_tool(self, name: str) -> BaseTool:
|
||||||
"""
|
"""
|
||||||
Get a new instance of a tool by name.
|
Get a new instance of a tool by name.
|
||||||
|
|||||||
Reference in New Issue
Block a user