From 96b1fccf76c00b10aadecc251f3b72ce524b8cf4 Mon Sep 17 00:00:00 2001 From: fengyl07 <1059732235@qq.com> Date: Tue, 7 Jul 2026 16:00:31 +0800 Subject: [PATCH] feat(mcp): add stateless on-demand tool retrieval module --- agent/tools/mcp/tool_retrieval.py | 159 ++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 agent/tools/mcp/tool_retrieval.py diff --git a/agent/tools/mcp/tool_retrieval.py b/agent/tools/mcp/tool_retrieval.py new file mode 100644 index 00000000..1a9c8c0c --- /dev/null +++ b/agent/tools/mcp/tool_retrieval.py @@ -0,0 +1,159 @@ +# encoding:utf-8 +""" +On-demand MCP tool retrieval. + +Pure, stateless selection helpers used by the streaming executor to decide +which MCP tools to inject into a given LLM turn. Vector precompute + caching +live in ToolManager (the tool-lifecycle owner, a process-wide singleton); +only the context-aware selection lives here, because only the executor knows +the conversation context. + +Invariants (per maintainer review of the feature proposal): + * Built-in tools are never handled here — the caller injects them in full. + * Any failure / missing input returns None so the caller falls back to + full injection; tools must never be silently dropped. + * Selection is union-accumulated across turns by the caller (only-grows), + so a tool that already produced a tool_use in the message history can + never disappear from the schema mid-run (which would make Claude/MiniMax + raise a message-format error). +""" +import math +from typing import Dict, List, Optional, Sequence, Set + +try: + import numpy as np + _HAS_NUMPY = True +except ImportError: + _HAS_NUMPY = False + +# How many trailing messages to concatenate into the retrieval query. Tool +# needs drift across a multi-turn tool-call loop, so a single (initial) user +# query is not enough; a short recent window captures the drift without +# bloating the query with stale context. +DEFAULT_QUERY_MESSAGES = 5 + + +def build_retrieval_query(messages: list, max_messages: int = DEFAULT_QUERY_MESSAGES) -> str: + """Concatenate the text of the most recent messages into a retrieval query. + + Only ``text`` content blocks are kept; ``tool_use`` / ``tool_result`` blocks + are skipped so the query stays short and focused on natural-language intent + rather than large serialized tool payloads. + + Args: + messages: Claude-style message list, each ``{"role", "content"}`` where + content is either a string or a list of typed blocks. + max_messages: Size of the trailing window to consider. + + Returns: + A single string (possibly empty if no text is found). + """ + if not messages: + return "" + + parts: List[str] = [] + for message in messages[-max_messages:]: + content = message.get("content") if isinstance(message, dict) else None + if isinstance(content, str): + if content.strip(): + parts.append(content.strip()) + continue + if isinstance(content, list): + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "text": + text = block.get("text", "") + if isinstance(text, str) and text.strip(): + parts.append(text.strip()) + return "\n".join(parts) + + +def cosine_similarity(a: Sequence[float], b: Sequence[float]) -> float: + """Cosine similarity of two equal-length vectors; 0.0 on degenerate input.""" + if not a or not b or len(a) != len(b): + return 0.0 + dot = sum(x * y for x, y in zip(a, b)) + norm_a = math.sqrt(sum(x * x for x in a)) + norm_b = math.sqrt(sum(y * y for y in b)) + if norm_a == 0 or norm_b == 0: + return 0.0 + return dot / (norm_a * norm_b) + + +def select_mcp_tools( + query_vector: Optional[Sequence[float]], + tool_vectors: Dict[str, Sequence[float]], + top_k: int, + already_selected: Optional[Set[str]] = None, +) -> Optional[Set[str]]: + """Return the accumulated set of MCP tool names to inject this turn. + + Computes cosine similarity between ``query_vector`` and each candidate + tool vector, keeps the ``top_k`` best, and unions them with + ``already_selected`` so the injected set only ever grows within a run. + + Args: + query_vector: Embedding of the current retrieval query, or None. + tool_vectors: ``{mcp_tool_name: vector}`` for candidate MCP tools. + top_k: Max number of tools to add from this turn's ranking. + already_selected: Names accumulated in previous turns of this run. + + Returns: + The union set of tool names to inject, or None to signal + "fall back to full injection" (no query vector, empty/invalid index, + or any unexpected error). This function never raises. + """ + accumulated: Set[str] = set(already_selected) if already_selected else set() + + if not query_vector or not tool_vectors or top_k <= 0: + return None + + try: + expected_dim = len(query_vector) + # Only rank candidates whose vector dimensionality matches the query. + # A dimension mismatch means the index was built with a different + # embedding model; ranking across dims is meaningless. + candidates = { + name: vec + for name, vec in tool_vectors.items() + if vec and len(vec) == expected_dim + } + if not candidates: + return None + + ranked = _rank_by_similarity(query_vector, candidates) + for name, _score in ranked[:top_k]: + accumulated.add(name) + return accumulated + except Exception: + # Selection must never break the agent — fall back to full injection. + return None + + +def _rank_by_similarity( + query_vector: Sequence[float], + candidates: Dict[str, Sequence[float]], +) -> List[tuple]: + """Return ``[(name, score), ...]`` sorted by descending cosine similarity. + + Uses numpy when available (vectorized, matching the memory-search path), + with a pure-Python fallback so the feature works without numpy installed. + """ + names = list(candidates.keys()) + + if _HAS_NUMPY: + matrix = np.array([candidates[n] for n in names], dtype=np.float32) # (N, D) + q_vec = np.array(query_vector, dtype=np.float32) # (D,) + dots = matrix @ q_vec # (N,) + row_norms = np.linalg.norm(matrix, axis=1) # (N,) + q_norm = float(np.linalg.norm(q_vec)) + denominators = row_norms * q_norm + np.maximum(denominators, 1e-10, out=denominators) # avoid div-by-zero + sims = dots / denominators + order = np.argsort(sims)[::-1] + return [(names[i], float(sims[i])) for i in order] + + scored = [(n, cosine_similarity(query_vector, candidates[n])) for n in names] + scored.sort(key=lambda x: x[1], reverse=True) + return scored