diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 644253ceac7..b50fb9cd716 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1326,6 +1326,7 @@ _MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS = ( "/batches", "/skills", "/evals", + "/vector_stores", ) _MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS = ( "/files", diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 4e7890e6ed8..d8f8d349614 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -6,6 +6,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i LiteLLM_ManagedVectorStore, ) from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import MODEL_ROUTING_HEADER_NAME from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import jsonify_object @@ -88,6 +89,35 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( return data +_EXPLICIT_ROUTING_KEYS = ( + "model", + "custom_llm_provider", + "litellm_credential_name", + "api_key", + "api_base", +) + + +def _apply_model_routing_hint(data: Dict, request: Request) -> Dict: + """ + Resolve ``?model=`` / ``x-litellm-model`` into ``data["model"]`` so the router + picks the requested deployment. + + Routing already established by the request body or by the managed vector store + registry wins, so the hint only applies when nothing else selected a deployment. + The hint is authorized against the key/team model allowlists by + ``get_model_from_request``, which reads the same query param and header. + """ + if any(data.get(key) for key in _EXPLICIT_ROUTING_KEYS): + return data + + model_hint = request.query_params.get("model") or request.headers.get(MODEL_ROUTING_HEADER_NAME) + if not model_hint: + return data + + return {**data, "model": model_hint} + + @router.post( "/v1/vector_stores/{vector_store_id:path}/search", dependencies=[Depends(user_api_key_auth)], @@ -130,6 +160,7 @@ async def vector_store_search( 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 = _apply_model_routing_hint(data=data, request=request) # The managed_vector_stores pre-call hook will handle: # 1. Decoding managed vector store IDs @@ -244,6 +275,8 @@ async def vector_store_create( return response + data = _apply_model_routing_hint(data=data, request=request) + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -306,6 +339,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 = _apply_model_routing_hint(data=data, request=request) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -376,6 +410,7 @@ async def vector_store_list( data["limit"] = limit if order is not None: data["order"] = order + data = _apply_model_routing_hint(data=data, request=request) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -442,6 +477,7 @@ async def vector_store_update( 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 = _apply_model_routing_hint(data=data, request=request) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -505,6 +541,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 = _apply_model_routing_hint(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 1610d76efb7..d9649506ed0 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -466,6 +466,25 @@ def test_get_model_from_request_authorizes_all_file_routing_model_sources(): } +def test_get_model_from_request_includes_vector_store_query_and_header_model(): + assert ( + get_model_from_request( + request_data={}, + route="/v1/vector_stores", + request_query_params={"model": "restricted-model"}, + ) + == "restricted-model" + ) + assert ( + get_model_from_request( + request_data={}, + route="/v1/vector_stores/vs_123", + request_headers={"x-litellm-model": "restricted-model"}, + ) + == "restricted-model" + ) + + def test_get_model_from_request_extracts_simple_encoded_file_id_model(): from litellm.proxy.openai_files_endpoints.common_utils import ( encode_file_id_with_model, 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 e7de8b54e4e..3f86ded0774 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 @@ -18,6 +18,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.vector_store_endpoints.endpoints import ( + _apply_model_routing_hint, _update_request_data_with_litellm_managed_vector_store_registry, index_create, ) @@ -592,6 +593,116 @@ async def test_update_request_data_passes_through_legacy_embedding_config(): resolve_mock.assert_not_awaited() +def _model_hint_request(query_params=None, headers=None): + request = MagicMock(spec=Request) + request.query_params = query_params or {} + request.headers = headers or {} + return request + + +def test_apply_model_routing_hint_uses_query_param(): + result = _apply_model_routing_hint( + data={"vector_store_id": "vs_123"}, + request=_model_hint_request(query_params={"model": "azure-vector-store"}), + ) + + assert result["model"] == "azure-vector-store" + + +def test_apply_model_routing_hint_uses_header(): + result = _apply_model_routing_hint( + data={"vector_store_id": "vs_123"}, + request=_model_hint_request(headers={"x-litellm-model": "azure-vector-store"}), + ) + + assert result["model"] == "azure-vector-store" + + +def test_apply_model_routing_hint_prefers_body_model(): + result = _apply_model_routing_hint( + data={"vector_store_id": "vs_123", "model": "body-model"}, + request=_model_hint_request(query_params={"model": "query-model"}), + ) + + assert result["model"] == "body-model" + + +def test_apply_model_routing_hint_preserves_registry_routing(): + """A managed vector store row already selected the provider/credentials, so a + caller-supplied hint must not repoint the request at another deployment.""" + data = { + "vector_store_id": "vs_123", + "custom_llm_provider": "azure_ai", + "api_key": "registry-key", + } + + result = _apply_model_routing_hint( + data=data, + request=_model_hint_request(query_params={"model": "other-deployment"}), + ) + + assert "model" not in result + assert result == data + + +def test_apply_model_routing_hint_without_hint_is_noop(): + data = {"vector_store_id": "vs_123"} + + assert _apply_model_routing_hint(data=data, request=_model_hint_request()) == data + + +@pytest.mark.asyncio +async def test_vector_store_retrieve_routes_by_model_query_param(): + from litellm.proxy.vector_store_endpoints import endpoints as vector_store_endpoints + + processor = MagicMock() + processor.base_process_llm_request = AsyncMock(return_value={"id": "vs_123"}) + + with ( + patch.object( + vector_store_endpoints, + "get_litellm_managed_vector_store", + AsyncMock(return_value=None), + ), + patch.object( + vector_store_endpoints, + "ProxyBaseLLMRequestProcessing", + MagicMock(return_value=processor), + ) as mock_processor_cls, + ): + response = await vector_store_endpoints.vector_store_retrieve( + request=_model_hint_request(query_params={"model": "azure-vector-store"}), + vector_store_id="vs_123", + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response == {"id": "vs_123"} + assert mock_processor_cls.call_args.kwargs["data"]["model"] == "azure-vector-store" + + +@pytest.mark.asyncio +async def test_vector_store_list_routes_by_model_header(): + from litellm.proxy.vector_store_endpoints import endpoints as vector_store_endpoints + + processor = MagicMock() + processor.base_process_llm_request = AsyncMock(return_value={"data": []}) + + with patch.object( + vector_store_endpoints, + "ProxyBaseLLMRequestProcessing", + MagicMock(return_value=processor), + ) as mock_processor_cls: + response = await vector_store_endpoints.vector_store_list( + request=_model_hint_request(headers={"x-litellm-model": "azure-vector-store"}), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert response == {"data": []} + assert mock_processor_cls.call_args.kwargs["data"]["model"] == "azure-vector-store" + + class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function."""