mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
fix(mcp): leave out tool texts too long to embed instead of failing tool search (#43554)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
1999751bae
commit
058c84476a
13 changed files with 637 additions and 29 deletions
|
|
@ -107,7 +107,11 @@ class SkillSearchIndex:
|
|||
if isinstance(scores, EmbeddingFailed):
|
||||
return SkillSearchEmbeddingFailed(reason=scores.reason)
|
||||
ranked: Final = sorted(
|
||||
(SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)),
|
||||
(
|
||||
SkillSearchHit(skill=skill, score=score)
|
||||
for skill, score in zip(skills, scores, strict=True)
|
||||
if score is not None
|
||||
),
|
||||
key=lambda hit: hit.score,
|
||||
reverse=True,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -258,9 +258,11 @@ def _split_core_tools(tools: Sequence[Tool], core_tools: Sequence[str]) -> tuple
|
|||
|
||||
|
||||
def _top_hits(
|
||||
tools: Sequence[Tool], scores: Sequence[float], minimum: float, limit: int
|
||||
tools: Sequence[Tool], scores: Sequence[float | None], minimum: float, limit: int
|
||||
) -> tuple[tuple[float, Tool], ...]:
|
||||
hits: Final = ((score, tool) for score, tool in zip(scores, tools, strict=True) if score >= minimum)
|
||||
hits: Final = (
|
||||
(score, tool) for score, tool in zip(scores, tools, strict=True) if score is not None and score >= minimum
|
||||
)
|
||||
return tuple(sorted(hits, key=lambda hit: hit[0], reverse=True)[:limit])
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -116,7 +116,11 @@ class AgentSearchIndex:
|
|||
if isinstance(scores, EmbeddingFailed):
|
||||
return AgentSearchEmbeddingFailed(reason=scores.reason)
|
||||
ranked: Final = sorted(
|
||||
(AgentSearchHit(agent=agent, score=score) for agent, score in zip(agents, scores, strict=True)),
|
||||
(
|
||||
AgentSearchHit(agent=agent, score=score)
|
||||
for agent, score in zip(agents, scores, strict=True)
|
||||
if score is not None
|
||||
),
|
||||
key=lambda hit: hit.score,
|
||||
reverse=True,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain, islice
|
||||
from types import MappingProxyType
|
||||
|
|
@ -13,7 +14,7 @@ from fastapi import HTTPException
|
|||
from openai import OpenAIError
|
||||
from pydantic import ConfigDict
|
||||
|
||||
from litellm.exceptions import BudgetExceededError
|
||||
from litellm.exceptions import BudgetExceededError, ContextWindowExceededError
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -25,6 +26,7 @@ Vector: TypeAlias = tuple[float, ...]
|
|||
|
||||
DEFAULT_MAX_CACHED_VECTORS: Final = 5000
|
||||
"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first."""
|
||||
DEFAULT_TOO_LONG_TTL_SECONDS: Final = 3600.0
|
||||
|
||||
|
||||
class Embedder(Protocol):
|
||||
|
|
@ -36,6 +38,18 @@ class EmbeddingFailed:
|
|||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _InputTooLong:
|
||||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TextVectors:
|
||||
vectors: Mapping[str, Vector]
|
||||
too_long: frozenset[str]
|
||||
failure: EmbeddingFailed | None = None
|
||||
|
||||
|
||||
class _EmbeddingItem(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
|
|
@ -102,11 +116,13 @@ def router_embedder(
|
|||
_CacheKey: TypeAlias = tuple[str, str]
|
||||
|
||||
|
||||
async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed:
|
||||
async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed | _InputTooLong:
|
||||
try:
|
||||
vectors: Final = tuple(await embed(texts))
|
||||
except HTTPException:
|
||||
raise
|
||||
except ContextWindowExceededError as exc:
|
||||
return _InputTooLong(reason=f"embedding the search query failed: {exc}")
|
||||
except (OpenAIError, ValueError, BudgetExceededError) as exc:
|
||||
return EmbeddingFailed(reason=f"embedding the search query failed: {exc}")
|
||||
if len(vectors) != len(texts):
|
||||
|
|
@ -114,32 +130,97 @@ async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...
|
|||
return vectors
|
||||
|
||||
|
||||
async def _embed_texts(embed: Embedder, texts: Sequence[str]) -> _TextVectors:
|
||||
if not texts:
|
||||
return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset())
|
||||
embedded: Final = await _embed_all(embed, texts)
|
||||
if isinstance(embedded, EmbeddingFailed):
|
||||
return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset(), failure=embedded)
|
||||
if not isinstance(embedded, _InputTooLong):
|
||||
return _TextVectors(vectors=MappingProxyType(dict(zip(texts, embedded, strict=True))), too_long=frozenset())
|
||||
if len(texts) == 1:
|
||||
return _TextVectors(vectors=MappingProxyType({}), too_long=frozenset({texts[0]}))
|
||||
middle: Final = len(texts) // 2
|
||||
left: Final = await _embed_texts(embed, texts[:middle])
|
||||
if left.failure is not None:
|
||||
return left
|
||||
right: Final = await _embed_texts(embed, texts[middle:])
|
||||
return _TextVectors(
|
||||
vectors=MappingProxyType({**left.vectors, **right.vectors}),
|
||||
too_long=left.too_long | right.too_long,
|
||||
failure=right.failure,
|
||||
)
|
||||
|
||||
|
||||
async def _embed_query_with(
|
||||
embed: Embedder, query: str, texts: Sequence[str]
|
||||
) -> tuple[Vector, _TextVectors] | EmbeddingFailed:
|
||||
if not texts:
|
||||
embedded_query: Final = await _embed_all(embed, (query,))
|
||||
if isinstance(embedded_query, EmbeddingFailed):
|
||||
return embedded_query
|
||||
if isinstance(embedded_query, _InputTooLong):
|
||||
return EmbeddingFailed(reason=embedded_query.reason)
|
||||
return embedded_query[0], _TextVectors(vectors=MappingProxyType({}), too_long=frozenset())
|
||||
|
||||
embedded: Final = await _embed_all(embed, (query, *texts))
|
||||
if isinstance(embedded, EmbeddingFailed):
|
||||
return embedded
|
||||
if not isinstance(embedded, _InputTooLong):
|
||||
return embedded[0], _TextVectors(
|
||||
vectors=MappingProxyType(dict(zip(texts, embedded[1:], strict=True))), too_long=frozenset()
|
||||
)
|
||||
query_embedding: Final = await _embed_all(embed, (query,))
|
||||
if isinstance(query_embedding, EmbeddingFailed):
|
||||
return query_embedding
|
||||
if isinstance(query_embedding, _InputTooLong):
|
||||
return EmbeddingFailed(reason=query_embedding.reason)
|
||||
vectors: Final = await _embed_texts(embed, texts)
|
||||
return query_embedding[0], vectors
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Embedded:
|
||||
query_vector: Vector
|
||||
vectors: Mapping[str, Vector]
|
||||
too_long: frozenset[str]
|
||||
failure: EmbeddingFailed | None
|
||||
|
||||
|
||||
def _same_dimension(query_vector: Vector, vectors: Mapping[str, Vector], texts: Sequence[str]) -> bool:
|
||||
return all(len(vectors[text]) == len(query_vector) for text in texts)
|
||||
return all(len(vectors[text]) == len(query_vector) for text in texts if text in vectors)
|
||||
|
||||
|
||||
async def _embed_query_and_texts(
|
||||
embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector]
|
||||
embed: Embedder, query: str, texts: Sequence[str], cached: Mapping[str, Vector], too_long: frozenset[str]
|
||||
) -> _Embedded | EmbeddingFailed:
|
||||
missing: Final = tuple(dict.fromkeys(text for text in texts if text not in cached))
|
||||
embedded: Final = await _embed_all(embed, (query, *missing))
|
||||
if isinstance(embedded, EmbeddingFailed):
|
||||
pending: Final = (text for text in texts if text not in cached and text not in too_long)
|
||||
missing: Final = tuple(dict.fromkeys(pending))
|
||||
query_embedding: Final = await _embed_query_with(embed, query, missing)
|
||||
if isinstance(query_embedding, EmbeddingFailed):
|
||||
return query_embedding
|
||||
query_vector, new_text_vectors = query_embedding
|
||||
vectors: Final = MappingProxyType(dict(chain(cached.items(), new_text_vectors.vectors.items())))
|
||||
embedded: Final = _Embedded(
|
||||
query_vector=query_vector,
|
||||
vectors=vectors,
|
||||
too_long=new_text_vectors.too_long,
|
||||
failure=new_text_vectors.failure,
|
||||
)
|
||||
if embedded.failure is not None:
|
||||
return embedded
|
||||
vectors: Final = MappingProxyType(dict(chain(cached.items(), zip(missing, embedded[1:], strict=True))))
|
||||
if _same_dimension(embedded[0], vectors, texts):
|
||||
return _Embedded(query_vector=embedded[0], vectors=vectors)
|
||||
unique: Final = tuple(dict.fromkeys(texts))
|
||||
reembedded: Final = await _embed_all(embed, (query, *unique))
|
||||
if _same_dimension(query_vector, vectors, texts):
|
||||
return embedded
|
||||
unique: Final = tuple(dict.fromkeys(text for text in texts if text in vectors))
|
||||
reembedded: Final = await _embed_query_with(embed, query, unique)
|
||||
if isinstance(reembedded, EmbeddingFailed):
|
||||
return reembedded
|
||||
reembedded_query, reembedded_text_vectors = reembedded
|
||||
return _Embedded(
|
||||
query_vector=reembedded[0], vectors=MappingProxyType(dict(zip(unique, reembedded[1:], strict=True)))
|
||||
query_vector=reembedded_query,
|
||||
vectors=reembedded_text_vectors.vectors,
|
||||
too_long=new_text_vectors.too_long | reembedded_text_vectors.too_long,
|
||||
failure=reembedded_text_vectors.failure,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -148,18 +229,49 @@ class SemanticTextIndex:
|
|||
|
||||
Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first."""
|
||||
|
||||
def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
max_entries: int = DEFAULT_MAX_CACHED_VECTORS,
|
||||
too_long_ttl_seconds: float = DEFAULT_TOO_LONG_TTL_SECONDS,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._max_entries: Final = max_entries
|
||||
self._too_long_ttl_seconds: Final = too_long_ttl_seconds
|
||||
self._clock: Final = clock
|
||||
self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({})
|
||||
self._too_long: Mapping[_CacheKey, float] = MappingProxyType({})
|
||||
|
||||
def _cached(self, embedding_model: str) -> Mapping[str, Vector]:
|
||||
return MappingProxyType(
|
||||
{text: vector for (model, text), vector in self._vectors.items() if model == embedding_model}
|
||||
)
|
||||
|
||||
def _too_long_for(self, embedding_model: str, now: float) -> frozenset[str]:
|
||||
return frozenset(
|
||||
text
|
||||
for (model, text), seen_at in self._too_long.items()
|
||||
if model == embedding_model and now - seen_at < self._too_long_ttl_seconds
|
||||
)
|
||||
|
||||
def _merged_too_long(
|
||||
self, embedding_model: str, embedded: _Embedded, texts: Sequence[str], now: float
|
||||
) -> Mapping[_CacheKey, float]:
|
||||
fresh: Final = MappingProxyType(
|
||||
{key: seen_at for key, seen_at in self._too_long.items() if now - seen_at < self._too_long_ttl_seconds}
|
||||
)
|
||||
newly_too_long: Final = (
|
||||
(embedding_model, text)
|
||||
for text in texts
|
||||
if text in embedded.too_long and (embedding_model, text) not in fresh
|
||||
)
|
||||
ordered: Final = MappingProxyType({**fresh, **dict.fromkeys(newly_too_long, now)})
|
||||
return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None)))
|
||||
|
||||
def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]:
|
||||
dimension: Final = len(embedded.query_vector)
|
||||
touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts})
|
||||
touched: Final = MappingProxyType(
|
||||
{(embedding_model, text): embedded.vectors[text] for text in texts if text in embedded.vectors}
|
||||
)
|
||||
untouched: Final = MappingProxyType(
|
||||
{
|
||||
key: vector
|
||||
|
|
@ -175,14 +287,27 @@ class SemanticTextIndex:
|
|||
|
||||
async def scores(
|
||||
self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str
|
||||
) -> tuple[float, ...] | EmbeddingFailed:
|
||||
"""Cosine similarity of `query` to each entry of `texts`, in the same order."""
|
||||
) -> tuple[float | None, ...] | EmbeddingFailed:
|
||||
"""Cosine similarity for embeddable entries of `texts`, with None for omitted entries."""
|
||||
if not texts:
|
||||
return ()
|
||||
embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model))
|
||||
now: Final = self._clock()
|
||||
embedded: Final = await _embed_query_and_texts(
|
||||
embed, query, texts, self._cached(embedding_model), self._too_long_for(embedding_model, now)
|
||||
)
|
||||
if isinstance(embedded, EmbeddingFailed):
|
||||
return embedded
|
||||
if not _same_dimension(embedded.query_vector, embedded.vectors, texts):
|
||||
same_dimension: Final = _same_dimension(embedded.query_vector, embedded.vectors, texts)
|
||||
if embedded.failure is not None:
|
||||
if same_dimension:
|
||||
self._vectors = self._merged(embedding_model, embedded, texts)
|
||||
self._too_long = self._merged_too_long(embedding_model, embedded, texts, now)
|
||||
return embedded.failure
|
||||
if not same_dimension:
|
||||
return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions")
|
||||
self._vectors = self._merged(embedding_model, embedded, texts)
|
||||
return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts)
|
||||
self._too_long = self._merged_too_long(embedding_model, embedded, texts, now)
|
||||
return tuple(
|
||||
cosine_similarity(embedded.query_vector, embedded.vectors[text]) if text in embedded.vectors else None
|
||||
for text in texts
|
||||
)
|
||||
|
|
|
|||
|
|
@ -39,6 +39,14 @@
|
|||
assertions: [denied_without_permission]
|
||||
source: "rest_endpoints.py:305-386"
|
||||
rationale: Tool-level permission guard; multi-tenant safety
|
||||
- id: mcp.tool_search.api_key.oversized_description_omitted
|
||||
module: mcp
|
||||
tier: P0
|
||||
operation: tool_search
|
||||
auth_family: api_key
|
||||
assertions: [oversized_description_omitted]
|
||||
source: "test_mcp_tool_search_e2e.py"
|
||||
rationale: Native semantic search omits tools whose descriptions exceed the embedding input limit
|
||||
- id: mcp.list_tools.bearer.succeeds
|
||||
module: mcp
|
||||
tier: P1
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Protocol
|
||||
|
||||
from e2e_config import datadog_mcp_url, unique_marker
|
||||
|
|
@ -44,6 +44,7 @@ def register_datadog_mcp(
|
|||
*,
|
||||
mcp_access_groups: list[str] | None = None,
|
||||
allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,),
|
||||
tool_name_to_description: Mapping[str, str] | None = None,
|
||||
) -> str:
|
||||
"""Register the core Datadog toolset with its credentials from the env. By default
|
||||
the server exposes only `search_datadog_logs`; pass `allowed_tools=None` to expose
|
||||
|
|
@ -61,6 +62,7 @@ def register_datadog_mcp(
|
|||
},
|
||||
allowed_tools=None if allowed_tools is None else list(allowed_tools),
|
||||
mcp_access_groups=mcp_access_groups,
|
||||
tool_name_to_description=tool_name_to_description,
|
||||
)
|
||||
resources.defer(lambda: client.delete_server(server_id))
|
||||
return server_id
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ class McpServerNewBody(BaseModel):
|
|||
static_headers: dict[str, str] | None = None
|
||||
allowed_tools: list[str] | None = None
|
||||
mcp_access_groups: list[str] | None = None
|
||||
tool_name_to_description: dict[str, str] | None = None
|
||||
|
||||
|
||||
class McpServerNewResponse(BaseModel):
|
||||
|
|
@ -131,6 +132,34 @@ class McpCallToolBody(BaseModel):
|
|||
server_id: str
|
||||
|
||||
|
||||
class McpVirtualToolCallBody(BaseModel):
|
||||
name: str
|
||||
arguments: dict[str, McpToolArg]
|
||||
|
||||
|
||||
class McpToolsListParams(BaseModel):
|
||||
server_id: str | None = None
|
||||
|
||||
|
||||
class McpToolSearchSettings(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
embedding_model: str | None = None
|
||||
top_k: int = 5
|
||||
similarity_threshold: float = 0.0
|
||||
core_tools: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class McpToolSearchSettingsResponse(BaseModel):
|
||||
values: McpToolSearchSettings
|
||||
|
||||
|
||||
class McpToolSearchSettingsUpdateResponse(BaseModel):
|
||||
message: str
|
||||
status: str
|
||||
settings: McpToolSearchSettings
|
||||
|
||||
|
||||
class McpCallContent(BaseModel):
|
||||
type: str | None = None
|
||||
text: str | None = None
|
||||
|
|
@ -166,6 +195,7 @@ class McpClient:
|
|||
static_headers: dict[str, str] | None = None,
|
||||
allowed_tools: list[str] | None = None,
|
||||
mcp_access_groups: list[str] | None = None,
|
||||
tool_name_to_description: Mapping[str, str] | None = None,
|
||||
) -> str:
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
|
|
@ -180,6 +210,9 @@ class McpClient:
|
|||
static_headers=static_headers,
|
||||
allowed_tools=allowed_tools,
|
||||
mcp_access_groups=mcp_access_groups,
|
||||
tool_name_to_description=(
|
||||
dict(tool_name_to_description) if tool_name_to_description is not None else None
|
||||
),
|
||||
),
|
||||
response_type=McpServerNewResponse,
|
||||
)
|
||||
|
|
@ -244,14 +277,21 @@ class McpClient:
|
|||
mcp_access_groups: list[str] | None = None,
|
||||
mcp_toolsets: list[str] | None = None,
|
||||
models: list[str] | None = None,
|
||||
mcp_tool_search_enabled: bool | None = None,
|
||||
) -> str:
|
||||
object_permission = (
|
||||
ObjectPermission(
|
||||
mcp_servers=mcp_servers,
|
||||
mcp_access_groups=mcp_access_groups,
|
||||
mcp_toolsets=mcp_toolsets,
|
||||
mcp_tool_search_enabled=mcp_tool_search_enabled,
|
||||
)
|
||||
if (
|
||||
mcp_servers is not None
|
||||
or mcp_access_groups is not None
|
||||
or mcp_toolsets is not None
|
||||
or mcp_tool_search_enabled is not None
|
||||
)
|
||||
if mcp_servers is not None or mcp_access_groups is not None or mcp_toolsets is not None
|
||||
else None
|
||||
)
|
||||
return self.proxy.generate_key(
|
||||
|
|
@ -263,14 +303,37 @@ class McpClient:
|
|||
)
|
||||
|
||||
@step("List the MCP tools the key can see from /mcp-rest/tools/list")
|
||||
def list_tools(self, key: str) -> Result[McpToolsListResponse]:
|
||||
def list_tools(self, key: str, server_id: str | None = None) -> Result[McpToolsListResponse]:
|
||||
return self.proxy.transport.get(
|
||||
"/mcp-rest/tools/list",
|
||||
headers=ApiKeyHeaders(x_litellm_api_key=key),
|
||||
params=NoBody(),
|
||||
params=McpToolsListParams(server_id=server_id),
|
||||
response_type=McpToolsListResponse,
|
||||
)
|
||||
|
||||
@step("Read the MCP tool search settings from /get/mcp_tool_search_settings")
|
||||
def get_mcp_tool_search_settings(self) -> McpToolSearchSettings:
|
||||
return unwrap(
|
||||
self.proxy.transport.get(
|
||||
"/get/mcp_tool_search_settings",
|
||||
headers=self.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=McpToolSearchSettingsResponse,
|
||||
)
|
||||
).values
|
||||
|
||||
@step("Update the MCP tool search settings with /update/mcp_tool_search_settings")
|
||||
def update_mcp_tool_search_settings(self, settings: McpToolSearchSettings) -> None:
|
||||
_ = unwrap(
|
||||
self.proxy.transport.patch(
|
||||
"/update/mcp_tool_search_settings",
|
||||
headers=self.proxy.transport.master,
|
||||
json=settings,
|
||||
response_type=McpToolSearchSettingsUpdateResponse,
|
||||
)
|
||||
)
|
||||
settle_propagation(time.monotonic())
|
||||
|
||||
@step('Wait for /mcp-rest/tools/list to show the MCP server\'s tool matching "{needle}"')
|
||||
def await_tool(self, key: str, server_id: str, needle: str) -> str:
|
||||
"""Poll tools/list until `server_id` serves a tool matching `needle`, and
|
||||
|
|
@ -419,6 +482,21 @@ class McpClient:
|
|||
response_type=McpCallToolResponse,
|
||||
)
|
||||
|
||||
@step("Call the MCP virtual tool {name} through /mcp-rest/tools/call")
|
||||
def call_virtual_tool(
|
||||
self,
|
||||
key: str,
|
||||
*,
|
||||
name: str,
|
||||
arguments: McpToolArguments,
|
||||
) -> Result[McpCallToolResponse]:
|
||||
return self.proxy.transport.post(
|
||||
"/mcp-rest/tools/call",
|
||||
headers=ApiKeyHeaders(x_litellm_api_key=key),
|
||||
json=McpVirtualToolCallBody(name=name, arguments=dict(arguments)),
|
||||
response_type=McpCallToolResponse,
|
||||
)
|
||||
|
||||
|
||||
def _is_mcp_not_synced(
|
||||
result: Result[McpCallToolResponse],
|
||||
|
|
|
|||
86
tests/e2e/mcp/test_mcp_tool_search_e2e.py
Normal file
86
tests/e2e/mcp/test_mcp_tool_search_e2e.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp
|
||||
from e2e_config import provider_edge_base, unique_marker
|
||||
from e2e_http import unwrap
|
||||
from lifecycle import ResourceManager
|
||||
from mcp_client import McpClient, McpToolSearchSettings
|
||||
from models import LiteLLMParamsBody
|
||||
from proxy_client import ProxyClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MCP_TOOL_SEARCH_TOOL_NAME: Final = "mcp_tool_search"
|
||||
OVERSIZED_DESCRIPTION: Final = f"OVERSIZED {'weather ' * 12_000}"
|
||||
|
||||
|
||||
class TestMcpToolSearch:
|
||||
@pytest.mark.covers("mcp.tool_search.api_key.oversized_description_omitted")
|
||||
def test_native_tool_search_omits_oversized_description(
|
||||
self, proxy: ProxyClient, client: McpClient, resources: ResourceManager
|
||||
) -> None:
|
||||
"""Exercise shared ranking through the native virtual tool because proxy mode does not apply description overrides."""
|
||||
model_name: Final = f"e2e-mcp-tool-search-{unique_marker()}"
|
||||
embedding_api_base: Final = provider_edge_base("openai")
|
||||
model_id: Final = proxy.create_model(
|
||||
model_name,
|
||||
LiteLLMParamsBody(
|
||||
model="openai/text-embedding-3-large",
|
||||
api_key="os.environ/OPENAI_API_KEY",
|
||||
api_base=None if embedding_api_base is None else f"{embedding_api_base}/v1",
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
|
||||
previous_settings: Final = client.get_mcp_tool_search_settings()
|
||||
resources.defer(lambda: client.update_mcp_tool_search_settings(previous_settings))
|
||||
client.update_mcp_tool_search_settings(
|
||||
McpToolSearchSettings(
|
||||
embedding_model=model_name,
|
||||
top_k=100,
|
||||
similarity_threshold=0.0,
|
||||
core_tools=previous_settings.core_tools,
|
||||
)
|
||||
)
|
||||
|
||||
server_id: Final = register_datadog_mcp(
|
||||
client,
|
||||
resources,
|
||||
allowed_tools=None,
|
||||
tool_name_to_description={SEARCH_LOGS_TOOL: OVERSIZED_DESCRIPTION},
|
||||
)
|
||||
client.await_registered(server_id)
|
||||
key: Final = client.generate_key(
|
||||
user_id=f"e2e-mcp-tool-search-{unique_marker()}",
|
||||
mcp_servers=[server_id],
|
||||
models=[model_name],
|
||||
mcp_tool_search_enabled=True,
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_key(key))
|
||||
|
||||
tools: Final = unwrap(client.list_tools(key, server_id=server_id))
|
||||
assert len(tools.tool_names_for_server(server_id)) >= 2
|
||||
oversized_tool_name: Final = tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL)
|
||||
assert oversized_tool_name is not None
|
||||
oversized_tool: Final = next((tool for tool in tools.tools if tool.name == oversized_tool_name), None)
|
||||
assert oversized_tool is not None
|
||||
assert oversized_tool.description == OVERSIZED_DESCRIPTION
|
||||
normal_tool_name: Final = next(
|
||||
(name for name in tools.tool_names_for_server(server_id) if name != oversized_tool_name), None
|
||||
)
|
||||
assert normal_tool_name is not None
|
||||
|
||||
result: Final = unwrap(
|
||||
client.call_virtual_tool(
|
||||
key,
|
||||
name=MCP_TOOL_SEARCH_TOOL_NAME,
|
||||
arguments={"query": normal_tool_name, "top_k": 100},
|
||||
)
|
||||
)
|
||||
|
||||
assert result.is_error is not True, result.all_text
|
||||
assert normal_tool_name in result.all_text
|
||||
assert oversized_tool_name not in result.all_text
|
||||
|
|
@ -67,6 +67,7 @@ class ObjectPermission(BaseModel):
|
|||
mcp_servers: list[str] | None = None
|
||||
mcp_access_groups: list[str] | None = None
|
||||
mcp_toolsets: list[str] | None = None
|
||||
mcp_tool_search_enabled: bool | None = None
|
||||
|
||||
|
||||
class KeyGenerateBody(BaseModel):
|
||||
|
|
|
|||
|
|
@ -89,6 +89,19 @@ class FixedDimensionEmbedder:
|
|||
return tuple((1.0,) * self.dimensions for _ in texts)
|
||||
|
||||
|
||||
class OversizedAwareEmbedder:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.calls.append(tuple(texts))
|
||||
if any("OVERSIZED" in text for text in texts):
|
||||
raise litellm.ContextWindowExceededError(
|
||||
message="input is too long", model="text-embedding-3-large", llm_provider="openai"
|
||||
)
|
||||
return tuple(VECTORS[text] for text in texts)
|
||||
|
||||
|
||||
class TestSkillSearchText:
|
||||
def test_joins_title_description_and_instructions(self) -> None:
|
||||
assert skill_search_text(TRANSLATOR) == (
|
||||
|
|
@ -133,6 +146,19 @@ class TestSkillSearchIndex:
|
|||
assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file", "trip-planner"]
|
||||
assert outcome.hits[0].score > outcome.hits[1].score
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_omits_skills_with_oversized_search_text(self) -> None:
|
||||
oversized: Final = LiteLLM_SkillsTable(
|
||||
skill_id="oversized", display_title="Oversized", instructions="OVERSIZED instructions"
|
||||
)
|
||||
|
||||
outcome: Final = await SkillSearchIndex().search(
|
||||
"language translation", (*SKILLS, oversized), top_k=5, embed=OversizedAwareEmbedder(), embedding_model="m"
|
||||
)
|
||||
|
||||
assert isinstance(outcome, SkillSearchHits)
|
||||
assert {hit.skill.skill_id for hit in outcome.hits} == {skill.skill_id for skill in SKILLS}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_search_only_embeds_the_query(self) -> None:
|
||||
index = SkillSearchIndex()
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ from litellm.proxy._experimental.mcp_server.tool_search import (
|
|||
coerce_top_k,
|
||||
get_virtual_tool_definitions,
|
||||
handle_mcp_tool_search,
|
||||
rank_mcp_tools,
|
||||
search_mcp_tools,
|
||||
search_tools,
|
||||
)
|
||||
|
|
@ -193,6 +194,19 @@ class RecordingEmbedder:
|
|||
return tuple(FAKE_VECTORS[text] for text in texts)
|
||||
|
||||
|
||||
class OversizedRecordingEmbedder:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.calls.append(tuple(texts))
|
||||
if any("OVERSIZED" in text for text in texts):
|
||||
raise litellm.ContextWindowExceededError(
|
||||
message="input is too long", model="text-embedding-3-large", llm_provider="openai"
|
||||
)
|
||||
return tuple((1.0, 0.0) if "weather" in text.lower() else (0.0, 1.0) for text in texts)
|
||||
|
||||
|
||||
def _ranker(embedder: RecordingEmbedder | None = None) -> SemanticToolRanker:
|
||||
return SemanticToolRanker(embed=embedder or RecordingEmbedder(), embedding_model="emb", index=SemanticTextIndex())
|
||||
|
||||
|
|
@ -203,6 +217,26 @@ def _names(results: Sequence[ToolSearchResult] | EmbeddingFailed) -> list[str]:
|
|||
|
||||
|
||||
class TestSearchMcpTools:
|
||||
@pytest.mark.asyncio
|
||||
async def test_rank_mcp_tools_omits_oversized_description(self) -> None:
|
||||
tools: Final = _make_tools(
|
||||
[
|
||||
("get_weather", "Get the current weather forecast for a city"),
|
||||
("send_email", "Send an email to a recipient"),
|
||||
("list_files", "List files in a directory"),
|
||||
("huge_tool", "OVERSIZED description"),
|
||||
]
|
||||
)
|
||||
embedder = OversizedRecordingEmbedder()
|
||||
ranker = SemanticToolRanker(embed=embedder, embedding_model="emb", index=SemanticTextIndex())
|
||||
|
||||
hits: Final = await rank_mcp_tools("weather", tools, 5, MCPToolSearchSettings(embedding_model="emb"), ranker)
|
||||
|
||||
assert not isinstance(hits, EmbeddingFailed)
|
||||
names: Final = [hit.tool.name for hit in hits]
|
||||
assert names[0] == "get_weather"
|
||||
assert "huge_tool" not in names
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_mode_finds_foreign_exchange_tool_for_fx(self) -> None:
|
||||
keyword_only = await search_mcp_tools("FX", CATALOG, 5, MCPToolSearchSettings(), ranker=None)
|
||||
|
|
|
|||
|
|
@ -88,6 +88,19 @@ class FixedDimensionEmbedder:
|
|||
return tuple((1.0,) * self.dimensions for _ in texts)
|
||||
|
||||
|
||||
class OversizedAwareEmbedder:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.calls.append(tuple(texts))
|
||||
if any("OVERSIZED" in text for text in texts):
|
||||
raise litellm.ContextWindowExceededError(
|
||||
message="input is too long", model="text-embedding-3-large", llm_provider="openai"
|
||||
)
|
||||
return tuple(VECTORS[text] for text in texts)
|
||||
|
||||
|
||||
class TestAgentSearchText:
|
||||
def test_joins_name_description_and_skills_with_tags(self) -> None:
|
||||
assert agent_search_text(TRANSLATOR) == (
|
||||
|
|
@ -127,6 +140,21 @@ class TestAgentSearchIndex:
|
|||
assert [hit.agent.agent_id for hit in outcome.hits] == ["translator", "trip"]
|
||||
assert outcome.hits[0].score > outcome.hits[1].score
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_omits_agents_with_oversized_search_text(self) -> None:
|
||||
oversized: Final = AgentResponse(
|
||||
agent_id="oversized",
|
||||
agent_name="oversized-agent",
|
||||
agent_card_params={"description": "OVERSIZED description"},
|
||||
)
|
||||
|
||||
outcome: Final = await AgentSearchIndex().search(
|
||||
"language translation", (*AGENTS, oversized), top_k=5, embed=OversizedAwareEmbedder(), embedding_model="m"
|
||||
)
|
||||
|
||||
assert isinstance(outcome, AgentSearchHits)
|
||||
assert {hit.agent.agent_id for hit in outcome.hits} == {agent.agent_id for agent in AGENTS}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_second_search_only_embeds_the_query(self) -> None:
|
||||
index = AgentSearchIndex()
|
||||
|
|
|
|||
210
tests/unit/proxy/common_utils/test_semantic_text_index.py
Normal file
210
tests/unit/proxy/common_utils/test_semantic_text_index.py
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.common_utils.semantic_text_index import EmbeddingFailed, SemanticTextIndex, Vector
|
||||
|
||||
QUERY: Final = "weather query"
|
||||
TEXTS: Final = (
|
||||
"valid one",
|
||||
"OVERSIZED first",
|
||||
"valid two",
|
||||
"valid three",
|
||||
"OVERSIZED second",
|
||||
"valid four",
|
||||
"valid five",
|
||||
)
|
||||
|
||||
|
||||
class ContextWindowEmbedder:
|
||||
def __init__(self, dimensions: int = 2) -> None:
|
||||
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
|
||||
self.vector: Final = (1.0,) * dimensions
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.calls.append(tuple(texts))
|
||||
if any("OVERSIZED" in text for text in texts):
|
||||
raise litellm.ContextWindowExceededError(
|
||||
message="input is too long", model="text-embedding-3-large", llm_provider="openai"
|
||||
)
|
||||
return tuple(self.vector for _ in texts)
|
||||
|
||||
|
||||
class RateLimitedEmbedder:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.calls.append(tuple(texts))
|
||||
raise litellm.RateLimitError(message="rate limited", model="text-embedding-3-large", llm_provider="openai")
|
||||
|
||||
|
||||
class RateLimitedAfterSplitEmbedder:
|
||||
def __init__(self, rate_limited: bool = True) -> None:
|
||||
self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs
|
||||
self.rate_limited = rate_limited
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.calls.append(tuple(texts))
|
||||
if any("OVERSIZED" in text for text in texts):
|
||||
raise litellm.ContextWindowExceededError(
|
||||
message="input is too long", model="text-embedding-3-large", llm_provider="openai"
|
||||
)
|
||||
if self.rate_limited and any("RATE LIMITED" in text for text in texts):
|
||||
raise litellm.RateLimitError(message="rate limited", model="text-embedding-3-large", llm_provider="openai")
|
||||
return tuple((1.0, 1.0) for _ in texts)
|
||||
|
||||
|
||||
class ConcurrencyTrackingEmbedder:
|
||||
def __init__(self) -> None:
|
||||
self.in_flight = 0
|
||||
self.max_in_flight = 0
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
self.in_flight += 1
|
||||
self.max_in_flight = max(self.max_in_flight, self.in_flight)
|
||||
await asyncio.sleep(0)
|
||||
self.in_flight -= 1
|
||||
if any("OVERSIZED" in text for text in texts):
|
||||
raise litellm.ContextWindowExceededError(
|
||||
message="input is too long", model="text-embedding-3-large", llm_provider="openai"
|
||||
)
|
||||
return tuple((1.0, 1.0) for _ in texts)
|
||||
|
||||
|
||||
class FakeClock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
class FixedDimensionEmbedder:
|
||||
def __init__(self, dimensions: int) -> None:
|
||||
self.vector: Final = (1.0,) * dimensions
|
||||
|
||||
async def __call__(self, texts: Sequence[str]) -> Sequence[Vector]:
|
||||
return tuple(self.vector for _ in texts)
|
||||
|
||||
|
||||
class TestSemanticTextIndex:
|
||||
@pytest.mark.asyncio
|
||||
async def test_omits_multiple_oversized_texts_and_preserves_score_order(self) -> None:
|
||||
embedder = ContextWindowEmbedder()
|
||||
scores: Final = await SemanticTextIndex().scores(QUERY, TEXTS, embedder, "emb")
|
||||
|
||||
assert not isinstance(scores, EmbeddingFailed)
|
||||
assert len(scores) == len(TEXTS)
|
||||
assert scores[0] == pytest.approx(1.0)
|
||||
assert scores[1] is None
|
||||
assert scores[2] == pytest.approx(1.0)
|
||||
assert scores[3] == pytest.approx(1.0)
|
||||
assert scores[4] is None
|
||||
assert scores[5] == pytest.approx(1.0)
|
||||
assert scores[6] == pytest.approx(1.0)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeat_search_only_embeds_the_query(self) -> None:
|
||||
index = SemanticTextIndex()
|
||||
embedder = ContextWindowEmbedder()
|
||||
first: Final = await index.scores(QUERY, TEXTS, embedder, "emb")
|
||||
assert not isinstance(first, EmbeddingFailed)
|
||||
first_call_count: Final = len(embedder.calls)
|
||||
|
||||
second: Final = await index.scores(QUERY, TEXTS, embedder, "emb")
|
||||
|
||||
assert not isinstance(second, EmbeddingFailed)
|
||||
assert embedder.calls[first_call_count:] == [(QUERY,)]
|
||||
assert second == first
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_query_returns_embedding_failed(self) -> None:
|
||||
result: Final = await SemanticTextIndex().scores(
|
||||
"OVERSIZED query", ("valid text",), ContextWindowEmbedder(), "emb"
|
||||
)
|
||||
|
||||
assert isinstance(result, EmbeddingFailed)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_context_provider_error_fails_after_one_embed_call(self) -> None:
|
||||
embedder = RateLimitedEmbedder()
|
||||
|
||||
result: Final = await SemanticTextIndex().scores(QUERY, ("valid text",), embedder, "emb")
|
||||
|
||||
assert isinstance(result, EmbeddingFailed)
|
||||
assert len(embedder.calls) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_text_is_retried_once_the_exclusion_expires(self) -> None:
|
||||
clock = FakeClock()
|
||||
index = SemanticTextIndex(too_long_ttl_seconds=60.0, clock=clock)
|
||||
embedder = ContextWindowEmbedder()
|
||||
texts: Final = ("valid one", "OVERSIZED first")
|
||||
_ = await index.scores(QUERY, texts, embedder, "emb")
|
||||
clock.now = 59.0
|
||||
before_expiry: Final = len(embedder.calls)
|
||||
_ = await index.scores(QUERY, texts, embedder, "emb")
|
||||
assert embedder.calls[before_expiry:] == [(QUERY,)]
|
||||
|
||||
clock.now = 61.0
|
||||
after_expiry: Final = len(embedder.calls)
|
||||
scores: Final = await index.scores(QUERY, texts, embedder, "emb")
|
||||
|
||||
assert scores == (pytest.approx(1.0), None)
|
||||
assert embedder.calls[after_expiry] == (QUERY, "OVERSIZED first")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_splitting_sends_one_embedding_request_at_a_time(self) -> None:
|
||||
embedder = ConcurrencyTrackingEmbedder()
|
||||
|
||||
scores: Final = await SemanticTextIndex().scores(QUERY, TEXTS, embedder, "emb")
|
||||
|
||||
assert not isinstance(scores, EmbeddingFailed)
|
||||
assert embedder.max_in_flight == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_context_provider_error_while_splitting_fails_the_search(self) -> None:
|
||||
result: Final = await SemanticTextIndex().scores(
|
||||
QUERY, ("RATE LIMITED text", "OVERSIZED text", "valid text"), RateLimitedAfterSplitEmbedder(), "emb"
|
||||
)
|
||||
|
||||
assert isinstance(result, EmbeddingFailed)
|
||||
assert "rate limited" in result.reason
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_split_keeps_vectors_and_too_long_texts_for_the_next_search(self) -> None:
|
||||
index = SemanticTextIndex()
|
||||
embedder = RateLimitedAfterSplitEmbedder()
|
||||
texts: Final = ("valid text", "OVERSIZED text", "RATE LIMITED text")
|
||||
|
||||
first: Final = await index.scores(QUERY, texts, embedder, "emb")
|
||||
assert isinstance(first, EmbeddingFailed)
|
||||
assert "rate limited" in first.reason
|
||||
|
||||
embedder.rate_limited = False
|
||||
before: Final = len(embedder.calls)
|
||||
second: Final = await index.scores(QUERY, texts, embedder, "emb")
|
||||
|
||||
assert second == (pytest.approx(1.0), None, pytest.approx(1.0))
|
||||
assert embedder.calls[before:] == [(QUERY, "RATE LIMITED text")]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_dimension_reembedding_omits_oversized_text(self) -> None:
|
||||
index = SemanticTextIndex()
|
||||
seeded: Final = await index.scores(QUERY, ("cached text", "OVERSIZED text"), FixedDimensionEmbedder(3), "emb")
|
||||
assert not isinstance(seeded, EmbeddingFailed)
|
||||
embedder = ContextWindowEmbedder(2)
|
||||
|
||||
scores: Final = await index.scores(QUERY, ("cached text", "OVERSIZED text"), embedder, "emb")
|
||||
|
||||
assert not isinstance(scores, EmbeddingFailed)
|
||||
assert scores[0] == pytest.approx(1.0)
|
||||
assert scores[1] is None
|
||||
assert (QUERY, "cached text", "OVERSIZED text") in embedder.calls
|
||||
assert ("OVERSIZED text",) in embedder.calls
|
||||
Loading…
Add table
Reference in a new issue