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