feat(mcp): add stateless on-demand tool retrieval module

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