mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(mcp): index authed request-time tools missing from the semantic filter startup index
The semantic tool filter builds its index by listing every MCP server without per-user credentials, so servers needing per-user auth (interactive OAuth tokens, user-scoped env vars) reject the anonymous tools/list and contribute zero routes. Request-time expansion resolves that auth, so filter_tools received tools the router could never select: an empty index failed open N->N (past 128 tools OpenAI rejects the request outright) and a partial index matched only unavailable tools, stripping every tool from the request. filter_tools now syncs missing tools into the router before matching (building the router when absent) behind an asyncio lock so each tool embeds once, and falls back to the full tool list when matches map to no available tool, consistent with the zero-match fallback. Context-window overflows keep failing closed.
This commit is contained in:
Tin Chi Lo
parent
b2202cb1aa
commit
ae2462b369
2 changed files with 219 additions and 7 deletions
|
|
@ -4,6 +4,7 @@ Semantic MCP Tool Filtering using semantic-router
|
|||
Filters MCP tools semantically for /chat/completions and /responses endpoints.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -76,6 +77,7 @@ class SemanticMCPToolFilter:
|
|||
self.tool_router: Optional["SemanticRouter"] = None
|
||||
self.context_window_error: Optional[str] = None
|
||||
self._tool_map: Dict[str, Any] = {} # MCPTool objects or OpenAI function dicts
|
||||
self._index_sync_lock = asyncio.Lock()
|
||||
|
||||
async def build_router_from_mcp_registry(self) -> None:
|
||||
"""Build semantic router from all MCP tools in the registry (no auth checks)."""
|
||||
|
|
@ -182,6 +184,55 @@ class SemanticMCPToolFilter:
|
|||
return
|
||||
raise
|
||||
|
||||
def _tools_missing_from_index(self, tools: list[Any]) -> dict[str, Any]:
|
||||
"""Map name -> tool for every named tool not yet in the semantic index."""
|
||||
return {
|
||||
name: tool
|
||||
for name, tool in ((self._extract_tool_info(t)[0], t) for t in tools)
|
||||
if name and name not in self._tool_map
|
||||
}
|
||||
|
||||
async def _ensure_tools_indexed(self, available_tools: list[Any]) -> None:
|
||||
"""
|
||||
Index request-time tools the startup build never saw.
|
||||
|
||||
The startup index lists every registered MCP server WITHOUT per-user
|
||||
credentials, so servers requiring per-user auth (interactive OAuth
|
||||
tokens, user-scoped env vars) contribute zero routes. Tools reaching
|
||||
the filter came through an authenticated expansion; without indexing
|
||||
them here they can never be selected, so requests either bypass
|
||||
filtering entirely (N->N) or lose every tool to unrelated matches.
|
||||
"""
|
||||
from semantic_router.routers.base import Route
|
||||
|
||||
if not self._tools_missing_from_index(available_tools):
|
||||
return
|
||||
|
||||
async with self._index_sync_lock:
|
||||
missing = self._tools_missing_from_index(available_tools)
|
||||
if not missing:
|
||||
return
|
||||
|
||||
if self.tool_router is None:
|
||||
self._build_router(list(missing.values()))
|
||||
return
|
||||
|
||||
descriptions = {name: self._extract_tool_info(tool)[1] for name, tool in missing.items()}
|
||||
routes = [
|
||||
Route(
|
||||
name=name,
|
||||
description=description,
|
||||
utterances=[description],
|
||||
score_threshold=self.similarity_threshold,
|
||||
)
|
||||
for name, description in descriptions.items()
|
||||
]
|
||||
await self.tool_router.aadd(routes)
|
||||
self._tool_map.update(missing)
|
||||
verbose_logger.info(
|
||||
f"Semantic tool filter indexed {len(routes)} request-time tools missing from the startup index"
|
||||
)
|
||||
|
||||
async def filter_tools(
|
||||
self,
|
||||
query: str,
|
||||
|
|
@ -216,13 +267,21 @@ class SemanticMCPToolFilter:
|
|||
if not query or not query.strip():
|
||||
return available_tools
|
||||
|
||||
# Router should be built on startup - if not, something went wrong
|
||||
if self.tool_router is None:
|
||||
verbose_logger.warning("Router not initialized - was build_router_from_mcp_registry() called on startup?")
|
||||
return available_tools
|
||||
|
||||
# Run semantic filtering
|
||||
try:
|
||||
await self._ensure_tools_indexed(available_tools)
|
||||
|
||||
if self.context_window_error is not None:
|
||||
raise SemanticToolFilterContextWindowError(
|
||||
embedding_model=self.embedding_model,
|
||||
stage="the MCP tool descriptions during semantic router build",
|
||||
original_error=self.context_window_error,
|
||||
)
|
||||
|
||||
if self.tool_router is None:
|
||||
verbose_logger.warning("Semantic router could not be built from the request's tools")
|
||||
return available_tools
|
||||
|
||||
limit = top_k or self.top_k
|
||||
matches = self.tool_router(text=query, limit=limit)
|
||||
matched_tool_names = self._extract_tool_names_from_matches(matches)
|
||||
|
|
@ -230,8 +289,13 @@ class SemanticMCPToolFilter:
|
|||
if not matched_tool_names:
|
||||
return available_tools
|
||||
|
||||
return self._get_tools_by_names(matched_tool_names, available_tools)
|
||||
filtered_tools = self._get_tools_by_names(matched_tool_names, available_tools)
|
||||
if not filtered_tools:
|
||||
return available_tools
|
||||
return filtered_tools
|
||||
|
||||
except SemanticToolFilterContextWindowError:
|
||||
raise
|
||||
except Exception as e:
|
||||
if _is_context_window_error(e):
|
||||
verbose_logger.error(
|
||||
|
|
@ -240,7 +304,7 @@ class SemanticMCPToolFilter:
|
|||
)
|
||||
raise SemanticToolFilterContextWindowError(
|
||||
embedding_model=self.embedding_model,
|
||||
stage="the user query",
|
||||
stage="the user query or the MCP tool descriptions being indexed",
|
||||
original_error=str(e),
|
||||
) from e
|
||||
verbose_logger.error(f"Semantic tool filter failed: {e}", exc_info=True)
|
||||
|
|
|
|||
|
|
@ -1664,3 +1664,151 @@ def test_is_context_window_error_detection_variants():
|
|||
assert _is_context_window_error(ValueError("Invalid 'input[0]': maximum input length is 8192 tokens."))
|
||||
assert not _is_context_window_error(ValueError("A generic API error occurred."))
|
||||
assert not _is_context_window_error(None)
|
||||
|
||||
|
||||
def _make_keyword_embedding_router(recorded_inputs):
|
||||
"""
|
||||
Mock litellm Router whose embeddings are deterministic keyword one-hots:
|
||||
texts mentioning linear/issue/ticket embed to [1, 0], everything else to
|
||||
[0, 1]. Lets tests assert real similarity ranking through the actual
|
||||
semantic-router index. Every embedding input batch is appended to
|
||||
recorded_inputs.
|
||||
"""
|
||||
from litellm.types.utils import Embedding, EmbeddingResponse
|
||||
|
||||
def _vector(text):
|
||||
lowered = text.lower()
|
||||
if "linear" in lowered or "issue" in lowered or "ticket" in lowered:
|
||||
return [1.0, 0.0]
|
||||
return [0.0, 1.0]
|
||||
|
||||
def mock_embedding_sync(*args, **kwargs):
|
||||
texts = kwargs["input"]
|
||||
recorded_inputs.append(list(texts))
|
||||
return EmbeddingResponse(
|
||||
data=[Embedding(embedding=_vector(t), index=i, object="embedding") for i, t in enumerate(texts)],
|
||||
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(*args, **kwargs)
|
||||
|
||||
mock_router = Mock()
|
||||
mock_router.embedding = mock_embedding_sync
|
||||
mock_router.aembedding = mock_embedding_async
|
||||
return mock_router
|
||||
|
||||
|
||||
def _make_keyword_filter(recorded_inputs, top_k: int = 3):
|
||||
from litellm.proxy._experimental.mcp_server.semantic_tool_filter import (
|
||||
SemanticMCPToolFilter,
|
||||
)
|
||||
|
||||
return SemanticMCPToolFilter(
|
||||
embedding_model="text-embedding-3-small",
|
||||
litellm_router_instance=_make_keyword_embedding_router(recorded_inputs),
|
||||
top_k=top_k,
|
||||
similarity_threshold=0.3,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
def _linear_issue_tool():
|
||||
return MCPTool(
|
||||
name="linear_stub-get_issue",
|
||||
description="Get a Linear issue (ticket) by its identifier such as LIT-1234",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
|
||||
|
||||
def _linear_list_tool():
|
||||
return MCPTool(
|
||||
name="linear_stub-list_issues",
|
||||
description="List Linear issues (tickets) in the workspace",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
|
||||
|
||||
def _weather_tool():
|
||||
return MCPTool(
|
||||
name="weather_stub-get_weather",
|
||||
description="Get the current weather conditions for a city",
|
||||
inputSchema={"type": "object"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_indexes_request_tools_when_startup_index_is_empty():
|
||||
"""
|
||||
Regression test: the startup index is built by listing every MCP server
|
||||
WITHOUT per-user credentials, so a gateway whose servers all require
|
||||
per-user auth (e.g. interactive OAuth) starts with an empty index
|
||||
(tool_router is None). filter_tools then failed open and returned all N
|
||||
tools unfiltered (customer-visible as an N->N header and, past 128 tools,
|
||||
an OpenAI 400 "tools array too long"). The authed request-time tools must
|
||||
instead be indexed on first sight so filtering actually runs.
|
||||
"""
|
||||
filter_instance = _make_keyword_filter([])
|
||||
assert filter_instance.tool_router is None
|
||||
|
||||
tools = [_linear_issue_tool(), _weather_tool()]
|
||||
filtered = await filter_instance.filter_tools(
|
||||
query="what is Linear ticket LIT-3794 about",
|
||||
available_tools=tools,
|
||||
)
|
||||
|
||||
assert [t.name for t in filtered] == ["linear_stub-get_issue"]
|
||||
print("✅ Empty startup index is built from authed request-time tools")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_indexes_tools_missing_from_partial_index():
|
||||
"""
|
||||
Regression test: servers whose tools/list needs per-user auth contribute
|
||||
zero routes to the startup index while anonymously listable servers are
|
||||
indexed. Tools reaching the filter through the authed request-time
|
||||
expansion must be added to the existing router (and only embedded once;
|
||||
repeat requests embed just the query).
|
||||
"""
|
||||
recorded_inputs = []
|
||||
filter_instance = _make_keyword_filter(recorded_inputs)
|
||||
filter_instance._build_router([_weather_tool()])
|
||||
assert filter_instance.tool_router is not None
|
||||
|
||||
tools = [_linear_issue_tool(), _weather_tool()]
|
||||
query = "what is Linear ticket LIT-3794 about"
|
||||
|
||||
filtered = await filter_instance.filter_tools(query=query, available_tools=tools)
|
||||
assert [t.name for t in filtered] == ["linear_stub-get_issue"]
|
||||
|
||||
calls_after_first = len(recorded_inputs)
|
||||
filtered_again = await filter_instance.filter_tools(query=query, available_tools=tools)
|
||||
assert [t.name for t in filtered_again] == ["linear_stub-get_issue"]
|
||||
assert len(recorded_inputs) == calls_after_first + 1
|
||||
|
||||
print("✅ Partial startup index is completed from request-time tools, embedding each tool once")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_fails_open_when_matches_are_not_in_available_tools():
|
||||
"""
|
||||
Regression test: when the semantic router's matches are all tools that are
|
||||
NOT in the request's available_tools (an index/request mismatch), the
|
||||
filter returned an empty list, stripping every tool from the request and
|
||||
breaking it outright (observed live as a 3->0 header followed by a
|
||||
provider 400). It must fail open with the full tool list instead, matching
|
||||
the zero-match fallback.
|
||||
"""
|
||||
filter_instance = _make_keyword_filter([])
|
||||
filter_instance._build_router([_weather_tool()])
|
||||
|
||||
tools = [_linear_issue_tool(), _linear_list_tool()]
|
||||
filtered = await filter_instance.filter_tools(
|
||||
query="current weather in San Francisco",
|
||||
available_tools=tools,
|
||||
)
|
||||
|
||||
assert [t.name for t in filtered] == ["linear_stub-get_issue", "linear_stub-list_issues"]
|
||||
print("✅ Matches outside available_tools fail open instead of dropping every tool")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue