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..3cf1d911d7f 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -87,6 +87,18 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() +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: def __init__( self, @@ -99,6 +111,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 +219,7 @@ 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(kwargs=kwargs): # LOG SUCCESS self._async_log_cache_hit_on_callbacks( logging_obj=logging_obj, @@ -215,11 +228,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 +279,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 +294,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 +339,19 @@ 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(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 +705,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 +734,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 +861,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/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3a83162fb20..ba840bc3d89 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4582,6 +4582,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 +4869,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 +5004,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/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/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6cba6a3e96b..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,78 +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() - - 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 - for model_name in team_object.models: - _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 - db_models = await prisma_client.db.litellm_proxymodeltable.find_many( - where={"model_name": {"in": team_object.models}} - ) - 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 145ec3a641a..da8da1b486f 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,26 @@ 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.llms.openai import ResponsesAPIStreamEvents 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 +56,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 +68,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 +115,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 +136,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 +213,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 +238,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 +250,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 +267,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 +290,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 +373,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 +639,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 +679,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 +753,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 +777,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 +1396,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 +1530,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..3799a0b9121 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,768 @@ 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: + # 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 + + 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..2b6712cbaa3 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,339 @@ 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_defers_streaming_completion_hit_callbacks(): + 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_not_called() + + +def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request(): + assert ( + _should_defer_streaming_cache_hit_callbacks( + kwargs={"stream": True}, + ) + is True + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + kwargs={"stream": False}, + ) + is False + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + kwargs={}, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_async_get_cache_defers_streaming_completion_hit_callbacks(): + 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_not_called() + + +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 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..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 @@ -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, @@ -2485,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", @@ -2505,157 +2503,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"] == pdf_b64 -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" + ) 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" 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"