scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(proxy): expose native streaming timing headers

Browse source
This commit is contained in:
Yucheng Zhu 2026-08-29 11:24:46 -07:00
parent 36ea28b092
commit 034d5fe62b
20 changed files with 1780 additions and 175 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/_internal_context.py
View file

@ -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)

48
litellm/caching/caching_handler.py
View file

@ -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(

7
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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,

130
litellm/litellm_core_utils/llm_response_utils/response_metadata.py
View file

@ -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)

5
litellm/llms/anthropic/experimental_pass_through/adapters/handler.py
View file

@ -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(

32
litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py
View file

@ -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

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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,
)

80
litellm/proxy/common_request_processing.py
View file

@ -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,

4
litellm/responses/litellm_completion_transformation/handler.py
View file

@ -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] = (

4
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -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

63
litellm/utils.py
View file

@ -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()

280
tests/test_litellm/caching/test_caching_handler.py
View file

@ -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()),

379
tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py
View file

@ -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 = {

76
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -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

50
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -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,

160
tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_output_config_passthrough.py
View file

@ -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

42
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -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

313
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -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()

195
tests/test_litellm/responses/litellm_completion_transformation/test_handler.py
View file

@ -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()

83
tests/test_litellm/test_utils.py
View file

@ -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