diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 8cb2ee09370..3e97b480779 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -29,12 +29,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: litellm_logging_obj: LiteLLMLoggingObj, request_body: dict, model: str, + hidden_params: Optional[Dict[str, Any]] = None, ): self.litellm_logging_obj = litellm_logging_obj self.request_body = request_body self.start_time = datetime.now() self.collected_chunks: List[bytes] = [] self.model = model + self._hidden_params: Dict[str, Any] = hidden_params or {} async def _handle_async_streaming_logging( self, @@ -76,11 +78,13 @@ class GoogleGenAIGenerateContentStreamingIterator( litellm_metadata: dict, custom_llm_provider: str, request_body: Optional[dict] = None, + hidden_params: Optional[Dict[str, Any]] = None, ): super().__init__( litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + hidden_params=hidden_params, ) self.response = response self.model = model @@ -130,11 +134,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator( litellm_metadata: dict, custom_llm_provider: str, request_body: Optional[dict] = None, + hidden_params: Optional[Dict[str, Any]] = None, ): super().__init__( litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + hidden_params=hidden_params, ) self.response = response self.model = model diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c9abaad908..ce2ea9cec39 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -150,6 +150,30 @@ else: LiteLLMLoggingObj = Any +def _google_genai_streaming_hidden_params( + *, + api_base: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + response_headers: httpx.Headers, +) -> Dict[str, Any]: + """Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params).""" + from litellm.litellm_core_utils.core_helpers import process_response_headers + + _model_info: Dict[str, Any] = dict( + getattr(litellm_params, "model_info", None) or {} + ) + _raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or "" + _model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id) + return { + "model_id": _model_id, + "api_base": api_base, + "cache_key": "", + "response_cost": "", + "additional_headers": process_response_headers(response_headers), + } + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -4495,9 +4519,9 @@ class BaseLLMHTTPHandler: # Second: Execute agentic loop # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider[ - "custom_llm_provider" - ] = custom_llm_provider + kwargs_with_provider["custom_llm_provider"] = ( + custom_llm_provider + ) agentic_response = await callback.async_run_agentic_loop( tools=tool_calls, model=model, @@ -4613,9 +4637,9 @@ class BaseLLMHTTPHandler: # Second: Execute agentic loop # Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider[ - "custom_llm_provider" - ] = custom_llm_provider + kwargs_with_provider["custom_llm_provider"] = ( + custom_llm_provider + ) agentic_response = ( await callback.async_run_chat_completion_agentic_loop( tools=tool_calls, @@ -5099,7 +5123,10 @@ class BaseLLMHTTPHandler: _is_async: bool = False, fake_stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, - ) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]: + ) -> Union[ + ImageResponse, + Coroutine[Any, Any, ImageResponse], + ]: """ Handles image edit requests. @@ -5311,7 +5338,10 @@ class BaseLLMHTTPHandler: fake_stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, api_key: Optional[str] = None, - ) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]: + ) -> Union[ + ImageResponse, + Coroutine[Any, Any, ImageResponse], + ]: """ Handles image generation requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -5551,7 +5581,10 @@ class BaseLLMHTTPHandler: fake_stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, api_key: Optional[str] = None, - ) -> Union[VideoObject, Coroutine[Any, Any, VideoObject],]: + ) -> Union[ + VideoObject, + Coroutine[Any, Any, VideoObject], + ]: """ Handles video generation requests. When _is_async=True, returns a coroutine instead of making the call directly. @@ -10031,6 +10064,12 @@ class BaseLLMHTTPHandler: litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, request_body=data, + hidden_params=_google_genai_streaming_hidden_params( + api_base=api_base, + litellm_params=litellm_params, + logging_obj=logging_obj, + response_headers=response.headers, + ), ) else: response = sync_httpx_client.post( @@ -10140,6 +10179,12 @@ class BaseLLMHTTPHandler: litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, request_body=data, + hidden_params=_google_genai_streaming_hidden_params( + api_base=api_base, + litellm_params=litellm_params, + logging_obj=logging_obj, + response_headers=response.headers, + ), ) else: response = await async_httpx_client.post( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d86872302a5..cb2b08199b9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -576,7 +576,6 @@ class ProxyBaseLLMRequestProcessing: logging_obj: LiteLLMLoggingObj, version: Optional[str], proxy_logging_obj: ProxyLogging, - llm_router: Optional[Router] = None, ) -> dict[str, str]: """ Build LiteLLM proxy response headers for routes that call the LLM directly @@ -601,9 +600,6 @@ class ProxyBaseLLMRequestProcessing: ) additional_headers = hidden_params.get("additional_headers", {}) or {} - if llm_router is not None: - request_data["deployment"] = llm_router.get_deployment(model_id=model_id) - custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, call_id=logging_obj.litellm_call_id, @@ -874,20 +870,18 @@ class ProxyBaseLLMRequestProcessing: return _payload_str = json.dumps(self.data, default=str) if len(_payload_str) > MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG: + _key_count = len(self.data) if isinstance(self.data, dict) else 0 verbose_proxy_logger.debug( - "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Keys: %s", + "Request received by LiteLLM: payload too large to log (%d bytes, limit %d). Dict key count: %d; type: %s", len(_payload_str), MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, - ( - list(self.data.keys()) - if isinstance(self.data, dict) - else type(self.data).__name__ - ), + _key_count, + type(self.data).__name__, ) else: verbose_proxy_logger.debug( "Request received by LiteLLM:\n%s", - json.dumps(self.data, indent=4, default=str), + _payload_str, ) async def base_process_llm_request( # noqa: PLR0915 diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 9eca414614c..6ada8f58783 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -82,7 +82,6 @@ async def google_generate_content( logging_obj=logging_obj, version=version, proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, ) fastapi_response.headers.update(success_headers) return response @@ -158,7 +157,6 @@ async def google_stream_generate_content( logging_obj=logging_obj, version=version, proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, ) # Check if response is an async iterator (streaming response) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index a3512bc6e7b..39a1193635f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2,12 +2,16 @@ import os import sys from unittest.mock import AsyncMock, Mock, patch +import httpx import pytest sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import ( + BaseLLMHTTPHandler, + _google_genai_streaming_hidden_params, +) from litellm.types.router import GenericLiteLLMParams @@ -81,7 +85,7 @@ async def test_async_anthropic_messages_handler_extra_headers(): extra_headers from kwargs with proper priority. """ handler = BaseLLMHTTPHandler() - + # Mock the config mock_config = Mock() mock_config.validate_anthropic_messages_environment = Mock( @@ -90,7 +94,7 @@ async def test_async_anthropic_messages_handler_extra_headers(): mock_config.transform_anthropic_messages_request = Mock( return_value={"model": "claude-3-opus-20240229", "messages": []} ) - + # Mock the client mock_client = AsyncMock() mock_response = Mock() @@ -104,13 +108,13 @@ async def test_async_anthropic_messages_handler_extra_headers(): "stop_reason": "end_turn", } mock_client.post = AsyncMock(return_value=mock_response) - + # Mock logging object mock_logging_obj = Mock() mock_logging_obj.update_environment_variables = Mock() mock_logging_obj.model_call_details = {} mock_logging_obj.stream = False - + # Test case 1: Only extra_headers in kwargs kwargs = { "extra_headers": { @@ -118,20 +122,21 @@ async def test_async_anthropic_messages_handler_extra_headers(): "X-Auth-Token": "token123", } } - + with patch( "litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers" ) as mock_provider_headers: mock_provider_headers.return_value = None - + # Capture what headers are passed to validate_anthropic_messages_environment captured_headers = {} + def capture_validate(*args, **kwargs): captured_headers.update(kwargs.get("headers", {})) return ({"x-api-key": "test-key"}, "https://api.anthropic.com") - + mock_config.validate_anthropic_messages_environment = capture_validate - + try: await handler.async_anthropic_messages_handler( model="claude-3-opus-20240229", @@ -146,7 +151,7 @@ async def test_async_anthropic_messages_handler_extra_headers(): ) except Exception: pass # We're testing header extraction, not the full flow - + # Verify extra_headers were extracted and merged assert "X-Custom-Header" in captured_headers assert captured_headers["X-Custom-Header"] == "from-kwargs" @@ -219,9 +224,11 @@ async def test_async_anthropic_messages_handler_passes_litellm_metadata(): mock_logging_obj.update_from_kwargs.assert_called_once() call_kwargs = mock_logging_obj.update_from_kwargs.call_args - kwargs_arg = call_kwargs.kwargs.get( - "kwargs", call_kwargs[1].get("kwargs", {}) - ) if call_kwargs.kwargs else call_kwargs[1].get("kwargs", {}) + kwargs_arg = ( + call_kwargs.kwargs.get("kwargs", call_kwargs[1].get("kwargs", {})) + if call_kwargs.kwargs + else call_kwargs[1].get("kwargs", {}) + ) assert "litellm_metadata" in kwargs_arg assert kwargs_arg["litellm_metadata"]["model_info"] == custom_model_info @@ -234,7 +241,7 @@ async def test_async_anthropic_messages_handler_header_priority(): forwarded < extra_headers < provider_specific """ handler = BaseLLMHTTPHandler() - + # Mock the config mock_config = Mock() mock_client = AsyncMock() @@ -242,31 +249,32 @@ async def test_async_anthropic_messages_handler_header_priority(): mock_logging_obj.update_environment_variables = Mock() mock_logging_obj.model_call_details = {} mock_logging_obj.stream = False - + # Test with all three header sources kwargs = { "headers": {"X-Priority": "forwarded", "X-Forwarded-Only": "keep"}, "extra_headers": {"X-Priority": "extra", "X-Extra-Only": "also-keep"}, } - + with patch( "litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers" ) as mock_provider_headers: mock_provider_headers.return_value = { "X-Priority": "provider", - "X-Provider-Only": "keep-this-too" + "X-Provider-Only": "keep-this-too", } - + captured_headers = {} + def capture_validate(*args, **kwargs): captured_headers.update(kwargs.get("headers", {})) return ({"x-api-key": "test-key"}, "https://api.anthropic.com") - + mock_config.validate_anthropic_messages_environment = capture_validate mock_config.transform_anthropic_messages_request = Mock( return_value={"model": "claude-3-opus-20240229", "messages": []} ) - + try: await handler.async_anthropic_messages_handler( model="claude-3-opus-20240229", @@ -281,10 +289,36 @@ async def test_async_anthropic_messages_handler_header_priority(): ) except Exception: pass - + # Verify priority: provider_specific should win assert captured_headers["X-Priority"] == "provider" # Verify all unique headers from different sources are present assert captured_headers["X-Forwarded-Only"] == "keep" assert captured_headers["X-Extra-Only"] == "also-keep" assert captured_headers["X-Provider-Only"] == "keep-this-too" + + +def test_google_genai_streaming_hidden_params_model_info_and_router_fallback(): + logging_obj = Mock() + logging_obj.get_router_model_id = Mock(return_value="router-model-id") + + from_model_info = _google_genai_streaming_hidden_params( + api_base="https://generativelanguage.googleapis.com/v1beta", + litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}), + logging_obj=logging_obj, + response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}), + ) + assert from_model_info["model_id"] == "info-id" + assert ( + from_model_info["api_base"] + == "https://generativelanguage.googleapis.com/v1beta" + ) + assert isinstance(from_model_info["additional_headers"], dict) + + from_router = _google_genai_streaming_hidden_params( + api_base="https://x", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + response_headers=httpx.Headers({}), + ) + assert from_router["model_id"] == "router-model-id" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 124bd43cc8a..c00a6fba58c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -251,9 +251,6 @@ class TestProxyBaseLLMRequestProcessing: return_value={"x-ratelimit-remaining-requests": "999"} ) - llm_router = MagicMock() - llm_router.get_deployment.return_value = {"litellm_params": {}} - headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( response=_FakeGenaiResponse(), request_data={"model": "gemini/gemini-1.5-flash"}, @@ -262,7 +259,6 @@ class TestProxyBaseLLMRequestProcessing: logging_obj=logging_obj, version="9.9.9", proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, ) assert headers["x-litellm-call-id"] == "call-id-test" @@ -271,10 +267,92 @@ class TestProxyBaseLLMRequestProcessing: assert headers["llm_provider-ratelimit-requests"] == "1000" assert headers["x-ratelimit-remaining-requests"] == "999" proxy_logging_obj.post_call_response_headers_hook.assert_awaited_once() - llm_router.get_deployment.assert_called_once_with( - model_id="deployment-model-id" + + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_streaming_style_iterator(self): + """AsyncGoogleGenAIGenerateContentStreamingIterator sets _hidden_params at init; headers must propagate.""" + + class _FakeStreamLike: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + _hidden_params = { + "model_id": "stream-model-id", + "api_base": "https://generativelanguage.googleapis.com/v1beta", + "cache_key": "", + "response_cost": "", + "additional_headers": {"llm_provider-x": "y"}, + } + + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + logging_obj = MagicMock() + logging_obj.litellm_call_id = "cid-stream" + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_FakeStreamLike(), + request_data={"model": "gemini/gemini-2.0-flash"}, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="1.0.0", + proxy_logging_obj=proxy_logging_obj, ) + assert headers["x-litellm-model-id"] == "stream-model-id" + assert headers["x-litellm-model-api-base"] == ( + "https://generativelanguage.googleapis.com/v1beta" + ) + assert headers["llm_provider-x"] == "y" + + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_no_hidden_params_metadata_fallback( + self, + ): + """When response has no _hidden_params, model_id can still come from litellm_metadata.""" + + class _BareResponse: + pass + + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + logging_obj = MagicMock() + logging_obj.litellm_call_id = "cid-meta" + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_BareResponse(), + request_data={ + "model": "gemini/gemini-1.5-flash", + "litellm_metadata": {"model_info": {"id": "meta-model-id"}}, + }, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="1.0.0", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-model-id"] == "meta-model-id" + @pytest.mark.asyncio async def test_add_litellm_data_to_request_with_stream_timeout_header(self): """