mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(mcp): preserve native tools in semantic filter hook (#26650)
* fix(mcp): preserve native tools in semantic filter hook The SemanticToolFilterHook.async_pre_call_hook passed ALL tools (MCP + native) to filter_tools(), which only knows MCP-registered tool names. Native tools silently failed the name match in _get_tools_by_names() and were dropped from the request. Fix: partition tools into native and MCP-registered before filtering. Run the semantic filter only on MCP tools, then merge native tools back unconditionally. Changes: - Robust _is_mcp_tool() using shape-based detection for OpenAI-format dicts, safe regardless of future _extract_tool_info changes - Single-pass partition loop (no double _is_mcp_tool calls) - Preserve native tools in MCP expansion path (mixed requests) - Track MCP expansion to prevent expanded tools bypassing filtering - filter_stats reports MCP-only counts for accurate metrics - Extracted _emit_filter_metadata() helper - Skip spurious filter headers for all-native tool requests Closes #26212 * remove stale docstring note referencing tools_expanded_from_mcp * fix: handle Responses API name collision and preserve tool ordering - Classify Responses API tools ({type: 'function', name: '...'}) as native to prevent name collisions with MCP canonical names - Preserve original request tool ordering using id()-based merge instead of naive native+mcp concatenation - Add 2 regression tests: name collision and ordering preservation * style: apply black formatting * fix(mcp): harden semantic filter — preserve all native tool formats, safe metadata access, graceful expansion failure, name-based merge * lint: suppress PLR0915 on async_pre_call_hook (matches codebase convention) * ci: retrigger checks after rebase onto litellm_internal_staging
This commit is contained in:
parent
82aa40b7b7
commit
438c825bd4
2 changed files with 588 additions and 52 deletions
|
|
@ -123,12 +123,78 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
|
||||
return openai_tools_as_dicts
|
||||
|
||||
def _is_mcp_tool(self, tool: Any) -> bool:
|
||||
"""
|
||||
Check whether *tool* is registered in the MCP semantic router.
|
||||
|
||||
Classification strategy:
|
||||
1. Standard OpenAI Chat Completions-format dicts
|
||||
(``{"type": "function", "function": {"name": ...}}``) are
|
||||
**always** classified as native — detected by shape (both
|
||||
``"function"`` and ``"type"`` keys present).
|
||||
2. Responses API function tools
|
||||
(``{"type": "function", "name": ..., ...}``) are **always**
|
||||
classified as native — detected by ``type == "function"``
|
||||
with a top-level string ``name`` but no nested ``"function"``
|
||||
dict. This prevents name collisions with MCP canonicals.
|
||||
3. Everything else (MCP ``Tool`` objects, flat dicts, etc.) is
|
||||
looked up by name in ``self.filter._tool_map``.
|
||||
"""
|
||||
# Chat Completions format: {"type": "function", "function": {...}}
|
||||
if isinstance(tool, dict) and "function" in tool and "type" in tool:
|
||||
return False # Standard OpenAI function tool — always native
|
||||
# Responses API format: {"type": "function", "name": "...", ...}
|
||||
if (
|
||||
isinstance(tool, dict)
|
||||
and tool.get("type") == "function"
|
||||
and isinstance(tool.get("name"), str)
|
||||
):
|
||||
return False # Responses API function tool — always native
|
||||
name, _ = self.filter._extract_tool_info(tool)
|
||||
return bool(name) and name in self.filter._tool_map
|
||||
|
||||
def _get_metadata_variable_name(self, data: dict) -> str:
|
||||
if "litellm_metadata" in data:
|
||||
return "litellm_metadata"
|
||||
return "metadata"
|
||||
|
||||
async def async_pre_call_hook(
|
||||
def _emit_filter_metadata(
|
||||
self,
|
||||
data: dict,
|
||||
mcp_tools: list,
|
||||
filtered_mcp_tools: list,
|
||||
native_tools: list,
|
||||
filtered_tools: list,
|
||||
) -> None:
|
||||
"""
|
||||
Emit response-header metadata when MCP tools were filtered.
|
||||
|
||||
Stats report MCP-only counts so downstream consumers see accurate
|
||||
semantic filter metrics (not inflated by native tools that always
|
||||
pass through). Skips metadata entirely for purely-native requests
|
||||
to avoid spurious ``N->N`` headers.
|
||||
"""
|
||||
if mcp_tools:
|
||||
filter_stats = f"{len(mcp_tools)}->{len(filtered_mcp_tools)}"
|
||||
tool_names_csv = self._get_tool_names_csv(filtered_mcp_tools)
|
||||
|
||||
_metadata_variable_name = self._get_metadata_variable_name(data)
|
||||
metadata = data.setdefault(_metadata_variable_name, {})
|
||||
metadata["litellm_semantic_filter_stats"] = filter_stats
|
||||
metadata["litellm_semantic_filter_tools"] = tool_names_csv
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Semantic tool filter: {filter_stats} MCP tools "
|
||||
f"({len(native_tools)} native preserved, "
|
||||
f"{len(filtered_tools)} total)"
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.info(
|
||||
f"Semantic tool filter: all {len(native_tools)} tools "
|
||||
f"are native, no MCP filtering applied"
|
||||
)
|
||||
|
||||
async def async_pre_call_hook( # noqa: PLR0915
|
||||
self,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
cache: "DualCache",
|
||||
|
|
@ -163,30 +229,54 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
verbose_proxy_logger.debug("No tools in request, skipping semantic filter")
|
||||
return None
|
||||
|
||||
original_tool_count = len(tools)
|
||||
|
||||
# Check for MCP references (server_url="litellm_proxy") and expand them
|
||||
# Check for MCP references (server_url="litellm_proxy") and expand them.
|
||||
# Expanded MCP tools are in OpenAI nested format which
|
||||
# filter_tools/_extract_tool_info cannot name-match, so we skip
|
||||
# semantic filtering and return early — expanded tools already
|
||||
# went through access-control filtering in
|
||||
# _process_mcp_tools_to_openai_format.
|
||||
if self._should_expand_mcp_tools(tools):
|
||||
verbose_proxy_logger.debug(
|
||||
"Detected litellm_proxy MCP references, expanding before semantic filtering"
|
||||
)
|
||||
|
||||
try:
|
||||
# Preserve all non-MCP-reference tools before expansion.
|
||||
# _expand_mcp_tools only returns expanded MCP tools;
|
||||
# non-MCP tools in the original list are discarded by
|
||||
# _parse_mcp_tools. Without this, native tools in
|
||||
# mixed requests (MCP refs + native) are silently lost.
|
||||
# We detect MCP references by type=="mcp" (the only
|
||||
# shape _should_expand_mcp_tools triggers on), keeping
|
||||
# both Chat Completions and Responses API native formats.
|
||||
native_tools_before_expand = [
|
||||
t
|
||||
for t in tools
|
||||
if not (isinstance(t, dict) and t.get("type") == "mcp")
|
||||
]
|
||||
|
||||
expanded_tools = await self._expand_mcp_tools(tools, user_api_key_dict)
|
||||
|
||||
if not expanded_tools:
|
||||
if native_tools_before_expand:
|
||||
data["tools"] = native_tools_before_expand
|
||||
verbose_proxy_logger.warning(
|
||||
"No MCP tools expanded, preserving "
|
||||
f"{len(native_tools_before_expand)} native tools"
|
||||
)
|
||||
return data
|
||||
verbose_proxy_logger.warning(
|
||||
"No tools expanded from MCP references"
|
||||
)
|
||||
return None
|
||||
|
||||
data["tools"] = native_tools_before_expand + expanded_tools
|
||||
verbose_proxy_logger.info(
|
||||
f"Expanded {len(tools)} MCP reference(s) to {len(expanded_tools)} tools"
|
||||
f"Expanded MCP references to {len(expanded_tools)} tools "
|
||||
f"({len(native_tools_before_expand)} native preserved), "
|
||||
f"skipping semantic filter (OpenAI nested format)"
|
||||
)
|
||||
|
||||
# Update tools for filtering
|
||||
tools = expanded_tools
|
||||
original_tool_count = len(tools)
|
||||
return data
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -218,33 +308,71 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
)
|
||||
return None
|
||||
|
||||
# ── Partition tools into MCP-registered and native ──
|
||||
native_tools: List[Any] = []
|
||||
mcp_tools: List[Any] = []
|
||||
mcp_indices: set = set()
|
||||
for i, t in enumerate(tools):
|
||||
if self._is_mcp_tool(t):
|
||||
mcp_tools.append(t)
|
||||
mcp_indices.add(i)
|
||||
else:
|
||||
native_tools.append(t)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Applying semantic filter to {len(tools)} tools "
|
||||
f"with query: '{user_query[:50]}...'"
|
||||
f"Applying semantic filter: {len(mcp_tools)} MCP tools, "
|
||||
f"{len(native_tools)} native tools, "
|
||||
f"query: '{user_query[:50]}...'"
|
||||
)
|
||||
|
||||
# Filter tools semantically
|
||||
filtered_tools = await self.filter.filter_tools(
|
||||
query=user_query,
|
||||
available_tools=tools, # type: ignore
|
||||
)
|
||||
# ── Filter only MCP tools semantically ──
|
||||
if mcp_tools:
|
||||
filtered_mcp_tools = await self.filter.filter_tools(
|
||||
query=user_query,
|
||||
available_tools=mcp_tools, # type: ignore
|
||||
)
|
||||
else:
|
||||
filtered_mcp_tools = []
|
||||
|
||||
# Always update tools and emit header (even if count unchanged)
|
||||
# Build set of surviving MCP tool names for order-preserving
|
||||
# merge. Uses names (not id()) so the merge stays correct
|
||||
# even if filter_tools internally copies tool objects.
|
||||
filtered_mcp_names: set = set()
|
||||
for t in filtered_mcp_tools:
|
||||
name, _ = self.filter._extract_tool_info(t)
|
||||
if name:
|
||||
filtered_mcp_names.add(name)
|
||||
|
||||
# Merge preserving original request order — some LLMs are
|
||||
# sensitive to the position tools appear in.
|
||||
filtered_tools: List[Any] = []
|
||||
for i, t in enumerate(tools):
|
||||
if i in mcp_indices:
|
||||
name, _ = self.filter._extract_tool_info(t)
|
||||
if name in filtered_mcp_names:
|
||||
filtered_tools.append(t)
|
||||
else:
|
||||
filtered_tools.append(t)
|
||||
|
||||
# Always update tools
|
||||
data["tools"] = filtered_tools
|
||||
|
||||
# Store filter stats and tool names for response header
|
||||
filter_stats = f"{original_tool_count}->{len(filtered_tools)}"
|
||||
tool_names_csv = self._get_tool_names_csv(filtered_tools)
|
||||
|
||||
_metadata_variable_name = self._get_metadata_variable_name(data)
|
||||
data[_metadata_variable_name][
|
||||
"litellm_semantic_filter_stats"
|
||||
] = filter_stats
|
||||
data[_metadata_variable_name][
|
||||
"litellm_semantic_filter_tools"
|
||||
] = tool_names_csv
|
||||
|
||||
verbose_proxy_logger.info(f"Semantic tool filter: {filter_stats} tools")
|
||||
# Emit metadata — wrapped separately so a metadata error
|
||||
# does not cause the outer except to return None after
|
||||
# data["tools"] has already been mutated.
|
||||
try:
|
||||
self._emit_filter_metadata(
|
||||
data=data,
|
||||
mcp_tools=mcp_tools,
|
||||
filtered_mcp_tools=filtered_mcp_tools,
|
||||
native_tools=native_tools,
|
||||
filtered_tools=filtered_tools,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to emit semantic filter metadata: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
|
|
@ -266,7 +394,7 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
from litellm.constants import MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH
|
||||
|
||||
_metadata_variable_name = self._get_metadata_variable_name(data)
|
||||
metadata = data[_metadata_variable_name]
|
||||
metadata = data.get(_metadata_variable_name, {})
|
||||
|
||||
filter_stats = metadata.get("litellm_semantic_filter_stats")
|
||||
if not filter_stats:
|
||||
|
|
|
|||
|
|
@ -452,6 +452,428 @@ async def test_semantic_filter_hook_skips_no_tools():
|
|||
print("✅ Hook correctly skips requests without tools")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_preserves_native_tools():
|
||||
"""
|
||||
Regression test: mixed MCP + native tools.
|
||||
|
||||
Given: 5 MCP tools (registered in _tool_map) + 2 native OpenAI-format
|
||||
function tools (not in _tool_map)
|
||||
When: The hook filters tools
|
||||
Then: The native tools must survive unconditionally, and only MCP
|
||||
tools go through the semantic filter.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
)
|
||||
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse
|
||||
|
||||
mock_router = Mock()
|
||||
|
||||
def mock_embedding_sync(*args, **kwargs):
|
||||
return EmbeddingResponse(
|
||||
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
|
||||
model="text-embedding-3-small",
|
||||
object="list",
|
||||
usage={"prompt_tokens": 10, "total_tokens": 10},
|
||||
)
|
||||
|
||||
async def mock_embedding_async(*args, **kwargs):
|
||||
return mock_embedding_sync()
|
||||
|
||||
mock_router.embedding = mock_embedding_sync
|
||||
mock_router.aembedding = mock_embedding_async
|
||||
|
||||
filter_instance = SemanticMCPToolFilter(
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=mock_router,
|
||||
top_k=2,
|
||||
similarity_threshold=0.3,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
# --- MCP tools (registered in the semantic router) ---
|
||||
mcp_tools = [
|
||||
MCPTool(
|
||||
name=f"mcp_tool_{i}",
|
||||
description=f"MCP tool {i}",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
for i in range(5)
|
||||
]
|
||||
filter_instance._build_router(mcp_tools)
|
||||
|
||||
# --- Native OpenAI-format function tools (NOT in _tool_map) ---
|
||||
native_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search_web",
|
||||
"description": "Search the web",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
# Combine: MCP tools + native tools
|
||||
all_tools = list(mcp_tools) + native_tools
|
||||
|
||||
hook = SemanticToolFilterHook(filter_instance)
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "What is the weather?"}],
|
||||
"tools": all_tools,
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
result = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert result is not None, "Hook should return modified data"
|
||||
filtered = result["tools"]
|
||||
|
||||
# Native tools must survive
|
||||
native_in_result = [
|
||||
t for t in filtered if isinstance(t, dict) and t.get("type") == "function"
|
||||
]
|
||||
assert (
|
||||
len(native_in_result) == 2
|
||||
), f"Both native tools must survive, got {len(native_in_result)}"
|
||||
|
||||
# MCP tools should be filtered (top_k=2)
|
||||
mcp_in_result = [t for t in filtered if not isinstance(t, dict)]
|
||||
assert (
|
||||
len(mcp_in_result) <= 2
|
||||
), f"MCP tools should be filtered to top_k=2, got {len(mcp_in_result)}"
|
||||
|
||||
# Total should be native + filtered MCP
|
||||
assert len(filtered) <= 4, f"Expected at most 4 tools, got {len(filtered)}"
|
||||
|
||||
# Filter stats should be emitted (MCP tools were present)
|
||||
assert "litellm_semantic_filter_stats" in result["metadata"]
|
||||
|
||||
# Stats should report MCP-only counts, not inflated with native tools
|
||||
stats = result["metadata"]["litellm_semantic_filter_stats"]
|
||||
mcp_before, mcp_after = stats.split("->")
|
||||
assert (
|
||||
int(mcp_before) == 5
|
||||
), f"Stats 'from' should be MCP count (5), got {mcp_before}"
|
||||
assert int(mcp_after) == len(
|
||||
mcp_in_result
|
||||
), f"Stats 'to' should match filtered MCP count, got {mcp_after}"
|
||||
|
||||
print(
|
||||
f"✅ Hook preserves native tools: {len(all_tools)} -> {len(filtered)} "
|
||||
f"({len(native_in_result)} native + {len(mcp_in_result)} MCP), "
|
||||
f"stats={stats}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_all_native_tools():
|
||||
"""
|
||||
Regression test: all-native request.
|
||||
|
||||
Given: Only native OpenAI-format function tools (none registered in
|
||||
the MCP semantic router)
|
||||
When: The hook processes the request
|
||||
Then: All tools pass through, and NO spurious semantic filter response
|
||||
headers are emitted (no litellm_semantic_filter_stats in metadata).
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
)
|
||||
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
|
||||
|
||||
mock_router = Mock()
|
||||
filter_instance = SemanticMCPToolFilter(
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=mock_router,
|
||||
top_k=3,
|
||||
similarity_threshold=0.3,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
# Build router with some MCP tools (so tool_router is not None)
|
||||
mcp_tools = [
|
||||
MCPTool(
|
||||
name="some_mcp_tool",
|
||||
description="An MCP tool",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
]
|
||||
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse
|
||||
|
||||
def mock_embedding_sync(*args, **kwargs):
|
||||
return EmbeddingResponse(
|
||||
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
|
||||
model="text-embedding-3-small",
|
||||
object="list",
|
||||
usage={"prompt_tokens": 10, "total_tokens": 10},
|
||||
)
|
||||
|
||||
async def mock_embedding_async(*args, **kwargs):
|
||||
return mock_embedding_sync()
|
||||
|
||||
mock_router.embedding = mock_embedding_sync
|
||||
mock_router.aembedding = mock_embedding_async
|
||||
|
||||
filter_instance._build_router(mcp_tools)
|
||||
|
||||
# --- Only native tools in the request ---
|
||||
native_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": f"native_func_{i}",
|
||||
"description": f"Native function {i}",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
hook = SemanticToolFilterHook(filter_instance)
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"tools": native_tools,
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
result = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert result is not None, "Hook should return modified data"
|
||||
filtered = result["tools"]
|
||||
|
||||
# All native tools must pass through
|
||||
assert (
|
||||
len(filtered) == 3
|
||||
), f"All 3 native tools must pass through, got {len(filtered)}"
|
||||
|
||||
# No spurious semantic filter stats (P2 fix)
|
||||
assert (
|
||||
"litellm_semantic_filter_stats" not in result["metadata"]
|
||||
), "Should NOT emit semantic filter stats for all-native-tool requests"
|
||||
|
||||
print(
|
||||
f"✅ Hook passes through all {len(filtered)} native tools, "
|
||||
f"no spurious filter headers emitted"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_responses_api_name_collision():
|
||||
"""
|
||||
Regression test: Responses API native tool with MCP-matching name.
|
||||
|
||||
Given: A Responses-API native tool whose top-level ``name`` collides
|
||||
with an MCP canonical name in ``_tool_map``
|
||||
When: The hook classifies tools
|
||||
Then: The native tool must NOT be sent to the semantic filter, even
|
||||
though its name matches an MCP canonical.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
)
|
||||
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse
|
||||
|
||||
mock_router = Mock()
|
||||
|
||||
def mock_embedding_sync(*args, **kwargs):
|
||||
return EmbeddingResponse(
|
||||
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
|
||||
model="text-embedding-3-small",
|
||||
object="list",
|
||||
usage={"prompt_tokens": 10, "total_tokens": 10},
|
||||
)
|
||||
|
||||
async def mock_embedding_async(*args, **kwargs):
|
||||
return mock_embedding_sync()
|
||||
|
||||
mock_router.embedding = mock_embedding_sync
|
||||
mock_router.aembedding = mock_embedding_async
|
||||
|
||||
filter_instance = SemanticMCPToolFilter(
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=mock_router,
|
||||
top_k=2,
|
||||
similarity_threshold=0.3,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
# Register an MCP tool with name "github-search"
|
||||
mcp_tools = [
|
||||
MCPTool(
|
||||
name="github-search",
|
||||
description="Search GitHub repos",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
]
|
||||
filter_instance._build_router(mcp_tools)
|
||||
|
||||
# Responses API native tool with SAME name as MCP canonical
|
||||
responses_api_tool = {
|
||||
"type": "function",
|
||||
"name": "github-search",
|
||||
"description": "Caller-owned search tool",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
|
||||
hook = SemanticToolFilterHook(filter_instance)
|
||||
|
||||
# Verify classification: should be native, not MCP
|
||||
assert not hook._is_mcp_tool(responses_api_tool), (
|
||||
"Responses API tool with type=function + top-level name "
|
||||
"should be classified as native, not MCP"
|
||||
)
|
||||
|
||||
# Full hook test: all-native request should preserve tools
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Search GitHub"}],
|
||||
"tools": [responses_api_tool],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
result = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
# All tools are native → hook returns data with all tools preserved
|
||||
filtered = (result or data)["tools"]
|
||||
assert len(filtered) == 1, f"Native tool must survive, got {len(filtered)}"
|
||||
assert filtered[0]["name"] == "github-search"
|
||||
|
||||
print("✅ Responses API tool with MCP-matching name correctly classified as native")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_preserves_tool_order():
|
||||
"""
|
||||
Regression test: tool ordering preservation.
|
||||
|
||||
Given: An interleaved request [mcp_tool_A, native_tool, mcp_tool_B]
|
||||
When: The hook filters tools (all MCP tools survive)
|
||||
Then: The output order must match the original request order,
|
||||
NOT native-first.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
)
|
||||
from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse
|
||||
|
||||
mock_router = Mock()
|
||||
|
||||
def mock_embedding_sync(*args, **kwargs):
|
||||
return EmbeddingResponse(
|
||||
data=[Embedding(embedding=[0.1] * 1536, index=0, object="embedding")],
|
||||
model="text-embedding-3-small",
|
||||
object="list",
|
||||
usage={"prompt_tokens": 10, "total_tokens": 10},
|
||||
)
|
||||
|
||||
async def mock_embedding_async(*args, **kwargs):
|
||||
return mock_embedding_sync()
|
||||
|
||||
mock_router.embedding = mock_embedding_sync
|
||||
mock_router.aembedding = mock_embedding_async
|
||||
|
||||
filter_instance = SemanticMCPToolFilter(
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=mock_router,
|
||||
top_k=5,
|
||||
similarity_threshold=0.3,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
# Register MCP tools
|
||||
mcp_tool_a = MCPTool(
|
||||
name="github-search",
|
||||
description="Search GitHub",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
mcp_tool_b = MCPTool(
|
||||
name="github-issue",
|
||||
description="Create GitHub issue",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
filter_instance._build_router([mcp_tool_a, mcp_tool_b])
|
||||
|
||||
# Mock filter_tools to return both MCP tools (deterministic)
|
||||
filter_instance.filter_tools = AsyncMock( # type: ignore[method-assign]
|
||||
return_value=[mcp_tool_a, mcp_tool_b]
|
||||
)
|
||||
|
||||
# Native tool (interleaved between MCP tools)
|
||||
native_tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "weather_lookup",
|
||||
"description": "Look up weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
|
||||
# Original order: [mcp_A, native, mcp_B]
|
||||
original_tools = [mcp_tool_a, native_tool, mcp_tool_b]
|
||||
|
||||
hook = SemanticToolFilterHook(filter_instance)
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Search GitHub and check weather"}],
|
||||
"tools": original_tools,
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
result = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
assert result is not None, "Hook should return modified data"
|
||||
filtered = result["tools"]
|
||||
|
||||
# All tools should survive
|
||||
assert len(filtered) == 3, f"Expected 3 tools, got {len(filtered)}"
|
||||
|
||||
# Order must be preserved: [mcp_A, native, mcp_B]
|
||||
assert filtered[0] is mcp_tool_a, "First tool should be mcp_tool_a"
|
||||
assert filtered[1] is native_tool, "Second tool should be native_tool"
|
||||
assert filtered[2] is mcp_tool_b, "Third tool should be mcp_tool_b"
|
||||
|
||||
print("✅ Tool ordering preserved: [mcp_A, native, mcp_B] maintained after filtering")
|
||||
|
||||
|
||||
class TestGetToolsByNames:
|
||||
"""
|
||||
Regression coverage for SemanticMCPToolFilter._get_tools_by_names
|
||||
|
|
@ -489,9 +911,7 @@ class TestGetToolsByNames:
|
|||
{"name": "send_email", "description": "send mail"},
|
||||
]
|
||||
|
||||
matched = filter_instance._get_tools_by_names(
|
||||
["send_email"], available_tools
|
||||
)
|
||||
matched = filter_instance._get_tools_by_names(["send_email"], available_tools)
|
||||
|
||||
assert len(matched) == 1
|
||||
assert matched[0]["name"] == "send_email"
|
||||
|
|
@ -503,9 +923,7 @@ class TestGetToolsByNames:
|
|||
client_name = "litellm_" + canonical
|
||||
available_tools = [{"name": client_name, "description": "scrape"}]
|
||||
|
||||
matched = filter_instance._get_tools_by_names(
|
||||
[canonical], available_tools
|
||||
)
|
||||
matched = filter_instance._get_tools_by_names([canonical], available_tools)
|
||||
|
||||
assert len(matched) == 1
|
||||
# Must return the incoming tool unchanged so the client-facing
|
||||
|
|
@ -516,13 +934,9 @@ class TestGetToolsByNames:
|
|||
"""Some clients use dash as alias separator; accept that too."""
|
||||
filter_instance = self._make_filter()
|
||||
canonical = "weather_svc-get_weather"
|
||||
available_tools = [
|
||||
{"name": "mcp-" + canonical, "description": "weather"}
|
||||
]
|
||||
available_tools = [{"name": "mcp-" + canonical, "description": "weather"}]
|
||||
|
||||
matched = filter_instance._get_tools_by_names(
|
||||
[canonical], available_tools
|
||||
)
|
||||
matched = filter_instance._get_tools_by_names([canonical], available_tools)
|
||||
|
||||
assert len(matched) == 1
|
||||
assert matched[0]["name"] == "mcp-" + canonical
|
||||
|
|
@ -552,9 +966,7 @@ class TestGetToolsByNames:
|
|||
{"name": "litellm_" + canonical, "description": "wrapped"},
|
||||
]
|
||||
|
||||
matched = filter_instance._get_tools_by_names(
|
||||
[canonical], available_tools
|
||||
)
|
||||
matched = filter_instance._get_tools_by_names([canonical], available_tools)
|
||||
|
||||
assert len(matched) == 1
|
||||
assert matched[0]["name"] == canonical
|
||||
|
|
@ -567,9 +979,7 @@ class TestGetToolsByNames:
|
|||
separator-anchored suffixes of ``litellm_api-fs-read_file``.
|
||||
"""
|
||||
filter_instance = self._make_filter()
|
||||
available_tools = [
|
||||
{"name": "litellm_api-fs-read_file", "description": "read"}
|
||||
]
|
||||
available_tools = [{"name": "litellm_api-fs-read_file", "description": "read"}]
|
||||
|
||||
matched = filter_instance._get_tools_by_names(
|
||||
["fs-read_file", "api-fs-read_file"], available_tools
|
||||
|
|
@ -590,9 +1000,7 @@ class TestGetToolsByNames:
|
|||
{"name": "my_" + canonical, "description": "plain search"},
|
||||
]
|
||||
|
||||
matched = filter_instance._get_tools_by_names(
|
||||
[canonical], available_tools
|
||||
)
|
||||
matched = filter_instance._get_tools_by_names([canonical], available_tools)
|
||||
|
||||
assert len(matched) == 1
|
||||
assert matched[0]["name"] == "my_" + canonical
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue