mirror of
https://github.com/zhayujie/chatgpt-on-wechat.git
synced 2026-07-17 11:07:11 +08:00
Merge pull request #2940 from fengyl07/feat/mcp-tool-retrieval
feat(mcp): on-demand vector retrieval for large MCP tool sets
This commit is contained in:
@@ -379,6 +379,12 @@ class AgentStreamExecutor:
|
|||||||
|
|
||||||
self._emit_event("agent_start")
|
self._emit_event("agent_start")
|
||||||
|
|
||||||
|
# Reset the run-scoped MCP tool-retrieval accumulator. On-demand tool
|
||||||
|
# retrieval only grows this set within a run, so a tool that already
|
||||||
|
# produced a tool_use never disappears from the schema mid-run (which
|
||||||
|
# would make Claude/MiniMax raise a message-format error).
|
||||||
|
self._retrieved_mcp_names = set()
|
||||||
|
|
||||||
final_response = ""
|
final_response = ""
|
||||||
turn = 0
|
turn = 0
|
||||||
|
|
||||||
@@ -702,6 +708,70 @@ class AgentStreamExecutor:
|
|||||||
|
|
||||||
return final_response
|
return final_response
|
||||||
|
|
||||||
|
def _select_tools_for_injection(self) -> list:
|
||||||
|
"""Decide which tools to inject into the current LLM turn.
|
||||||
|
|
||||||
|
Built-in tools are ALWAYS injected in full (skills and core flows hard
|
||||||
|
depend on them). MCP tools are also injected in full UNLESS on-demand
|
||||||
|
retrieval is enabled AND the MCP tool count exceeds the configured
|
||||||
|
threshold — then only the most relevant MCP tools are injected, unioned
|
||||||
|
with those already selected earlier in this run (only-grows, so a tool
|
||||||
|
that already produced a tool_use never vanishes from the schema).
|
||||||
|
|
||||||
|
Degrades safely: disabled feature, no embedding provider, embedding
|
||||||
|
failure, count below threshold, or any error → inject all tools. Tools
|
||||||
|
are never silently dropped.
|
||||||
|
"""
|
||||||
|
all_tools = list(self.tools.values())
|
||||||
|
try:
|
||||||
|
from config import conf
|
||||||
|
if not conf().get("mcp_tool_retrieval_enabled", False):
|
||||||
|
return all_tools
|
||||||
|
|
||||||
|
from agent.tools.mcp.mcp_tool import McpTool
|
||||||
|
mcp_tools = [t for t in all_tools if isinstance(t, McpTool)]
|
||||||
|
builtin_tools = [t for t in all_tools if not isinstance(t, McpTool)]
|
||||||
|
|
||||||
|
threshold = int(conf().get("mcp_tool_retrieval_threshold", 20) or 20)
|
||||||
|
if len(mcp_tools) <= threshold:
|
||||||
|
return all_tools
|
||||||
|
|
||||||
|
top_k = int(conf().get("mcp_tool_retrieval_top_k", 10) or 10)
|
||||||
|
|
||||||
|
from agent.tools import ToolManager
|
||||||
|
from agent.tools.mcp.tool_retrieval import (
|
||||||
|
build_retrieval_query,
|
||||||
|
select_mcp_tools,
|
||||||
|
)
|
||||||
|
|
||||||
|
tm = ToolManager()
|
||||||
|
tool_vectors = tm.get_mcp_tool_vectors()
|
||||||
|
query = build_retrieval_query(self.messages)
|
||||||
|
query_vector = tm.embed_query(query)
|
||||||
|
|
||||||
|
selected = select_mcp_tools(
|
||||||
|
query_vector,
|
||||||
|
tool_vectors,
|
||||||
|
top_k,
|
||||||
|
getattr(self, "_retrieved_mcp_names", set()),
|
||||||
|
)
|
||||||
|
if selected is None:
|
||||||
|
# No provider / empty index / error → full injection.
|
||||||
|
return all_tools
|
||||||
|
|
||||||
|
# Persist the accumulated selection for subsequent turns.
|
||||||
|
self._retrieved_mcp_names = selected
|
||||||
|
|
||||||
|
selected_mcp = [t for t in mcp_tools if t.name in selected]
|
||||||
|
logger.info(
|
||||||
|
f"[ToolRetrieval] Injecting {len(builtin_tools)} built-in + "
|
||||||
|
f"{len(selected_mcp)}/{len(mcp_tools)} MCP tool(s) (top_k={top_k})"
|
||||||
|
)
|
||||||
|
return builtin_tools + selected_mcp
|
||||||
|
except Exception as e:
|
||||||
|
logger.debug(f"[ToolRetrieval] full injection (retrieval skipped): {e}")
|
||||||
|
return all_tools
|
||||||
|
|
||||||
def _call_llm_stream(self, retry_on_empty=True, retry_count=0, max_retries=3,
|
def _call_llm_stream(self, retry_on_empty=True, retry_count=0, max_retries=3,
|
||||||
_overflow_retry: bool = False) -> Tuple[str, List[Dict]]:
|
_overflow_retry: bool = False) -> Tuple[str, List[Dict]]:
|
||||||
"""
|
"""
|
||||||
@@ -742,7 +812,7 @@ class AgentStreamExecutor:
|
|||||||
tools_schema = None
|
tools_schema = None
|
||||||
if self.tools:
|
if self.tools:
|
||||||
tools_schema = []
|
tools_schema = []
|
||||||
for tool in self.tools.values():
|
for tool in self._select_tools_for_injection():
|
||||||
input_schema = tool.params
|
input_schema = tool.params
|
||||||
try:
|
try:
|
||||||
dynamic = (tool.get_json_schema() or {}).get("parameters") or {}
|
dynamic = (tool.get_json_schema() or {}).get("parameters") or {}
|
||||||
|
|||||||
159
agent/tools/mcp/tool_retrieval.py
Normal file
159
agent/tools/mcp/tool_retrieval.py
Normal file
@@ -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
|
||||||
@@ -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.
|
||||||
|
|||||||
@@ -41,5 +41,8 @@
|
|||||||
"enable_thinking": false,
|
"enable_thinking": false,
|
||||||
"reasoning_effort": "high",
|
"reasoning_effort": "high",
|
||||||
"knowledge": true,
|
"knowledge": true,
|
||||||
"self_evolution_enabled": true
|
"self_evolution_enabled": true,
|
||||||
|
"mcp_tool_retrieval_enabled": false,
|
||||||
|
"mcp_tool_retrieval_threshold": 20,
|
||||||
|
"mcp_tool_retrieval_top_k": 10
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -269,6 +269,13 @@ available_setting = {
|
|||||||
"deep_dream_enabled": True, # scheduled deep dream switch; manual /memory dream is unaffected
|
"deep_dream_enabled": True, # scheduled deep dream switch; manual /memory dream is unaffected
|
||||||
"skill": {}, # Per-skill runtime config; nested keys flatten to SKILL_<NAME>_<KEY> env vars at startup
|
"skill": {}, # Per-skill runtime config; nested keys flatten to SKILL_<NAME>_<KEY> env vars at startup
|
||||||
"mcp_servers": [], # MCP server list; each entry supports type "stdio" (local process) or "sse" (remote URL)
|
"mcp_servers": [], # MCP server list; each entry supports type "stdio" (local process) or "sse" (remote URL)
|
||||||
|
# On-demand MCP tool retrieval: when many MCP tools are connected, inject
|
||||||
|
# only the most query-relevant ones instead of all of them. Built-in tools
|
||||||
|
# are always injected in full; degrades to full injection when disabled,
|
||||||
|
# below threshold, or when no embedding provider is available.
|
||||||
|
"mcp_tool_retrieval_enabled": False, # switch for on-demand MCP tool retrieval
|
||||||
|
"mcp_tool_retrieval_threshold": 20, # only retrieve when MCP tool count exceeds this
|
||||||
|
"mcp_tool_retrieval_top_k": 10, # max relevant MCP tools injected per turn
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
253
tests/test_mcp_tool_retrieval.py
Normal file
253
tests/test_mcp_tool_retrieval.py
Normal file
@@ -0,0 +1,253 @@
|
|||||||
|
# encoding:utf-8
|
||||||
|
"""
|
||||||
|
Unit tests for on-demand MCP tool retrieval.
|
||||||
|
|
||||||
|
Covers the invariants the maintainer asked for in the feature review:
|
||||||
|
* below threshold / degrade paths behave exactly like today (full injection),
|
||||||
|
* the injected MCP tool set only ever grows within a run (never shrinks),
|
||||||
|
plus the pure selection helpers (query building, cosine, top-k, fallbacks).
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||||
|
|
||||||
|
from agent.tools.mcp.tool_retrieval import (
|
||||||
|
build_retrieval_query,
|
||||||
|
cosine_similarity,
|
||||||
|
select_mcp_tools,
|
||||||
|
)
|
||||||
|
from agent.tools.mcp.mcp_tool import McpTool
|
||||||
|
|
||||||
|
|
||||||
|
def _mcp_tool(name, description=""):
|
||||||
|
"""Build a real McpTool without needing a live MCP client."""
|
||||||
|
return McpTool(
|
||||||
|
client=None,
|
||||||
|
tool_schema={"name": name, "description": description, "inputSchema": {}},
|
||||||
|
server_name="test-server",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeBuiltinTool:
|
||||||
|
"""Minimal stand-in for a built-in BaseTool."""
|
||||||
|
|
||||||
|
def __init__(self, name):
|
||||||
|
self.name = name
|
||||||
|
self.description = f"builtin {name}"
|
||||||
|
self.params = {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
def get_json_schema(self):
|
||||||
|
return {}
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeToolManager:
|
||||||
|
"""Controls the vectors/query embeddings seen by the executor."""
|
||||||
|
|
||||||
|
def __init__(self, tool_vectors, query_vectors):
|
||||||
|
self._tool_vectors = tool_vectors
|
||||||
|
self._query_vectors = list(query_vectors)
|
||||||
|
self._call = 0
|
||||||
|
|
||||||
|
def get_mcp_tool_vectors(self):
|
||||||
|
return dict(self._tool_vectors)
|
||||||
|
|
||||||
|
def embed_query(self, text):
|
||||||
|
if self._call < len(self._query_vectors):
|
||||||
|
vec = self._query_vectors[self._call]
|
||||||
|
else:
|
||||||
|
vec = self._query_vectors[-1] if self._query_vectors else None
|
||||||
|
self._call += 1
|
||||||
|
return vec
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
# Pure helper: build_retrieval_query
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestBuildRetrievalQuery(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_empty_messages(self):
|
||||||
|
self.assertEqual(build_retrieval_query([]), "")
|
||||||
|
|
||||||
|
def test_extracts_text_blocks(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": [{"type": "text", "text": "hello"}]},
|
||||||
|
{"role": "assistant", "content": [{"type": "text", "text": "world"}]},
|
||||||
|
]
|
||||||
|
self.assertEqual(build_retrieval_query(messages), "hello\nworld")
|
||||||
|
|
||||||
|
def test_skips_tool_result_and_tool_use(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": [{"type": "text", "text": "do it"}]},
|
||||||
|
{"role": "assistant", "content": [
|
||||||
|
{"type": "tool_use", "name": "read", "input": {}},
|
||||||
|
]},
|
||||||
|
{"role": "user", "content": [
|
||||||
|
{"type": "tool_result", "content": "huge payload " * 100},
|
||||||
|
]},
|
||||||
|
]
|
||||||
|
self.assertEqual(build_retrieval_query(messages), "do it")
|
||||||
|
|
||||||
|
def test_string_content_supported(self):
|
||||||
|
messages = [{"role": "user", "content": "plain string"}]
|
||||||
|
self.assertEqual(build_retrieval_query(messages), "plain string")
|
||||||
|
|
||||||
|
def test_respects_recent_window(self):
|
||||||
|
messages = [
|
||||||
|
{"role": "user", "content": [{"type": "text", "text": f"m{i}"}]}
|
||||||
|
for i in range(10)
|
||||||
|
]
|
||||||
|
# Only the last 3 messages should be kept.
|
||||||
|
self.assertEqual(build_retrieval_query(messages, max_messages=3), "m7\nm8\nm9")
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
# Pure helper: cosine_similarity
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestCosineSimilarity(unittest.TestCase):
|
||||||
|
|
||||||
|
def test_identical_vectors(self):
|
||||||
|
self.assertAlmostEqual(cosine_similarity([1.0, 0.0], [1.0, 0.0]), 1.0)
|
||||||
|
|
||||||
|
def test_orthogonal_vectors(self):
|
||||||
|
self.assertAlmostEqual(cosine_similarity([1.0, 0.0], [0.0, 1.0]), 0.0)
|
||||||
|
|
||||||
|
def test_degenerate_inputs(self):
|
||||||
|
self.assertEqual(cosine_similarity([], [1.0]), 0.0)
|
||||||
|
self.assertEqual(cosine_similarity([0.0, 0.0], [1.0, 1.0]), 0.0)
|
||||||
|
self.assertEqual(cosine_similarity([1.0], [1.0, 0.0]), 0.0)
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
# Pure helper: select_mcp_tools
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestSelectMcpTools(unittest.TestCase):
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.vectors = {
|
||||||
|
"a": [1.0, 0.0, 0.0],
|
||||||
|
"b": [0.0, 1.0, 0.0],
|
||||||
|
"c": [0.0, 0.0, 1.0],
|
||||||
|
"d": [0.9, 0.1, 0.0],
|
||||||
|
}
|
||||||
|
|
||||||
|
def test_returns_top_k(self):
|
||||||
|
selected = select_mcp_tools([1.0, 0.0, 0.0], self.vectors, top_k=2,
|
||||||
|
already_selected=set())
|
||||||
|
self.assertEqual(selected, {"a", "d"})
|
||||||
|
|
||||||
|
def test_union_only_grows_across_turns(self):
|
||||||
|
"""The core invariant: a later turn never drops earlier selections."""
|
||||||
|
first = select_mcp_tools([1.0, 0.0, 0.0], self.vectors, top_k=2,
|
||||||
|
already_selected=set())
|
||||||
|
second = select_mcp_tools([0.0, 1.0, 0.0], self.vectors, top_k=1,
|
||||||
|
already_selected=first)
|
||||||
|
self.assertTrue(first.issubset(second))
|
||||||
|
self.assertIn("b", second)
|
||||||
|
|
||||||
|
def test_none_query_vector_falls_back(self):
|
||||||
|
self.assertIsNone(select_mcp_tools(None, self.vectors, top_k=2,
|
||||||
|
already_selected=set()))
|
||||||
|
|
||||||
|
def test_empty_index_falls_back(self):
|
||||||
|
self.assertIsNone(select_mcp_tools([1.0, 0.0, 0.0], {}, top_k=2,
|
||||||
|
already_selected=set()))
|
||||||
|
|
||||||
|
def test_dimension_mismatch_falls_back(self):
|
||||||
|
mismatched = {"a": [1.0, 0.0]} # dim 2 vs query dim 3
|
||||||
|
self.assertIsNone(select_mcp_tools([1.0, 0.0, 0.0], mismatched, top_k=2,
|
||||||
|
already_selected=set()))
|
||||||
|
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
# Executor integration: AgentStream._select_tools_for_injection
|
||||||
|
# --------------------------------------------------------------------------
|
||||||
|
|
||||||
|
class TestSelectToolsForInjection(unittest.TestCase):
|
||||||
|
"""Exercise the executor decision without spinning up a real agent."""
|
||||||
|
|
||||||
|
def _make_self(self, mcp_count, builtins=("read", "write", "bash")):
|
||||||
|
from types import SimpleNamespace
|
||||||
|
tools = {}
|
||||||
|
for name in builtins:
|
||||||
|
tools[name] = _FakeBuiltinTool(name)
|
||||||
|
for i in range(mcp_count):
|
||||||
|
name = f"mcp_{i}"
|
||||||
|
tools[name] = _mcp_tool(name, f"tool number {i}")
|
||||||
|
return SimpleNamespace(
|
||||||
|
tools=tools,
|
||||||
|
messages=[{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
|
||||||
|
_retrieved_mcp_names=set(),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _call(self, fake_self):
|
||||||
|
from agent.protocol.agent_stream import AgentStreamExecutor
|
||||||
|
return AgentStreamExecutor._select_tools_for_injection(fake_self)
|
||||||
|
|
||||||
|
def _conf(self, **overrides):
|
||||||
|
cfg = {
|
||||||
|
"mcp_tool_retrieval_enabled": True,
|
||||||
|
"mcp_tool_retrieval_threshold": 20,
|
||||||
|
"mcp_tool_retrieval_top_k": 2,
|
||||||
|
}
|
||||||
|
cfg.update(overrides)
|
||||||
|
return cfg
|
||||||
|
|
||||||
|
def test_disabled_returns_all_tools(self):
|
||||||
|
fake = self._make_self(mcp_count=50)
|
||||||
|
with patch("config.conf", return_value=self._conf(mcp_tool_retrieval_enabled=False)):
|
||||||
|
result = self._call(fake)
|
||||||
|
self.assertEqual(len(result), len(fake.tools))
|
||||||
|
|
||||||
|
def test_below_threshold_returns_all_tools(self):
|
||||||
|
"""Maintainer scenario 1: below threshold → behavior unchanged."""
|
||||||
|
fake = self._make_self(mcp_count=5) # <= threshold 20
|
||||||
|
with patch("config.conf", return_value=self._conf()):
|
||||||
|
result = self._call(fake)
|
||||||
|
self.assertEqual(len(result), len(fake.tools))
|
||||||
|
|
||||||
|
def test_degrade_no_provider_returns_all_tools(self):
|
||||||
|
"""Maintainer scenario 2: no embedding provider → full injection."""
|
||||||
|
fake = self._make_self(mcp_count=25) # > threshold
|
||||||
|
fake_tm = _FakeToolManager(tool_vectors={}, query_vectors=[None])
|
||||||
|
with patch("config.conf", return_value=self._conf()), \
|
||||||
|
patch("agent.tools.ToolManager", return_value=fake_tm):
|
||||||
|
result = self._call(fake)
|
||||||
|
self.assertEqual(len(result), len(fake.tools))
|
||||||
|
|
||||||
|
def test_builtins_always_injected_and_set_grows(self):
|
||||||
|
"""Maintainer scenario 3: multi-turn MCP set only grows; builtins stay."""
|
||||||
|
fake = self._make_self(mcp_count=25)
|
||||||
|
# Deterministic vectors: mcp_0 wins turn 1, mcp_1 wins turn 2.
|
||||||
|
tool_vectors = {f"mcp_{i}": [0.1, 0.1, 0.1] for i in range(25)}
|
||||||
|
tool_vectors["mcp_0"] = [1.0, 0.0, 0.0]
|
||||||
|
tool_vectors["mcp_1"] = [0.0, 1.0, 0.0]
|
||||||
|
fake_tm = _FakeToolManager(
|
||||||
|
tool_vectors=tool_vectors,
|
||||||
|
query_vectors=[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]],
|
||||||
|
)
|
||||||
|
with patch("config.conf", return_value=self._conf(mcp_tool_retrieval_top_k=1)), \
|
||||||
|
patch("agent.tools.ToolManager", return_value=fake_tm):
|
||||||
|
result1 = self._call(fake)
|
||||||
|
names1 = {t.name for t in result1}
|
||||||
|
result2 = self._call(fake)
|
||||||
|
names2 = {t.name for t in result2}
|
||||||
|
|
||||||
|
# Built-in tools present in both turns.
|
||||||
|
for b in ("read", "write", "bash"):
|
||||||
|
self.assertIn(b, names1)
|
||||||
|
self.assertIn(b, names2)
|
||||||
|
# Turn 1 selected mcp_0; turn 2 must still contain it (only-grows).
|
||||||
|
self.assertIn("mcp_0", names1)
|
||||||
|
self.assertIn("mcp_0", names2)
|
||||||
|
self.assertIn("mcp_1", names2)
|
||||||
|
self.assertTrue(fake._retrieved_mcp_names >= {"mcp_0", "mcp_1"})
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user