scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-22 00:31:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 96

Merge 483b282ba6 into 252c71c0b2

Browse source
This commit is contained in:
Kent 2026-09-15 20:27:09 -04:00 • committed by GitHub
commit 4f83101915
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 69 additions and 62 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/caching/_embedding_router.py
View file

@ -5,9 +5,8 @@ configured embedding model is a proxy Router deployment, embeddings must run
through the Router so per-deployment auth (e.g. Bedrock aws_role_name) is
applied. Otherwise fall back to a direct litellm embedding call.
This module is dependency-injected: callers pass the proxy ``llm_router`` and
``llm_model_list`` in, so the decision logic is unit-testable without importing
``litellm.proxy.proxy_server``.
This module is dependency-injected: callers pass the proxy ``llm_router`` in, so the
decision logic is unit-testable without importing ``litellm.proxy.proxy_server``.
"""
from __future__ import annotations
@ -25,15 +24,11 @@ if TYPE_CHECKING:
def resolve_embedding_router(
embedding_model: str,
llm_router: Router | None,
llm_model_list: list[dict[str, Any]] | None,
) -> Router | None:
"""Return ``llm_router`` iff it serves ``embedding_model`` as a deployment."""
if llm_router is None:
return None
router_model_names: Final[list[str]] = (
[m["model_name"] for m in llm_model_list if "model_name" in m] if llm_model_list is not None else []
)
if embedding_model in router_model_names:
if llm_router.get_model_list(model_name=embedding_model):
return llm_router
return None

10
litellm/caching/qdrant_semantic_cache.py
View file

@ -223,12 +223,11 @@ class QdrantSemanticCache(BaseCache):
def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
"""Embed via the proxy Router when it serves the model, else direct."""
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
from litellm.proxy.proxy_server import llm_router
except ImportError:
llm_model_list = None
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
router: Final = resolve_embedding_router(self.embedding_model, llm_router)
embedding_input: Final = self._embedding_input(prompt, router)
if router is not None:
return router.embedding(
@ -249,12 +248,11 @@ class QdrantSemanticCache(BaseCache):
async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
from litellm.proxy.proxy_server import llm_router
except ImportError:
llm_model_list = None
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
router: Final = resolve_embedding_router(self.embedding_model, llm_router)
embedding_input: Final = self._embedding_input(prompt, router)
embedding_call: Final = (
router.aembedding(

10
litellm/caching/redis_semantic_cache.py
View file

@ -341,12 +341,11 @@ class RedisSemanticCache(BaseCache):
mirroring ``_get_async_embedding``; otherwise embeds directly.
"""
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
from litellm.proxy.proxy_server import llm_router
except ImportError:
llm_model_list = None
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
router: Final = resolve_embedding_router(self.embedding_model, llm_router)
embedding_input: Final = self._embedding_input(prompt, router)
if router is not None:
embedding_response = cast(
@ -516,12 +515,11 @@ class RedisSemanticCache(BaseCache):
List[float]: The embedding vector
"""
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
from litellm.proxy.proxy_server import llm_router
except ImportError:
llm_model_list = None
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
router: Final = resolve_embedding_router(self.embedding_model, llm_router)
embedding_input: Final = self._embedding_input(prompt, router)
embedding_call: Final = (
router.aembedding(

73
tests/test_litellm/caching/test_embedding_router.py
View file

@ -10,41 +10,68 @@ from litellm.caching._embedding_router import (
)
def test_resolve_returns_router_when_model_is_a_deployment():
router = MagicMock()
def test_resolve_routes_exact_name_model_via_real_router():
router = litellm.Router(
model_list=[
{
"model_name": "sem-embed",
"litellm_params": {"model": "text-embedding-3-small"},
}
]
)
assert resolve_embedding_router("sem-embed", router) is router
def test_resolve_routes_provider_prefixed_wildcard_via_real_router():
router = litellm.Router(
model_list=[
{"model_name": "bedrock/*", "litellm_params": {"model": "bedrock/*"}}
]
)
# bedrock/amazon.titan-embed-text-v2:0 is NOT an exact model_name; only the
# bedrock/* pattern serves it. The old exact-name code returned None here.
assert (
resolve_embedding_router("sem-embed", router, [{"model_name": "sem-embed"}])
resolve_embedding_router("bedrock/amazon.titan-embed-text-v2:0", router)
is router
)
def test_resolve_returns_none_when_model_not_in_router():
router = MagicMock()
assert (
resolve_embedding_router("sem-embed", router, [{"model_name": "other"}]) is None
def test_resolve_routes_visible_model_group_alias_via_real_router():
router = litellm.Router(
model_list=[
{
"model_name": "real-embed",
"litellm_params": {"model": "text-embedding-3-small"},
}
],
model_group_alias={"aliased-embed": "real-embed"},
)
# aliased-embed is only reachable through model_group_alias.
assert resolve_embedding_router("aliased-embed", router) is router
def test_resolve_returns_none_when_real_router_does_not_serve_model():
router = litellm.Router(
model_list=[
{
"model_name": "other-embed",
"litellm_params": {"model": "text-embedding-3-small"},
}
]
)
assert resolve_embedding_router("sem-embed", router) is None
def test_resolve_returns_none_when_router_is_none():
assert (
resolve_embedding_router("sem-embed", None, [{"model_name": "sem-embed"}])
is None
)
assert resolve_embedding_router("sem-embed", None) is None
def test_resolve_returns_none_when_model_list_is_none():
def test_resolve_returns_none_when_get_model_list_returns_none():
# get_model_list is annotated Optional[List]; in practice it returns [],
# but pin the falsy-None path so the `if ...:` gate stays correct.
router = MagicMock()
assert resolve_embedding_router("sem-embed", router, None) is None
def test_resolve_skips_entries_missing_model_name():
router = MagicMock()
model_list = [
{"litellm_params": {"model": "bedrock/x"}},
{"model_name": "sem-embed"},
]
assert resolve_embedding_router("sem-embed", router, model_list) is router
assert resolve_embedding_router("other", router, [{"litellm_params": {}}]) is None
router.get_model_list = MagicMock(return_value=None)
assert resolve_embedding_router("sem-embed", router) is None
def test_build_metadata_preserves_request_fields_and_adds_flag():

19
tests/test_litellm/caching/test_qdrant_semantic_cache.py
View file

@ -18,7 +18,6 @@ def test_qdrant_semantic_cache_initialization(monkeypatch):
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -80,7 +79,6 @@ def test_qdrant_semantic_cache_get_cache_hit():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -160,7 +158,6 @@ def test_qdrant_semantic_cache_rejects_unscoped_cache_hit():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"result": {"exists": True}}
@ -321,7 +318,6 @@ def test_qdrant_semantic_cache_get_cache_miss():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -377,7 +373,6 @@ async def test_qdrant_semantic_cache_async_get_cache_hit():
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -468,7 +463,6 @@ async def test_qdrant_semantic_cache_async_get_cache_miss():
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -528,7 +522,6 @@ def test_qdrant_semantic_cache_set_cache():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -594,7 +587,6 @@ async def test_qdrant_semantic_cache_async_set_cache():
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client"
) as mock_async_client,
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -664,7 +656,6 @@ def test_qdrant_semantic_cache_custom_vector_size():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection does NOT exist (so it will be created)
mock_exists_response = MagicMock()
mock_exists_response.status_code = 200
@ -725,7 +716,6 @@ def test_qdrant_semantic_cache_default_vector_size():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection exists check
mock_response = MagicMock()
mock_response.status_code = 200
@ -761,7 +751,6 @@ def test_qdrant_semantic_cache_large_vector_size():
) as mock_sync_client,
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client"),
):
# Mock the collection does NOT exist (so it will be created)
mock_exists_response = MagicMock()
mock_exists_response.status_code = 200
@ -807,10 +796,10 @@ def test_qdrant_semantic_cache_large_vector_size():
assert create_payload["vectors"]["size"] == 4096
def _router_proxy_module(router, model_name):
def _router_proxy_module(router, model_name="sem-embed"):
router.get_model_list = MagicMock(return_value=[{"model_name": model_name}])
mod = types.ModuleType("litellm.proxy.proxy_server")
mod.llm_router = router
mod.llm_model_list = [{"model_name": model_name}]
return mod
@ -837,7 +826,7 @@ def test_qdrant_sync_get_cache_routes_through_router(monkeypatch):
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
_router_proxy_module(router, "sem-embed"),
_router_proxy_module(router),
)
with patch("litellm.embedding") as direct_embed:
@ -896,7 +885,7 @@ async def test_qdrant_async_embedding_forwards_full_metadata(monkeypatch):
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
_router_proxy_module(router, "sem-embed"),
_router_proxy_module(router),
)
await cache._get_async_embedding(

8
tests/test_litellm/caching/test_redis_semantic_cache.py
View file

@ -829,11 +829,11 @@ def test_redis_get_embedding_routes_through_router(monkeypatch):
cache.embedding_model = "sem-embed"
router = MagicMock()
router.get_model_list = MagicMock(return_value=[{"model_name": "sem-embed"}])
router.get_configured_token_limits.return_value = (None, None)
router.embedding = MagicMock(return_value={"data": [{"embedding": [0.5, 0.6]}]})
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
fake_proxy.llm_router = router
fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
with patch("litellm.embedding") as direct_embed:
@ -1072,11 +1072,11 @@ async def test_redis_async_embedding_forwards_full_metadata(monkeypatch):
cache.embedding_model = "sem-embed"
router = MagicMock()
router.get_model_list = MagicMock(return_value=[{"model_name": "sem-embed"}])
router.get_configured_token_limits.return_value = (None, None)
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
fake_proxy.llm_router = router
fake_proxy.llm_model_list = [{"model_name": "sem-embed"}]
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
await cache._get_async_embedding(
@ -1096,9 +1096,9 @@ LONG_PROMPT = " ".join(f"token{i}" for i in range(300))
def _proxy_with_router(monkeypatch: pytest.MonkeyPatch, router: MagicMock, model_name: str) -> None:
import types
router.get_model_list = MagicMock(return_value=[{"model_name": model_name}])
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
fake_proxy.llm_router = router
fake_proxy.llm_model_list = [{"model_name": model_name}]
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy)
@ -1242,9 +1242,9 @@ def test_redis_llmcache_setter_supported():
def _router_proxy_module(router, model_name):
import types
router.get_model_list = MagicMock(return_value=[{"model_name": model_name}])
fake_proxy = types.ModuleType("litellm.proxy.proxy_server")
fake_proxy.llm_router = router
fake_proxy.llm_model_list = [{"model_name": model_name}]
return fake_proxy