diff --git a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py index 9833891f1de..e992c47f283 100644 --- a/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py +++ b/litellm/proxy/openai_files_endpoints/file_content_streaming_handler.py @@ -64,14 +64,16 @@ class FileContentStreamingHandler: ) -> StreamingResponse: effective_custom_llm_provider = custom_llm_provider if should_route: + if credentials is None or credentials.get("custom_llm_provider") is None: + raise ValueError( + "Model-based file routing requires credentials with custom_llm_provider" + ) prepare_data_with_credentials( data=data, - credentials=credentials, # type: ignore[arg-type] + credentials=credentials, file_id=original_file_id, ) - effective_custom_llm_provider = cast( - str, credentials["custom_llm_provider"] - ) + effective_custom_llm_provider = cast(str, credentials["custom_llm_provider"]) stream_result = cast( FileContentStreamingResult, diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index aeca75e9365..66afded70a4 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -5,6 +5,7 @@ from unittest.mock import ANY, AsyncMock import pytest import respx +import httpx from fastapi.testclient import TestClient from pytest_mock import MockerFixture @@ -22,7 +23,7 @@ from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import FileContentStreamingHandler, ) from litellm.proxy.proxy_server import app -from litellm.types.llms.openai import OpenAIFileObject +from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIFileObject client = TestClient(app) from litellm.caching.caching import DualCache @@ -1713,3 +1714,67 @@ def test_get_file_content_streams_with_routed_provider( assert captured_kwargs["stream"] is True proxy_logging_obj.update_request_status.assert_awaited_once() proxy_logging_obj.post_call_failure_hook.assert_not_called() + + +def test_get_file_content_non_openai_provider_skips_streaming_handler( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs = {} + + async def _mock_afile_content(**kwargs): + captured_kwargs.update(kwargs) + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=b"azure-bytes", + headers={ + "content-type": "application/octet-stream", + "content-length": "11", + }, + ) + ) + + mock_streaming_response = mocker.AsyncMock() + + monkeypatch.setattr(litellm, "afile_content", _mock_afile_content) + monkeypatch.setattr( + FileContentStreamingHandler, + "get_streaming_file_content_response", + mock_streaming_response, + ) + monkeypatch.setattr( + "litellm.proxy.openai_files_endpoints.files_endpoints.handle_model_based_routing", + lambda **kwargs: (False, None, None, None), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + + try: + response = client.get( + "/v1/files/file-abc123/content", + headers={ + "Authorization": "Bearer test-key", + "custom-llm-provider": "azure", + }, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert response.content == b"azure-bytes" + assert captured_kwargs["custom_llm_provider"] == "azure" + assert "stream" not in captured_kwargs + mock_streaming_response.assert_not_awaited() + proxy_logging_obj.post_call_failure_hook.assert_not_called()