From 1f7355db0c2c806a5cee0278648ddcac266a83e8 Mon Sep 17 00:00:00 2001 From: pragnyanramtha Date: Wed, 20 May 2026 05:42:08 +0000 Subject: [PATCH] fix(proxy): route vector store retrieve/delete by request model --- litellm/proxy/auth/auth_utils.py | 1 + .../proxy/vector_store_endpoints/endpoints.py | 27 ++++ .../proxy/auth/test_auth_utils.py | 22 +++ .../test_vector_store_tenant_guard.py | 151 ++++++++++++++++++ 4 files changed, 201 insertions(+) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 637a4a070c4..e3fef557c3f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1120,6 +1120,7 @@ _MODEL_ROUTING_ROUTE_MARKERS = ( _MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS = ( "/files", "/batches", + "/vector_stores", "/skills", "/evals", ) diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index ccf15c206b0..86b8444df07 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -87,6 +87,31 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( return data +def _update_request_data_with_model_from_request( + data: Dict, + request: Request, +) -> Dict: + """ + Add model routing from request query/header when no body or registry model exists. + + The vector store retrieve/delete endpoints do not have request bodies, but + clients can still pass the proxy model through query params or + x-litellm-model. Preserve registry-derived routing first. + """ + if ( + data.get("model") + or data.get("custom_llm_provider") + or data.get("litellm_credential_name") + ): + return data + + model = request.query_params.get("model") or request.headers.get("x-litellm-model") + if model: + data["model"] = model + + return data + + @router.post( "/v1/vector_stores/{vector_store_id:path}/search", dependencies=[Depends(user_api_key_auth)], @@ -311,6 +336,7 @@ async def vector_store_retrieve( data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) + data = _update_request_data_with_model_from_request(data=data, request=request) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -518,6 +544,7 @@ async def vector_store_delete( data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) + data = _update_request_data_with_model_from_request(data=data, request=request) processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 08035fb7173..9f6299fac49 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -282,6 +282,28 @@ def test_get_model_from_request_includes_file_endpoint_header_model(): ) +def test_get_model_from_request_includes_vector_store_query_model(): + assert ( + get_model_from_request( + request_data={"vector_store_id": "vs_plain_openai"}, + route="/v1/vector_stores/{vector_store_id}", + request_query_params={"model": "restricted-model"}, + ) + == "restricted-model" + ) + + +def test_get_model_from_request_includes_vector_store_header_model(): + assert ( + get_model_from_request( + request_data={"vector_store_id": "vs_plain_openai"}, + route="/v1/vector_stores/{vector_store_id}", + request_headers={"x-litellm-model": "restricted-model"}, + ) + == "restricted-model" + ) + + def test_get_model_from_request_ignores_routing_header_on_standard_llm_routes(): assert ( get_model_from_request( diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py index 48262afd363..ccc8b55e59c 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py @@ -55,6 +55,157 @@ async def test_vector_store_search_forces_path_id_over_body_id(): assert captured_data["vector_store_id"] == "vs_path_allowed" +@pytest.mark.asyncio +async def test_vector_store_retrieve_uses_query_model_for_routing(): + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_retrieve + + captured_data = {} + + async def fake_base_process(self, **kwargs): + captured_data.update(self.data) + return {"ok": True} + + request = _mock_request() + request.method = "GET" + request.query_params = {"model": "gpt-5.2"} + request.headers = {"x-litellm-model": "header-model"} + + with ( + patch( + "litellm.proxy.vector_store_endpoints.endpoints.get_litellm_managed_vector_store", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.vector_store_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=fake_base_process, + ), + ): + response = await vector_store_retrieve( + request=request, + vector_store_id="vs_plain_openai", + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert response == {"ok": True} + assert captured_data["vector_store_id"] == "vs_plain_openai" + assert captured_data["model"] == "gpt-5.2" + + +@pytest.mark.asyncio +async def test_vector_store_delete_uses_header_model_for_routing(): + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_delete + + captured_data = {} + + async def fake_base_process(self, **kwargs): + captured_data.update(self.data) + return {"ok": True} + + request = _mock_request() + request.method = "DELETE" + request.headers = {"x-litellm-model": "gpt-5.2"} + + with ( + patch( + "litellm.proxy.vector_store_endpoints.endpoints.get_litellm_managed_vector_store", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.vector_store_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=fake_base_process, + ), + ): + response = await vector_store_delete( + request=request, + vector_store_id="vs_plain_openai", + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert response == {"ok": True} + assert captured_data["vector_store_id"] == "vs_plain_openai" + assert captured_data["model"] == "gpt-5.2" + + +@pytest.mark.asyncio +async def test_vector_store_delete_uses_query_model_for_routing(): + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_delete + + captured_data = {} + + async def fake_base_process(self, **kwargs): + captured_data.update(self.data) + return {"ok": True} + + request = _mock_request() + request.method = "DELETE" + request.query_params = {"model": "gpt-5.2"} + request.headers = {"x-litellm-model": "header-model"} + + with ( + patch( + "litellm.proxy.vector_store_endpoints.endpoints.get_litellm_managed_vector_store", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.vector_store_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=fake_base_process, + ), + ): + response = await vector_store_delete( + request=request, + vector_store_id="vs_plain_openai", + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert response == {"ok": True} + assert captured_data["vector_store_id"] == "vs_plain_openai" + assert captured_data["model"] == "gpt-5.2" + + +def test_vector_store_model_from_request_preserves_existing_model(): + from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_model_from_request, + ) + + request = _mock_request() + request.query_params = {"model": "query-model"} + request.headers = {"x-litellm-model": "header-model"} + data = {"model": "registry-model"} + + updated_data = _update_request_data_with_model_from_request( + data=data, request=request + ) + + assert updated_data is data + assert updated_data["model"] == "registry-model" + + +def test_vector_store_model_from_request_preserves_registry_routing(): + from litellm.proxy.vector_store_endpoints.endpoints import ( + _update_request_data_with_model_from_request, + ) + + request = _mock_request() + request.query_params = {"model": "query-model"} + request.headers = {"x-litellm-model": "header-model"} + data = { + "custom_llm_provider": "openai", + "litellm_credential_name": "vector-store-credential", + } + + updated_data = _update_request_data_with_model_from_request( + data=data, request=request + ) + + assert updated_data is data + assert "model" not in updated_data + assert updated_data["custom_llm_provider"] == "openai" + assert updated_data["litellm_credential_name"] == "vector-store-credential" + + @pytest.mark.asyncio async def test_vector_store_file_create_forces_path_id_over_body_id(): from litellm.proxy.vector_store_files_endpoints.endpoints import (