diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1269f58213a..a562a4cd101 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -237,6 +237,12 @@ async def create_response( with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): yield chunk + # Tell nginx/CDNs not to buffer streaming responses — without this, + # intermediary layers hold the entire response before forwarding it, + # making TTFB equal to total latency. + if media_type == "text/event-stream": + headers["X-Accel-Buffering"] = "no" + return StreamingResponse( combined_generator(), media_type=media_type, @@ -1318,7 +1324,7 @@ class ProxyBaseLLMRequestProcessing: request_data=request_data, ): verbose_proxy_logger.debug( - "async_data_generator: received streaming chunk - {}".format(chunk) + "async_data_generator: received streaming chunk - %s", chunk ) chunk = await proxy_logging_obj.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, @@ -1340,12 +1346,11 @@ class ProxyBaseLLMRequestProcessing: elif isinstance(chunk, dict): str_so_far += str(chunk.get("content", "")) - model_name = request_data.get("model", "") - chunk = ( - ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + if litellm.include_cost_in_streaming_usage: + model_name = request_data.get("model", "") + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( chunk, model_name ) - ) yield serialize_chunk(chunk) except Exception as e: verbose_proxy_logger.exception( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6796af09146..5e53dde1fb6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5278,6 +5278,11 @@ def _restamp_streaming_chunk_model( ) model_mismatch_logged = True + # Short-circuit: model already matches — avoid the setattr and keep the object + # "clean" (not mutated), which is a prerequisite for any future JSON caching. + if downstream_model == requested_model_from_client: + return chunk, model_mismatch_logged + if isinstance(chunk, dict): chunk["model"] = requested_model_from_client return chunk, model_mismatch_logged @@ -5302,8 +5307,7 @@ async def async_data_generator( ): verbose_proxy_logger.debug("inside generator") try: - # Use a list to accumulate response segments to avoid O(n^2) string concatenation - str_so_far_parts: list[str] = [] + str_so_far: str = "" error_message: Optional[str] = None requested_model_from_client = _get_client_requested_model_for_streaming( request_data=request_data @@ -5315,21 +5319,22 @@ async def async_data_generator( request_data=request_data, ): ### CALL HOOKS ### - modify outgoing data - # Only compute str_so_far when callbacks are registered — joining - # the accumulated parts on every chunk is O(n²) otherwise. + # Extract chunk text once; pass into hook so it doesn't re-extract. + # Accumulate str_so_far AFTER the hook so the hook receives text + # from *previous* chunks (not the current one) in str_so_far. if litellm.callbacks: + chunk_str: Optional[str] = None + if isinstance(chunk, (ModelResponse, ModelResponseStream)): + chunk_str = litellm.get_response_string(response_obj=chunk) chunk = await proxy_logging_obj.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, response=chunk, data=request_data, - str_so_far="".join(str_so_far_parts), + str_so_far=str_so_far, + response_str=chunk_str, ) - - if litellm.callbacks and isinstance( - chunk, (ModelResponse, ModelResponseStream) - ): - response_str = litellm.get_response_string(response_obj=chunk) - str_so_far_parts.append(response_str) + if chunk_str is not None: + str_so_far += chunk_str chunk, model_mismatch_logged = _restamp_streaming_chunk_model( chunk=chunk, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2c1a9beceff..23a49789230 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2011,6 +2011,7 @@ class ProxyLogging: ], user_api_key_dict: UserAPIKeyAuth, str_so_far: Optional[str] = None, + response_str: Optional[str] = None, ): """ Allow user to modify outgoing streaming data -> per chunk @@ -2022,12 +2023,11 @@ class ProxyLogging: if not litellm.callbacks: return response - from litellm.proxy.proxy_server import llm_router - - response_str: Optional[str] = None - if isinstance(response, (ModelResponse, ModelResponseStream)): - response_str = litellm.get_response_string(response_obj=response) - elif isinstance(response, dict) and self.is_a2a_streaming_response(response): + # Use pre-extracted string if caller already computed it to avoid re-extracting. + if response_str is None: + if isinstance(response, (ModelResponse, ModelResponseStream)): + response_str = litellm.get_response_string(response_obj=response) + if response_str is None and isinstance(response, dict) and self.is_a2a_streaming_response(response): from litellm.llms.a2a.common_utils import extract_text_from_a2a_response response_str = extract_text_from_a2a_response(response) @@ -2037,6 +2037,10 @@ class ProxyLogging: _callback: Optional[CustomLogger] = None if isinstance(callback, CustomGuardrail): # Main - V2 Guardrails implementation + # llm_router import is deferred here to avoid a module-level + # circular import AND to skip the lookup on every chunk when + # no guardrails are registered. + from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks ## CHECK FOR MODEL-LEVEL GUARDRAILS @@ -2059,10 +2063,18 @@ class ProxyLogging: else: _callback = callback # type: ignore if _callback is not None and isinstance(_callback, CustomLogger): - if str_so_far is not None: - complete_response = str_so_far + response_str - else: - complete_response = response_str + # Skip the await entirely when the callback doesn't override + # async_post_call_streaming_hook — the base-class default just + # returns None, so calling it is pure overhead. This is the + # common case for logging-only integrations (Datadog, OTEL, …). + if ( + "async_post_call_streaming_hook" + not in type(_callback).__dict__ + ): + continue + complete_response = ( + str_so_far + response_str if str_so_far else response_str + ) callback_response = ( await _callback.async_post_call_streaming_hook( user_api_key_dict=user_api_key_dict, @@ -2095,6 +2107,20 @@ class ProxyLogging: yield chunk return + # Fast path: if no callback overrides the iterator hook or applies a + # guardrail, skip the wrapping entirely — most callbacks only use the + # per-chunk hook, not the full-iterator hook. + _needs_iterator_wrap = any( + "async_post_call_streaming_iterator_hook" in type(cb).__dict__ + or "apply_guardrail" in type(cb).__dict__ + for cb in litellm.callbacks + if isinstance(cb, CustomLogger) or isinstance(cb, CustomGuardrail) + ) + if not _needs_iterator_wrap: + async for chunk in response: + yield chunk + return + current_response = response for callback in litellm.callbacks: diff --git a/litellm/utils.py b/litellm/utils.py index 81e772b1765..6714c97b025 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4967,7 +4967,18 @@ def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) response_obj.choices ) - # Use list accumulation to avoid O(n^2) string concatenation across choices + # Fast path: single choice (overwhelmingly common for streaming proxy traffic). + # Avoids creating a list and calling join. + if len(_choices) == 1: + choice = _choices[0] + if isinstance(choice, StreamingChoices): + content = choice.delta.content + return str(content) if content is not None else "" + if isinstance(choice, Choices): + content = choice.message.content + return str(content) if content is not None else "" + + # General path for multi-choice responses response_parts: List[str] = [] for choice in _choices: if isinstance(choice, Choices): diff --git a/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py b/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py index d9e882072cd..825fbd4ca68 100644 --- a/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py +++ b/tests/test_litellm/proxy/hooks/test_post_call_streaming_hook_integration.py @@ -304,6 +304,114 @@ async def test_streaming_iterator_hook_fast_path_no_callbacks(): assert chunk.id == f"chunk-{i}" +@pytest.mark.asyncio +async def test_streaming_hook_skips_await_for_non_overriding_callback(): + """ + Callbacks that don't override async_post_call_streaming_hook should be + skipped (no await) — the base-class default just returns None, so calling + it is pure overhead. This covers logging-only integrations. + """ + + class LoggingOnlyCallback(CustomLogger): + """Overrides only log_success, not the per-chunk hook.""" + + def __init__(self): + self.hook_called = False + + async def async_log_success_event( + self, kwargs, response_obj, start_time, end_time + ): + pass + + logger = LoggingOnlyCallback() + assert "async_post_call_streaming_hook" not in type(logger).__dict__ + + with patch("litellm.callbacks", [logger]): + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="Hello", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + ) + + # Original response must be unchanged (hook was not awaited) + assert result is original_response + assert logger.hook_called is False + + +@pytest.mark.asyncio +async def test_streaming_hook_uses_pre_extracted_response_str(): + """ + When response_str is supplied by the caller, the hook must use it without + calling get_response_string again — the received text in the callback should + match what the caller passed in. + """ + + class CapturingLogger(CustomLogger): + def __init__(self): + self.received_text: str = "" + + async def async_post_call_streaming_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: str, + ): + self.received_text = response + return None + + logger = CapturingLogger() + + with patch("litellm.callbacks", [logger]): + from litellm.caching.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + original_response = ModelResponseStream( + id="original-stream", + choices=[ + StreamingChoices( + delta=Delta(content="actual chunk text", role="assistant"), + index=0, + ) + ], + model="test-model", + ) + + data = {"model": "test-model"} + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + # Pass pre-extracted response_str and a non-empty str_so_far + await proxy_logging.async_post_call_streaming_hook( + data=data, + response=original_response, + user_api_key_dict=user_api_key_dict, + str_so_far="previous text ", + response_str="actual chunk text", + ) + + # The callback receives str_so_far + response_str + assert logger.received_text == "previous text actual chunk text" + + @pytest.mark.asyncio async def test_streaming_hook_handles_exceptions(): """ diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 35cb290fccd..573be7a795a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -3599,3 +3599,43 @@ class TestValidateAndFixThinkingParam: validate_and_fix_thinking_param(thinking=thinking) assert "budgetTokens" in thinking assert "budget_tokens" not in thinking + + +class TestGetResponseString: + """Tests for get_response_string, including the single-choice fast path.""" + + def test_single_streaming_choice(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + from litellm.utils import get_response_string + + chunk = ModelResponseStream( + id="chatcmpl-1", + choices=[StreamingChoices(delta=Delta(content="hello"), index=0)], + model="gpt-4.1-mini", + ) + assert get_response_string(chunk) == "hello" + + def test_single_streaming_choice_none_content(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + from litellm.utils import get_response_string + + chunk = ModelResponseStream( + id="chatcmpl-1", + choices=[StreamingChoices(delta=Delta(content=None), index=0)], + model="gpt-4.1-mini", + ) + assert get_response_string(chunk) == "" + + def test_multi_choice_streaming(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + from litellm.utils import get_response_string + + chunk = ModelResponseStream( + id="chatcmpl-1", + choices=[ + StreamingChoices(delta=Delta(content="foo"), index=0), + StreamingChoices(delta=Delta(content="bar"), index=1), + ], + model="gpt-4.1-mini", + ) + assert get_response_string(chunk) == "foobar"