From a8051270085e4e666f7dbe9a18562cb424195215 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 5 Sep 2026 19:41:30 +0000 Subject: [PATCH] fix(skills): enforce key limits on search embeddings and bound the semantic index Semantic search embeddings now run the same pre_call_hook the /embeddings route runs, so key rate limits, budgets and guardrails apply before the embedding model is called. The shared SemanticTextIndex caps cached vectors and evicts the least recently searched entries, and each skill's embedded text is capped so one skill cannot inflate the embedding batch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/litellm_proxy/skills/skill_search.py | 19 ++- .../_experimental/mcp_server/tool_search.py | 10 +- litellm/proxy/agent_endpoints/agent_search.py | 7 +- litellm/proxy/agent_endpoints/endpoints.py | 3 +- .../anthropic_endpoints/skills_endpoints.py | 3 +- .../proxy/common_utils/semantic_text_index.py | 62 +++++-- .../litellm_proxy/skills/test_skill_search.py | 155 +++++++++++++++--- .../mcp_server/test_mcp_tool_search.py | 7 + .../agent_endpoints/test_agent_search.py | 27 ++- 9 files changed, 233 insertions(+), 60 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/skill_search.py b/litellm/llms/litellm_proxy/skills/skill_search.py index d466f49fdfb..37701aaf923 100644 --- a/litellm/llms/litellm_proxy/skills/skill_search.py +++ b/litellm/llms/litellm_proxy/skills/skill_search.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Final, TypeAlias from pydantic import BaseModel, ConfigDict +from litellm.llms.litellm_proxy.skills.constants import MAX_SKILLS_PER_SEARCH from litellm.proxy.common_utils.semantic_text_index import ( Embedder, EmbeddingFailed, @@ -17,12 +18,17 @@ from litellm.proxy.common_utils.semantic_text_index import ( if TYPE_CHECKING: from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router DEFAULT_SKILL_SEARCH_TOP_K: Final = 5 MAX_SKILL_SEARCH_TOP_K: Final = 100 """Matches the ``le=100`` bound GET /v1/skills?query= enforces via FastAPI's Query validation, so the MCP tool can't return a larger payload than the REST endpoint allows.""" +MAX_SKILL_SEARCH_TEXT_CHARS: Final = 4000 +"""Per-skill cap on the title + description + instructions text that gets embedded, so one +search embeds at most ``MAX_SKILLS_PER_SEARCH * MAX_SKILL_SEARCH_TEXT_CHARS`` characters no +matter how long the stored instructions are.""" @dataclass(frozen=True, slots=True) @@ -59,7 +65,8 @@ class SkillSearchResult(BaseModel): def skill_search_text(skill: LiteLLM_SkillsTable) -> str: - return "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part) + joined: Final = "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part) + return joined[:MAX_SKILL_SEARCH_TEXT_CHARS] def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult: @@ -74,8 +81,8 @@ def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult: class SkillSearchIndex: """Caches one vector per distinct skill text per embedding model, so repeat searches only embed the query.""" - def __init__(self) -> None: - self._index: Final = SemanticTextIndex() + def __init__(self, max_entries: int = MAX_SKILLS_PER_SEARCH) -> None: + self._index: Final = SemanticTextIndex(max_entries=max_entries) async def search( self, @@ -108,6 +115,7 @@ async def search_skills( embedding_model: str | None, index: SkillSearchIndex, user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, ) -> SkillSearchOutcome: if embedding_model is None: return SkillSearchNotConfigured( @@ -115,6 +123,5 @@ async def search_skills( ) if router is None: return SkillSearchNotConfigured(reason="skill search needs a model_list so the embedding model can be called") - return await index.search( - query, skills, top_k, router_embedder(router, embedding_model, user_api_key_dict), embedding_model - ) + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, skills, top_k, embed, embedding_model) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 7f9e8dbfff9..f19340d30cb 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -247,7 +247,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI ) from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj await check_feature_access_for_user(user_api_key_dict, "agents") outcome: Final = await search_agents( @@ -258,6 +258,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): @@ -280,7 +281,7 @@ async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPI search_skills, skill_search_result, ) - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj outcome: Final = await search_skills( query=query, @@ -290,6 +291,7 @@ async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPI embedding_model=litellm.skill_search_embedding_model, index=global_skill_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case SkillSearchHits(hits): @@ -313,7 +315,7 @@ async def handle_mcp_tool_search( raw_headers: dict[str, str] | None = None, ) -> CallToolResult: from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj settings: Final = mcp_tool_search_settings() if isinstance(settings, ValidationError): @@ -327,7 +329,7 @@ async def handle_mcp_tool_search( ) ranker: Final = ( SemanticToolRanker( - embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict), + embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj), embedding_model=settings.embedding_model, index=global_mcp_tool_search_index, ) diff --git a/litellm/proxy/agent_endpoints/agent_search.py b/litellm/proxy/agent_endpoints/agent_search.py index 76e3fe6c5ad..65a89bb2c7a 100644 --- a/litellm/proxy/agent_endpoints/agent_search.py +++ b/litellm/proxy/agent_endpoints/agent_search.py @@ -18,6 +18,7 @@ from litellm.types.agents import AgentResponse if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router DEFAULT_AGENT_SEARCH_TOP_K: Final = 5 @@ -132,6 +133,7 @@ async def search_agents( embedding_model: str | None, index: AgentSearchIndex, user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, ) -> AgentSearchOutcome: if embedding_model is None: return AgentSearchNotConfigured( @@ -139,6 +141,5 @@ async def search_agents( ) if router is None: return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called") - return await index.search( - query, agents, top_k, router_embedder(router, embedding_model, user_api_key_dict), embedding_model - ) + embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj) + return await index.search(query, agents, top_k, embed, embedding_model) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index cc17672553b..aa8979a73c6 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -249,7 +249,7 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep async def _rank_agents_by_query( query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth ) -> tuple[AgentResponse, ...]: - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj outcome: Final = await search_agents( query=query, @@ -259,6 +259,7 @@ async def _rank_agents_by_query( embedding_model=litellm.agent_search_embedding_model, index=global_agent_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) match outcome: case AgentSearchHits(hits): diff --git a/litellm/proxy/anthropic_endpoints/skills_endpoints.py b/litellm/proxy/anthropic_endpoints/skills_endpoints.py index a997f3f71e6..fb4d13a9383 100644 --- a/litellm/proxy/anthropic_endpoints/skills_endpoints.py +++ b/litellm/proxy/anthropic_endpoints/skills_endpoints.py @@ -50,7 +50,7 @@ async def _search_litellm_skills(query: str, top_k: int, user_api_key_dict: User from litellm.llms.litellm_proxy.skills.transformation import ( LiteLLMSkillsTransformationHandler, ) - from litellm.proxy.proxy_server import llm_router + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj db_skills: Final = await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict=user_api_key_dict) outcome: Final = await search_skills( @@ -61,6 +61,7 @@ async def _search_litellm_skills(query: str, top_k: int, user_api_key_dict: User embedding_model=litellm.skill_search_embedding_model, index=global_skill_search_index, user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) to_response: Final = LiteLLMSkillsTransformationHandler().db_skill_to_response match outcome: diff --git a/litellm/proxy/common_utils/semantic_text_index.py b/litellm/proxy/common_utils/semantic_text_index.py index 0820459af49..4f56ab13af4 100644 --- a/litellm/proxy/common_utils/semantic_text_index.py +++ b/litellm/proxy/common_utils/semantic_text_index.py @@ -5,7 +5,7 @@ from __future__ import annotations import math from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass -from itertools import chain +from itertools import chain, islice from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias @@ -16,10 +16,14 @@ from litellm.exceptions import BudgetExceededError if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import ProxyLogging from litellm.router import Router 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.""" + class Embedder(Protocol): def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ... @@ -57,18 +61,27 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj } -def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder: +def router_embedder( + router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging +) -> Embedder: + """Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs.""" + async def embed(texts: Sequence[str]) -> Sequence[Vector]: batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input - response: Final = await router.aembedding( - model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict) - ) + metadata: Final = embedding_spend_metadata(user_api_key_dict) + request: Final = { # mutable-ok: pre_call_hook mutates the request dict in place + "model": embedding_model, + "input": batch, + "metadata": metadata, + } + await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding") + response: Final = await router.aembedding(model=embedding_model, input=batch, metadata=metadata) return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data) return embed -_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({}) +_CacheKey: TypeAlias = tuple[str, str] async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed: @@ -111,20 +124,34 @@ async def _embed_query_and_texts( class SemanticTextIndex: - """Caches one vector per distinct text per embedding model, so repeat searches only embed the query.""" + """Caches one vector per distinct text per embedding model, so repeat searches only embed the query. - def __init__(self) -> None: - self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({}) + Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first.""" - def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]: - kept: Final = MappingProxyType( + def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None: + self._max_entries: Final = max_entries + self._vectors: Mapping[_CacheKey, Vector] = 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 _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}) + untouched: Final = MappingProxyType( { - text: vector - for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items() - if len(vector) == len(embedded.query_vector) + key: vector + for key, vector in chain( + self._vectors.items(), + (((embedding_model, text), vector) for text, vector in embedded.vectors.items()), + ) + if key not in touched and (key[0] != embedding_model or len(vector) == dimension) } ) - return MappingProxyType({**kept, **embedded.vectors}) + ordered: Final = MappingProxyType({**untouched, **touched}) + return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None))) async def scores( self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str @@ -132,11 +159,10 @@ class SemanticTextIndex: """Cosine similarity of `query` to each entry of `texts`, in the same order.""" if not texts: return () - cached: Final = self._vectors.get(embedding_model, _NO_VECTORS) - embedded: Final = await _embed_query_and_texts(embed, query, texts, cached) + embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model)) if isinstance(embedded, EmbeddingFailed): return embedded if not _same_dimension(embedded.query_vector, embedded.vectors, texts): return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions") - self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)}) + self._vectors = self._merged(embedding_model, embedded, texts) return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts) diff --git a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py index 2a21ef6803a..908279e6184 100644 --- a/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py +++ b/tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py @@ -6,12 +6,13 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from openai import APIConnectionError import litellm from litellm.llms.litellm_proxy.skills.skill_search import ( + MAX_SKILL_SEARCH_TEXT_CHARS, SkillSearchEmbeddingFailed, SkillSearchHits, SkillSearchIndex, @@ -54,6 +55,23 @@ VECTORS: Final = MappingProxyType( ) +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + +def _embedding_router() -> MagicMock: + router = MagicMock() + router.aembedding = AsyncMock( + side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], + ) + ) + return router + + class FakeEmbedder: def __init__(self) -> None: self.calls: list[tuple[str, ...]] = [] # mutable-ok: test spy recording embed inputs @@ -88,6 +106,14 @@ class TestSkillSearchText: def test_all_fields_absent_is_an_empty_string(self) -> None: assert skill_search_text(LiteLLM_SkillsTable(skill_id="empty")) == "" + def test_oversized_instructions_are_cut_so_one_skill_cannot_blow_up_the_embedding_batch(self) -> None: + bloated = LiteLLM_SkillsTable( + skill_id="bloated", display_title="Bloated", instructions="x" * (MAX_SKILL_SEARCH_TEXT_CHARS * 3) + ) + text = skill_search_text(bloated) + assert len(text) == MAX_SKILL_SEARCH_TEXT_CHARS + assert text.startswith("Bloated\n") + class TestCosineSimilarity: def test_identical_direction_scores_one(self) -> None: @@ -160,6 +186,28 @@ class TestSkillSearchIndex: await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") assert embedder.calls[-1] == ("q",) + @pytest.mark.asyncio + async def test_least_recently_searched_skills_are_evicted_once_the_index_is_full(self) -> None: + index = SkillSearchIndex(max_entries=len(SKILLS)) + embedder = FixedDimensionEmbedder(3) + newcomer = LiteLLM_SkillsTable(skill_id="newcomer", display_title="Newcomer") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS[:1], top_k=5, embed=embedder, embedding_model="m") + await index.search("q", (newcomer,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[1])) + + @pytest.mark.asyncio + async def test_deleted_skills_stop_occupying_the_index_after_enough_new_ones(self) -> None: + index = SkillSearchIndex(max_entries=2) + embedder = FixedDimensionEmbedder(3) + for generation in range(50): + skill = LiteLLM_SkillsTable(skill_id=f"gen-{generation}", display_title=f"Generation {generation}") + await index.search("q", (skill,), top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + await index.search("q", SKILLS, top_k=5, embed=embedder, embedding_model="m") + assert embedder.calls[-1] == ("q", skill_search_text(SKILLS[0])) + @pytest.mark.asyncio async def test_mixed_dimensions_in_one_batch_become_embedding_failed(self) -> None: async def mixed(texts: Sequence[str]) -> Sequence[Vector]: @@ -198,7 +246,14 @@ class TestSearchSkills: @pytest.mark.asyncio async def test_no_embedding_model_is_not_configured(self) -> None: outcome = await search_skills( - "q", SKILLS, 5, router=MagicMock(), embedding_model=None, index=SkillSearchIndex(), user_api_key_dict=CALLER + "q", + SKILLS, + 5, + router=MagicMock(), + embedding_model=None, + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, SkillSearchNotConfigured) assert "skill_search_embedding_model" in outcome.reason @@ -206,19 +261,20 @@ class TestSearchSkills: @pytest.mark.asyncio async def test_no_router_is_not_configured(self) -> None: outcome = await search_skills( - "q", SKILLS, 5, router=None, embedding_model="m", index=SkillSearchIndex(), user_api_key_dict=CALLER + "q", + SKILLS, + 5, + router=None, + embedding_model="m", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, SkillSearchNotConfigured) @pytest.mark.asyncio async def test_router_embeddings_are_read_from_the_response(self) -> None: - router = MagicMock() - router.aembedding = AsyncMock( - side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( - model=model, - data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], - ) - ) + router = _embedding_router() outcome = await search_skills( "language translation", SKILLS, @@ -227,6 +283,7 @@ class TestSearchSkills: embedding_model="text-embedding-3-small", index=SkillSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, SkillSearchHits) assert [hit.skill.skill_id for hit in outcome.hits] == ["translate-file"] @@ -234,13 +291,7 @@ class TestSearchSkills: @pytest.mark.asyncio async def test_embedding_spend_is_attributed_to_the_calling_key(self) -> None: - router = MagicMock() - router.aembedding = AsyncMock( - side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( - model=model, - data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], - ) - ) + router = _embedding_router() await search_skills( "language translation", SKILLS, @@ -249,12 +300,53 @@ class TestSearchSkills: embedding_model="text-embedding-3-small", index=SkillSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) metadata = router.aembedding.await_args.kwargs["metadata"] assert metadata["user_api_key"] == "hashed-caller-key" assert metadata["user_api_key_team_id"] == "team-1" assert metadata["user_api_key_user_id"] == "user-1" + @pytest.mark.asyncio + async def test_key_limits_are_checked_against_the_real_embedding_call_before_it_runs(self) -> None: + router = _embedding_router() + key_limits = _pass_through_key_limits() + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + checked = key_limits.pre_call_hook.await_args.kwargs + assert checked["user_api_key_dict"] is CALLER + assert checked["call_type"] == "aembedding" + assert checked["data"]["model"] == "text-embedding-3-small" + assert checked["data"]["input"] == router.aembedding.await_args.kwargs["input"] + assert checked["data"]["metadata"]["user_api_key"] == "hashed-caller-key" + + @pytest.mark.asyncio + async def test_a_key_over_its_limit_never_reaches_the_embedding_model(self) -> None: + router = _embedding_router() + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=HTTPException(status_code=429, detail="rpm exceeded")) + with pytest.raises(HTTPException) as raised: + await search_skills( + "language translation", + SKILLS, + 1, + router=router, + embedding_model="text-embedding-3-small", + index=SkillSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=key_limits, + ) + assert raised.value.status_code == 429 + router.aembedding.assert_not_awaited() + def _client(role: LitellmUserRoles) -> TestClient: app = FastAPI() @@ -273,14 +365,15 @@ def accessible_skills(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: @pytest.fixture -def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock: - router = MagicMock() - router.aembedding = AsyncMock( - side_effect=lambda model, input, metadata: litellm.EmbeddingResponse( - model=model, - data=[{"object": "embedding", "index": i, "embedding": list(VECTORS[t])} for i, t in enumerate(input)], - ) - ) +def key_limits(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + limits = _pass_through_key_limits() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", limits) + return limits + + +@pytest.fixture +def embedding_router(monkeypatch: pytest.MonkeyPatch, key_limits: MagicMock) -> MagicMock: + router = _embedding_router() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) monkeypatch.setattr(litellm, "skill_search_embedding_model", "text-embedding-3-small") return router @@ -360,6 +453,18 @@ class TestGetSkillsQuery: assert response.status_code == 503 assert response.json()["detail"]["error"] == "skill_search_unavailable" + def test_a_key_over_its_rate_limit_gets_a_429_without_embedding( + self, accessible_skills: AsyncMock, embedding_router: MagicMock, key_limits: MagicMock + ) -> None: + key_limits.pre_call_hook = AsyncMock(side_effect=HTTPException(status_code=429, detail="rpm exceeded")) + response = _client(LitellmUserRoles.PROXY_ADMIN).get( + "/v1/skills", + params={"custom_llm_provider": "litellm_proxy", "query": "language translation"}, + headers={"Authorization": "Bearer k"}, + ) + assert response.status_code == 429 + embedding_router.aembedding.assert_not_awaited() + def test_top_k_is_validated(self, accessible_skills: AsyncMock, embedding_router: MagicMock) -> None: response = _client(LitellmUserRoles.PROXY_ADMIN).get( "/v1/skills", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 3f8b2ad5fff..798b0001af1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -814,10 +814,15 @@ class TestCallToolRestApiVirtualTools: router = MagicMock() router.aembedding = AsyncMock(side_effect=fake_aembedding) + key_limits = MagicMock() + key_limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) with ( patch( # test-quality-ok: the proxy's router is a module global; the handler reaches it the way production does "litellm.proxy.proxy_server.llm_router", router ), + patch( # test-quality-ok: the proxy's key-limit hooks are a module global; the embedding call runs them like /embeddings does + "litellm.proxy.proxy_server.proxy_logging_obj", key_limits + ), patch( # test-quality-ok: the authorized catalog is the seam every virtual tool shares; the ranking under test stays real "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", new_callable=AsyncMock, @@ -827,6 +832,8 @@ class TestCallToolRestApiVirtualTools: result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict) assert mock_list.await_args.kwargs["user_api_key_auth"] is user_api_key_dict + assert key_limits.pre_call_hook.await_args.kwargs["call_type"] == "aembedding" + assert key_limits.pre_call_hook.await_args.kwargs["data"]["model"] == "emb" assert result.isError is False assert [t["name"] for t in json.loads(result.content[0].text)] == [FX_TOOL.name] diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py index daca244c0a1..c02ed1f37e5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_search.py @@ -211,11 +211,24 @@ class TestAgentSearchIndex: assert isinstance(outcome, AgentSearchEmbeddingFailed) +def _pass_through_key_limits() -> MagicMock: + limits = MagicMock() + limits.pre_call_hook = AsyncMock(side_effect=lambda user_api_key_dict, data, call_type: data) + return limits + + class TestSearchAgents: @pytest.mark.asyncio async def test_no_embedding_model_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=MagicMock(), embedding_model=None, index=AgentSearchIndex(), user_api_key_dict=CALLER + "q", + AGENTS, + 5, + router=MagicMock(), + embedding_model=None, + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) assert "agent_search_embedding_model" in outcome.reason @@ -223,7 +236,14 @@ class TestSearchAgents: @pytest.mark.asyncio async def test_no_router_is_not_configured(self) -> None: outcome = await search_agents( - "q", AGENTS, 5, router=None, embedding_model="m", index=AgentSearchIndex(), user_api_key_dict=CALLER + "q", + AGENTS, + 5, + router=None, + embedding_model="m", + index=AgentSearchIndex(), + user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchNotConfigured) @@ -244,6 +264,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) assert isinstance(outcome, AgentSearchHits) assert [hit.agent.agent_id for hit in outcome.hits] == ["translator"] @@ -266,6 +287,7 @@ class TestSearchAgents: embedding_model="text-embedding-3-small", index=AgentSearchIndex(), user_api_key_dict=CALLER, + proxy_logging_obj=_pass_through_key_limits(), ) metadata = router.aembedding.await_args.kwargs["metadata"] assert metadata["user_api_key"] == "hashed-caller-key" @@ -302,6 +324,7 @@ def embedding_router(monkeypatch: pytest.MonkeyPatch) -> MagicMock: ) ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", _pass_through_key_limits()) monkeypatch.setattr(litellm, "agent_search_embedding_model", "text-embedding-3-small") return router