diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index e7c549b7358..3547bf81448 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -86,9 +86,12 @@ class RouterVectorStoreEmbeddingExecutor: ) return bool(resolved) or model in deployment_models + def _embeds_through_sdk(self, model: str, configuration: Mapping[str, object]) -> bool: + return bool(configuration) and not self._router_serves(model) + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if not self._router_serves(model): + if self._embeds_through_sdk(model, configuration): return LiteLLMVectorStoreEmbeddingExecutor().embed(model, query, embedding_kwargs) return self.router.embedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, @@ -98,7 +101,7 @@ class RouterVectorStoreEmbeddingExecutor: async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: embedding_kwargs: Final = self._embedding_kwargs(configuration) - if not self._router_serves(model): + if self._embeds_through_sdk(model, configuration): return await LiteLLMVectorStoreEmbeddingExecutor().aembed(model, query, embedding_kwargs) return await self.router.aembedding( # pyright: ignore[reportUnknownMemberType] # Router embedding input retains a legacy untyped list model=model, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index ade4d19815b..d530a1c9d52 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -16,7 +16,7 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.vector_store.transformation import ( - LiteLLMVectorStoreEmbeddingExecutor, + BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -42,12 +42,10 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -def _direct_vector_store_embedding_executor(value: object) -> VectorStoreEmbeddingExecutor: - if value is None: - return LiteLLMVectorStoreEmbeddingExecutor() - if isinstance(value, VectorStoreEmbeddingExecutor): - return value - raise TypeError("Invalid direct vector store embedding executor") +def _direct_vector_store_embedding_executor(value: object, router: "Router | None") -> VectorStoreEmbeddingExecutor: + if value is not None and not isinstance(value, VectorStoreEmbeddingExecutor): + raise TypeError("Invalid direct vector store embedding executor") + return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor(value, router) def mock_vector_store_search_response( @@ -302,7 +300,7 @@ async def asearch( Async: Search a vector store for relevant chunks based on a query and file attributes filter. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None) + kwargs.pop("_direct_vector_store_embedding_executor", None), router ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" @@ -388,7 +386,7 @@ def search( VectorStoreSearchResponse containing the search results. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None) + kwargs.pop("_direct_vector_store_embedding_executor", None), router ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" diff --git a/tests/router_unit_tests/test_router_embedding_integration.py b/tests/router_unit_tests/test_router_embedding_integration.py index 96c6ce3708e..2cc9914c9b3 100644 --- a/tests/router_unit_tests/test_router_embedding_integration.py +++ b/tests/router_unit_tests/test_router_embedding_integration.py @@ -186,6 +186,25 @@ class TestRouterEmbeddingIntegration: assert _sent(store_route, 0) == ("Bearer store-key", "text-embedding-3-large", ["sync query"]) assert _sent(store_route, 1) == ("Bearer store-key", "text-embedding-3-large", ["async query"]) + @pytest.mark.asyncio + async def test_router_executor_rejects_unserved_models_without_explicit_config( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + openai_route = _mock_embedding_route(respx_mock, OPENAI_EMBEDDINGS_URL) + executor = RouterVectorStoreEmbeddingExecutor( + router=_alias_router(), + metadata={"user_api_key_team_id": "team-a"}, + ) + + with pytest.raises(litellm.BadRequestError): + executor.embed("openai/text-embedding-3-large", "sync query", {}) + with pytest.raises(litellm.BadRequestError): + await executor.aembed("openai/text-embedding-3-large", "async query", {}) + + assert openai_route.call_count == 0 + def test_router_executor_routes_deployment_model_names_through_the_router( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 854ce351f2d..a956aef51d6 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -52,7 +52,7 @@ def _serialize_litellm_params(litellm_params): def test_direct_vector_store_embedding_executor_rejects_invalid_value(): with pytest.raises(TypeError, match="Invalid direct vector store embedding executor"): - _direct_vector_store_embedding_executor(object()) + _direct_vector_store_embedding_executor(object(), None) def test_router_vector_store_search_injects_executor_and_request_metadata(): diff --git a/tests/vector_store_tests/test_milvus_vector_store.py b/tests/vector_store_tests/test_milvus_vector_store.py index ea3c1883e46..3b3bd444b20 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -575,6 +575,32 @@ async def test_router_search_resolves_bare_embedding_alias_async( _assert_alias_resolved(embedding_route, search_route, response) +def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_sync( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = litellm.vector_stores.search(router=_alias_router(), **ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_async( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + embedding_route = _mock_embedding_route(respx_mock) + search_route = _mock_search_route(respx_mock) + + response = await litellm.vector_stores.asearch(router=_alias_router(), **ALIAS_SEARCH_KWARGS) + + _assert_alias_resolved(embedding_route, search_route, response) + + @pytest.mark.asyncio async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE)