From 1df402b80d2d3f4fed8d504a5a73946162c133b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:53:05 -0700 Subject: [PATCH] fix(vector-store): carry request metadata into the Router executor built from the router kwarg --- .../base_llm/vector_store/transformation.py | 16 ++++++- litellm/router.py | 11 +---- litellm/vector_stores/main.py | 13 ++++-- .../test_vector_store_endpoints.py | 2 +- .../test_milvus_vector_store.py | 46 +++++++++++++++++++ 5 files changed, 73 insertions(+), 15 deletions(-) diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 3547bf81448..9624a721870 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -62,6 +62,19 @@ class LiteLLMVectorStoreEmbeddingExecutor: ) +_REQUEST_METADATA: Final = TypeAdapter(dict[str, object]) + + +def vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: + litellm_metadata: Final = kwargs.get("litellm_metadata") + if isinstance(litellm_metadata, dict): + return _REQUEST_METADATA.validate_python(litellm_metadata) + metadata: Final = kwargs.get("metadata") + if isinstance(metadata, dict): + return _REQUEST_METADATA.validate_python(metadata) + return MappingProxyType({}) + + @dataclass(frozen=True, slots=True) class RouterVectorStoreEmbeddingExecutor: router: Router @@ -312,11 +325,12 @@ class BaseQueryEmbeddingVectorStoreConfig(BaseVectorStoreConfig): def query_embedding_executor( embedding_executor: VectorStoreEmbeddingExecutor | None, router: Router | None, + request_metadata: Mapping[str, object] = MappingProxyType({}), ) -> VectorStoreEmbeddingExecutor: if embedding_executor is not None: return embedding_executor if router is not None: - return RouterVectorStoreEmbeddingExecutor(router=router, metadata=MappingProxyType({})) + return RouterVectorStoreEmbeddingExecutor(router=router, metadata=request_metadata) return LiteLLMVectorStoreEmbeddingExecutor() def embed_query( diff --git a/litellm/router.py b/litellm/router.py index 991f1bc2828..8b510e748f5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -87,6 +87,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( ) from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, + vector_store_request_metadata, ) from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -6697,15 +6698,7 @@ class Router: @staticmethod def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: - litellm_metadata: Final = kwargs.get("litellm_metadata") - if isinstance(litellm_metadata, dict): - return cast( # cast-ok: isinstance validates the runtime dict boundary - "dict[str, object]", litellm_metadata - ) - metadata: Final = kwargs.get("metadata") - if isinstance(metadata, dict): - return cast("dict[str, object]", metadata) # cast-ok: isinstance validates the runtime dict boundary - return MappingProxyType({}) + return vector_store_request_metadata(kwargs) async def _init_vector_store_api_endpoints( self, diff --git a/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index d530a1c9d52..636bdd4b52e 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, + vector_store_request_metadata, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -42,10 +43,14 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# -def _direct_vector_store_embedding_executor(value: object, router: "Router | None") -> VectorStoreEmbeddingExecutor: +def _direct_vector_store_embedding_executor( + value: object, router: "Router | None", request_kwargs: Mapping[str, object] +) -> 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) + return BaseQueryEmbeddingVectorStoreConfig.query_embedding_executor( + value, router, vector_store_request_metadata(request_kwargs) + ) def mock_vector_store_search_response( @@ -300,7 +305,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), router + kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) local_vars: Final = { # mutable-ok: exception logging requires a sanitized mutable snapshot key: value for key, value in locals().items() if key != "embedding_executor" @@ -386,7 +391,7 @@ def search( VectorStoreSearchResponse containing the search results. """ embedding_executor: Final = _direct_vector_store_embedding_executor( - kwargs.pop("_direct_vector_store_embedding_executor", None), router + kwargs.pop("_direct_vector_store_embedding_executor", None), router, kwargs ) 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/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 a956aef51d6..1abbbe91e97 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(), None) + _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 3b3bd444b20..2ba9168b49f 100644 --- a/tests/vector_store_tests/test_milvus_vector_store.py +++ b/tests/vector_store_tests/test_milvus_vector_store.py @@ -601,6 +601,52 @@ async def test_sdk_search_with_router_kwarg_resolves_bare_embedding_alias_async( _assert_alias_resolved(embedding_route, search_route, response) +def _team_alias_router(): + return Router( + model_list=[ + { + "model_name": "team-a-embedder", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "deployment-key", + }, + "model_info": {"team_id": "team-a", "team_public_model_name": "multilingual-e5-large"}, + } + ] + ) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_resolves_team_alias_from_request_metadata( + 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=_team_alias_router(), metadata={"user_api_key_team_id": "team-a"}, **ALIAS_SEARCH_KWARGS + ) + + _assert_alias_resolved(embedding_route, search_route, response) + + +@pytest.mark.asyncio +async def test_sdk_search_with_router_kwarg_rejects_team_alias_without_team_metadata( + 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) + _mock_search_route(respx_mock) + + with pytest.raises(litellm.APIConnectionError): + await litellm.vector_stores.asearch(router=_team_alias_router(), **ALIAS_SEARCH_KWARGS) + + assert embedding_route.call_count == 0 + + @pytest.mark.asyncio async def test_transform_uses_injected_executor_without_embedding_config(respx_mock: respx.MockRouter): executor = RecordingEmbeddingExecutor(ALIAS_EMBEDDING_RESPONSE)