mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #33584 from BerriAI/litellm_lit4451_semfilter_chat_tool_format
fix(mcp): keep the MCP reference intact when the semantic filter narrows tools
This commit is contained in:
tin-berri
GitHub
commit
68f0fb0346
2 changed files with 303 additions and 28 deletions
|
|
@ -155,6 +155,40 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
|
||||
return await self.filter.filter_tools(query=user_query, available_tools=expanded_tools)
|
||||
|
||||
def _selected_tool_names(self, filtered_tools: list[dict[str, Any]]) -> list[str]:
|
||||
"""Names of the semantically selected tools, as produced by the MCP expansion."""
|
||||
names = (self.filter._extract_tool_info(tool)[0] for tool in filtered_tools)
|
||||
return [name for name in names if name]
|
||||
|
||||
@staticmethod
|
||||
def _narrow_mcp_references(tools: list[Any], selected_tool_names: list[str]) -> list[Any]:
|
||||
"""
|
||||
Restrict each litellm_proxy MCP reference to the semantically selected tools.
|
||||
|
||||
The reference block is preserved rather than replaced with expanded tools, so the
|
||||
MCP gateway still performs the expansion. That keeps the per-endpoint tool shape
|
||||
and tool auto-execution intact. Expansion already applied any caller-supplied
|
||||
allowed_tools, so this selection can only narrow a block further.
|
||||
|
||||
Whether an undecidable selection exposes every tool or none is owned by
|
||||
SemanticMCPToolFilter.filter_tools, which returns the full set when nothing
|
||||
matches; the same policy therefore governs references and plain tools. Passing an
|
||||
empty selection through is safe rather than a hidden allow-all: the gateway reads
|
||||
the union of every reference's allowed_tools and treats an empty union as unset.
|
||||
"""
|
||||
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
|
||||
LiteLLM_Proxy_MCP_Handler,
|
||||
)
|
||||
|
||||
return [
|
||||
(
|
||||
{**tool, "allowed_tools": selected_tool_names}
|
||||
if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool])
|
||||
else tool
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
def _is_mcp_tool(self, tool: object) -> bool:
|
||||
"""
|
||||
Check whether *tool* is registered in the MCP semantic router.
|
||||
|
|
@ -261,36 +295,30 @@ class SemanticToolFilterHook(CustomLogger):
|
|||
if self._should_expand_mcp_tools(tools):
|
||||
verbose_proxy_logger.debug("Detected litellm_proxy MCP references, expanding before semantic filtering")
|
||||
|
||||
if not self.filter.enabled:
|
||||
verbose_proxy_logger.debug("Semantic filter disabled, leaving MCP references untouched")
|
||||
return None
|
||||
|
||||
try:
|
||||
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(
|
||||
f"No MCP tools expanded, preserving {len(native_tools_before_expand)} native tools"
|
||||
)
|
||||
return data
|
||||
verbose_proxy_logger.warning("No tools expanded from MCP references")
|
||||
return None
|
||||
|
||||
if not self.filter.enabled:
|
||||
data["tools"] = native_tools_before_expand + expanded_tools
|
||||
verbose_proxy_logger.debug("Semantic filter disabled, forwarding expanded MCP tools unfiltered")
|
||||
return data
|
||||
|
||||
filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools)
|
||||
|
||||
combined_tools = native_tools_before_expand + filtered_expanded_tools
|
||||
data["tools"] = combined_tools
|
||||
selected_tool_names = self._selected_tool_names(filtered_expanded_tools)
|
||||
narrowed_tools = self._narrow_mcp_references(tools, selected_tool_names)
|
||||
data["tools"] = narrowed_tools
|
||||
self._emit_filter_metadata_safe(
|
||||
data=data,
|
||||
mcp_tools=expanded_tools,
|
||||
filtered_mcp_tools=filtered_expanded_tools,
|
||||
native_tools=native_tools_before_expand,
|
||||
filtered_tools=combined_tools,
|
||||
filtered_tools=narrowed_tools,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
f"Expanded MCP references to {len(expanded_tools)} tools "
|
||||
|
|
|
|||
|
|
@ -865,14 +865,17 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools():
|
|||
)
|
||||
|
||||
assert result is not None, "Hook should return modified data"
|
||||
filtered = result["tools"]
|
||||
mcp_references = [tool for tool in result["tools"] if tool.get("type") == "mcp"]
|
||||
assert len(mcp_references) == 1, "The litellm_proxy MCP reference must be preserved for the MCP gateway to expand"
|
||||
|
||||
assert len(filtered) <= 2, f"Expanded tools should be filtered to top_k=2, got {len(filtered)}"
|
||||
assert len(filtered) < len(expanded_tools), (
|
||||
f"Hook must not forward all {len(expanded_tools)} expanded tools unfiltered, got {len(filtered)}"
|
||||
allowed_tools = mcp_references[0]["allowed_tools"]
|
||||
assert len(allowed_tools) <= 2, f"Expanded tools should be filtered to top_k=2, got {len(allowed_tools)}"
|
||||
assert len(allowed_tools) < len(expanded_tools), (
|
||||
f"Hook must not forward all {len(expanded_tools)} expanded tools unfiltered, got {len(allowed_tools)}"
|
||||
)
|
||||
for tool in filtered:
|
||||
assert tool in expanded_tools, "Filtered tools must be the original expanded tool dicts"
|
||||
expanded_names = {tool["name"] for tool in expanded_tools}
|
||||
for name in allowed_tools:
|
||||
assert name in expanded_names, "Selected tool names must come from the expanded tools"
|
||||
|
||||
assert (
|
||||
"litellm_semantic_filter_stats" in result["metadata"]
|
||||
|
|
@ -880,9 +883,246 @@ async def test_semantic_filter_hook_filters_expanded_litellm_proxy_tools():
|
|||
stats = result["metadata"]["litellm_semantic_filter_stats"]
|
||||
total, selected = stats.split("->")
|
||||
assert int(total) == 5, f"Stats 'from' should be pre-filter expanded count (5), got {total}"
|
||||
assert int(selected) == len(filtered), f"Stats 'to' should match post-filter count, got {selected}"
|
||||
assert int(selected) == len(allowed_tools), f"Stats 'to' should match post-filter count, got {selected}"
|
||||
|
||||
print(f"✅ Expanded litellm_proxy tools filtered: {len(expanded_tools)} -> {len(filtered)}, stats={stats}")
|
||||
print(f"✅ Expanded litellm_proxy tools filtered: {len(expanded_tools)} -> {len(allowed_tools)}, stats={stats}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_narrows_mcp_reference_for_chat_completions():
|
||||
"""
|
||||
Regression test (LIT-4451): the hook must narrow the litellm_proxy MCP
|
||||
reference instead of replacing it with expanded tool definitions.
|
||||
|
||||
Given: A /chat/completions request whose tools are a single
|
||||
{"type": "mcp", "server_url": "litellm_proxy"} reference that
|
||||
expands to 5 tools, with the semantic filter selecting top_k=2
|
||||
When: The hook processes the request with call_type="acompletion"
|
||||
Then: The MCP reference survives in data["tools"], carrying the selected
|
||||
tools in allowed_tools, and no expanded function definitions are
|
||||
written into the request.
|
||||
|
||||
Replacing the reference made the hook write Responses-API-shaped tools
|
||||
({"type": "function", "name": ...}) into /chat/completions, which expects
|
||||
{"type": "function", "function": {...}}. The provider transformation then
|
||||
rejected every MCP tool (Anthropic raised KeyError: 'function') or dropped
|
||||
it silently (Bedrock), so the model saw no MCP tools at all. Replacing the
|
||||
reference also removed the marker the MCP gateway matches on, which
|
||||
disabled tool auto-execution for require_approval="never".
|
||||
"""
|
||||
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,
|
||||
)
|
||||
|
||||
registry_tools = [
|
||||
MCPTool(
|
||||
name=f"srv-tool_{i}",
|
||||
description=f"Registry tool {i}",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
for i in range(5)
|
||||
]
|
||||
filter_instance._build_router(registry_tools)
|
||||
|
||||
expanded_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": f"srv-tool_{i}",
|
||||
"description": f"Registry tool {i}",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
for i in range(5)
|
||||
]
|
||||
|
||||
hook = SemanticToolFilterHook(filter_instance)
|
||||
hook._expand_mcp_tools = AsyncMock( # type: ignore[method-assign]
|
||||
return_value=expanded_tools
|
||||
)
|
||||
|
||||
mcp_reference = {
|
||||
"type": "mcp",
|
||||
"server_url": "litellm_proxy",
|
||||
"require_approval": "never",
|
||||
}
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Send an email"}],
|
||||
"tools": [mcp_reference],
|
||||
"metadata": {},
|
||||
}
|
||||
|
||||
result = await hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=data,
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
assert result is not None, "Hook should return modified data"
|
||||
forwarded = result["tools"]
|
||||
|
||||
assert [tool.get("type") for tool in forwarded] == ["mcp"], (
|
||||
"The MCP reference must be the only forwarded tool; writing expanded function "
|
||||
f"definitions into a chat completion loses every MCP tool. Got: {forwarded}"
|
||||
)
|
||||
assert forwarded[0]["server_url"] == "litellm_proxy", "The MCP reference must keep routing to the gateway"
|
||||
assert forwarded[0]["require_approval"] == "never", "The MCP reference must keep its auto-execute marker"
|
||||
|
||||
allowed_tools = forwarded[0]["allowed_tools"]
|
||||
assert allowed_tools, "The narrowed reference must still carry the selected tools"
|
||||
assert len(allowed_tools) <= 2, f"Selection must narrow the reference to top_k=2, got {allowed_tools}"
|
||||
assert len(allowed_tools) < len(expanded_tools), (
|
||||
f"Hook must not forward all {len(expanded_tools)} expanded tools unfiltered, got {allowed_tools}"
|
||||
)
|
||||
assert set(allowed_tools) <= {tool["name"] for tool in expanded_tools}, (
|
||||
f"Selected names must come from the expanded tools, got {allowed_tools}"
|
||||
)
|
||||
|
||||
print(f"✅ chat completions: MCP reference preserved, narrowed to {allowed_tools}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_filter_hook_zero_matches_exposes_all_tools_on_both_paths():
|
||||
"""
|
||||
A query that matches nothing must expose every MCP tool, whether the request
|
||||
carries a litellm_proxy MCP reference or plain MCP tool objects.
|
||||
|
||||
Given: A router that returns no matches for the query
|
||||
When: The hook processes an MCP reference request and a plain MCP tool request
|
||||
Then: Both expose all 3 tools, because filter_tools owns the undecidable-selection
|
||||
policy and returns the full set rather than an empty one
|
||||
|
||||
The two paths narrow through different mechanisms (allowed_tools on the reference
|
||||
versus dropping unmatched entries), so they could drift into opposite fail
|
||||
behaviours. Pinning both here keeps that single policy honest: flipping
|
||||
filter_tools to fail closed must fail this test on both paths at once, instead of
|
||||
silently hard-limiting one surface and not the other.
|
||||
"""
|
||||
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
|
||||
|
||||
registry_tools = [
|
||||
MCPTool(
|
||||
name=f"srv-tool_{i}",
|
||||
description=f"Registry tool {i}",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
def build_hook():
|
||||
filter_instance = SemanticMCPToolFilter(
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=mock_router,
|
||||
top_k=2,
|
||||
similarity_threshold=0.3,
|
||||
enabled=True,
|
||||
)
|
||||
filter_instance._build_router(registry_tools)
|
||||
zero_match_router = Mock(return_value=[])
|
||||
zero_match_router.top_k = 2
|
||||
filter_instance.tool_router = zero_match_router
|
||||
return SemanticToolFilterHook(filter_instance)
|
||||
|
||||
expanded_tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": f"srv-tool_{i}",
|
||||
"description": f"Registry tool {i}",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
reference_hook = build_hook()
|
||||
reference_hook._expand_mcp_tools = AsyncMock( # type: ignore[method-assign]
|
||||
return_value=expanded_tools
|
||||
)
|
||||
reference_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "something entirely unrelated"}],
|
||||
"tools": [{"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"}],
|
||||
"metadata": {},
|
||||
}
|
||||
reference_result = await reference_hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=reference_data,
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
reference_tools = (reference_result or reference_data)["tools"]
|
||||
mcp_references = [tool for tool in reference_tools if tool.get("type") == "mcp"]
|
||||
assert len(mcp_references) == 1, "The MCP reference must survive a zero-match query"
|
||||
assert set(mcp_references[0].get("allowed_tools") or []) == {tool["name"] for tool in expanded_tools}, (
|
||||
"A zero-match query must leave every expanded tool reachable through the reference"
|
||||
)
|
||||
|
||||
plain_hook = build_hook()
|
||||
plain_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "something entirely unrelated"}],
|
||||
"tools": list(registry_tools),
|
||||
"metadata": {},
|
||||
}
|
||||
plain_result = await plain_hook.async_pre_call_hook(
|
||||
user_api_key_dict=Mock(),
|
||||
cache=Mock(),
|
||||
data=plain_data,
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
plain_tools = (plain_result or plain_data)["tools"]
|
||||
assert len(plain_tools) == len(registry_tools), (
|
||||
f"A zero-match query must not drop plain MCP tools, got {len(plain_tools)} of {len(registry_tools)}"
|
||||
)
|
||||
|
||||
print("✅ zero matches: both the MCP reference path and the plain tool path expose every tool")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -958,8 +1198,9 @@ async def test_semantic_filter_hook_filters_expanded_tools_with_string_input():
|
|||
async def test_semantic_filter_hook_expansion_skips_filter_when_disabled():
|
||||
"""
|
||||
When the filter is disabled at runtime (e.g. via the UI toggle), the
|
||||
expansion path must forward all expanded tools and emit NO filter
|
||||
stats, mirroring the generic path's enabled guard.
|
||||
expansion path must leave the MCP reference untouched and emit NO filter
|
||||
stats, mirroring the generic path's enabled guard. The MCP gateway then
|
||||
expands the reference itself, so no tool is narrowed away.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
|
|
@ -1009,13 +1250,19 @@ async def test_semantic_filter_hook_expansion_skips_filter_when_disabled():
|
|||
call_type="aresponses",
|
||||
)
|
||||
|
||||
assert result is not None, "Hook should still expand MCP references when the filter is disabled"
|
||||
assert len(result["tools"]) == 5, f"All expanded tools must be forwarded when disabled, got {len(result['tools'])}"
|
||||
assert result is None, "Hook must not modify the request when the filter is disabled"
|
||||
assert data["tools"] == [
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_url": "litellm_proxy",
|
||||
"require_approval": "never",
|
||||
}
|
||||
], "The MCP reference must be left intact for the MCP gateway to expand"
|
||||
assert (
|
||||
"litellm_semantic_filter_stats" not in result["metadata"]
|
||||
"litellm_semantic_filter_stats" not in data["metadata"]
|
||||
), "No filter stats may be emitted when the filter is disabled"
|
||||
|
||||
print("✅ Disabled filter: expansion preserved, no spurious stats")
|
||||
print("✅ Disabled filter: MCP reference untouched, no spurious stats")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue