From 8947a74e13e27257cee26a30293f9ce0b5e5d399 Mon Sep 17 00:00:00 2001 From: Noah <46087295+ohnoah@users.noreply.github.com> Date: Tue, 7 Apr 2026 05:28:21 +0100 Subject: [PATCH 1/9] fix(cache): persist and replay streamed Responses API requests (#24580) * fix(cache): persist and replay streamed Responses API requests * Add focused coverage for streamed responses cache * Cover streamed responses cache helper branches * Exercise streamed responses cache edge branches --- litellm/caching/caching.py | 3 +- litellm/caching/caching_handler.py | 106 ++- litellm/responses/streaming_iterator.py | 738 ++++++++++++--- litellm/types/llms/openai.py | 5 +- .../test_responses_hooks.py | 888 +++++++++++++++++- tests/local_testing/test_caching_handler.py | 399 +++++++- .../test_responses_stream_cache_keys.py | 141 +++ 7 files changed, 2105 insertions(+), 175 deletions(-) create mode 100644 tests/local_testing/test_responses_stream_cache_keys.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index ce1bc26c5e0..11733ce4cee 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -432,9 +432,10 @@ class Cache: str: The final hashed cache key with the redis namespace. """ dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {}) + metadata = kwargs.get("metadata") or {} namespace = ( dynamic_cache_control.get("namespace") - or kwargs.get("metadata", {}).get("redis_namespace") + or metadata.get("redis_namespace") or self.namespace ) if namespace: diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7d514e648fe..02207a207ed 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -85,6 +85,19 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() +_RESPONSES_STREAMING_CALLBACK_CALL_TYPES = { + CallTypes.aresponses.value, + CallTypes.responses.value, +} + + +def _should_defer_streaming_cache_hit_callbacks( + *, call_type: str, kwargs: Dict[str, Any] +) -> bool: + return ( + kwargs.get("stream", False) is True + and call_type in _RESPONSES_STREAMING_CALLBACK_CALL_TYPES + ) class LLMCachingHandler: @@ -99,6 +112,7 @@ class LLMCachingHandler: self.async_streaming_chunks: List[ModelResponse] = [] self.sync_streaming_chunks: List[ModelResponse] = [] self.request_kwargs = request_kwargs + self.preset_cache_key: Optional[str] = None self.original_function = original_function self.start_time = start_time if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache): @@ -206,7 +220,10 @@ class LLMCachingHandler: custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) - if kwargs.get("stream", False) is False: + if not _should_defer_streaming_cache_hit_callbacks( + call_type=call_type, + kwargs=kwargs, + ): # LOG SUCCESS self._async_log_cache_hit_on_callbacks( logging_obj=logging_obj, @@ -215,11 +232,12 @@ class LLMCachingHandler: end_time=end_time, cache_hit=cache_hit, ) - cache_key = litellm.cache.get_cache_key(**kwargs) - if ( - isinstance(cached_result, BaseModel) - or isinstance(cached_result, CustomStreamWrapper) - ) and hasattr(cached_result, "_hidden_params"): + cache_key = ( + self.preset_cache_key + or self.request_kwargs.get("cache_key") + or litellm.cache.get_cache_key(**self.request_kwargs) + ) + if hasattr(cached_result, "_hidden_params"): cached_result._hidden_params["cache_key"] = cache_key # type: ignore return CachingHandlerResponse(cached_result=cached_result) elif ( @@ -265,8 +283,6 @@ class LLMCachingHandler: kwargs: Dict[str, Any], args: Optional[Tuple[Any, ...]] = None, ) -> CachingHandlerResponse: - from litellm.utils import CustomStreamWrapper - cached_result: Optional[Any] = None # Check if caching should be performed BEFORE doing expensive kwargs copy @@ -282,6 +298,11 @@ class LLMCachingHandler: args, ) ) + if new_kwargs.get("metadata") is None: + new_kwargs.pop("metadata", None) + if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: + new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) + self.request_kwargs = new_kwargs print_verbose("Checking Sync Cache") cached_result = litellm.cache.get_cache(**new_kwargs) if cached_result is not None: @@ -322,17 +343,22 @@ class LLMCachingHandler: is_async=False, ) - logging_obj.handle_sync_success_callbacks_for_async_calls( - result=cached_result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, + if not _should_defer_streaming_cache_hit_callbacks( + call_type=call_type, + kwargs=kwargs, + ): + logging_obj.handle_sync_success_callbacks_for_async_calls( + result=cached_result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + ) + cache_key = ( + self.preset_cache_key + or self.request_kwargs.get("cache_key") + or litellm.cache.get_cache_key(**self.request_kwargs) ) - cache_key = litellm.cache.get_cache_key(**kwargs) - if ( - isinstance(cached_result, BaseModel) - or isinstance(cached_result, CustomStreamWrapper) - ) and hasattr(cached_result, "_hidden_params"): + if hasattr(cached_result, "_hidden_params"): cached_result._hidden_params["cache_key"] = cache_key # type: ignore return CachingHandlerResponse(cached_result=cached_result) return CachingHandlerResponse(cached_result=cached_result) @@ -686,6 +712,11 @@ class LLMCachingHandler: args, ) ) + if new_kwargs.get("metadata") is None: + new_kwargs.pop("metadata", None) + if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: + new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) + self.request_kwargs = new_kwargs cached_result: Optional[Any] = None if call_type == CallTypes.aembedding.value: if isinstance(new_kwargs["input"], str): @@ -710,14 +741,26 @@ class LLMCachingHandler: if all(result is None for result in cached_result): cached_result = None else: + request_kwargs = new_kwargs.copy() + request_cache_key = request_kwargs.pop("cache_key", None) if litellm.cache._supports_async() is True: ## check if dual cache is supported ## + self.preset_cache_key = ( + request_cache_key or litellm.cache.get_cache_key(**request_kwargs) + ) cached_result = await litellm.cache.async_get_cache( - dynamic_cache_object=self.dual_cache, **new_kwargs + dynamic_cache_object=self.dual_cache, + cache_key=self.preset_cache_key, + **request_kwargs, ) else: # fallback for caches that don't support async + self.preset_cache_key = ( + request_cache_key or litellm.cache.get_cache_key(**request_kwargs) + ) cached_result = litellm.cache.get_cache( - dynamic_cache_object=self.dual_cache, **new_kwargs + dynamic_cache_object=self.dual_cache, + cache_key=self.preset_cache_key, + **request_kwargs, ) return cached_result @@ -825,8 +868,27 @@ class LLMCachingHandler: elif (call_type == "aresponses" or call_type == "responses") and isinstance( cached_result, dict ): - # Convert cached dict back to ResponsesAPIResponse object - cached_result = ResponsesAPIResponse(**cached_result) + from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + ) + + response_obj = ResponsesAPIResponse(**cached_result) + if ( + hasattr(response_obj, "_hidden_params") + and response_obj._hidden_params is not None + and isinstance(response_obj._hidden_params, dict) + ): + response_obj._hidden_params["cache_hit"] = True + + if kwargs.get("stream", False) is True: + cached_result = CachedResponsesAPIStreamingIterator( + response=response_obj, + logging_obj=logging_obj, + request_data=kwargs, + call_type=call_type, + ) + else: + cached_result = response_obj if ( hasattr(cached_result, "_hidden_params") diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 145ec3a641a..1902a347215 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,9 +1,12 @@ +from __future__ import annotations + import asyncio import json import time import traceback from datetime import datetime -from typing import Any, Dict, List, Optional +from functools import lru_cache +from typing import Any, Dict, List, Literal, Optional import httpx @@ -22,19 +25,25 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.llms.openai import ( - OutputTextDeltaEvent, - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIRequestParams, - ResponsesAPIResponse, - ResponsesAPIStreamEvents, - ResponsesAPIStreamingResponse, -) from litellm.types.utils import CallTypes from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook +@lru_cache(maxsize=1) +def _get_openai_response_types(): + from litellm.types.llms import openai as openai_types + + return openai_types + + +def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -> None: + if task.cancelled(): + return + exception = task.exception() + if exception is not None: + verbose_logger.error("%s failed: %s", task_name, exception) + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -46,7 +55,7 @@ class BaseResponsesAPIStreamingIterator: self, response: httpx.Response, model: str, - responses_api_provider_config: BaseResponsesAPIConfig, + responses_api_provider_config: Optional[BaseResponsesAPIConfig], logging_obj: LiteLLMLoggingObj, litellm_metadata: Optional[Dict[str, Any]] = None, custom_llm_provider: Optional[str] = None, @@ -58,9 +67,13 @@ class BaseResponsesAPIStreamingIterator: self.logging_obj = logging_obj self.finished = False self.responses_api_provider_config = responses_api_provider_config - self.completed_response: Optional[ResponsesAPIStreamingResponse] = None + self.completed_response: Optional[Any] = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called + self._completed_response_cached = False + self._completed_response_logged = False + self._completed_response_cache_hit: Optional[bool] = None + self._persist_completed_response_before_logging = True self._stream_created_time: float = time.time() # track request context for hooks @@ -101,7 +114,7 @@ class BaseResponsesAPIStreamingIterator: llm_provider=self.custom_llm_provider or "", ) - def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]: + def _process_chunk(self, chunk) -> Optional[Any]: """Process a single chunk of data from the stream""" if not chunk: return None @@ -122,6 +135,10 @@ class BaseResponsesAPIStreamingIterator: # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): + if self.responses_api_provider_config is None: + raise ValueError( + "responses_api_provider_config is required to process live streaming chunks" + ) openai_responses_api_chunk = ( self.responses_api_provider_config.transform_streaming_response( model=self.model, @@ -195,10 +212,11 @@ class BaseResponsesAPIStreamingIterator: if self.litellm_metadata and self.litellm_metadata.get( "encrypted_content_affinity_enabled" ): + openai_types = _get_openai_response_types() event_type = getattr(openai_responses_api_chunk, "type", None) if event_type in ( - ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): item = getattr(openai_responses_api_chunk, "item", None) if item: @@ -219,10 +237,11 @@ class BaseResponsesAPIStreamingIterator: # Store the completed response (also for incomplete/failed so logging still fires) _chunk_type = getattr(openai_responses_api_chunk, "type", None) + openai_types = _get_openai_response_types() if openai_responses_api_chunk and _chunk_type in ( - ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - ResponsesAPIStreamEvents.RESPONSE_FAILED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True @@ -230,11 +249,11 @@ class BaseResponsesAPIStreamingIterator: litellm.include_cost_in_streaming_usage and self.logging_obj is not None ): - response_obj: Optional[ResponsesAPIResponse] = getattr( + response_obj: Optional[Any] = getattr( openai_responses_api_chunk, "response", None ) if response_obj: - usage_obj: Optional[ResponseAPIUsage] = getattr( + usage_obj: Optional[Any] = getattr( response_obj, "usage", None ) if usage_obj is not None: @@ -247,9 +266,13 @@ class BaseResponsesAPIStreamingIterator: if cost is not None: setattr(usage_obj, "cost", cost) except Exception: + # Best-effort usage cost annotation should not break stream replay. pass - if _chunk_type == ResponsesAPIStreamEvents.RESPONSE_FAILED: + if ( + _chunk_type + == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED + ): self._handle_logging_failed_response() else: self._handle_logging_completed_response() @@ -266,6 +289,59 @@ class BaseResponsesAPIStreamingIterator: self._handle_failure(e) raise + def _log_completed_response(self, *, is_async: bool) -> None: + if self._completed_response_logged: + return + self._completed_response_logged = True + + if self._persist_completed_response_before_logging: + self._persist_completed_response_to_cache(is_async=is_async) + + # Create a copy for logging to avoid modifying the response object that will be returned to the user + # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) + # to chat completion format (prompt_tokens/completion_tokens) for internal logging + # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with + # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) + logging_response = self.completed_response + if self.completed_response is not None and hasattr( + self.completed_response, "model_dump" + ): + try: + logging_response = type(self.completed_response).model_validate( + self.completed_response.model_dump() + ) + except Exception: + # Fallback to original if serialization fails + pass + + end_time = datetime.now() + if is_async: + asyncio.create_task( + self.logging_obj.async_success_handler( + result=logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + ) + ) + else: + run_async_function( + async_function=self.logging_obj.async_success_handler, + result=logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + ) + + executor.submit( + self.logging_obj.success_handler, + result=logging_response, + cache_hit=self._completed_response_cache_hit, + start_time=self.start_time, + end_time=end_time, + ) + self._run_post_success_hooks(end_time=end_time) + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" pass @@ -296,6 +372,88 @@ class BaseResponsesAPIStreamingIterator: ) self._handle_failure(exception) + def _get_completed_response_object(self) -> Optional[Any]: + openai_types = _get_openai_response_types() + completed_response = self.completed_response + if isinstance(completed_response, openai_types.ResponsesAPIResponse): + return completed_response + + response_obj = getattr(completed_response, "response", None) + if isinstance(response_obj, openai_types.ResponsesAPIResponse): + return response_obj + + return None + + def _persist_completed_response_to_cache(self, *, is_async: bool) -> None: + if self._completed_response_cached: + return + + completed_response = self.completed_response + openai_types = _get_openai_response_types() + if ( + getattr(completed_response, "type", None) + != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + return + + response_obj = self._get_completed_response_object() + if response_obj is None: + return + + caching_handler = getattr(self.logging_obj, "_llm_caching_handler", None) + if caching_handler is None: + return + + request_kwargs = getattr(caching_handler, "request_kwargs", None) + if ( + not isinstance(request_kwargs, dict) + or request_kwargs.get("stream") is not True + ): + return + request_kwargs = request_kwargs.copy() + preset_cache_key = getattr(caching_handler, "preset_cache_key", None) + request_cache_key = request_kwargs.pop("cache_key", None) + if preset_cache_key is None: + preset_cache_key = request_cache_key + if request_kwargs.get("metadata") is None: + request_kwargs.pop("metadata", None) + request_kwargs.pop("custom_llm_provider", None) + if preset_cache_key is not None: + request_kwargs["cache_key"] = preset_cache_key + + if not caching_handler._should_store_result_in_cache( + original_function=caching_handler.original_function, + kwargs=request_kwargs, + ): + return + + if litellm.cache is None: + return + + cached_response = response_obj.model_dump_json() + if is_async: + cache_write_task = asyncio.create_task( + litellm.cache.async_add_cache( + cached_response, + dynamic_cache_object=getattr(caching_handler, "dual_cache", None), + **request_kwargs, + ) + ) + cache_write_task.add_done_callback( + lambda task: _log_background_task_failure( + task, + task_name="Responses stream cache write", + ) + ) + else: + litellm.cache.add_cache( + cached_response, + dynamic_cache_object=getattr(caching_handler, "dual_cache", None), + **request_kwargs, + ) + + self._completed_response_cached = True + async def _call_post_streaming_deployment_hook(self, chunk): """ Allow callbacks to modify streaming chunks before returning (parity with chat). @@ -480,7 +638,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> ResponsesAPIStreamingResponse: + async def __anext__(self) -> Any: try: self._check_max_streaming_duration() while True: @@ -520,40 +678,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in async context""" - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr( - self.completed_response, "model_dump" - ): - try: - logging_response = type(self.completed_response).model_validate( - self.completed_response.model_dump() - ) - except Exception: - # Fallback to original if serialization fails - pass - - asyncio.create_task( - self.logging_obj.async_success_handler( - result=logging_response, - start_time=self.start_time, - end_time=datetime.now(), - cache_hit=None, - ) - ) - - executor.submit( - self.logging_obj.success_handler, - result=logging_response, - cache_hit=None, - start_time=self.start_time, - end_time=datetime.now(), - ) - self._run_post_success_hooks(end_time=datetime.now()) + self._log_completed_response(is_async=True) class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -627,39 +752,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in sync context""" - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr( - self.completed_response, "model_dump" - ): - try: - logging_response = type(self.completed_response).model_validate( - self.completed_response.model_dump() - ) - except Exception: - # Fallback to original if serialization fails - pass - - run_async_function( - async_function=self.logging_obj.async_success_handler, - result=logging_response, - start_time=self.start_time, - end_time=datetime.now(), - cache_hit=None, - ) - - executor.submit( - self.logging_obj.success_handler, - result=logging_response, - cache_hit=None, - start_time=self.start_time, - end_time=datetime.now(), - ) - self._run_post_success_hooks(end_time=datetime.now()) + self._log_completed_response(is_async=False) class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -683,90 +776,441 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data: Optional[Dict[str, Any]] = None, call_type: Optional[str] = None, ): - super().__init__( - response=response, + transformed = responses_api_provider_config.transform_response_api_response( model=model, - responses_api_provider_config=responses_api_provider_config, + raw_response=response, + logging_obj=logging_obj, + ) + super().__init__( + response=httpx.Response(200), + model=model, + responses_api_provider_config=None, logging_obj=logging_obj, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, request_data=request_data, call_type=call_type, ) + self._set_events_from_response(transformed=transformed, logging_obj=logging_obj) - # one-time transform - transformed = ( - self.responses_api_provider_config.transform_response_api_response( - model=self.model, - raw_response=response, - logging_obj=logging_obj, - ) + def _set_events_from_response( + self, + transformed: Any, + logging_obj: LiteLLMLoggingObj, + ) -> None: + self._events = _build_synthetic_response_events( + transformed=transformed, + logging_obj=logging_obj, + chunk_size=self.CHUNK_SIZE, ) - full_text = self._collect_text(transformed) - - # build a list of 5‑char delta events - deltas = [ - OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - delta=full_text[i : i + self.CHUNK_SIZE], - item_id=transformed.id, - output_index=0, - content_index=0, - ) - for i in range(0, len(full_text), self.CHUNK_SIZE) - ] - - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Optional[ResponseAPIUsage] = getattr(transformed, "usage", None) - if usage_obj is not None: - try: - cost: Optional[float] = logging_obj._response_cost_calculator( - result=transformed - ) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - # If cost calculation fails, continue without cost - pass - - # append the completed event - self._events = deltas + [ - ResponseCompletedEvent( - type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - response=transformed, - ) - ] self._idx = 0 + self.completed_response = self._events[-1] def __aiter__(self): return self - async def __anext__(self) -> ResponsesAPIStreamingResponse: + async def __anext__(self) -> Any: if self._idx >= len(self._events): raise StopAsyncIteration evt = self._events[self._idx] self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=True) return evt def __iter__(self): return self - def __next__(self) -> ResponsesAPIStreamingResponse: + def __next__(self) -> Any: if self._idx >= len(self._events): raise StopIteration evt = self._events[self._idx] self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=False) return evt - def _collect_text(self, resp: ResponsesAPIResponse) -> str: - out = "" - for out_item in resp.output: - item_type = getattr(out_item, "type", None) - if item_type == "message": - for c in getattr(out_item, "content", []): - out += c.text - return out + +class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): + def __init__( + self, + response: Any, + logging_obj: LiteLLMLoggingObj, + request_data: Optional[Dict[str, Any]] = None, + call_type: Optional[str] = None, + ): + BaseResponsesAPIStreamingIterator.__init__( + self, + response=httpx.Response(200), + model=getattr(response, "model", ""), + responses_api_provider_config=None, + logging_obj=logging_obj, + litellm_metadata=None, + custom_llm_provider="cached_response", + request_data=request_data, + call_type=call_type, + ) + self._completed_response_cache_hit = True + self._persist_completed_response_before_logging = False + self._events: List[Any] = [] + self._idx = 0 + self._set_events_from_response(transformed=response, logging_obj=logging_obj) + + def _set_events_from_response( + self, + transformed: Any, + logging_obj: LiteLLMLoggingObj, + ) -> None: + self._events = _build_synthetic_response_events( + transformed=transformed, + logging_obj=logging_obj, + chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE, + ) + self._idx = 0 + self.completed_response = self._events[-1] + + def __aiter__(self): + return self + + async def __anext__(self) -> Any: + if self._idx >= len(self._events): + raise StopAsyncIteration + evt = self._events[self._idx] + self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=True) + return evt + + def __iter__(self): + return self + + def __next__(self) -> Any: + if self._idx >= len(self._events): + raise StopIteration + evt = self._events[self._idx] + self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=False) + return evt + + +def _dump_response_object(obj: Any) -> Dict[str, Any]: + if hasattr(obj, "model_dump"): + return obj.model_dump() + if isinstance(obj, dict): + return obj + return {} + + +def _build_response_status_event( + event_type: Literal[ + "response.created", + "response.in_progress", + ], + transformed: Any, +) -> Any: + openai_types = _get_openai_response_types() + in_progress_response = transformed.model_copy( + deep=True, + update={"status": "in_progress", "output": []}, + ) + if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED: + return openai_types.ResponseCreatedEvent( + type=event_type, response=in_progress_response + ) + return openai_types.ResponseInProgressEvent( + type=event_type, response=in_progress_response + ) + + +def _build_content_part_done_event( + *, + item_id: str, + output_index: int, + content_index: int, + part_payload: Dict[str, Any], +) -> Optional[Any]: + openai_types = _get_openai_response_types() + part_type = part_payload.get("type") + part: Any + if part_type == "output_text": + annotations = [ + openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) + for annotation in part_payload.get("annotations", []) or [] + ] + part = openai_types.ContentPartDonePartOutputText( + type="output_text", + text=str(part_payload.get("text") or ""), + annotations=annotations, + logprobs=part_payload.get("logprobs"), + ) + elif part_type == "refusal": + part = openai_types.ContentPartDonePartRefusal( + type="refusal", + refusal=str(part_payload.get("refusal") or ""), + ) + elif part_type == "reasoning_text": + part = openai_types.ContentPartDonePartReasoningText( + type="reasoning_text", + reasoning=str(part_payload.get("reasoning") or ""), + ) + else: + return None + + return openai_types.ContentPartDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=item_id, + output_index=output_index, + content_index=content_index, + part=part, + ) + + +def _add_text_like_part_events( + *, + events: List[Any], + item_id: str, + output_index: int, + content_index: int, + part_payload: Dict[str, Any], + chunk_size: int, +) -> None: + openai_types = _get_openai_response_types() + part_type = part_payload.get("type") + if part_type == "output_text": + text = str(part_payload.get("text") or "") + for i in range(0, len(text), chunk_size): + events.append( + openai_types.OutputTextDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id=item_id, + output_index=output_index, + content_index=content_index, + delta=text[i : i + chunk_size], + ) + ) + for annotation_index, annotation in enumerate( + part_payload.get("annotations", []) or [] + ): + events.append( + openai_types.OutputTextAnnotationAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, + item_id=item_id, + output_index=output_index, + content_index=content_index, + annotation_index=annotation_index, + annotation=annotation, + ) + ) + events.append( + openai_types.OutputTextDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=item_id, + output_index=output_index, + content_index=content_index, + text=text, + ) + ) + elif part_type == "refusal": + refusal = str(part_payload.get("refusal") or "") + for i in range(0, len(refusal), chunk_size): + events.append( + openai_types.RefusalDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DELTA, + item_id=item_id, + output_index=output_index, + content_index=content_index, + delta=refusal[i : i + chunk_size], + ) + ) + events.append( + openai_types.RefusalDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DONE, + item_id=item_id, + output_index=output_index, + content_index=content_index, + refusal=refusal, + ) + ) + + +def _build_synthetic_response_events( + *, + transformed: Any, + logging_obj: LiteLLMLoggingObj, + chunk_size: int, +) -> List[Any]: + openai_types = _get_openai_response_types() + if litellm.include_cost_in_streaming_usage and logging_obj is not None: + usage_obj: Optional[Any] = getattr(transformed, "usage", None) + if usage_obj is not None: + try: + cost: Optional[float] = logging_obj._response_cost_calculator( + result=transformed + ) + if cost is not None: + setattr(usage_obj, "cost", cost) + except Exception: + pass + + events: List[Any] = [ + _build_response_status_event( + openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed + ), + _build_response_status_event( + openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed + ), + ] + + sequence_number = 0 + for output_index, output_item in enumerate( + getattr(transformed, "output", []) or [] + ): + output_item_payload = _dump_response_object(output_item) + item_id = str(output_item_payload.get("id") or transformed.id) + item_type = output_item_payload.get("type") + + events.append( + openai_types.OutputItemAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=output_index, + item=openai_types.BaseLiteLLMOpenAIResponseObject( + **output_item_payload + ), + ) + ) + + if item_type == "message": + for content_index, part in enumerate( + output_item_payload.get("content", []) or [] + ): + part_payload = _dump_response_object(part) + events.append( + openai_types.ContentPartAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id=item_id, + output_index=output_index, + content_index=content_index, + part=openai_types.BaseLiteLLMOpenAIResponseObject( + **part_payload + ), + ) + ) + _add_text_like_part_events( + events=events, + item_id=item_id, + output_index=output_index, + content_index=content_index, + part_payload=part_payload, + chunk_size=chunk_size, + ) + done_event = _build_content_part_done_event( + item_id=item_id, + output_index=output_index, + content_index=content_index, + part_payload=part_payload, + ) + if done_event is not None: + events.append(done_event) + elif item_type == "function_call": + arguments = str(output_item_payload.get("arguments") or "") + for i in range(0, len(arguments), chunk_size): + events.append( + openai_types.FunctionCallArgumentsDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + item_id=item_id, + output_index=output_index, + delta=arguments[i : i + chunk_size], + ) + ) + events.append( + openai_types.FunctionCallArgumentsDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, + item_id=item_id, + output_index=output_index, + arguments=arguments, + ) + ) + elif item_type == "reasoning": + for summary_index, summary in enumerate( + output_item_payload.get("summary", []) or [] + ): + summary_payload = _dump_response_object(summary) + summary_text = str(summary_payload.get("text") or "") + for i in range(0, len(summary_text), chunk_size): + events.append( + openai_types.ReasoningSummaryTextDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id=item_id, + output_index=output_index, + summary_index=summary_index, + delta=summary_text[i : i + chunk_size], + ) + ) + sequence_number += 1 + events.append( + openai_types.ReasoningSummaryTextDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, + item_id=item_id, + output_index=output_index, + sequence_number=sequence_number, + summary_index=summary_index, + text=summary_text, + ) + ) + sequence_number += 1 + events.append( + openai_types.ReasoningSummaryPartDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, + item_id=item_id, + output_index=output_index, + sequence_number=sequence_number, + summary_index=summary_index, + part=openai_types.BaseLiteLLMOpenAIResponseObject( + **summary_payload + ), + ) + ) + + sequence_number += 1 + events.append( + openai_types.OutputItemDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + sequence_number=sequence_number, + item=openai_types.BaseLiteLLMOpenAIResponseObject( + **output_item_payload + ), + ) + ) + + events.append( + openai_types.ResponseCompletedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=transformed, + ) + ) + return events # --------------------------------------------------------------------------- @@ -951,8 +1395,8 @@ class ResponsesWebSocketStreaming: # --------------------------------------------------------------------------- _RESPONSE_CREATE_PARAMS: frozenset = ( - ResponsesAPIRequestParams.__required_keys__ - | ResponsesAPIRequestParams.__optional_keys__ + _get_openai_response_types().ResponsesAPIRequestParams.__required_keys__ + | _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__ ) _MANAGED_WS_SKIP_KWARGS: frozenset = frozenset( @@ -1085,7 +1529,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: Dict[str, Any] + completed_event: Dict[str, Any], ) -> List[Dict[str, Any]]: """ Convert the output items in a ``response.completed`` event into diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 2fd0c4ea970..986ec39f3bb 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1482,6 +1482,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA] item_id: str output_index: int + summary_index: int = 0 delta: str @@ -1490,7 +1491,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): item_id: str output_index: int sequence_number: int - summary_index: int + summary_index: int = 0 text: str @@ -1499,7 +1500,7 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): item_id: str output_index: int sequence_number: int - summary_index: int + summary_index: int = 0 part: BaseLiteLLMOpenAIResponseObject diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 3227fecdfb2..8933ec8f3c5 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -1,6 +1,9 @@ import asyncio +from contextlib import suppress from datetime import datetime +import json from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -8,8 +11,17 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.responses import streaming_iterator as streaming_module -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + MockResponsesAPIStreamingIterator, + ResponsesAPIStreamingIterator, + SyncResponsesAPIStreamingIterator, +) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) from litellm.types.utils import CallTypes @@ -19,15 +31,19 @@ class _FakeLoggingObj: self.async_success_calls = 0 self.failure_calls = 0 self.async_failure_calls = 0 + self.last_success_kwargs = None + self.last_async_success_kwargs = None self.start_time = datetime.now() self.model_call_details = {"litellm_params": {}} # Signature alignment with Logging handlers def success_handler(self, *args, **kwargs): self.success_calls += 1 + self.last_success_kwargs = kwargs async def async_success_handler(self, *args, **kwargs): self.async_success_calls += 1 + self.last_async_success_kwargs = kwargs def failure_handler(self, *args, **kwargs): self.failure_calls += 1 @@ -36,6 +52,115 @@ class _FakeLoggingObj: self.async_failure_calls += 1 +def _make_completed_response(response_id: str = "resp_test") -> ResponseCompletedEvent: + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id=response_id, + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[ + { + "type": "message", + "id": f"msg_{response_id}", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "cached streamed response", + "annotations": [], + } + ], + } + ], + ), + ) + + +@pytest.mark.asyncio +async def test_log_background_task_failure_logs_task_exceptions(monkeypatch): + error_logger = MagicMock() + monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger) + + async def _boom(): + raise RuntimeError("boom") + + task = asyncio.create_task(_boom()) + with suppress(RuntimeError): + await task + + streaming_module._log_background_task_failure(task, task_name="cache write") + + error_logger.assert_called_once() + assert error_logger.call_args.args == ( + "%s failed: %s", + "cache write", + task.exception(), + ) + + +@pytest.mark.asyncio +async def test_log_background_task_failure_ignores_cancelled_tasks(monkeypatch): + error_logger = MagicMock() + monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger) + + task = asyncio.create_task(asyncio.sleep(1)) + task.cancel() + with suppress(asyncio.CancelledError): + await task + + streaming_module._log_background_task_failure(task, task_name="cache write") + + error_logger.assert_not_called() + + +def test_content_part_done_event_supports_refusal_and_reasoning_text(): + refusal_event = streaming_module._build_content_part_done_event( + item_id="msg_1", + output_index=0, + content_index=0, + part_payload={"type": "refusal", "refusal": "no"}, + ) + reasoning_event = streaming_module._build_content_part_done_event( + item_id="msg_1", + output_index=0, + content_index=1, + part_payload={"type": "reasoning_text", "reasoning": "because"}, + ) + unsupported_event = streaming_module._build_content_part_done_event( + item_id="msg_1", + output_index=0, + content_index=2, + part_payload={"type": "image"}, + ) + + assert refusal_event.part.type == "refusal" + assert refusal_event.part.refusal == "no" + assert reasoning_event.part.type == "reasoning_text" + assert reasoning_event.part.reasoning == "because" + assert unsupported_event is None + + +def test_dump_response_object_handles_model_and_unknown_values(): + response = ResponsesAPIResponse( + id="resp_dump", + created_at=int(datetime.now().timestamp()), + status="completed", + model="gpt-4.1-mini", + object="response", + output=[], + ) + + assert streaming_module._dump_response_object(response)["id"] == "resp_dump" + assert streaming_module._dump_response_object({"type": "message"}) == { + "type": "message" + } + assert streaming_module._dump_response_object(object()) == {} + + @pytest.mark.asyncio async def test_responses_streaming_triggers_hooks(monkeypatch): """ @@ -167,3 +292,762 @@ async def test_responses_streaming_failure_triggers_failure_handlers(): await asyncio.sleep(0.2) assert logging_obj.failure_calls >= 1 assert logging_obj.async_failure_calls >= 1 + + +def test_process_chunk_requires_provider_config(): + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=None, + logging_obj=_FakeLoggingObj(), + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + with pytest.raises(ValueError, match="responses_api_provider_config is required"): + iterator._process_chunk(json.dumps({"type": "response.completed"})) + + +def test_process_chunk_wraps_encrypted_content_with_model_id(): + openai_types = streaming_module._get_openai_response_types() + + class _EncryptedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.OutputItemAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=openai_types.BaseLiteLLMOpenAIResponseObject( + id="rs_123", + type="reasoning", + encrypted_content="ciphertext", + ), + ) + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_EncryptedConfig(), + logging_obj=_FakeLoggingObj(), + litellm_metadata={ + "encrypted_content_affinity_enabled": True, + "model_info": {"id": "model-123"}, + }, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + event = iterator._process_chunk(json.dumps({"type": "response.output_item.added"})) + + assert event.item.encrypted_content.startswith("litellm_enc:") + assert event.item.encrypted_content.endswith(";ciphertext") + + +def test_process_chunk_completed_response_updates_id_and_usage_cost(monkeypatch): + original_include_cost = litellm.include_cost_in_streaming_usage + litellm.include_cost_in_streaming_usage = True + openai_types = streaming_module._get_openai_response_types() + + class _CompletedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.ResponseCompletedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_live", + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[], + usage=openai_types.ResponseAPIUsage( + input_tokens=1, + output_tokens=2, + total_tokens=3, + ), + ), + ) + + logging_obj = _FakeLoggingObj() + logging_obj._response_cost_calculator = MagicMock(return_value=1.23) + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_CompletedConfig(), + logging_obj=logging_obj, + litellm_metadata={"model_info": {"id": "model-123"}}, + custom_llm_provider="openai", + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + completion_handler = MagicMock() + monkeypatch.setattr( + iterator, "_handle_logging_completed_response", completion_handler + ) + + try: + event = iterator._process_chunk(json.dumps({"type": "response.completed"})) + finally: + litellm.include_cost_in_streaming_usage = original_include_cost + + assert iterator.completed_response is event + assert event.response.id != "resp_live" + assert event.response.id.startswith("resp_") + assert event.response.usage.cost == 1.23 + completion_handler.assert_called_once() + + +def test_process_chunk_failed_response_triggers_failure_logging(monkeypatch): + openai_types = streaming_module._get_openai_response_types() + + class _FailedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.ResponseFailedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, + response=ResponsesAPIResponse( + id="resp_failed", + created_at=int(datetime.now().timestamp()), + status="failed", + model="test-model", + object="response", + output=[], + error={"message": "provider failed"}, + ), + ) + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_FailedConfig(), + logging_obj=_FakeLoggingObj(), + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + failure_handler = MagicMock() + monkeypatch.setattr(iterator, "_handle_logging_failed_response", failure_handler) + + event = iterator._process_chunk(json.dumps({"type": "response.failed"})) + + assert iterator.completed_response is event + failure_handler.assert_called_once() + + +@pytest.mark.asyncio +async def test_handle_logging_failed_response_uses_response_error_message(): + openai_types = streaming_module._get_openai_response_types() + logging_obj = _FakeLoggingObj() + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + iterator.completed_response = openai_types.ResponseFailedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, + response=ResponsesAPIResponse( + id="resp_failed_real", + created_at=int(datetime.now().timestamp()), + status="failed", + model="test-model", + object="response", + output=[], + error={"message": "provider failed"}, + ), + ) + + iterator._handle_logging_failed_response() + await asyncio.sleep(0.2) + + assert logging_obj.failure_calls == 1 + assert logging_obj.async_failure_calls == 1 + + +def test_process_chunk_returns_none_for_invalid_json_and_non_dict_payload(): + class _NoopConfig: + def transform_streaming_response(self, **kwargs): + raise AssertionError("should not be called") + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_NoopConfig(), + logging_obj=_FakeLoggingObj(), + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + assert iterator._process_chunk("not-json") is None + assert iterator._process_chunk(json.dumps(["not", "a", "dict"])) is None + + +def test_process_chunk_cost_annotation_failure_is_nonfatal(monkeypatch): + original_include_cost = litellm.include_cost_in_streaming_usage + litellm.include_cost_in_streaming_usage = True + openai_types = streaming_module._get_openai_response_types() + + class _CompletedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.ResponseCompletedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_cost_failure", + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[], + usage=openai_types.ResponseAPIUsage( + input_tokens=1, + output_tokens=2, + total_tokens=3, + ), + ), + ) + + logging_obj = _FakeLoggingObj() + logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom")) + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_CompletedConfig(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + completion_handler = MagicMock() + monkeypatch.setattr( + iterator, "_handle_logging_completed_response", completion_handler + ) + + try: + event = iterator._process_chunk(json.dumps({"type": "response.completed"})) + finally: + litellm.include_cost_in_streaming_usage = original_include_cost + + assert iterator.completed_response is event + assert event.response.usage.cost is None + completion_handler.assert_called_once() + + +def test_get_completed_response_object_accepts_direct_response(): + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + direct_response = _make_completed_response("resp_direct").response + iterator.completed_response = direct_response + + assert iterator._get_completed_response_object() is direct_response + + +@pytest.mark.asyncio +async def test_responses_streaming_completed_event_persists_async_cache(): + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": True, + "caching": True, + "cache_key": "stale-request-cache-key", + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.aresponses, + async_set_cache=AsyncMock(), + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + logging_obj._llm_caching_handler = caching_handler + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data=caching_handler.request_kwargs, + call_type=CallTypes.aresponses.value, + ) + iterator.completed_response = _make_completed_response() + + iterator._handle_logging_completed_response() + await asyncio.sleep(0.2) + + litellm.cache.async_add_cache.assert_called_once() + assert litellm.cache.async_add_cache.call_args.kwargs["stream"] is True + assert ( + litellm.cache.async_add_cache.call_args.kwargs["cache_key"] + == "responses-stream-cache-key" + ) + assert "metadata" not in litellm.cache.async_add_cache.call_args.kwargs + assert "custom_llm_provider" not in litellm.cache.async_add_cache.call_args.kwargs + assert ( + json.loads(litellm.cache.async_add_cache.call_args.args[0])["id"] + == iterator.completed_response.response.id + ) + litellm.cache = original_cache + + +def test_responses_streaming_completed_event_persists_sync_cache(): + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": True, + "caching": True, + "cache_key": "stale-request-cache-key", + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.responses, + sync_set_cache=MagicMock(), + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + logging_obj._llm_caching_handler = caching_handler + + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data=caching_handler.request_kwargs, + call_type=CallTypes.responses.value, + ) + iterator.completed_response = _make_completed_response("resp_sync") + + iterator._handle_logging_completed_response() + + litellm.cache.add_cache.assert_called_once() + assert litellm.cache.add_cache.call_args.kwargs["stream"] is True + assert ( + litellm.cache.add_cache.call_args.kwargs["cache_key"] + == "responses-stream-cache-key" + ) + assert "metadata" not in litellm.cache.add_cache.call_args.kwargs + assert "custom_llm_provider" not in litellm.cache.add_cache.call_args.kwargs + assert ( + json.loads(litellm.cache.add_cache.call_args.args[0])["id"] + == iterator.completed_response.response.id + ) + litellm.cache = original_cache + + +def test_log_completed_response_sync_direct_path(monkeypatch): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + iterator._persist_completed_response_before_logging = False + iterator.completed_response = _make_completed_response("resp_log_sync") + + iterator._log_completed_response(is_async=False) + asyncio.run(asyncio.sleep(0.2)) + + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + + +def test_log_completed_response_falls_back_when_model_validate_fails(monkeypatch): + class _BadSerializableResponse: + @classmethod + def model_validate(cls, value): + raise RuntimeError("nope") + + def model_dump(self): + return {"id": "bad"} + + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + iterator._persist_completed_response_before_logging = False + iterator.completed_response = _BadSerializableResponse() + monkeypatch.setattr(iterator, "_run_post_success_hooks", MagicMock()) + + iterator._log_completed_response(is_async=False) + asyncio.run(asyncio.sleep(0.2)) + + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + + +@pytest.mark.parametrize( + "scenario", + [ + "already_cached", + "not_completed", + "missing_caching_handler", + "not_streaming", + "store_disabled", + "missing_cache_backend", + ], +) +def test_persist_completed_response_to_cache_guard_branches(monkeypatch, scenario): + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + openai_types = streaming_module._get_openai_response_types() + completed_event = _make_completed_response("resp_guard") + iterator.completed_response = completed_event + + if scenario == "already_cached": + iterator._completed_response_cached = True + elif scenario == "not_completed": + iterator.completed_response = openai_types.ResponseIncompleteEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + response=completed_event.response, + ) + elif scenario == "missing_caching_handler": + logging_obj._llm_caching_handler = None + else: + logging_obj._llm_caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": scenario != "not_streaming", + "cache_key": "request-cache-key", + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key=None, + original_function=litellm.responses, + dual_cache=None, + _should_store_result_in_cache=lambda original_function, kwargs: ( + scenario != "store_disabled" + ), + ) + if scenario == "missing_cache_backend": + monkeypatch.setattr(streaming_module.litellm, "cache", None) + else: + monkeypatch.setattr( + streaming_module.litellm, + "cache", + SimpleNamespace(add_cache=MagicMock(), async_add_cache=AsyncMock()), + ) + + iterator._persist_completed_response_to_cache(is_async=False) + + expected_cached_flag = scenario == "already_cached" + assert iterator._completed_response_cached is expected_cached_flag + + +def test_build_synthetic_response_events_covers_annotations_function_calls_and_refusals(): + original_include_cost = litellm.include_cost_in_streaming_usage + litellm.include_cost_in_streaming_usage = True + logging_obj = _FakeLoggingObj() + logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom")) + transformed = ResponsesAPIResponse( + id="resp_events", + created_at=int(datetime.now().timestamp()), + status="completed", + model="gpt-4.1-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_events", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "hello world", + "annotations": [{"type": "file_citation", "file_id": "file_1"}], + }, + { + "type": "refusal", + "refusal": "no thanks", + }, + ], + }, + { + "type": "function_call", + "id": "fc_events", + "call_id": "call_123", + "name": "lookup", + "arguments": '{"id":1}', + }, + ], + ) + + try: + events = streaming_module._build_synthetic_response_events( + transformed=transformed, + logging_obj=logging_obj, + chunk_size=5, + ) + finally: + litellm.include_cost_in_streaming_usage = original_include_cost + + event_types = [ + event.type.value if hasattr(event.type, "value") else str(event.type) + for event in events + ] + + assert "response.output_text.annotation.added" in event_types + assert "response.refusal.delta" in event_types + assert "response.refusal.done" in event_types + assert "response.function_call_arguments.delta" in event_types + assert "response.function_call_arguments.done" in event_types + assert event_types[-1] == "response.completed" + + +@pytest.mark.asyncio +async def test_mock_responses_streaming_iterator_async_iteration_logs_completion( + monkeypatch, +): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + class _MockTransformConfig: + def transform_response_api_response(self, **kwargs): + return _make_completed_response("resp_mock").response + + logging_obj = _FakeLoggingObj() + + iterator = MockResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_MockTransformConfig(), + logging_obj=logging_obj, + request_data={"model": "test-model", "stream": True}, + call_type=CallTypes.responses.value, + ) + + streamed_events = [event async for event in iterator] + await asyncio.sleep(0.2) + + assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + + +def test_mock_responses_streaming_iterator_sync_iteration_logs_completion(monkeypatch): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + class _MockTransformConfig: + def transform_response_api_response(self, **kwargs): + return _make_completed_response("resp_mock_sync").response + + logging_obj = _FakeLoggingObj() + iterator = MockResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_MockTransformConfig(), + logging_obj=logging_obj, + request_data={"model": "test-model", "stream": True}, + call_type=CallTypes.responses.value, + ) + + streamed_events = list(iterator) + asyncio.run(asyncio.sleep(0.2)) + + assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + + +@pytest.mark.asyncio +async def test_cached_responses_stream_async_hit_triggers_success_callbacks( + monkeypatch, +): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + logging_obj._llm_caching_handler = SimpleNamespace( + request_kwargs={"model": "test-model", "input": "hello", "stream": True}, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.aresponses, + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + + iterator = CachedResponsesAPIStreamingIterator( + response=_make_completed_response("resp_cached_async").response, + logging_obj=logging_obj, + request_data={"model": "test-model", "input": "hello", "stream": True}, + call_type=CallTypes.aresponses.value, + ) + + streamed_events = [event async for event in iterator] + await asyncio.sleep(0.2) + + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert logging_obj.last_success_kwargs["cache_hit"] is True + assert logging_obj.last_async_success_kwargs["cache_hit"] is True + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + litellm.cache.async_add_cache.assert_not_called() + litellm.cache.add_cache.assert_not_called() + litellm.cache = original_cache + + +def test_cached_responses_stream_sync_hit_triggers_success_callbacks(monkeypatch): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + logging_obj._llm_caching_handler = SimpleNamespace( + request_kwargs={"model": "test-model", "input": "hello", "stream": True}, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.responses, + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + + iterator = CachedResponsesAPIStreamingIterator( + response=_make_completed_response("resp_cached_sync").response, + logging_obj=logging_obj, + request_data={"model": "test-model", "input": "hello", "stream": True}, + call_type=CallTypes.responses.value, + ) + + streamed_events = list(iterator) + asyncio.run(asyncio.sleep(0.2)) + + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert logging_obj.last_success_kwargs["cache_hit"] is True + assert logging_obj.last_async_success_kwargs["cache_hit"] is True + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + litellm.cache.async_add_cache.assert_not_called() + litellm.cache.add_cache.assert_not_called() + litellm.cache = original_cache diff --git a/tests/local_testing/test_caching_handler.py b/tests/local_testing/test_caching_handler.py index 806f72bfde8..0ce9938e710 100644 --- a/tests/local_testing/test_caching_handler.py +++ b/tests/local_testing/test_caching_handler.py @@ -19,9 +19,14 @@ import pytest import litellm from litellm import aembedding, completion, embedding, aresponses, responses from litellm.caching.caching import Cache +from litellm.responses.streaming_iterator import CachedResponsesAPIStreamingIterator from unittest.mock import AsyncMock, patch, MagicMock -from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse +from litellm.caching.caching_handler import ( + LLMCachingHandler, + CachingHandlerResponse, + _should_defer_streaming_cache_hit_callbacks, +) from litellm.caching.caching import LiteLLMCacheType from litellm.types.utils import CallTypes from litellm.types.rerank import RerankResponse @@ -627,6 +632,55 @@ async def test_async_responses_api_caching(): assert cached_response.cached_result._hidden_params["cache_hit"] == True +@pytest.mark.asyncio +async def test_async_get_cache_updates_request_kwargs_for_streaming_responses(): + """ + Ensure streamed responses retain the normalized lookup kwargs so a later + cache write can reuse the exact cache key from the read path. + """ + setup_cache() + + caching_handler = LLMCachingHandler( + original_function=aresponses, + request_kwargs={"stale": True}, + start_time=datetime.now(), + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-4o", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + kwargs = { + "model": "gpt-4o", + "input": "hello", + "stream": True, + "caching": True, + } + + await caching_handler._async_get_cache( + model="gpt-4o", + original_function=aresponses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aresponses.value, + kwargs=kwargs, + ) + + assert "stale" not in caching_handler.request_kwargs + assert caching_handler.request_kwargs["model"] == "gpt-4o" + assert caching_handler.request_kwargs["input"] == "hello" + assert caching_handler.request_kwargs["stream"] is True + assert caching_handler.request_kwargs["cache_key"] == litellm.cache.get_cache_key( + **caching_handler.request_kwargs + ) + + def test_sync_responses_api_caching(): """ Test that synchronous responses API calls are properly cached and retrieved. @@ -769,6 +823,349 @@ def test_convert_cached_responses_api_result_to_model_response(): assert len(result.output) == 1 +def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): + litellm.set_verbose = True + setup_cache() + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + + original_model = "gpt-4o" + responses_api_response = ResponsesAPIResponse( + id="resp_stream_sync_hit", + created_at=int(time.time()), + status="completed", + model=original_model, + object="response", + output=[ + { + "type": "message", + "id": "msg_stream_sync_hit", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Sync streamed cache hit response.", + "annotations": [], + } + ], + } + ], + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.responses.value, + model=original_model, + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + + kwargs = { + "model": original_model, + "input": "Tell me a cached story", + "stream": True, + "caching": True, + } + + caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs) + time.sleep(0.2) + + cached_response = caching_handler._sync_get_cache( + model=original_model, + original_function=responses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.responses.value, + kwargs=kwargs, + ) + + assert cached_response.cached_result is not None + assert isinstance( + cached_response.cached_result, CachedResponsesAPIStreamingIterator + ) + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + + +def test_sync_get_cache_still_eagerly_logs_streaming_completion_hits(): + litellm.set_verbose = True + setup_cache() + caching_handler = LLMCachingHandler( + original_function=completion, request_kwargs={}, start_time=datetime.now() + ) + + original_model = "gpt-4o" + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.completion.value, + model=original_model, + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + + kwargs = { + "model": original_model, + "messages": [{"role": "user", "content": "Tell me a cached joke"}], + "stream": True, + "caching": True, + } + + caching_handler.sync_set_cache(result=chat_completion_response, kwargs=kwargs) + time.sleep(0.2) + + cached_response = caching_handler._sync_get_cache( + model=original_model, + original_function=completion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.completion.value, + kwargs=kwargs, + ) + + assert cached_response.cached_result is not None + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() + + +def test_should_defer_streaming_cache_hit_callbacks_only_for_responses_streams(): + assert ( + _should_defer_streaming_cache_hit_callbacks( + call_type=CallTypes.responses.value, + kwargs={"stream": True}, + ) + is True + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + call_type=CallTypes.aresponses.value, + kwargs={"stream": True}, + ) + is True + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + call_type=CallTypes.completion.value, + kwargs={"stream": True}, + ) + is False + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + call_type=CallTypes.responses.value, + kwargs={"stream": False}, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_async_get_cache_still_eagerly_logs_streaming_completion_hits(): + litellm.set_verbose = True + setup_cache() + caching_handler = LLMCachingHandler( + original_function=completion, request_kwargs={}, start_time=datetime.now() + ) + + original_model = "gpt-4o" + kwargs = { + "model": original_model, + "messages": [{"role": "user", "content": "Tell me a cached joke"}], + "stream": True, + "caching": True, + } + + await caching_handler.async_set_cache( + result=chat_completion_response, + original_function=litellm.acompletion, + kwargs=kwargs, + ) + await asyncio.sleep(0.2) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.acompletion.value, + model=original_model, + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + caching_handler._async_log_cache_hit_on_callbacks = MagicMock() + + cached_response = await caching_handler._async_get_cache( + model=original_model, + original_function=litellm.acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + ) + + assert cached_response is not None + assert cached_response.cached_result is not None + caching_handler._async_log_cache_hit_on_callbacks.assert_called_once() + + +def test_convert_cached_streaming_responses_result_to_iterator(): + """ + Test that cached streaming Responses results are replayed through a synthetic + streaming iterator instead of being returned as a full response object. + """ + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.responses.value, + model="gpt-4o", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + cached_result = { + "id": "resp_stream_cache_test", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_stream_cache_test", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Streaming cache replay test.", + "annotations": [], + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "test", "stream": True}, + logging_obj=logging_obj, + model="gpt-4o", + args=(), + ) + + assert isinstance(result, CachedResponsesAPIStreamingIterator) + assert result.completed_response is not None + assert result.completed_response.response.id == cached_result["id"] + + streamed_events = list(result) + assert streamed_events[0].type == "response.created" + assert streamed_events[1].type == "response.in_progress" + assert streamed_events[2].type == "response.output_item.added" + assert streamed_events[3].type == "response.content_part.added" + assert streamed_events[-4].type == "response.output_text.done" + assert streamed_events[-3].type == "response.content_part.done" + assert streamed_events[-2].type == "response.output_item.done" + assert streamed_events[-1].type == "response.completed" + assert streamed_events[-1].response.id == cached_result["id"] + assert streamed_events[-1].response.output[0].content[0].text == ( + "Streaming cache replay test." + ) + + +def test_convert_cached_streaming_reasoning_result_to_iterator(): + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.responses.value, + model="gpt-4o", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + cached_result = { + "id": "resp_stream_reasoning_cache_test", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "reasoning", + "id": "rs_stream_cache_test", + "summary": [ + { + "type": "summary_text", + "text": "Cached reasoning summary.", + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "test", "stream": True}, + logging_obj=logging_obj, + model="gpt-4o", + args=(), + ) + + assert isinstance(result, CachedResponsesAPIStreamingIterator) + + streamed_events = list(result) + streamed_event_types = [ + event.type.value if hasattr(event.type, "value") else str(event.type) + for event in streamed_events + ] + + assert streamed_event_types[:3] == [ + "response.created", + "response.in_progress", + "response.output_item.added", + ] + assert streamed_event_types[-4:] == [ + "response.reasoning_summary_text.done", + "response.reasoning_summary_part.done", + "response.output_item.done", + "response.completed", + ] + assert streamed_event_types.count("response.reasoning_summary_text.delta") >= 1 + + delta_events = [ + event + for event in streamed_events + if (event.type.value if hasattr(event.type, "value") else str(event.type)) + == "response.reasoning_summary_text.delta" + ] + text_done_event = streamed_events[-4] + part_done_event = streamed_events[-3] + output_item_done_event = streamed_events[-2] + + assert all(delta_event.summary_index == 0 for delta_event in delta_events) + assert text_done_event.text == "Cached reasoning summary." + assert text_done_event.summary_index == 0 + assert part_done_event.part.type == "summary_text" + assert part_done_event.part.text == "Cached reasoning summary." + assert output_item_done_event.item.type == "reasoning" + assert output_item_done_event.item.summary[0]["text"] == "Cached reasoning summary." + + @pytest.mark.asyncio async def test_responses_api_cache_with_different_inputs(): """ diff --git a/tests/local_testing/test_responses_stream_cache_keys.py b/tests/local_testing/test_responses_stream_cache_keys.py new file mode 100644 index 00000000000..5637028f550 --- /dev/null +++ b/tests/local_testing/test_responses_stream_cache_keys.py @@ -0,0 +1,141 @@ +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import aresponses +from litellm._uuid import uuid +from litellm.caching.caching_handler import LLMCachingHandler +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.llms import openai as openai_types +from litellm.types.utils import CallTypes + + +@pytest.mark.asyncio +async def test_async_get_cache_reuses_preset_cache_key_for_responses(): + caching_handler = LLMCachingHandler( + original_function=aresponses, + request_kwargs={}, + start_time=datetime.now(), + ) + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-4.1-mini", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + original_cache = litellm.cache + mock_cache = MagicMock() + mock_cache.supported_call_types = [CallTypes.aresponses.value] + mock_cache._supports_async.return_value = True + mock_cache.get_cache_key.return_value = "responses-stream-cache-key" + mock_cache.async_get_cache = AsyncMock(return_value=None) + litellm.cache = mock_cache + + kwargs = { + "model": "gpt-4.1-mini", + "input": "hello", + "stream": True, + "litellm_params": {}, + } + await caching_handler._async_get_cache( + model="gpt-4.1-mini", + original_function=aresponses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aresponses.value, + kwargs=kwargs, + ) + + assert caching_handler.preset_cache_key == "responses-stream-cache-key" + mock_cache.async_get_cache.assert_awaited_once() + assert ( + mock_cache.async_get_cache.call_args.kwargs["cache_key"] + == "responses-stream-cache-key" + ) + + litellm.cache = original_cache + + +@pytest.mark.asyncio +async def test_async_get_cache_falls_back_to_sync_cache_for_responses(): + caching_handler = LLMCachingHandler( + original_function=aresponses, + request_kwargs={}, + start_time=datetime.now(), + ) + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-4.1-mini", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + original_cache = litellm.cache + mock_cache = MagicMock() + mock_cache.supported_call_types = [CallTypes.aresponses.value] + mock_cache._supports_async.return_value = False + mock_cache.get_cache_key.return_value = "responses-stream-cache-key" + mock_cache.get_cache.return_value = None + litellm.cache = mock_cache + + kwargs = { + "model": "gpt-4.1-mini", + "input": "hello", + "stream": True, + "litellm_params": {}, + } + await caching_handler._async_get_cache( + model="gpt-4.1-mini", + original_function=aresponses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aresponses.value, + kwargs=kwargs, + ) + + assert caching_handler.preset_cache_key == "responses-stream-cache-key" + mock_cache.get_cache.assert_called_once() + assert mock_cache.get_cache.call_args.kwargs["cache_key"] == ( + "responses-stream-cache-key" + ) + + litellm.cache = original_cache + + +def test_reasoning_summary_events_default_summary_index(): + delta_event = openai_types.ReasoningSummaryTextDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id="rs_1", + output_index=0, + delta="abc", + ) + text_done_event = openai_types.ReasoningSummaryTextDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, + item_id="rs_1", + output_index=0, + sequence_number=1, + text="abc", + ) + part_done_event = openai_types.ReasoningSummaryPartDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, + item_id="rs_1", + output_index=0, + sequence_number=2, + part=openai_types.BaseLiteLLMOpenAIResponseObject( + type="summary_text", + text="abc", + ), + ) + + assert delta_event.summary_index == 0 + assert text_done_event.summary_index == 0 + assert part_done_event.summary_index == 0 From 262782f5b225a0d4bc5308af8818b947c6096894 Mon Sep 17 00:00:00 2001 From: Gowtham Raj Date: Tue, 7 Apr 2026 10:01:29 +0530 Subject: [PATCH 2/9] fix: resolve access group names in _filter_models_by_team_id (#25224) The /v2/model/info endpoint (used by the UI's Models + Endpoints page) was not resolving access group names when filtering models by team. When a team has models: ["Group-A"] where "Group-A" is an access group, _filter_models_by_team_id() passed it as a literal model name to get_model_list(), which found no deployments with that name. This caused the UI to show all models instead of only team-accessible ones. The request-time auth path (model_in_access_group in auth_checks.py) correctly resolves access groups via get_model_access_groups(). This fix applies the same resolution in _filter_models_by_team_id() for both the in-memory router lookup and the database fallback query. Tests added: - test_filter_resolves_access_group_names - test_filter_resolves_mix_of_access_groups_and_literal_names - test_filter_excludes_models_from_other_access_group - test_filter_db_fallback_receives_resolved_model_names --- litellm/proxy/proxy_server.py | 20 +- ...test_filter_models_by_team_access_group.py | 236 ++++++++++++++++++ 2 files changed, 254 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/proxy/test_filter_models_by_team_access_group.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6cba6a3e96b..0668880f908 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10324,6 +10324,7 @@ async def _filter_models_by_team_id( # Get models accessible to this team (similar to _add_team_models_to_all_models) team_accessible_model_ids: Set[str] = set() + access_groups = llm_router.get_model_access_groups() if llm_router else {} if ( not team_object.models # empty list = all model access @@ -10347,8 +10348,16 @@ async def _filter_models_by_team_id( if can_add_model: team_accessible_model_ids.add(model_id) else: - # Team has access to specific models + # Team has access to specific models — resolve access group names first + resolved_model_names: Set[str] = set() for model_name in team_object.models: + if model_name in access_groups: + # This is an access group name — expand to member model names + resolved_model_names.update(access_groups[model_name]) + else: + resolved_model_names.add(model_name) + + for model_name in resolved_model_names: _models = ( llm_router.get_model_list(model_name=model_name, team_id=team_id) if llm_router @@ -10368,8 +10377,15 @@ async def _filter_models_by_team_id( and SpecialModelNames.all_proxy_models.value not in team_object.models ): # Team has specific models - check database for those model names + # Resolve access group names to actual model names for the DB query + _resolved_names: list[str] = [] + for _m in team_object.models: + if _m in access_groups: + _resolved_names.extend(access_groups[_m]) + else: + _resolved_names.append(_m) db_models = await prisma_client.db.litellm_proxymodeltable.find_many( - where={"model_name": {"in": team_object.models}} + where={"model_name": {"in": _resolved_names}} ) for db_model in db_models: model_id = db_model.model_id diff --git a/tests/test_litellm/proxy/test_filter_models_by_team_access_group.py b/tests/test_litellm/proxy/test_filter_models_by_team_access_group.py new file mode 100644 index 00000000000..2d8a9f30c1b --- /dev/null +++ b/tests/test_litellm/proxy/test_filter_models_by_team_access_group.py @@ -0,0 +1,236 @@ +""" +Tests for _filter_models_by_team_id resolving access group names. + +Verifies that when a team's `models` field contains an access group name +(e.g., "Group-A"), the filter resolves it to the member model names before +looking up deployments — matching the behavior of the auth path in +auth_checks.py:model_in_access_group(). +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.proxy_server import _filter_models_by_team_id + + +def _make_model(model_name: str, model_id: str, access_groups: list[str] = None): + """Helper to build a model dict matching the router's format.""" + return { + "model_name": model_name, + "litellm_params": {"model": model_name}, + "model_info": { + "id": model_id, + "access_groups": access_groups or [], + }, + } + + +def _make_team(models: list[str], team_id: str = "team_alpha"): + """Helper to build a mock team DB object.""" + mock = MagicMock() + mock.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Team Alpha", + "models": models, + "max_budget": None, + "spend": 0.0, + "blocked": False, + "members_with_roles": [], + "metadata": {}, + } + return mock + + +@pytest.mark.asyncio +async def test_filter_resolves_access_group_names(): + """ + When team.models contains an access group name, _filter_models_by_team_id + should resolve it to the member models and return only those deployments. + """ + # Models on the proxy + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + gpt5 = _make_model("gpt-5", "id-2", ["Group-A"]) + claude = _make_model("claude-3", "id-3", ["Group-B"]) + + all_models = [gpt4o, gpt5, claude] + + # Router mock + mock_router = MagicMock() + # get_model_access_groups returns {group_name: [model_names]} + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o", "gpt-5"], + "Group-B": ["claude-3"], + } + + # get_model_list returns deployments matching a model_name + def fake_get_model_list(model_name=None, team_id=None): + return [m for m in all_models if m["model_name"] == model_name] + + mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list) + + # Team has models: ["Group-A"] — an access group name, not a literal model + team_db = _make_team(models=["Group-A"]) + + # Prisma mock + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + result_ids = {m["model_info"]["id"] for m in result} + # Should include gpt-4o and gpt-5 (Group-A), but NOT claude-3 (Group-B) + assert result_ids == { + "id-1", + "id-2", + }, f"Expected Group-A models only, got {result_ids}" + + # Verify DB fallback query received resolved model names, not access group name + call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1] + assert set(call_kwargs["where"]["model_name"]["in"]) == { + "gpt-4o", + "gpt-5", + }, "find_many should receive resolved model names, not the access group name" + + +@pytest.mark.asyncio +async def test_filter_resolves_mix_of_access_groups_and_literal_names(): + """ + When team.models contains both an access group name and a literal model name, + both should be resolved correctly. + """ + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + gpt5 = _make_model("gpt-5", "id-2", ["Group-A"]) + claude = _make_model("claude-3", "id-3", ["Group-B"]) + mistral = _make_model("mistral-large", "id-4", []) # no access group + + all_models = [gpt4o, gpt5, claude, mistral] + + mock_router = MagicMock() + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o", "gpt-5"], + "Group-B": ["claude-3"], + } + + def fake_get_model_list(model_name=None, team_id=None): + return [m for m in all_models if m["model_name"] == model_name] + + mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list) + + # Team has access to Group-A (access group) + mistral-large (literal name) + team_db = _make_team(models=["Group-A", "mistral-large"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + result_ids = {m["model_info"]["id"] for m in result} + # Group-A models + mistral-large, but NOT claude-3 + assert result_ids == { + "id-1", + "id-2", + "id-4", + }, f"Expected Group-A + mistral-large, got {result_ids}" + + +@pytest.mark.asyncio +async def test_filter_excludes_models_from_other_access_group(): + """ + Models belonging only to a different access group must not appear in results. + """ + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + claude = _make_model("claude-3", "id-3", ["Group-B"]) + llama = _make_model("llama-4", "id-4", ["Group-B"]) + + all_models = [gpt4o, claude, llama] + + mock_router = MagicMock() + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o"], + "Group-B": ["claude-3", "llama-4"], + } + + def fake_get_model_list(model_name=None, team_id=None): + return [m for m in all_models if m["model_name"] == model_name] + + mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list) + + team_db = _make_team(models=["Group-A"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + result_names = {m["model_name"] for m in result} + assert "claude-3" not in result_names, "Group-B model should not be accessible" + assert "llama-4" not in result_names, "Group-B model should not be accessible" + assert "gpt-4o" in result_names, "Group-A model should be accessible" + + +@pytest.mark.asyncio +async def test_filter_db_fallback_receives_resolved_model_names(): + """ + When get_model_list returns no results (forcing the DB fallback path), + the DB query should receive resolved model names, not the raw access group name. + """ + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + all_models = [gpt4o] + + mock_router = MagicMock() + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o", "gpt-5"], + } + # get_model_list returns nothing — forces reliance on the DB fallback + mock_router.get_model_list = MagicMock(return_value=[]) + + team_db = _make_team(models=["Group-A"]) + + # DB returns a model that the router didn't find + mock_db_model = MagicMock() + mock_db_model.model_id = "id-db-1" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[mock_db_model] + ) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + # Verify DB query received resolved names, not "Group-A" + call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1] + queried_names = set(call_kwargs["where"]["model_name"]["in"]) + assert queried_names == { + "gpt-4o", + "gpt-5", + }, f"DB query should receive resolved model names, got {queried_names}" + assert "Group-A" not in queried_names, "Raw access group name should not be in DB query" From f49c91ea922f6ce7ee849a0571e9134ff167be40 Mon Sep 17 00:00:00 2001 From: shubham-arora-clear Date: Tue, 7 Apr 2026 10:03:41 +0530 Subject: [PATCH 3/9] fix(bedrock): handle document content blocks in Converse API message conversion (#24644) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(bedrock): handle document content blocks in Converse API message conversion Document content blocks (used for PDF support) were silently dropped during message conversion for Bedrock's Converse API. The content block processing loop only handled text, image_url, and file types — document blocks were skipped without warning, causing the model to respond as if no document was provided. Adds document block handling in three locations: - Sync user message processing (_bedrock_converse_messages_pt) - Async user message processing (_bedrock_converse_messages_pt_async) - Tool result conversion (_convert_to_bedrock_tool_call_result) Fixes #24641 Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use _validate_format for proper MIME type to Bedrock format mapping Address Greptile review: naive media_type.split("/")[1] produced invalid Bedrock format names for complex MIME types (e.g. OOXML → docx, text/plain → txt, text/markdown → md). Now reuses BedrockImageProcessor._validate_format which handles all MIME types correctly via mimetypes + fallback. Also fixes test assertions to expect correct Bedrock format values and adds text/plain and text/markdown test cases. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: reject non-base64 document sources with a clear error URL-type document sources (e.g. {"type": "url", "url": "..."}) would crash with an opaque KeyError on missing 'media_type'. Guard at the top of _process_document_message and raise a clear ValueError since Bedrock Converse only supports base64-encoded document sources. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- .../prompt_templates/factory.py | 59 ++++- ...llm_core_utils_prompt_templates_factory.py | 238 +++++++----------- 2 files changed, 149 insertions(+), 148 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3a83162fb20..4ef08b0a01f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4042,9 +4042,13 @@ def _convert_to_bedrock_tool_call_result( tool_result_content_blocks.append( BedrockToolResultContentBlock(image=_block["image"]) ) - elif "document" in _block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(document=_block["document"]) + elif content["type"] == "document": + _doc_block = BedrockConverseMessagesProcessor._process_document_message( + content + ) + if "document" in _doc_block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(document=_doc_block["document"]) ) else: verbose_logger.warning( @@ -4090,7 +4094,6 @@ def _convert_to_bedrock_tool_call_result( "%s for file tool-result block %s; dropping.", list(_file_block.keys()), content, - ) message.get("name", "") id = str(message.get("tool_call_id", str(uuid.uuid4()))) @@ -4582,6 +4585,11 @@ class BedrockConverseMessagesProcessor: message=cast(ChatCompletionFileObject, element) ) _parts.append(_part) + elif element["type"] == "document": + _part = BedrockConverseMessagesProcessor._process_document_message( + element + ) + _parts.append(_part) _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( message_block=cast( @@ -4864,6 +4872,44 @@ class BedrockConverseMessagesProcessor: image_url=cast(str, file_id or file_data), format=format ) + @staticmethod + def _process_document_message(element: dict) -> BedrockContentBlock: + """Convert a document content block to a Bedrock DocumentBlock. + + Handles the Anthropic-style document format: + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "..."}} + """ + source = element["source"] + source_type = source.get("type") + if source_type != "base64": + raise ValueError( + f"Bedrock Converse only supports base64-encoded document sources, got '{source_type}'. " + "Please convert the document to base64 before sending to Bedrock." + ) + media_type: str = source["media_type"] + data: str = source["data"] + doc_format = BedrockImageProcessor._validate_format( + mime_type=media_type, image_format=media_type.split("/")[1] + ) + + # Deterministic name using the same hashing pattern as _create_bedrock_block + HASH_SAMPLE_BYTES = 64 * 1024 + normalized = "".join(data.split()).encode("utf-8") + sample = normalized[:HASH_SAMPLE_BYTES] + hasher = hashlib.sha256() + hasher.update(sample) + hasher.update(str(len(normalized)).encode("utf-8")) + content_hash = hasher.hexdigest()[:16] + document_name = f"Document_{content_hash}_{doc_format}" + + return BedrockContentBlock( + document=BedrockDocumentBlock( + source=BedrockSourceBlock(bytes=data), + format=doc_format, + name=document_name, + ) + ) + @staticmethod def add_thinking_blocks_to_assistant_content( thinking_blocks: List[BedrockContentBlock], @@ -4961,6 +5007,11 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 ) ) _parts.append(_part) + elif element["type"] == "document": + _part = BedrockConverseMessagesProcessor._process_document_message( + element + ) + _parts.append(_part) _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( message_block=cast( diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index d424cd8599f..33bdf169ccd 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -9,8 +9,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, BedrockConverseMessagesProcessor, BedrockImageProcessor, - anthropic_messages_pt, + _bedrock_converse_messages_pt, _convert_to_bedrock_tool_call_invoke, + _convert_to_bedrock_tool_call_result, + anthropic_messages_pt, convert_to_gemini_tool_call_result, ollama_pt, sanitize_messages_for_tool_calling, @@ -2505,157 +2507,105 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): ], } - result = convert_to_anthropic_tool_result(message) + result = _convert_to_bedrock_tool_call_result(message) - assert result["type"] == "tool_result" - assert result["tool_use_id"] == "toolu_pdf_1" - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "document" - assert block["source"]["type"] == "base64" - assert block["source"]["media_type"] == "application/pdf" - assert block["source"]["data"] == pdf_b64 + tool_result = result["toolResult"] + assert len(tool_result["content"]) == 1 + assert "document" in tool_result["content"][0] + assert tool_result["content"][0]["document"]["format"] == "pdf" + assert tool_result["content"][0]["document"]["source"]["bytes"] == "dGVzdA==" -def test_convert_to_anthropic_tool_result_image_url_pdf_data_uri_becomes_document(): - """ - Regression: a PDF sent as an `image_url` data URI on the tool-result path - must translate to an Anthropic document block (not an image block — Anthropic - rejects image blocks whose media_type is a non-image like application/pdf). - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) +def test_bedrock_converse_messages_pt_document_various_formats(): + """Test that various document media types produce the correct format value.""" + test_cases = [ + ("application/pdf", "pdf"), + ("text/csv", "csv"), + ("text/html", "html"), + ("text/plain", "txt"), + ("text/markdown", "md"), + ( + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "docx", + ), + ] - pdf_b64 = "JVBERi0xLjQKJeLjz9MK" - message = { - "tool_call_id": "toolu_pdf_img_1", - "role": "tool", - "name": "fetch_document", - "content": [ + for media_type, expected_format in test_cases: + messages = [ { - "type": "image_url", - "image_url": { - "url": f"data:application/pdf;base64,{pdf_b64}", + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": media_type, + "data": "dGVzdA==", + }, + }, + ], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + doc_block = result[0]["content"][0] + assert doc_block["document"]["format"] == expected_format, ( + f"Expected format '{expected_format}' for media_type '{media_type}', " + f"got '{doc_block['document']['format']}'" + ) + + +def test_bedrock_converse_messages_pt_document_deterministic_name(): + """Test that the same document data always produces the same name.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, }, - }, - ], - } + ], + } + ] - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "document" - assert block["source"]["media_type"] == "application/pdf" - assert block["source"]["data"] == pdf_b64 - - -def test_convert_to_anthropic_tool_result_image_url_unsupported_mime_stays_image_path(): - """ - An `image_url` data URI whose mime is neither application/pdf nor text/plain - (e.g. application/json) must NOT be routed through the document path. Anthropic - only accepts application/pdf and text/plain as base64 document media_types — - anything else would produce a document block the API rejects. The old - (pre-fix) behavior was to wrap such data as an image block, which also - fails but stays on the image code path; preserve that failure mode rather - than switching to a document path that is equally broken. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, + result1 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + result2 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" ) - message = { - "tool_call_id": "toolu_json_1", - "role": "tool", - "name": "fetch_json", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:application/json;base64,eyJrIjoidiJ9", + name1 = result1[0]["content"][0]["document"]["name"] + name2 = result2[0]["content"][0]["document"]["name"] + assert name1 == name2 + + +def test_bedrock_converse_messages_pt_document_rejects_url_source(): + """Test that a URL-type document source raises a clear error instead of KeyError.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "url", + "url": "https://example.com/doc.pdf", + }, }, - }, - ], - } + ], + } + ] - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "image", ( - f"unsupported mime {block.get('source', {}).get('media_type')!r} " - f"should not be routed to document path; got {block}" - ) - - -def test_convert_to_anthropic_tool_result_image_url_text_plain_data_uri_becomes_document(): - """ - text/plain is one of the two mimes Anthropic accepts as a base64 document - media_type. Confirm it routes through the document path so tightening the - gate to {application/pdf, text/plain} (not "application/*") covers both. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) - - txt_b64 = "aGVsbG8=" # "hello" - message = { - "tool_call_id": "toolu_txt_1", - "role": "tool", - "name": "fetch_text", - "content": [ - { - "type": "image_url", - "image_url": { - "url": f"data:text/plain;base64,{txt_b64}", - }, - }, - ], - } - - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "document" - assert block["source"]["media_type"] == "text/plain" - assert block["source"]["data"] == txt_b64 - - -def test_convert_to_anthropic_tool_result_image_url_png_still_becomes_image(): - """ - Regression: image_url with a real image mime type must continue to translate - to an Anthropic image block. Locks in existing behavior after the - data-URI-mime-type branching for PDFs. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) - - png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg==" - message = { - "tool_call_id": "toolu_png_1", - "role": "tool", - "name": "fetch_image", - "content": [ - { - "type": "image_url", - "image_url": { - "url": f"data:image/png;base64,{png_b64}", - }, - }, - ], - } - - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "image" - assert block["source"]["media_type"] == "image/png" + with pytest.raises(ValueError, match="only supports base64-encoded"): + _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) From ec38f2b17beef24551339f04b282e1605da0515d Mon Sep 17 00:00:00 2001 From: Baqiao Date: Mon, 6 Apr 2026 21:35:28 -0700 Subject: [PATCH 4/9] feat(xai): add parallel_tool_calls to supported params (#25106) --- litellm/llms/xai/chat/transformation.py | 1 + .../llms/xai/test_xai_chat_transformation.py | 39 +++++++++++++++++++ 2 files changed, 40 insertions(+) create mode 100644 tests/test_litellm/llms/xai/test_xai_chat_transformation.py diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index bfa55105a6c..64b4a545acb 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -43,6 +43,7 @@ class XAIChatConfig(OpenAIGPTConfig): "logprobs", "max_tokens", "n", + "parallel_tool_calls", "presence_penalty", "response_format", "seed", diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py new file mode 100644 index 00000000000..5a236de900e --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -0,0 +1,39 @@ +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.xai.chat.transformation import XAIChatConfig + + +class TestXAIParallelToolCalls: + """Test suite for XAI parallel tool calls functionality.""" + + def test_get_supported_openai_params_includes_parallel_tool_calls(self): + """Test that parallel_tool_calls is in supported parameters.""" + config = XAIChatConfig() + supported_params = config.get_supported_openai_params( + "xai/grok-4.20" + ) + assert "parallel_tool_calls" in supported_params + + def test_transform_request_preserves_parallel_tool_calls(self): + """Test that transform_request preserves parallel_tool_calls parameter.""" + config = XAIChatConfig() + + messages = [{"role": "user", "content": "What's the weather like?"}] + optional_params = {"parallel_tool_calls": True} + + result = config.transform_request( + model="xai/grok-4.20", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result.get("parallel_tool_calls") is True + assert len(result["messages"]) == 1 + assert result["messages"][0]["role"] == "user" From d2346b525912a5c006c70bbf32b2cceeaba0c8bb Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 1 May 2026 12:09:19 +0530 Subject: [PATCH 5/9] Fix prompt factory --- .../litellm_core_utils/prompt_templates/factory.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 4ef08b0a01f..ba840bc3d89 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4042,13 +4042,9 @@ def _convert_to_bedrock_tool_call_result( tool_result_content_blocks.append( BedrockToolResultContentBlock(image=_block["image"]) ) - elif content["type"] == "document": - _doc_block = BedrockConverseMessagesProcessor._process_document_message( - content - ) - if "document" in _doc_block: - tool_result_content_blocks.append( - BedrockToolResultContentBlock(document=_doc_block["document"]) + elif "document" in _block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(document=_block["document"]) ) else: verbose_logger.warning( @@ -4094,6 +4090,7 @@ def _convert_to_bedrock_tool_call_result( "%s for file tool-result block %s; dropping.", list(_file_block.keys()), content, + ) message.get("name", "") id = str(message.get("tool_call_id", str(uuid.uuid4()))) From 94162916a81e37eaaa97dc21614d6b038e54f9a5 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 1 May 2026 12:14:53 +0530 Subject: [PATCH 6/9] Fix lint and ruff --- litellm/proxy/proxy_server.py | 186 +++++++++++++----------- litellm/responses/streaming_iterator.py | 1 + 2 files changed, 101 insertions(+), 86 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0668880f908..1c286b7f2e4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10285,6 +10285,101 @@ def _paginate_models_response( } +def _team_models_resolve_to_names( + team_models: List[str], access_groups: Dict[str, Any] +) -> List[str]: + """Expand team model entries (including access group names) to concrete model names.""" + resolved: List[str] = [] + for name in team_models: + if name in access_groups: + resolved.extend(access_groups[name]) + else: + resolved.append(name) + return resolved + + +async def _load_team_object_for_model_filter( + team_id: str, prisma_client: PrismaClient +) -> Optional[LiteLLM_TeamTable]: + """Load team row from DB; returns None if missing or on error.""" + try: + team_db_object = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + if team_db_object is None: + verbose_proxy_logger.warning(f"Team {team_id} not found in database") + return None + return LiteLLM_TeamTable(**team_db_object.model_dump()) + except Exception as e: + verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}") + return None + + +async def _gather_team_accessible_model_ids( + team_object: LiteLLM_TeamTable, + team_id: str, + prisma_client: PrismaClient, + llm_router: Router, +) -> Set[str]: + """Collect model IDs the team can use from router config and DB.""" + team_accessible_model_ids: Set[str] = set() + access_groups = llm_router.get_model_access_groups() if llm_router else {} + + if ( + not team_object.models + or SpecialModelNames.all_proxy_models.value in team_object.models + ): + model_list = llm_router.get_model_list() if llm_router else [] + if model_list is not None: + for model in model_list: + model_id = model.get("model_info", {}).get("id", None) + if model_id is None: + continue + team_model_id = model.get("model_info", {}).get("team_id", None) + if team_model_id is None or team_model_id == team_id: + team_accessible_model_ids.add(model_id) + else: + resolved_model_names: Set[str] = set() + for model_name in team_object.models: + if model_name in access_groups: + resolved_model_names.update(access_groups[model_name]) + else: + resolved_model_names.add(model_name) + + for model_name in resolved_model_names: + _models = ( + llm_router.get_model_list(model_name=model_name, team_id=team_id) + if llm_router + else [] + ) + if _models is not None: + for model in _models: + model_id = model.get("model_info", {}).get("id", None) + if model_id is not None: + team_accessible_model_ids.add(model_id) + + try: + if ( + team_object.models + and SpecialModelNames.all_proxy_models.value not in team_object.models + ): + _resolved_names = _team_models_resolve_to_names( + team_object.models, access_groups + ) + db_models = await prisma_client.db.litellm_proxymodeltable.find_many( + where={"model_name": {"in": _resolved_names}} + ) + for db_model in db_models: + if db_model.model_id: + team_accessible_model_ids.add(db_model.model_id) + except Exception as e: + verbose_proxy_logger.debug( + f"Error querying database models for team {team_id}: {str(e)}" + ) + + return team_accessible_model_ids + + async def _filter_models_by_team_id( all_models: List[Dict[str, Any]], team_id: str, @@ -10307,94 +10402,13 @@ async def _filter_models_by_team_id( Returns: Filtered list of models """ - # Get team from database - try: - team_db_object = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - if team_db_object is None: - verbose_proxy_logger.warning(f"Team {team_id} not found in database") - # If team doesn't exist, return empty list - return [] - - team_object = LiteLLM_TeamTable(**team_db_object.model_dump()) - except Exception as e: - verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}") + team_object = await _load_team_object_for_model_filter(team_id, prisma_client) + if team_object is None: return [] - # Get models accessible to this team (similar to _add_team_models_to_all_models) - team_accessible_model_ids: Set[str] = set() - access_groups = llm_router.get_model_access_groups() if llm_router else {} - - if ( - not team_object.models # empty list = all model access - or SpecialModelNames.all_proxy_models.value in team_object.models - ): - # Team has access to all models - model_list = llm_router.get_model_list() if llm_router else [] - if model_list is not None: - for model in model_list: - model_id = model.get("model_info", {}).get("id", None) - if model_id is None: - continue - # if team model id set, check if team id matches - team_model_id = model.get("model_info", {}).get("team_id", None) - can_add_model = False - if team_model_id is None: - can_add_model = True - elif team_model_id == team_id: - can_add_model = True - - if can_add_model: - team_accessible_model_ids.add(model_id) - else: - # Team has access to specific models — resolve access group names first - resolved_model_names: Set[str] = set() - for model_name in team_object.models: - if model_name in access_groups: - # This is an access group name — expand to member model names - resolved_model_names.update(access_groups[model_name]) - else: - resolved_model_names.add(model_name) - - for model_name in resolved_model_names: - _models = ( - llm_router.get_model_list(model_name=model_name, team_id=team_id) - if llm_router - else [] - ) - if _models is not None: - for model in _models: - model_id = model.get("model_info", {}).get("id", None) - if model_id is not None: - team_accessible_model_ids.add(model_id) - - # Also search database for models accessible to this team - # This complements the config search done above - try: - if ( - team_object.models - and SpecialModelNames.all_proxy_models.value not in team_object.models - ): - # Team has specific models - check database for those model names - # Resolve access group names to actual model names for the DB query - _resolved_names: list[str] = [] - for _m in team_object.models: - if _m in access_groups: - _resolved_names.extend(access_groups[_m]) - else: - _resolved_names.append(_m) - db_models = await prisma_client.db.litellm_proxymodeltable.find_many( - where={"model_name": {"in": _resolved_names}} - ) - for db_model in db_models: - model_id = db_model.model_id - if model_id: - team_accessible_model_ids.add(model_id) - except Exception as e: - verbose_proxy_logger.debug( - f"Error querying database models for team {team_id}: {str(e)}" - ) + team_accessible_model_ids = await _gather_team_accessible_model_ids( + team_object, team_id, prisma_client, llm_router + ) # Filter models based on direct_access or access_via_team_ids # Models are already enriched with these fields before this function is called diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 1902a347215..da8da1b486f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import CallTypes from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook From 900ef454c83097452884000ea3185520a07b4794 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 1 May 2026 15:35:30 +0530 Subject: [PATCH 7/9] test: fix Bedrock PDF tool-result bytes assertion in factory test The test supplies a minimal PDF base64 payload but expected the wrong constant (base64 for "test"). Assert against the same pdf_b64 value and drop the unused import. Made-with: Cursor --- .../test_litellm_core_utils_prompt_templates_factory.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 33bdf169ccd..27a3ddb553d 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -2487,10 +2487,6 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): inside the tool_result content. Reuses anthropic_process_openai_file_message, which already handles this for user messages. """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) - pdf_b64 = "JVBERi0xLjQKJeLjz9MK" message = { "tool_call_id": "toolu_pdf_1", @@ -2513,7 +2509,7 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): assert len(tool_result["content"]) == 1 assert "document" in tool_result["content"][0] assert tool_result["content"][0]["document"]["format"] == "pdf" - assert tool_result["content"][0]["document"]["source"]["bytes"] == "dGVzdA==" + assert tool_result["content"][0]["document"]["source"]["bytes"] == pdf_b64 def test_bedrock_converse_messages_pt_document_various_formats(): From 15288f3ae7d57177365ca37c0459532060f0fe23 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 1 May 2026 15:41:17 +0530 Subject: [PATCH 8/9] test: include response key in response.completed chunk for ID hook test _base_process_chunk only encodes response IDs when parsed_chunk contains a top-level "response" key. Align test_process_chunk_completed_response_ updates_id_and_usage_cost with that contract and test_base_responses_api_streaming_iterator. Made-with: Cursor --- tests/llm_responses_api_testing/test_responses_hooks.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 8933ec8f3c5..3799a0b9121 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -384,7 +384,13 @@ def test_process_chunk_completed_response_updates_id_and_usage_cost(monkeypatch) ) try: - event = iterator._process_chunk(json.dumps({"type": "response.completed"})) + # Chunk must include a top-level "response" key so BaseResponsesAPIStreamingIterator + # runs _update_responses_api_response_id_with_model_id (see streaming_iterator.py). + event = iterator._process_chunk( + json.dumps( + {"type": "response.completed", "response": {"id": "resp_live"}} + ) + ) finally: litellm.include_cost_in_streaming_usage = original_include_cost From e0398cade78ad10cb431b1c9a429d43b7133223f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 1 May 2026 17:03:32 +0530 Subject: [PATCH 9/9] fix(caching): defer streaming cache-hit callbacks for all stream=True Success handlers already run when CustomStreamWrapper or CachedResponsesAPIStreamingIterator finishes replay. Logging at cache-hit time for acompletion/completion streaming duplicated spend and callbacks. Align tests with deferred behavior. Made-with: Cursor --- litellm/caching/caching_handler.py | 31 ++++++++------------- tests/local_testing/test_caching_handler.py | 24 +++++----------- 2 files changed, 19 insertions(+), 36 deletions(-) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 02207a207ed..3cf1d911d7f 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -85,19 +85,18 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() -_RESPONSES_STREAMING_CALLBACK_CALL_TYPES = { - CallTypes.aresponses.value, - CallTypes.responses.value, -} -def _should_defer_streaming_cache_hit_callbacks( - *, call_type: str, kwargs: Dict[str, Any] -) -> bool: - return ( - kwargs.get("stream", False) is True - and call_type in _RESPONSES_STREAMING_CALLBACK_CALL_TYPES - ) +def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool: + """ + When stream=True, do not run success callbacks at cache-hit time. + + Cached chat/text completion replay uses CustomStreamWrapper; cached Responses + replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success + handlers when the stream finishes; firing them here too would double-count + spend and callback records. + """ + return kwargs.get("stream", False) is True class LLMCachingHandler: @@ -220,10 +219,7 @@ class LLMCachingHandler: custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) - if not _should_defer_streaming_cache_hit_callbacks( - call_type=call_type, - kwargs=kwargs, - ): + if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs): # LOG SUCCESS self._async_log_cache_hit_on_callbacks( logging_obj=logging_obj, @@ -343,10 +339,7 @@ class LLMCachingHandler: is_async=False, ) - if not _should_defer_streaming_cache_hit_callbacks( - call_type=call_type, - kwargs=kwargs, - ): + if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs): logging_obj.handle_sync_success_callbacks_for_async_calls( result=cached_result, start_time=start_time, diff --git a/tests/local_testing/test_caching_handler.py b/tests/local_testing/test_caching_handler.py index 0ce9938e710..2b6712cbaa3 100644 --- a/tests/local_testing/test_caching_handler.py +++ b/tests/local_testing/test_caching_handler.py @@ -891,7 +891,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() -def test_sync_get_cache_still_eagerly_logs_streaming_completion_hits(): +def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): litellm.set_verbose = True setup_cache() caching_handler = LLMCachingHandler( @@ -930,42 +930,32 @@ def test_sync_get_cache_still_eagerly_logs_streaming_completion_hits(): ) assert cached_response.cached_result is not None - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() -def test_should_defer_streaming_cache_hit_callbacks_only_for_responses_streams(): +def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request(): assert ( _should_defer_streaming_cache_hit_callbacks( - call_type=CallTypes.responses.value, kwargs={"stream": True}, ) is True ) assert ( _should_defer_streaming_cache_hit_callbacks( - call_type=CallTypes.aresponses.value, - kwargs={"stream": True}, - ) - is True - ) - assert ( - _should_defer_streaming_cache_hit_callbacks( - call_type=CallTypes.completion.value, - kwargs={"stream": True}, + kwargs={"stream": False}, ) is False ) assert ( _should_defer_streaming_cache_hit_callbacks( - call_type=CallTypes.responses.value, - kwargs={"stream": False}, + kwargs={}, ) is False ) @pytest.mark.asyncio -async def test_async_get_cache_still_eagerly_logs_streaming_completion_hits(): +async def test_async_get_cache_defers_streaming_completion_hit_callbacks(): litellm.set_verbose = True setup_cache() caching_handler = LLMCachingHandler( @@ -1009,7 +999,7 @@ async def test_async_get_cache_still_eagerly_logs_streaming_completion_hits(): assert cached_response is not None assert cached_response.cached_result is not None - caching_handler._async_log_cache_hit_on_callbacks.assert_called_once() + caching_handler._async_log_cache_hit_on_callbacks.assert_not_called() def test_convert_cached_streaming_responses_result_to_iterator():