scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(skills): enforce key limits on search embeddings and bound the semantic index

Browse source 
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>
This commit is contained in:
yassin 2026-09-05 19:41:30 +00:00
parent 3da44f43e6
commit a805127008
9 changed files with 233 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/llms/litellm_proxy/skills/skill_search.py
View file

@ -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)

10
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -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,
)

7
litellm/proxy/agent_endpoints/agent_search.py
View file

@ -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)

3
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -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):

3
litellm/proxy/anthropic_endpoints/skills_endpoints.py
View file

@ -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:

62
litellm/proxy/common_utils/semantic_text_index.py
View file

@ -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)

155
tests/test_litellm/llms/litellm_proxy/skills/test_skill_search.py
View file

@ -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",

7
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py
View file

@ -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]

27
tests/test_litellm/proxy/agent_endpoints/test_agent_search.py
View file

@ -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