diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index cec25634bb8..78ca0355bb8 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -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 diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index c5876e993d3..8fd58ba67d2 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -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( diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index 9a70bfc1418..6d4f3b0e6b6 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -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( diff --git a/tests/test_litellm/caching/test_embedding_router.py b/tests/test_litellm/caching/test_embedding_router.py index 00a80c63303..7b505368cd5 100644 --- a/tests/test_litellm/caching/test_embedding_router.py +++ b/tests/test_litellm/caching/test_embedding_router.py @@ -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(): diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index e07578dd7e5..36213de6fdc 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -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( diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 9884e9d9bc0..73758bdf62b 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -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