diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py index d8c1b6a368e..22d3f5b6dd1 100644 --- a/litellm/llms/litellm_proxy/skills/skill_search.py +++ b/litellm/llms/litellm_proxy/skills/skill_search.py @@ -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, ) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 23a63c00bd0..16ab77c2c2d 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -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]) diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 1d1778a960b..af2baa2bd89 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -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, ) diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 7450ef2b509..c2829f28e61 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -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 + ) diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index 9e25d78b6cd..cdd194b267f 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -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 diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index 63194af221e..b712f5ec45d 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -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 diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index e373a8e31ae..26a1264a0f1 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -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], diff --git a/tests/e2e/mcp/test_mcp_tool_search_e2e.py b/tests/e2e/mcp/test_mcp_tool_search_e2e.py new file mode 100644 index 00000000000..847239e8a66 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_tool_search_e2e.py @@ -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 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 5259166e9a9..7eb9252acfb 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -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): diff --git a/tests/unit/llms/litellm_proxy/skills/test_skill_search.py b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py index e28e927cf61..c0754f6f1af 100644 --- a/tests/unit/llms/litellm_proxy/skills/test_skill_search.py +++ b/tests/unit/llms/litellm_proxy/skills/test_skill_search.py @@ -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() diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 0d6b83b3275..fd94fccf8fd 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -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) diff --git a/tests/unit/proxy/agent_endpoints/test_agent_search.py b/tests/unit/proxy/agent_endpoints/test_agent_search.py index c02ed1f37e5..750fadadab0 100644 --- a/tests/unit/proxy/agent_endpoints/test_agent_search.py +++ b/tests/unit/proxy/agent_endpoints/test_agent_search.py @@ -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() diff --git a/tests/unit/proxy/common_utils/test_semantic_text_index.py b/tests/unit/proxy/common_utils/test_semantic_text_index.py new file mode 100644 index 00000000000..21fc33f839e --- /dev/null +++ b/tests/unit/proxy/common_utils/test_semantic_text_index.py @@ -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