diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index f856fe0f2b3..613e16925af 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -12,3 +12,5 @@ from typing import Final # When True, suppresses async logging and billing for internal sub-calls # (e.g., emulated file-search steps that make nested LLM calls). is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False) + +is_proxy_stream_header_prefetch: Final[ContextVar[bool]] = ContextVar("is_proxy_stream_header_prefetch", default=False) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 8fe60876b4e..d82a9ef2036 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -221,7 +221,6 @@ class LLMCachingHandler: # Init cache timing metrics ######################################################### cache_check_start_time: Final = time.perf_counter() - cache_check_end_time: float | None = None ######################################################### parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) kwargs["parent_otel_span"] = parent_otel_span @@ -233,19 +232,17 @@ class LLMCachingHandler: kwargs=kwargs, args=args, ) - cache_check_end_time = time.perf_counter() + cache_duration_ms: Final = (time.perf_counter() - cache_check_start_time) * 1000 if cached_result is not None and not isinstance(cached_result, list): verbose_logger.debug("Cache Hit!") cache_hit: Final = True - end_time: Final = datetime.datetime.now() model, custom_llm_provider, _, _ = litellm.get_llm_provider( model=model, custom_llm_provider=kwargs.get("custom_llm_provider", None), api_base=kwargs.get("api_base", None), api_key=kwargs.get("api_key", None), ) - cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000 self._update_litellm_logging_obj_environment( logging_obj=logging_obj, model=model, @@ -267,6 +264,7 @@ class LLMCachingHandler: custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) + end_time: Final = datetime.datetime.now() if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs): # LOG SUCCESS self._async_log_cache_hit_on_callbacks( @@ -344,44 +342,40 @@ class LLMCachingHandler: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) print_verbose("Checking Sync Cache") + cache_check_start_time: Final = time.perf_counter() cached_result = litellm.cache.get_cache(**new_kwargs) + cache_check_end_time: Final = time.perf_counter() if cached_result is not None: if "detail" in cached_result: # implies an error occurred pass else: - call_type = original_function.__name__ + cache_hit: Final = True + resolved_model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model or "", + custom_llm_provider=kwargs.get("custom_llm_provider", None), + api_base=kwargs.get("api_base", None), + api_key=kwargs.get("api_key", None), + ) + cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000 + self._update_litellm_logging_obj_environment( + logging_obj=logging_obj, + model=f"{custom_llm_provider}/{resolved_model}", + kwargs=kwargs, + cached_result=cached_result, + is_async=False, + cache_duration_ms=cache_duration_ms, + ) cached_result = self._convert_cached_result_to_model_response( cached_result=cached_result, - call_type=call_type, + call_type=original_function.__name__, kwargs=kwargs, logging_obj=logging_obj, model=model, custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) - - # LOG SUCCESS - cache_hit: Final = True end_time: Final = datetime.datetime.now() - ( - model, - custom_llm_provider, - dynamic_api_key, - api_base, - ) = litellm.get_llm_provider( - model=model or "", - custom_llm_provider=kwargs.get("custom_llm_provider", None), - api_base=kwargs.get("api_base", None), - api_key=kwargs.get("api_key", None), - ) - self._update_litellm_logging_obj_environment( - logging_obj=logging_obj, - model=f"{custom_llm_provider}/{model}", - kwargs=kwargs, - cached_result=cached_result, - is_async=False, - ) if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs): logging_obj.handle_sync_success_callbacks_for_async_calls( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ff9a5d41c2d..dbef0e780f8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -544,6 +544,7 @@ class Logging(LiteLLMLoggingBaseClass): # Init Caching related details self.caching_details: CachingDetails | None = None + self.response_timing_metrics: Mapping[str, float] = {} # Passthrough endpoint guardrails config for field targeting self.passthrough_guardrails_config: dict[str, Any] | None = None @@ -563,6 +564,9 @@ class Logging(LiteLLMLoggingBaseClass): self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None + def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: + self.response_timing_metrics = dict(timing_metrics) + def process_dynamic_callbacks(self): """ Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -5965,6 +5969,9 @@ def get_standard_logging_object_payload( clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params) if clean_hidden_params["response_cost"] is None and raw_response_cost is not None: clean_hidden_params["response_cost"] = llm_response_cost + request_overhead_ms: Final = logging_obj.response_timing_metrics.get("litellm_overhead_time_ms") + if clean_hidden_params["litellm_overhead_time_ms"] is None and request_overhead_ms is not None: + clean_hidden_params["litellm_overhead_time_ms"] = request_overhead_ms model_cost_information: Final = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index a375560288f..795d1a61bc2 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -1,10 +1,14 @@ import datetime +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final +from litellm._internal_context import is_internal_call, is_proxy_stream_header_prefetch from litellm.constants import LITELLM_DETAILED_TIMING from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.logging_utils import LiteLLMLoggingObject +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ( EmbeddingResponse, HiddenParams, @@ -13,6 +17,73 @@ from litellm.types.utils import ( ) +def request_timing_metrics( + start_time: datetime.datetime, + end_time: datetime.datetime, + logging_obj: LiteLLMLoggingObject, +) -> Mapping[str, float]: + """Total request duration and the LiteLLM overhead within it, in milliseconds.""" + total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000 + llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") + caching_details: Final = logging_obj.caching_details + cache_duration_ms: Final = ( + caching_details.get("cache_duration_ms") + if caching_details is not None and caching_details.get("cache_hit") is True + else None + ) + + if cache_duration_ms is not None: + overhead_ms: float | None = total_response_time_ms - cache_duration_ms + elif llm_api_duration_ms is not None: + overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) + else: + overhead_ms = None + + timing_items: Final = ( + ( + ("_response_ms", total_response_time_ms), + ("litellm_overhead_time_ms", overhead_ms), + ) + if overhead_ms is not None + else (("_response_ms", total_response_time_ms),) + ) + return MappingProxyType(dict(timing_items)) + + +def refresh_response_timing_metrics(logging_obj: LiteLLMLoggingObject) -> None: + start_time: Final = getattr(logging_obj, "start_time", None) + llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") + if ( + not is_proxy_stream_header_prefetch.get() + or is_internal_call.get() + or not isinstance(start_time, datetime.datetime) + or not isinstance(llm_api_duration_ms, (int, float)) + ): + return + logging_obj.set_response_timing_metrics( + request_timing_metrics( + start_time=start_time, + end_time=datetime.datetime.now(tz=start_time.tzinfo), + logging_obj=logging_obj, + ) + ) + + +async def prefetch_proxy_stream_for_timing(completion_response: object) -> None: + if ( + not is_proxy_stream_header_prefetch.get() + or is_internal_call.get() + or not isinstance(completion_response, CustomStreamWrapper) + or completion_response.custom_llm_provider != "vertex_ai_beta" + or completion_response.completion_stream is not None + or completion_response.make_call is None + ): + return + + await completion_response.fetch_stream() + refresh_response_timing_metrics(completion_response.logging_obj) + + class ResponseMetadata: """ Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses @@ -78,36 +149,15 @@ class ResponseMetadata: logging_obj: LiteLLMLoggingObject, ) -> None: """Set response timing metrics""" - total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000 + timing_metrics: Final = request_timing_metrics(start_time, end_time, logging_obj) + total_response_time_ms: Final = timing_metrics["_response_ms"] # Set total response time if supported if self.supports_response_time: self.result._response_ms = total_response_time_ms - ######################################################### - # 1. Add _response_ms total duration - ######################################################### - self._update_hidden_params( - { - "_response_ms": total_response_time_ms, - } - ) + self._update_hidden_params(dict(timing_metrics)) - ######################################################### - # 2. Add LiteLLM overhead duration - ######################################################### - llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") - if llm_api_duration_ms is not None: - overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4) - self._update_hidden_params( - { - "litellm_overhead_time_ms": overhead_ms, - } - ) - - ######################################################### - # 3. Add callback processing duration - ######################################################### callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None) if callback_duration_ms is not None: self._update_hidden_params( @@ -116,25 +166,7 @@ class ResponseMetadata: } ) - ######################################################### - # 4. Add duration for reading from cache - # In this case overhead from litellm is the difference between the cache read duration and the total response time - ######################################################### - if ( - logging_obj.caching_details is not None - and logging_obj.caching_details.get("cache_hit") is True - and (cache_duration_ms := logging_obj.caching_details.get("cache_duration_ms")) is not None - ): - overhead_ms = total_response_time_ms - cache_duration_ms - self._update_hidden_params( - { - "litellm_overhead_time_ms": overhead_ms, - } - ) - - ######################################################### - # 5. Detailed per-phase timing (opt-in via env var) - ######################################################### + llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms") if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None: detailed: Final[dict] = { "timing_llm_api_ms": round(llm_api_duration_ms, 4), @@ -178,7 +210,17 @@ def update_response_metadata( - response._hidden_params["litellm_overhead_time_ms"] - response.response_time_ms """ - if result is None or not hasattr(result, "_hidden_params"): + if result is None: + return + + if isinstance(result, dict): + if not is_internal_call.get(): + logging_obj.set_response_timing_metrics(request_timing_metrics(start_time, end_time, logging_obj)) + return + + if not hasattr(result, "_hidden_params"): + if getattr(logging_obj, "stream", False) is True and not is_internal_call.get(): + logging_obj.set_response_timing_metrics(request_timing_metrics(start_time, end_time, logging_obj)) return metadata: Final = ResponseMetadata(result) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 9d61701d26d..08f84062607 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -11,6 +11,9 @@ from typing_extensions import TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + prefetch_proxy_stream_for_timing, +) from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( AnthropicAdapter, ) @@ -613,6 +616,8 @@ class LiteLLMMessagesToCompletionTransformationHandler: ) completion_response: Final = await litellm.acompletion(**completion_kwargs) + if stream is True: + await prefetch_proxy_stream_for_timing(completion_response) if stream: transformed_stream: Final = ANTHROPIC_ADAPTER.translate_completion_output_params_streaming( diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index cc5879df56d..f15cfdf43a1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -254,6 +254,11 @@ class _CombinedChunkSplitter: ) return self._buffer.popleft() + async def aclose(self) -> None: + aclose: Final = getattr(self._stream, "aclose", None) + if aclose is not None: + await aclose() + class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): """ @@ -955,6 +960,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): return {"type": "message_stop"} raise StopAsyncIteration + async def aclose(self) -> None: + aclose: Final = getattr(self.completion_stream, "aclose", None) + if aclose is not None: + await aclose() + def anthropic_sse_wrapper(self) -> Iterator[bytes]: """ Convert AnthropicStreamWrapper dict chunks to Server-Sent Events format. @@ -972,18 +982,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): yield chunk async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]: - """ - Async version of anthropic_sse_wrapper. - Convert AnthropicStreamWrapper dict chunks to Server-Sent Events format. - """ - async for chunk in self: - if isinstance(chunk, dict): - event_type: str = str(chunk.get("type", "message")) - payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" - yield payload.encode() - else: - # For non-dict chunks, forward the original value unchanged - yield chunk + try: + async for chunk in self: + if isinstance(chunk, dict): + event_type: str = str(chunk.get("type", "message")) + payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" + yield payload.encode() + else: + yield chunk + finally: + await self.aclose() def _increment_content_block_index(self): self.current_content_block_index += 1 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 573ba85416f..834f7d564a2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2872,6 +2872,7 @@ class BaseLLMHTTPHandler: headers=headers, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + logging_obj=logging_obj, **body_kwargs, ) @@ -2903,6 +2904,7 @@ class BaseLLMHTTPHandler: url=api_base, headers=headers, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + logging_obj=logging_obj, **body_kwargs, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ff6c8d1b1f8..2cadb2fdb23 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -18,11 +18,13 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from starlette.types import Receive, Scope, Send import litellm +from litellm._internal_context import is_proxy_stream_header_prefetch from litellm._logging import _redact_string, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, DEFAULT_MAX_RECURSE_DEPTH, + EMPTY_MAPPING, LITELLM_DETAILED_TIMING, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, @@ -1481,8 +1483,31 @@ class ProxyBaseLLMRequestProcessing: **kwargs, ) -> dict: exclude_values: Final = {"", None, "None"} - hidden_params = hidden_params or {} - + response_hidden_params: Final = hidden_params or EMPTY_MAPPING + request_timing_metrics: Final[Mapping[str, float]] = ( + getattr(litellm_logging_obj, "response_timing_metrics", EMPTY_MAPPING) + if litellm_logging_obj is not None + else EMPTY_MAPPING + ) + response_duration_ms: Final = next( + ( + value + for value in (response_hidden_params.get("_response_ms"), request_timing_metrics.get("_response_ms")) + if value is not None + ), + None, + ) + overhead_duration_ms: Final = next( + ( + value + for value in ( + response_hidden_params.get("litellm_overhead_time_ms"), + request_timing_metrics.get("litellm_overhead_time_ms"), + ) + if value is not None + ), + None, + ) cost_breakdown: Final = _get_cost_breakdown_from_logging_obj( litellm_logging_obj=litellm_logging_obj, response_cost=response_cost ) @@ -1549,15 +1574,19 @@ class ProxyBaseLLMRequestProcessing: "x-litellm-key-rpm-limit": str(user_api_key_dict.rpm_limit), "x-litellm-key-max-budget": str(user_api_key_dict.max_budget), "x-litellm-key-spend": str(updated_spend), - "x-litellm-response-duration-ms": str(hidden_params.get("_response_ms", None)), - "x-litellm-overhead-duration-ms": str(hidden_params.get("litellm_overhead_time_ms", None)), - "x-litellm-callback-duration-ms": str(hidden_params.get("callback_duration_ms", None)), + "x-litellm-response-duration-ms": str(response_duration_ms), + "x-litellm-overhead-duration-ms": str(overhead_duration_ms), + "x-litellm-callback-duration-ms": str(response_hidden_params.get("callback_duration_ms")), **( { - "x-litellm-timing-pre-processing-ms": str(hidden_params.get("timing_pre_processing_ms", None)), - "x-litellm-timing-llm-api-ms": str(hidden_params.get("timing_llm_api_ms", None)), - "x-litellm-timing-post-processing-ms": str(hidden_params.get("timing_post_processing_ms", None)), - "x-litellm-timing-message-copy-ms": str(hidden_params.get("timing_message_copy_ms", None)), + "x-litellm-timing-pre-processing-ms": str( + response_hidden_params.get("timing_pre_processing_ms", None) + ), + "x-litellm-timing-llm-api-ms": str(response_hidden_params.get("timing_llm_api_ms", None)), + "x-litellm-timing-post-processing-ms": str( + response_hidden_params.get("timing_post_processing_ms", None) + ), + "x-litellm-timing-message-copy-ms": str(response_hidden_params.get("timing_message_copy_ms", None)), } if LITELLM_DETAILED_TIMING else {} @@ -2266,7 +2295,18 @@ class ProxyBaseLLMRequestProcessing: user_model=user_model, user_api_key_dict=user_api_key_dict, ) - llm_call_task: Final = asyncio.create_task(llm_call) + should_prefetch_stream_headers: Final = route_type in { + "anthropic_messages", + "aresponses", + } and self._is_streaming_request( + data=self.data, + is_streaming_request=is_streaming_request, + ) + prefetch_token: Final = is_proxy_stream_header_prefetch.set(should_prefetch_stream_headers) + try: + llm_call_task: Final = asyncio.create_task(llm_call) + finally: + is_proxy_stream_header_prefetch.reset(prefetch_token) tasks.append(llm_call_task) llm_responses: Final = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call @@ -2280,6 +2320,8 @@ class ProxyBaseLLMRequestProcessing: await _cancel_pending_gather_tasks(tasks) response = responses[1] + response_ownership_transferred = False + response_requires_cleanup = False _exception_raised = False try: @@ -2308,6 +2350,7 @@ class ProxyBaseLLMRequestProcessing: if self._is_streaming_request( data=self.data, is_streaming_request=is_streaming_request ) or self._is_streaming_response(response): # use generate_responses to stream responses + response_requires_cleanup = self._is_streaming_response(response) custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, call_id=logging_obj.litellm_call_id, @@ -2391,6 +2434,7 @@ class ProxyBaseLLMRequestProcessing: if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): + response_ownership_transferred = True if asyncio.iscoroutine(response): generator = await response else: @@ -2452,7 +2496,7 @@ class ProxyBaseLLMRequestProcessing: proxy_logging_obj=proxy_logging_obj, request=request, ) - return await create_response( + anthropic_stream_response: Final = await create_response( generator=wrap_sse_stream_with_keepalive_pings( stream=selected_data_generator, ping_interval_seconds=litellm.anthropic_sse_ping_interval_seconds, @@ -2461,6 +2505,8 @@ class ProxyBaseLLMRequestProcessing: headers=custom_headers, request=request, ) + response_ownership_transferred = True + return anthropic_stream_response # Non-streaming response - fall through to normal response handling elif select_data_generator: selected_data_generator = select_data_generator( @@ -2486,12 +2532,14 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, ) ) - return await create_response( + responses_stream_response: Final = await create_response( generator=selected_data_generator, media_type="text/event-stream", headers=custom_headers, request=request, ) + response_ownership_transferred = True + return responses_stream_response ### CALL HOOKS ### - modify outgoing data # If we reach here with a streaming closure still set, it means @@ -2536,6 +2584,14 @@ class ProxyBaseLLMRequestProcessing: _exception_raised = True raise finally: + if response_requires_cleanup and not response_ownership_transferred: + aclose = getattr(response, "aclose", None) + if aclose is not None: + with anyio.CancelScope(shield=True): + try: + await aclose() + except Exception as exc: # noqa: BLE001 + verbose_proxy_logger.debug("Error closing unowned response stream: %s", exc) ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( logging_obj=logging_obj, exception_raised=_exception_raised, diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index a0e8cd278e6..a2397963008 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -6,6 +6,9 @@ from collections.abc import Coroutine, Mapping from typing import Final import litellm +from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + prefetch_proxy_stream_for_timing, +) from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) @@ -110,6 +113,7 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( **acompletion_args, ) + await prefetch_proxy_stream_for_timing(litellm_completion_response) if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 8b1eeb30306..446517ec9e6 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -985,6 +985,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self.finished = True raise e + async def aclose(self) -> None: + self.finished = True + await self.litellm_custom_stream_wrapper.aclose() + def __iter__(self): return self diff --git a/litellm/utils.py b/litellm/utils.py index bb705e4ae9f..ce22ee1cd16 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1474,6 +1474,7 @@ def client(original_function): print_args_passed_to_litellm(original_function, args, kwargs) start_time: Final = datetime.datetime.now() result = None + _update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata") logging_obj: LiteLLMLoggingObject | None = kwargs.get("litellm_logging_obj", None) # only set litellm_call_id if its not in kwargs @@ -1601,8 +1602,7 @@ def client(original_function): return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") - update_response_metadata( + _update_response_metadata( result=result, logging_obj=logging_obj, model=model, @@ -1642,6 +1642,17 @@ def client(original_function): kwargs=kwargs, ) + returns_raw_dict: Final = isinstance(result, dict) + if returns_raw_dict: + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") # Copy the current context to propagate it to the background thread @@ -1655,16 +1666,15 @@ def client(original_function): start_time, end_time, ) - # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") - update_response_metadata( - result=result, - logging_obj=logging_obj, - model=model, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) + if not returns_raw_dict: + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) return result except Exception as e: call_type = original_function.__name__ @@ -1921,6 +1931,17 @@ def client(original_function): args=args, ) + returns_raw_dict: Final = isinstance(result, dict) + if returns_raw_dict: + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) + # LOG SUCCESS - handle streaming success logging in the _next_ object # Internal sub-calls (e.g. emulated file-search steps) share the # parent's logging obj; skip async logging here so only the outer call bills once. @@ -1970,15 +1991,15 @@ def client(original_function): end_time=end_time, ) - _update_response_metadata( - result=result, - logging_obj=logging_obj, - model=model, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - + if not returns_raw_dict: + _update_response_metadata( + result=result, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + start_time=start_time, + end_time=end_time, + ) return result except Exception as e: traceback_exception: Final = traceback.format_exc() diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 8e0bc200012..24f27cc3914 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -232,18 +232,10 @@ def test_combine_usage_handles_none_details(): def test_is_chat_completion_cached_dict(): from litellm.caching.caching_handler import _is_chat_completion_cached_dict - assert _is_chat_completion_cached_dict( - {"id": "chatcmpl-abc", "object": "chat.completion", "choices": []} - ) - assert _is_chat_completion_cached_dict( - {"id": "other", "object": "chat.completion.chunk", "choices": []} - ) - assert _is_chat_completion_cached_dict( - {"id": "no-object", "choices": [{"index": 0}]} - ) - assert not _is_chat_completion_cached_dict( - {"id": "resp_abc", "object": "response", "output": []} - ) + assert _is_chat_completion_cached_dict({"id": "chatcmpl-abc", "object": "chat.completion", "choices": []}) + assert _is_chat_completion_cached_dict({"id": "other", "object": "chat.completion.chunk", "choices": []}) + assert _is_chat_completion_cached_dict({"id": "no-object", "choices": [{"index": 0}]}) + assert not _is_chat_completion_cached_dict({"id": "resp_abc", "object": "response", "output": []}) def _build_logging_obj(call_type: str, stream: bool): @@ -262,15 +254,261 @@ def _build_logging_obj(call_type: str, stream: bool): ) +def test_sync_cached_anthropic_messages_records_cache_read_duration(monkeypatch): + import importlib + + import litellm + + caching_handler_module = importlib.import_module("litellm.caching.caching_handler") + + class Cache: + cache = None + supported_call_types = ["anthropic_messages"] + + def get_cache(self, **kwargs): + return { + "id": "msg_cached", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "cached"}], + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + + def get_cache_key(self, **kwargs): + return "cache-key" + + def _get_preset_cache_key_from_kwargs(self, **kwargs): + return None + + def anthropic_messages(): + return None + + cache = Cache() + logging_obj = MagicMock() + logging_obj.model_call_details = {} + start_time = datetime.now() + caching_handler = LLMCachingHandler( + original_function=anthropic_messages, + request_kwargs={"model": "openai/gpt-4o", "messages": []}, + start_time=start_time, + ) + + def resolve_model(**kwargs): + caching_handler_module.time.perf_counter() + return "gpt-4o", "openai", None, None + + monkeypatch.setattr(litellm, "cache", cache) + monkeypatch.setattr(litellm, "get_llm_provider", resolve_model) + monkeypatch.setattr( + caching_handler_module.time, + "perf_counter", + MagicMock(side_effect=(1.0, 1.25, 2.0)), + ) + + result = caching_handler._sync_get_cache( + model="openai/gpt-4o", + original_function=anthropic_messages, + logging_obj=logging_obj, + start_time=start_time, + call_type="anthropic_messages", + kwargs={"model": "openai/gpt-4o", "messages": []}, + args=(), + ) + + timing_metrics = logging_obj.set_response_timing_metrics.call_args.args[0] + assert result.cached_result is not None + assert "_hidden_params" not in result.cached_result + assert logging_obj.caching_details["cache_duration_ms"] == 250.0 + assert timing_metrics["_response_ms"] - timing_metrics["litellm_overhead_time_ms"] == pytest.approx(250.0) + + +def test_sync_cached_stream_response_preserves_model(monkeypatch): + import importlib + + import litellm + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + caching_handler_module = importlib.import_module("litellm.caching.caching_handler") + + class Cache: + cache = None + supported_call_types = ["completion"] + + def get_cache(self, **kwargs): + return { + "id": "chatcmpl_cached", + "object": "chat.completion", + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "cached"}}], + } + + def get_cache_key(self, **kwargs): + return "cache-key" + + def _get_preset_cache_key_from_kwargs(self, **kwargs): + return None + + def completion(): + return None + + cache = Cache() + logging_obj = MagicMock() + start_time = datetime(2025, 1, 1) + caching_handler = LLMCachingHandler( + original_function=completion, + request_kwargs={"model": "openai/gpt-4o", "messages": []}, + start_time=start_time, + ) + monkeypatch.setattr(litellm, "cache", cache) + monkeypatch.setattr( + caching_handler_module.time, + "perf_counter", + MagicMock(side_effect=(1.0, 1.25)), + ) + + result = caching_handler._sync_get_cache( + model="openai/gpt-4o", + original_function=completion, + logging_obj=logging_obj, + start_time=start_time, + call_type="completion", + kwargs={"model": "openai/gpt-4o", "messages": [], "stream": True}, + args=(), + ) + + assert isinstance(result.cached_result, CustomStreamWrapper) + assert result.cached_result.model == "openai/gpt-4o" + + +def test_sync_cached_response_records_callback_time_after_conversion(monkeypatch): + import importlib + + import litellm + + caching_handler_module = importlib.import_module("litellm.caching.caching_handler") + + class Cache: + cache = None + supported_call_types = ["completion"] + + def get_cache(self, **kwargs): + return {"id": "chatcmpl_cached", "object": "chat.completion", "choices": []} + + def get_cache_key(self, **kwargs): + return "cache-key" + + def _get_preset_cache_key_from_kwargs(self, **kwargs): + return None + + def completion(): + return None + + cache = Cache() + logging_obj = MagicMock() + start_time = datetime(2025, 1, 1) + callback_end_time = datetime(2025, 1, 1, 0, 0, 2) + conversion_completed = MagicMock() + + def convert_cached_result(**kwargs): + conversion_completed() + return MagicMock() + + def callback_now(): + conversion_completed.assert_called_once() + return callback_end_time + + caching_handler = LLMCachingHandler( + original_function=completion, + request_kwargs={"model": "openai/gpt-4o", "messages": []}, + start_time=start_time, + ) + monkeypatch.setattr(litellm, "cache", cache) + monkeypatch.setattr(caching_handler, "_convert_cached_result_to_model_response", convert_cached_result) + monkeypatch.setattr(caching_handler_module.datetime, "datetime", MagicMock(now=callback_now)) + + caching_handler._sync_get_cache( + model="openai/gpt-4o", + original_function=completion, + logging_obj=logging_obj, + start_time=start_time, + call_type="completion", + kwargs={"model": "openai/gpt-4o", "messages": []}, + args=(), + ) + + assert logging_obj.handle_sync_success_callbacks_for_async_calls.call_args.kwargs["end_time"] == callback_end_time + + +@pytest.mark.asyncio +async def test_async_cached_response_records_callback_time_after_conversion(monkeypatch): + import importlib + + import litellm + + caching_handler_module = importlib.import_module("litellm.caching.caching_handler") + + class Cache: + cache = None + supported_call_types = ["acompletion"] + + async def async_get_cache(self, **kwargs): + return {"id": "chatcmpl_cached", "object": "chat.completion", "choices": []} + + def get_cache_key(self, **kwargs): + return "cache-key" + + def _get_preset_cache_key_from_kwargs(self, **kwargs): + return None + + def _supports_async(self): + return True + + async def acompletion(): + return None + + cache = Cache() + logging_obj = MagicMock() + start_time = datetime(2025, 1, 1) + callback_end_time = datetime(2025, 1, 1, 0, 0, 2) + conversion_completed = MagicMock() + + def convert_cached_result(**kwargs): + conversion_completed() + return MagicMock() + + def callback_now(): + conversion_completed.assert_called_once() + return callback_end_time + + caching_handler = LLMCachingHandler( + original_function=acompletion, + request_kwargs={"model": "openai/gpt-4o", "messages": []}, + start_time=start_time, + ) + monkeypatch.setattr(litellm, "cache", cache) + monkeypatch.setattr(caching_handler, "_convert_cached_result_to_model_response", convert_cached_result) + monkeypatch.setattr(caching_handler_module.datetime, "datetime", MagicMock(now=callback_now)) + + await caching_handler._async_get_cache( + model="openai/gpt-4o", + original_function=acompletion, + logging_obj=logging_obj, + start_time=start_time, + call_type="acompletion", + kwargs={"model": "openai/gpt-4o", "messages": []}, + args=(), + ) + + assert logging_obj.async_success_handler.call_args.kwargs["end_time"] == callback_end_time + + def test_convert_cached_aresponses_bridge_chat_completion_stream(): """openai/responses chat-completions bridge: streaming cache hit replays as chat stream.""" from litellm import aresponses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import CallTypes - caching_handler = LLMCachingHandler( - original_function=aresponses, request_kwargs={}, start_time=datetime.now() - ) + caching_handler = LLMCachingHandler(original_function=aresponses, request_kwargs={}, start_time=datetime.now()) cached_result = { "id": "chatcmpl-bridge-cache-test", "object": "chat.completion", @@ -307,9 +545,7 @@ def test_convert_cached_responses_bridge_chat_completion_nonstream(): from litellm import responses from litellm.types.utils import CallTypes, ModelResponse - caching_handler = LLMCachingHandler( - original_function=responses, request_kwargs={}, start_time=datetime.now() - ) + caching_handler = LLMCachingHandler(original_function=responses, request_kwargs={}, start_time=datetime.now()) cached_result = { "id": "chatcmpl-bridge-nonstream", "object": "chat.completion", @@ -348,9 +584,7 @@ def test_convert_cached_responses_legacy_nonstream_path(): from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import CallTypes - caching_handler = LLMCachingHandler( - original_function=responses, request_kwargs={}, start_time=datetime.now() - ) + caching_handler = LLMCachingHandler(original_function=responses, request_kwargs={}, start_time=datetime.now()) cached_result = { "id": "resp_legacy_nonstream", "created_at": int(time.time()), @@ -395,9 +629,7 @@ def test_convert_cached_responses_legacy_stream_path(): ) from litellm.types.utils import CallTypes - caching_handler = LLMCachingHandler( - original_function=responses, request_kwargs={}, start_time=datetime.now() - ) + caching_handler = LLMCachingHandler(original_function=responses, request_kwargs={}, start_time=datetime.now()) cached_result = { "id": "resp_legacy_stream", "created_at": int(time.time()), diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index 996530daa2e..c146e41acd2 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -6,13 +6,19 @@ through _hidden_params to the x-litellm-callback-duration-ms response header. """ import datetime -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock +import pytest + +import litellm import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod import litellm.proxy.common_request_processing as common_request_processing_mod +from litellm._internal_context import is_internal_call, is_proxy_stream_header_prefetch from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ResponseMetadata, + prefetch_proxy_stream_for_timing, + refresh_response_timing_metrics, update_response_metadata, ) from litellm.proxy._types import UserAPIKeyAuth @@ -68,9 +74,7 @@ class TestCallbackDurationMs: def test_update_response_metadata_includes_callback_duration(self): """End-to-end: update_response_metadata should propagate callback_duration_ms.""" result = ModelResponse() - logging_obj = self._make_logging_obj( - callback_duration_ms=5.5, llm_api_duration_ms=800.0 - ) + logging_obj = self._make_logging_obj(callback_duration_ms=5.5, llm_api_duration_ms=800.0) logging_obj._response_cost_calculator = MagicMock(return_value=0.001) logging_obj.litellm_call_id = "test-call-id" @@ -92,25 +96,34 @@ class TestCallbackDurationMs: assert hidden.get("litellm_overhead_time_ms") is not None -class TestDictResultsSkipMetadataUpdate: - """Regression for /v1/messages cost-breakdown clobbering: AnthropicMessagesResponse - is a TypedDict, so apply() can never attach _hidden_params to it and the whole - metadata pass is discarded - except the cost recompute, whose only observable - effect was overwriting the logging object's already-correct cost breakdown with a - service-tier-less, reasoning-less recompute on the adapted response.""" +def _messages_response(): + return { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "usage": {"input_tokens": 7, "output_tokens": 320}, + } - def test_update_response_metadata_skips_cost_recompute_for_dict_results(self): - anthropic_response = { - "id": "msg_123", - "type": "message", - "role": "assistant", - "content": [{"type": "text", "text": "hi"}], - "usage": {"input_tokens": 7, "output_tokens": 320}, - } - logging_obj = MagicMock() - logging_obj.model_call_details = {} - logging_obj.caching_details = None - logging_obj.litellm_call_id = "test-call-id" + +def _messages_logging_obj(callback_duration_ms=None): + logging_obj = MagicMock() + logging_obj.model_call_details = {"llm_api_duration_ms": 900.0} + logging_obj.caching_details = None + logging_obj.litellm_call_id = "test-call-id" + if callback_duration_ms is None: + del logging_obj.callback_duration_ms + else: + logging_obj.callback_duration_ms = callback_duration_ms + return logging_obj + + +class TestDictResultsSkipCostRecompute: + """Regression coverage for metadata on Anthropic Messages dict responses.""" + + def test_update_response_metadata_adds_timing_without_recomputing_cost(self): + anthropic_response = _messages_response() + logging_obj = _messages_logging_obj() update_response_metadata( result=anthropic_response, @@ -123,6 +136,310 @@ class TestDictResultsSkipMetadataUpdate: logging_obj._response_cost_calculator.assert_not_called() assert "_hidden_params" not in anthropic_response + assert dict(logging_obj.set_response_timing_metrics.call_args.args[0]) == { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + + def test_request_timing_metrics_snapshot_is_read_only(self): + anthropic_response = _messages_response() + logging_obj = _messages_logging_obj() + + update_response_metadata( + result=anthropic_response, + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + snapshot = logging_obj.set_response_timing_metrics.call_args.args[0] + with pytest.raises(TypeError): + snapshot["_response_ms"] = 1.0 + + def test_cached_messages_response_uses_cache_read_duration(self): + anthropic_response = _messages_response() + logging_obj = _messages_logging_obj() + logging_obj.caching_details = {"cache_hit": True, "cache_duration_ms": 250.0} + + update_response_metadata( + result=anthropic_response, + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + assert dict(logging_obj.set_response_timing_metrics.call_args.args[0]) == { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 750.0, + } + + def test_internal_sub_call_preserves_outer_request_timing(self): + logging_obj = _messages_logging_obj() + outer_timing_metrics = {"_response_ms": 2500.0, "litellm_overhead_time_ms": 200.0} + logging_obj.response_timing_metrics = outer_timing_metrics + + token = is_internal_call.set(True) + try: + update_response_metadata( + result=_messages_response(), + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + finally: + is_internal_call.reset(token) + + logging_obj.set_response_timing_metrics.assert_not_called() + assert logging_obj.response_timing_metrics is outer_timing_metrics + + def test_streaming_opaque_result_uses_request_timing_without_mutation(self): + class OpaqueStream: + pass + + result = OpaqueStream() + logging_obj = _messages_logging_obj() + logging_obj.stream = True + + update_response_metadata( + result=result, + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + assert not hasattr(result, "_hidden_params") + assert dict(logging_obj.set_response_timing_metrics.call_args.args[0]) == { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + + @pytest.mark.parametrize("stream", [False, None, MagicMock()]) + def test_non_streaming_opaque_result_does_not_set_request_timing(self, stream): + class OpaqueResult: + pass + + logging_obj = _messages_logging_obj() + logging_obj.stream = stream + + update_response_metadata( + result=OpaqueResult(), + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + logging_obj.set_response_timing_metrics.assert_not_called() + + def test_internal_streaming_opaque_result_preserves_outer_request_timing(self): + class OpaqueStream: + pass + + logging_obj = _messages_logging_obj() + logging_obj.stream = True + outer_timing_metrics = {"_response_ms": 2500.0, "litellm_overhead_time_ms": 200.0} + logging_obj.response_timing_metrics = outer_timing_metrics + + token = is_internal_call.set(True) + try: + update_response_metadata( + result=OpaqueStream(), + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + finally: + is_internal_call.reset(token) + + logging_obj.set_response_timing_metrics.assert_not_called() + assert logging_obj.response_timing_metrics is outer_timing_metrics + + +class TestPrefetchedStreamTiming: + @pytest.mark.asyncio + async def test_prefetches_deferred_vertex_stream_and_refreshes_timing(self): + logging_obj = _messages_logging_obj() + logging_obj.start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=1000) + async def empty_stream(): + if False: + yield None + + stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=AsyncMock(return_value=empty_stream()), + ) + + token = is_proxy_stream_header_prefetch.set(True) + try: + await prefetch_proxy_stream_for_timing(stream) + finally: + is_proxy_stream_header_prefetch.reset(token) + + stream.make_call.assert_awaited_once() + assert dict(logging_obj.set_response_timing_metrics.call_args.args[0])["litellm_overhead_time_ms"] > 0 + + @pytest.mark.asyncio + async def test_prefetch_does_not_open_non_vertex_stream(self): + logging_obj = _messages_logging_obj() + logging_obj.start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=1000) + stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="other-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + make_call=AsyncMock(return_value=object()), + ) + + token = is_proxy_stream_header_prefetch.set(True) + try: + await prefetch_proxy_stream_for_timing(stream) + finally: + is_proxy_stream_header_prefetch.reset(token) + + stream.make_call.assert_not_awaited() + logging_obj.set_response_timing_metrics.assert_not_called() + + def test_refreshes_detached_timing_after_proxy_stream_prefetch(self): + logging_obj = _messages_logging_obj() + logging_obj.start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=1000) + + token = is_proxy_stream_header_prefetch.set(True) + try: + refresh_response_timing_metrics(logging_obj) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert dict(logging_obj.set_response_timing_metrics.call_args.args[0])["litellm_overhead_time_ms"] > 0 + + def test_does_not_refresh_timing_outside_proxy_stream_prefetch(self): + logging_obj = _messages_logging_obj() + logging_obj.start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=1000) + + refresh_response_timing_metrics(logging_obj) + + logging_obj.set_response_timing_metrics.assert_not_called() + + def test_internal_call_does_not_refresh_outer_request_timing(self): + logging_obj = _messages_logging_obj() + logging_obj.start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=1000) + logging_obj.response_timing_metrics = {"_response_ms": 2500.0, "litellm_overhead_time_ms": 200.0} + prefetch_token = is_proxy_stream_header_prefetch.set(True) + internal_token = is_internal_call.set(True) + try: + refresh_response_timing_metrics(logging_obj) + finally: + is_internal_call.reset(internal_token) + is_proxy_stream_header_prefetch.reset(prefetch_token) + + logging_obj.set_response_timing_metrics.assert_not_called() + assert logging_obj.response_timing_metrics == {"_response_ms": 2500.0, "litellm_overhead_time_ms": 200.0} + + +class TestDictResponseTimingHeaders: + """Regression coverage for Messages timing headers.""" + + def _headers_for(self, logging_obj, hidden_params): + return ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + ) + + def test_headers_use_request_timing_without_mutating_messages_response(self): + anthropic_response = _messages_response() + logging_obj = _messages_logging_obj(callback_duration_ms=7.25) + + update_response_metadata( + result=anthropic_response, + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + + logging_obj.response_timing_metrics = logging_obj.set_response_timing_metrics.call_args.args[0] + headers = self._headers_for(logging_obj, hidden_params={}) + + assert "_hidden_params" not in anthropic_response + assert headers["x-litellm-response-duration-ms"] == "1000.0" + assert headers["x-litellm-overhead-duration-ms"] == "100.0" + assert "x-litellm-callback-duration-ms" not in headers + + def test_headers_use_streaming_opaque_request_timing(self): + class OpaqueStream: + pass + + logging_obj = _messages_logging_obj() + logging_obj.stream = True + update_response_metadata( + result=OpaqueStream(), + logging_obj=logging_obj, + model="vertex_ai/gemini-3.5-flash", + kwargs={}, + start_time=datetime.datetime(2025, 1, 1, 0, 0, 0), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1), + ) + logging_obj.response_timing_metrics = logging_obj.set_response_timing_metrics.call_args.args[0] + + headers = self._headers_for(logging_obj, hidden_params={}) + + assert headers["x-litellm-response-duration-ms"] == "1000.0" + assert headers["x-litellm-overhead-duration-ms"] == "100.0" + assert "x-litellm-callback-duration-ms" not in headers + + def test_headers_accept_legacy_logging_object_without_timing_metrics(self): + class LegacyLoggingObject: + pass + + headers = self._headers_for(LegacyLoggingObject(), hidden_params={}) + + assert "x-litellm-response-duration-ms" not in headers + assert "x-litellm-overhead-duration-ms" not in headers + + def test_response_hidden_params_win_over_request_timing(self): + logging_obj = _messages_logging_obj() + logging_obj.response_timing_metrics = { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + + headers = self._headers_for( + logging_obj, + hidden_params={"_response_ms": 42.0, "litellm_overhead_time_ms": 4.0}, + ) + + assert headers["x-litellm-response-duration-ms"] == "42.0" + assert headers["x-litellm-overhead-duration-ms"] == "4.0" + + def test_request_timing_fills_explicit_none_response_timing(self): + logging_obj = _messages_logging_obj() + logging_obj.response_timing_metrics = { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + + headers = self._headers_for( + logging_obj, + hidden_params={"_response_ms": 42.0, "litellm_overhead_time_ms": None}, + ) + + assert headers["x-litellm-response-duration-ms"] == "42.0" + assert headers["x-litellm-overhead-duration-ms"] == "100.0" class TestCallbackDurationInCustomHeaders: @@ -221,11 +538,19 @@ class TestDetailedTiming: assert hidden.get("timing_llm_api_ms") is None assert hidden.get("timing_pre_processing_ms") is None + def test_detailed_timing_headers_accept_none_hidden_params(self, monkeypatch): + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True) + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + hidden_params=None, + ) + + assert "x-litellm-timing-llm-api-ms" not in headers + def test_detailed_timing_headers_in_custom_headers(self, monkeypatch): """When LITELLM_DETAILED_TIMING is true, headers flow to get_custom_headers.""" - monkeypatch.setattr( - common_request_processing_mod, "LITELLM_DETAILED_TIMING", True - ) + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") hidden_params = { @@ -248,9 +573,7 @@ class TestDetailedTiming: def test_detailed_timing_headers_absent_when_disabled(self, monkeypatch): """When LITELLM_DETAILED_TIMING is false, no timing headers emitted.""" - monkeypatch.setattr( - common_request_processing_mod, "LITELLM_DETAILED_TIMING", False - ) + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") hidden_params = { diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1c717228ecc..37f06092c68 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1,4 +1,5 @@ import contextlib +import copy import os import sys import asyncio @@ -4134,6 +4135,81 @@ def _anthropic_messages_logging_obj(): ) +def test_logging_response_timing_metrics_supports_deepcopy(): + logging_obj = _anthropic_messages_logging_obj() + timing_metrics = {"_response_ms": 1000.0} + logging_obj.set_response_timing_metrics(timing_metrics) + timing_metrics["_response_ms"] = 1.0 + logging_copy = copy.deepcopy(logging_obj) + + assert logging_copy.response_timing_metrics == {"_response_ms": 1000.0} + + +def test_anthropic_messages_standard_payload_uses_request_timing_without_mutating_response(): + from datetime import datetime + + logging_obj = LitellmLogging( + model="claude-haiku-4-5-20251001", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="lit-5466", + function_id="lit-5466", + ) + logging_obj.optional_params = {} + logging_obj.set_response_timing_metrics( + { + "_response_ms": 1000.0, + "litellm_overhead_time_ms": 100.0, + } + ) + response = { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5-20251001", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 7, "output_tokens": 320}, + } + now = datetime.now() + + logging_obj._success_handler_helper_fn(response, now, now) + + payload = logging_obj.model_call_details["standard_logging_object"] + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] == 100.0 + assert "_hidden_params" not in response + + +def test_standard_payload_prefers_response_overhead_to_request_timing(): + """A response carrying its own overhead keeps it; the request-scoped carrier only fills gaps.""" + from datetime import datetime + + logging_obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="lit-5466", + function_id="lit-5466", + ) + logging_obj.optional_params = {} + logging_obj.set_response_timing_metrics({"litellm_overhead_time_ms": 100.0}) + response = ModelResponse() + response._hidden_params = {"litellm_overhead_time_ms": 25.0} + now = datetime.now() + + logging_obj._success_handler_helper_fn(response, now, now) + + payload = logging_obj.model_call_details["standard_logging_object"] + assert payload is not None + assert payload["hidden_params"]["litellm_overhead_time_ms"] == 25.0 + + def _responses_api_response_with_text(text="hello world"): from openai.types.responses import ResponseOutputMessage, ResponseOutputText diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7f54fbfb4c2..cf2410db656 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -762,6 +762,31 @@ async def test_streaming_completion_start_time(logging_obj: Logging): ) +@pytest.mark.asyncio +async def test_vertex_fetch_stream_preserves_bad_request(logging_obj: Logging): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + expected_error = VertexAIError( + status_code=400, message="invalid maxOutputTokens", headers=None + ) + + async def _raise_bad_request(**kwargs): + raise expected_error + + response = CustomStreamWrapper( + completion_stream=None, + model="gemini-3-pro-preview", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=_raise_bad_request, + ) + + with pytest.raises(VertexAIError) as excinfo: + await response.fetch_stream() + + assert excinfo.value is expected_error + + @pytest.mark.asyncio async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging): """Ensure Vertex bad request errors surface as 400, not mid-stream fallbacks.""" @@ -787,6 +812,31 @@ async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging): assert "invalid maxOutputTokens" in str(excinfo.value) +@pytest.mark.asyncio +async def test_vertex_fetch_stream_preserves_rate_limit(logging_obj: Logging): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + expected_error = VertexAIError( + status_code=429, message="Resource exhausted.", headers=None + ) + + async def _raise_rate_limit(**kwargs): + raise expected_error + + response = CustomStreamWrapper( + completion_stream=None, + model="gemini-3-flash-preview", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=_raise_rate_limit, + ) + + with pytest.raises(VertexAIError) as excinfo: + await response.fetch_stream() + + assert excinfo.value is expected_error + + @pytest.mark.asyncio async def test_vertex_streaming_rate_limit_triggers_midstream_fallback( logging_obj: Logging, diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py index a944afc6152..c0e07ea0830 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py @@ -25,12 +25,16 @@ Tests cover (consolidating PRs #23706 and #22727): the fallback inference path from being exercised). """ +import datetime import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm +from litellm._internal_context import is_proxy_stream_header_prefetch + # Anchor sys.path to this file's location — not the working-directory-relative # pattern Greptile flagged on PR #23706. Resolves correctly regardless of # where pytest is invoked from. @@ -46,6 +50,33 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.handler import ( MESSAGES = [{"role": "user", "content": "hello"}] +class _EmptyAsyncStream: + def __init__(self): + self.aclosed = False + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def aclose(self): + self.aclosed = True + + +def _deferred_stream(provider: str = "vertex_ai_beta") -> litellm.CustomStreamWrapper: + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 40.0} + logging_obj.start_time = datetime.datetime.now() + return litellm.CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider=provider, + make_call=AsyncMock(return_value=_EmptyAsyncStream()), + ) + + def _call_prepare(extra_kwargs, model="gpt-4o", output_format=None, **overrides): """ Drive ``_prepare_completion_kwargs`` with the minimum scaffolding needed. @@ -76,6 +107,133 @@ def _call_prepare(extra_kwargs, model="gpt-4o", output_format=None, **overrides) ) +@pytest.mark.asyncio +async def test_async_messages_bridge_keeps_sdk_deferred_gemini_stream_lazy(): + deferred_stream = _deferred_stream() + + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=32, + messages=MESSAGES, + model="gemini-3.5-flash", + stream=True, + ) + + assert deferred_stream.make_call.await_count == 0 + first_sse_event = await anext(result) + assert b"event: message_start" in first_sse_event + assert deferred_stream.make_call.await_count == 0 + + +@pytest.mark.asyncio +async def test_async_messages_bridge_prefetches_without_consuming_source_event(): + deferred_stream = _deferred_stream() + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=32, + messages=MESSAGES, + model="gemini-3.5-flash", + stream=True, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + deferred_stream.make_call.assert_awaited_once() + first_sse_event = await anext(result) + assert b"event: message_start" in first_sse_event + remaining_events = [event async for event in result] + assert sum(b"event: message_start" in event for event in remaining_events) == 0 + + +@pytest.mark.asyncio +async def test_async_messages_bridge_does_not_prefetch_already_connected_gemini_stream_for_proxy(): + deferred_stream = _deferred_stream() + existing_stream = _EmptyAsyncStream() + deferred_stream.completion_stream = existing_stream + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=32, + messages=MESSAGES, + model="gemini-3.5-flash", + stream=True, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert deferred_stream.make_call.await_count == 0 + assert b"event: message_start" in await anext(result) + + +@pytest.mark.asyncio +async def test_async_messages_bridge_closes_prefetched_gemini_stream_on_early_close(): + upstream = _EmptyAsyncStream() + deferred_stream = _deferred_stream() + deferred_stream.make_call = AsyncMock(return_value=upstream) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=32, + messages=MESSAGES, + model="gemini-3.5-flash", + stream=True, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert b"event: message_start" in await anext(result) + await result.aclose() + + assert upstream.aclosed is True + assert deferred_stream.completion_stream is None + + +@pytest.mark.asyncio +async def test_async_messages_bridge_propagates_initial_fetch_failure(): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + expected_error = VertexAIError(status_code=429, message="Resource exhausted.", headers=None) + deferred_stream = _deferred_stream() + deferred_stream.make_call = AsyncMock(side_effect=expected_error) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + with pytest.raises(VertexAIError) as excinfo: + await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=32, + messages=MESSAGES, + model="gemini-3.5-flash", + stream=True, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert excinfo.value is expected_error + + +@pytest.mark.asyncio +async def test_async_messages_bridge_does_not_prefetch_non_gemini_stream_for_proxy(): + deferred_stream = _deferred_stream(provider="openai") + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await LiteLLMMessagesToCompletionTransformationHandler.async_anthropic_messages_handler( + max_tokens=32, + messages=MESSAGES, + model="other-model", + stream=True, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert deferred_stream.make_call.await_count == 0 + assert b"event: message_start" in await anext(result) + + class TestAnthropicOnlyRequestKeysExport: """The exclusion list must be a public, named constant for maintainability — Greptile P2 on PR #23706: ``excluded_keys`` was silently growing as a 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 b37c0f466d2..7dada07313e 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 @@ -18,6 +18,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -235,6 +236,45 @@ def test_response_api_handler_runs_responses_pre_call_hook_before_transform(): assert hook_litellm_params.get(_SANDBOX_KEY) +@pytest.mark.asyncio +async def test_async_response_api_handler_records_llm_api_duration(): + config = Mock(spec=BaseResponsesAPIConfig) + config.validate_environment.return_value = {} + config.get_complete_url.return_value = "https://responses.example.com/v1/responses" + config.transform_responses_api_request.return_value = {"model": "gpt-5.3-codex", "input": "hi"} + config.sign_request.return_value = ({}, None) + config.transform_response_api_response.return_value = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="gpt-5.3-codex", + ) + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={"id": "resp_1"}, request=request) + ) + ) + logging_obj = Mock() + logging_obj.dynamic_success_callbacks = [] + logging_obj.model_call_details = {} + + response = await BaseLLMHTTPHandler().async_response_api_handler( + model="gpt-5.3-codex", + input="hi", + responses_api_provider_config=config, + response_api_optional_request_params={}, + custom_llm_provider="chatgpt", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + client=client, + ) + + assert response.id == "resp_1" + assert logging_obj.model_call_details["llm_api_duration_ms"] >= 0 + + @pytest.mark.asyncio async def test_async_response_api_handler_streams_when_provider_transform_adds_stream(): handler = BaseLLMHTTPHandler() @@ -267,6 +307,8 @@ async def test_async_response_api_handler_streams_when_provider_transform_adds_s client=client, ) + assert client.post.call_args.kwargs["logging_obj"] is logging_obj + assert client.post.call_args.kwargs["stream"] is True assert client.post.call_args.kwargs["json"]["stream"] is True diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 71d4666416d..af195cc9d23 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -12,6 +12,7 @@ from fastapi import HTTPException, Request, Response, status from fastapi.responses import JSONResponse, StreamingResponse import litellm +from litellm._internal_context import is_proxy_stream_header_prefetch from litellm._uuid import uuid from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger @@ -2521,6 +2522,72 @@ class TestStreamingOverheadHeader: assert "x-litellm-overhead-duration-ms" in headers assert headers["x-litellm-overhead-duration-ms"] == "42.5" + def test_get_custom_headers_uses_detached_stream_timing_before_headers_commit(self): + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + logging_obj = MagicMock() + logging_obj.response_timing_metrics = { + "_response_ms": 500.0, + "litellm_overhead_time_ms": 42.5, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + hidden_params={}, + litellm_logging_obj=logging_obj, + ) + + assert headers["x-litellm-response-duration-ms"] == "500.0" + assert headers["x-litellm-overhead-duration-ms"] == "42.5" + + def test_get_custom_headers_uses_detached_timing_when_response_timing_is_none(self): + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + logging_obj = MagicMock() + logging_obj.response_timing_metrics = { + "_response_ms": 500.0, + "litellm_overhead_time_ms": 42.5, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + hidden_params={"_response_ms": None, "litellm_overhead_time_ms": None}, + litellm_logging_obj=logging_obj, + ) + + assert headers["x-litellm-response-duration-ms"] == "500.0" + assert headers["x-litellm-overhead-duration-ms"] == "42.5" + + def test_get_custom_headers_preserves_zero_response_timing_over_detached_timing(self): + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + logging_obj = MagicMock() + logging_obj.response_timing_metrics = { + "_response_ms": 500.0, + "litellm_overhead_time_ms": 42.5, + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + hidden_params={"_response_ms": 0, "litellm_overhead_time_ms": 0}, + litellm_logging_obj=logging_obj, + ) + + assert headers["x-litellm-response-duration-ms"] == "0" + assert headers["x-litellm-overhead-duration-ms"] == "0" + def test_get_custom_headers_omits_overhead_when_none(self): """ get_custom_headers() omits x-litellm-overhead-duration-ms @@ -4066,7 +4133,14 @@ class TestCancelOnDisconnect: await _await_llm_call_cancelling_on_disconnect(request, llm_call) async def _drive_base_process_llm_request( - self, monkeypatch, general_settings: dict, llm_call, request: Request + self, + monkeypatch, + general_settings: dict, + llm_call, + request: Request, + route_type: str = "acompletion", + data: dict | None = None, + is_streaming_request: bool = False, ): from litellm.proxy._types import UserAPIKeyAuth @@ -4076,9 +4150,12 @@ class TestCancelOnDisconnect: logging_obj._on_deferred_stream_complete = None logging_obj.cost_breakdown = None - processor = ProxyBaseLLMRequestProcessing( - data={"model": "fake-model", "litellm_logging_obj": logging_obj} + request_data = ( + {"model": "fake-model", "litellm_logging_obj": logging_obj} + if data is None + else {**data, "litellm_logging_obj": logging_obj} ) + processor = ProxyBaseLLMRequestProcessing(data=request_data) proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -4103,13 +4180,241 @@ class TestCancelOnDisconnect: request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), - route_type="acompletion", + route_type=route_type, proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, proxy_config=MagicMock(spec=ProxyConfig), + is_streaming_request=is_streaming_request, skip_pre_call_logic=True, ) + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_native_stream_task_inherits_prefetch_marker_only_for_child(self, monkeypatch, route_type): + seen_in_child: list[bool] = [] + + async def llm_call(): + seen_in_child.append(is_proxy_stream_header_prefetch.get()) + return litellm.ModelResponse() + + assert is_proxy_stream_header_prefetch.get() is False + await self._drive_base_process_llm_request( + monkeypatch, + general_settings={}, + llm_call=llm_call, + request=self._request([]), + route_type=route_type, + data={"model": "fake-model", "stream": True}, + is_streaming_request=True, + ) + + assert seen_in_child == [True] + assert is_proxy_stream_header_prefetch.get() is False + + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_native_stream_uses_detached_timing_for_headers_before_body_iteration(self, monkeypatch, route_type): + logging_obj = MagicMock() + logging_obj._hidden_params = {} + logging_obj.response_timing_metrics = { + "_response_ms": 500.0, + "litellm_overhead_time_ms": 42.5, + } + + body_started = asyncio.Event() + + async def response_stream(): + body_started.set() + yield "data: first event\n\n" + + async def llm_call(): + return response_stream() + + async def fake_route_request(**kwargs): + return llm_call() + + async def fake_create_response(generator, media_type, headers, **kwargs): + assert body_started.is_set() is False + assert headers["x-litellm-response-duration-ms"] == "500.0" + assert headers["x-litellm-overhead-duration-ms"] == "42.5" + return StreamingResponse(generator, media_type=media_type, headers=headers) + + from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "stream": True, "litellm_logging_obj": logging_obj} + ) + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + monkeypatch.setattr(litellm.proxy.common_request_processing, "create_response", fake_create_response) + + def select_data_generator(**kwargs): + return kwargs["response"] + + result = await processor.base_process_llm_request( + request=self._request([]), + fastapi_response=Response(), + user_api_key_dict=RealUserAPIKeyAuth(api_key="sk-test"), + route_type=route_type, + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, StreamingResponse) + assert body_started.is_set() is False + assert result.headers["x-litellm-response-duration-ms"] == "500.0" + assert result.headers["x-litellm-overhead-duration-ms"] == "42.5" + + @pytest.mark.parametrize("route_type", ["acompletion", "anthropic_messages", "aresponses"]) + async def test_non_streaming_task_does_not_inherit_prefetch_marker(self, monkeypatch, route_type): + seen_in_child: list[bool] = [] + + async def llm_call(): + seen_in_child.append(is_proxy_stream_header_prefetch.get()) + return litellm.ModelResponse() + + await self._drive_base_process_llm_request( + monkeypatch, + general_settings={}, + llm_call=llm_call, + request=self._request([]), + route_type=route_type, + data={"model": "fake-model", "stream": False}, + is_streaming_request=False, + ) + + assert seen_in_child == [False] + assert is_proxy_stream_header_prefetch.get() is False + + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_truthy_non_boolean_stream_flag_does_not_inherit_prefetch_marker(self, monkeypatch, route_type): + seen_in_child: list[bool] = [] + + async def llm_call(): + seen_in_child.append(is_proxy_stream_header_prefetch.get()) + return litellm.ModelResponse() + + await self._drive_base_process_llm_request( + monkeypatch, + general_settings={}, + llm_call=llm_call, + request=self._request([]), + route_type=route_type, + data={"model": "fake-model", "stream": "true"}, + is_streaming_request=False, + ) + + assert seen_in_child == [False] + assert is_proxy_stream_header_prefetch.get() is False + + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_native_stream_cancels_upstream_on_disconnect_when_enabled(self, monkeypatch, route_type): + upstream_cancelled = asyncio.Event() + + async def llm_call(): + try: + await asyncio.sleep(5) + return litellm.ModelResponse() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + with pytest.raises(HTTPException) as exc_info: + await self._drive_base_process_llm_request( + monkeypatch, + general_settings={"cancel_on_disconnect": True}, + llm_call=llm_call, + request=self._request([{"type": "http.disconnect"}]), + route_type=route_type, + data={"model": "fake-model", "stream": True}, + is_streaming_request=True, + ) + + assert exc_info.value.status_code == 499 + assert upstream_cancelled.is_set() + + async def test_cancellation_before_stream_response_transfers_ownership_closes_upstream(self, monkeypatch): + class Upstream: + def __init__(self): + self.aclosed = False + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def aclose(self): + self.aclosed = True + + upstream = Upstream() + logging_obj = MagicMock() + logging_obj.litellm_call_id = "test-prefetch-close" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + logging_obj.model_call_details = {"litellm_params": {}} + response = litellm.CustomStreamWrapper( + completion_stream=upstream, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + ) + header_hook_started = asyncio.Event() + + async def header_hook(**kwargs): + header_hook_started.set() + await asyncio.Event().wait() + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "stream": True, "litellm_logging_obj": logging_obj} + ) + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = header_hook + + async def fake_route_request(**kwargs): + async def llm_call(): + return response + + return llm_call() + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "route_request", + fake_route_request, + ) + + task = asyncio.create_task( + processor._process_llm_request( + request=self._request([]), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="anthropic_messages", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + is_streaming_request=True, + skip_pre_call_logic=True, + ) + ) + await header_hook_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert upstream.aclosed is True + async def test_disconnect_ignored_when_flag_disabled(self, monkeypatch): upstream_cancelled = asyncio.Event() model_response = litellm.ModelResponse() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py index 2cfec6a1844..bde3767512d 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_handler.py @@ -12,11 +12,13 @@ capture the forwarded kwargs; if the flag-setting line is removed the captured kwargs lack the flag and these tests fail. """ -from unittest.mock import patch +import datetime +from unittest.mock import AsyncMock, MagicMock, patch import pytest - +import litellm +from litellm._internal_context import is_proxy_stream_header_prefetch from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) @@ -68,3 +70,192 @@ async def test_async_fallback_tags_skip_responses_api_bridge(): await coro assert captured.get("_skip_responses_api_bridge") is True + + +class _EmptyAsyncStream: + def __init__(self): + self.aclosed = False + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + async def aclose(self): + self.aclosed = True + + +@pytest.mark.asyncio +async def test_async_responses_bridge_keeps_sdk_deferred_gemini_stream_lazy(): + handler = LiteLLMCompletionTransformationHandler() + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}} + deferred_stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=AsyncMock(return_value=_EmptyAsyncStream()), + ) + + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await handler.async_response_api_handler( + litellm_completion_request={"model": "gemini-3.5-flash", "stream": True}, + request_input="ping", + responses_api_request={}, + ) + + assert result.litellm_custom_stream_wrapper is deferred_stream + deferred_stream.make_call.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_async_responses_bridge_prefetches_deferred_gemini_stream_for_proxy(): + import datetime + + handler = LiteLLMCompletionTransformationHandler() + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 40.0} + logging_obj.start_time = datetime.datetime.now() + deferred_stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=AsyncMock(return_value=_EmptyAsyncStream()), + ) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await handler.async_response_api_handler( + litellm_completion_request={"model": "gemini-3.5-flash", "stream": True}, + request_input="ping", + responses_api_request={}, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert result.litellm_custom_stream_wrapper is deferred_stream + deferred_stream.make_call.assert_awaited_once() + logging_obj.set_response_timing_metrics.assert_called_once() + assert result._buffered_chunk is None + assert result._response_id_primed is False + + +@pytest.mark.asyncio +async def test_async_responses_bridge_does_not_prefetch_already_connected_gemini_stream_for_proxy(): + handler = LiteLLMCompletionTransformationHandler() + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 40.0} + deferred_stream = litellm.CustomStreamWrapper( + completion_stream=_EmptyAsyncStream(), + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=AsyncMock(return_value=_EmptyAsyncStream()), + ) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await handler.async_response_api_handler( + litellm_completion_request={"model": "gemini-3.5-flash", "stream": True}, + request_input="ping", + responses_api_request={}, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert result.litellm_custom_stream_wrapper is deferred_stream + assert deferred_stream.make_call.await_count == 0 + + +@pytest.mark.asyncio +async def test_async_responses_bridge_closes_prefetched_gemini_stream_on_early_close(): + handler = LiteLLMCompletionTransformationHandler() + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 40.0} + logging_obj.start_time = datetime.datetime.now() + upstream = _EmptyAsyncStream() + deferred_stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=AsyncMock(return_value=upstream), + ) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await handler.async_response_api_handler( + litellm_completion_request={"model": "gemini-3.5-flash", "stream": True}, + request_input="ping", + responses_api_request={}, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + await result.aclose() + + assert upstream.aclosed is True + assert deferred_stream.completion_stream is None + with pytest.raises(StopAsyncIteration): + await anext(result) + deferred_stream.make_call.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_responses_bridge_propagates_initial_fetch_failure(): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + handler = LiteLLMCompletionTransformationHandler() + expected_error = VertexAIError(status_code=500, message="upstream failed", headers=None) + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 40.0} + deferred_stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="gemini-3.5-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=AsyncMock(side_effect=expected_error), + ) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + with pytest.raises(VertexAIError) as excinfo: + await handler.async_response_api_handler( + litellm_completion_request={"model": "gemini-3.5-flash", "stream": True}, + request_input="ping", + responses_api_request={}, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert excinfo.value is expected_error + + +@pytest.mark.asyncio +async def test_async_responses_bridge_does_not_prefetch_non_gemini_stream_for_proxy(): + handler = LiteLLMCompletionTransformationHandler() + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}, "llm_api_duration_ms": 40.0} + deferred_stream = litellm.CustomStreamWrapper( + completion_stream=None, + model="other-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + make_call=AsyncMock(return_value=_EmptyAsyncStream()), + ) + token = is_proxy_stream_header_prefetch.set(True) + try: + with patch("litellm.acompletion", new=AsyncMock(return_value=deferred_stream)): # test-quality-ok: handler directly owns this module-level provider-call seam + result = await handler.async_response_api_handler( + litellm_completion_request={"model": "other-model", "stream": True}, + request_input="ping", + responses_api_request={}, + ) + finally: + is_proxy_stream_header_prefetch.reset(token) + + assert result.litellm_custom_stream_wrapper is deferred_stream + deferred_stream.make_call.assert_not_awaited() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 6524353aa48..c8c0e2c0bfe 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -5691,6 +5691,89 @@ async def test_wrapper_async_failure_hook_exception_snapshot_preserves_traceback assert received[0].__traceback__ is not None +def test_client_records_dict_timing_before_sync_success_handler() -> None: + events: list[str] = [] + response = {"id": "msg_123"} + logging_obj = MagicMock() + logging_obj.stream = False + logging_obj.success_handler.side_effect = lambda *args: events.append("success") + immediate_executor = MagicMock() + immediate_executor.submit.side_effect = lambda context_run, callback, *args: context_run(callback, *args) + + def capture_timing(**kwargs: object) -> None: + assert kwargs["result"] is response + events.append("metadata") + + @client + def messages(model: str, **kwargs: object) -> dict[str, str]: + return response + + with ( + patch("litellm.utils.executor", immediate_executor), # test-quality-ok: controls async success-handler scheduling + patch("litellm.utils.update_response_metadata", side_effect=capture_timing), # test-quality-ok: observes metadata-before-logging ordering + ): + result = messages(model="anthropic/claude-haiku-4-5-20251001", litellm_logging_obj=logging_obj) + + assert result is response + assert events == ["metadata", "success"] + + +@pytest.mark.asyncio +async def test_client_records_dict_timing_before_async_logging() -> None: + events: list[str] = [] + response = {"id": "msg_123"} + logging_obj = MagicMock() + logging_obj.stream = False + logging_obj._defer_async_logging = False + + async def capture_async_logging(**kwargs: object) -> None: + events.append("success") + + def capture_timing(**kwargs: object) -> None: + assert kwargs["result"] is response + events.append("metadata") + + @client + async def amessages(model: str, **kwargs: object) -> dict[str, str]: + return response + + with ( + patch("litellm.utils.update_response_metadata", side_effect=capture_timing), # test-quality-ok: observes metadata-before-logging ordering + patch("litellm.utils._client_async_logging_helper", side_effect=capture_async_logging), # test-quality-ok: observes async logging dispatch + ): + result = await amessages(model="anthropic/claude-haiku-4-5-20251001", litellm_logging_obj=logging_obj) + await asyncio.sleep(0) + + assert result is response + assert events == ["metadata", "success"] + + +@pytest.mark.asyncio +async def test_client_preserves_outer_timing_for_internal_dict_call() -> None: + response = {"id": "msg_123"} + outer_timing_metrics = {"_response_ms": 2500.0, "litellm_overhead_time_ms": 200.0} + logging_obj = MagicMock() + logging_obj.stream = False + logging_obj._defer_async_logging = False + logging_obj.model_call_details = {"llm_api_duration_ms": 900.0} + logging_obj.caching_details = None + logging_obj.response_timing_metrics = outer_timing_metrics + + @client + async def amessages(model: str, **kwargs: object) -> dict[str, str]: + return response + + token = is_internal_call.set(True) + try: + result = await amessages(model="anthropic/claude-haiku-4-5-20251001", litellm_logging_obj=logging_obj) + finally: + is_internal_call.reset(token) + + assert result is response + logging_obj.set_response_timing_metrics.assert_not_called() + assert logging_obj.response_timing_metrics is outer_timing_metrics + + def test_snapshot_exception_for_hook_preserves_suppress_context_flag() -> None: """Regression: setting __cause__ has a documented CPython side effect of implicitly forcing __suppress_context__ to True, even when the real exception's own