From 78997c2e35d10ae225cc49af612b2e415442dbff Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Thu, 31 Jul 2025 21:22:49 -0700 Subject: [PATCH] Anthropic - working mid-stream fallbacks (#13149) * fix(router.py): add acompletion_streaming_iterator inside router allows router to catch errors mid-stream for fallbacks Work for https://github.com/BerriAI/litellm/issues/6532 * fix(router.py): working mid-stream fallbacks * fix(router.py): more iterations * fix(router.py): working mid-stream fallbacks with fallbacks set on router * fix(router.py): pass prior content back in new request as assistant prefix message * fix(router.py): add a system prompt to help guide non-prefix supporting models to use the continued text correctly * fix(common_utils.py): support converting `prefix: true` for non-prefix supporting models * fix: reduce LOC in function * test(test_router.py): add unit tests for new function * test: add basic unit test * fix(router.py): ensure return type of fallback stream is compatible with CustomStreamWrapper prevent client code from breaking * fix: cleanup * test: update test * fix: fix linting error --- litellm/exceptions.py | 62 +++ .../prompt_templates/common_utils.py | 38 ++ .../litellm_core_utils/streaming_handler.py | 30 +- litellm/proxy/_new_secret_config.yaml | 2 +- litellm/router.py | 514 +++++++++++------- .../test_router_batch_completion.py | 1 + ...ore_utils_prompt_templates_common_utils.py | 18 + tests/test_litellm/test_router.py | 321 ++++++++++- 8 files changed, 781 insertions(+), 205 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 9f3411143a..153230518c 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -829,3 +829,65 @@ class BlockedPiiEntityError(Exception): self.guardrail_name = guardrail_name self.message = f"Blocked entity detected: {entity_type} by Guardrail: {guardrail_name}. This entity is not allowed to be used in this request." super().__init__(self.message) + + +class MidStreamFallbackError(ServiceUnavailableError): # type: ignore + def __init__( + self, + message: str, + model: str, + llm_provider: str, + original_exception: Optional[Exception] = None, + response: Optional[httpx.Response] = None, + litellm_debug_info: Optional[str] = None, + max_retries: Optional[int] = None, + num_retries: Optional[int] = None, + generated_content: str = "", + is_pre_first_chunk: bool = False, + ): + self.status_code = 503 # Service Unavailable + self.message = f"litellm.MidStreamFallbackError: {message}" + self.model = model + self.llm_provider = llm_provider + self.original_exception = original_exception + self.litellm_debug_info = litellm_debug_info + self.max_retries = max_retries + self.num_retries = num_retries + self.generated_content = generated_content + self.is_pre_first_chunk = is_pre_first_chunk + + # Create a response if one wasn't provided + if response is None: + self.response = httpx.Response( + status_code=self.status_code, + request=httpx.Request( + method="POST", + url=f"https://{llm_provider}.com/v1/", + ), + ) + else: + self.response = response + + # Call the parent constructor + super().__init__( + message=self.message, + llm_provider=llm_provider, + model=model, + response=self.response, + litellm_debug_info=self.litellm_debug_info, + max_retries=self.max_retries, + num_retries=self.num_retries, + ) + + def __str__(self): + _message = self.message + if self.num_retries: + _message += f" LiteLLM Retried: {self.num_retries} times" + if self.max_retries: + _message += f", LiteLLM Max Retries: {self.max_retries}" + if self.original_exception: + _message += f" Original exception: {type(self.original_exception).__name__}: {str(self.original_exception)}" + return _message + + def __repr__(self): + return self.__str__() diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 258601ff5a..827d28598e 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -822,3 +822,41 @@ def set_last_user_message( messages.reverse() messages.append({"role": "user", "content": content}) return messages + + +def convert_prefix_message_to_non_prefix_messages( + messages: List[AllMessageValues], +) -> List[AllMessageValues]: + """ + For models that don't support {prefix: true} in messages, we need to convert the prefix message to a non-prefix message. + + Use prompt: + + {"role": "assistant", "content": "value", "prefix": true} -> [ + { + "role": "system", + "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", + }, + { + "role": "assistant", + "content": message["content"], + }, + ] + + do this in place + """ + new_messages: List[AllMessageValues] = [] + for message in messages: + if message.get("prefix"): + new_messages.append( + { + "role": "system", + "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", + } + ) + new_messages.append( + {**{k: v for k, v in message.items() if k != "prefix"}} # type: ignore + ) + else: + new_messages.append(message) + return new_messages diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 17f5a8d1de..2e9e6770a1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -940,8 +940,8 @@ class CustomStreamWrapper: and not self.sent_last_thinking_block and model_response.choices[0].delta.content ): - model_response.choices[0].delta.content = ( - "" + (model_response.choices[0].delta.content or "") + model_response.choices[0].delta.content = "" + ( + model_response.choices[0].delta.content or "" ) self.sent_last_thinking_block = True @@ -1841,13 +1841,25 @@ class CustomStreamWrapper: self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore ) ## Map to OpenAI Exception - raise exception_type( - model=self.model, - custom_llm_provider=self.custom_llm_provider, - original_exception=e, - completion_kwargs={}, - extra_kwargs={}, - ) + try: + exception_type( + model=self.model, + custom_llm_provider=self.custom_llm_provider, + original_exception=e, + completion_kwargs={}, + extra_kwargs={}, + ) + except Exception as e: + from litellm.exceptions import MidStreamFallbackError + + raise MidStreamFallbackError( + message=str(e), + model=self.model, + llm_provider=self.custom_llm_provider or "anthropic", + original_exception=e, + generated_content=self.response_uptil_now, + is_pre_first_chunk=not self.sent_first_chunk, + ) @staticmethod def _strip_sse_data_from_chunk(chunk: Optional[str]) -> Optional[str]: diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 25843ec8fe..ed8fd02241 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -5,4 +5,4 @@ model_list: api_key: os.environ/OPENAI_API_KEY router_settings: - model_group_alias: {"gpt-4o": "gpt-4o-mini-openai"} \ No newline at end of file + model_group_alias: {"gpt-4o": "gpt-4o-mini-openai"} diff --git a/litellm/router.py b/litellm/router.py index 52d48fc71e..2673df1504 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1078,9 +1078,104 @@ class Router: ) raise e + async def _acompletion_streaming_iterator( + self, + model_response: CustomStreamWrapper, + messages: List[Dict[str, str]], + initial_kwargs: dict, + ) -> CustomStreamWrapper: + """ + Helper to iterate over a streaming response. + + Catches errors for fallbacks using the router's fallback system + """ + from litellm.exceptions import MidStreamFallbackError + + class FallbackStreamWrapper(CustomStreamWrapper): + def __init__(self, async_generator): + # Copy attributes from the original model_response + super().__init__( + completion_stream=async_generator, + model=model_response.model, + custom_llm_provider=model_response.custom_llm_provider, + logging_obj=model_response.logging_obj, + ) + self._async_generator = async_generator + + def __aiter__(self): + return self._async_generator + + async def __anext__(self): + return await self._async_generator.__anext__() + + async def stream_with_fallbacks(): + try: + async for item in model_response: + yield item + except MidStreamFallbackError as e: + try: + # Use the router's fallback system + model_group = cast(str, initial_kwargs.get("model")) + fallbacks: Optional[List] = initial_kwargs.get( + "fallbacks", self.fallbacks + ) + context_window_fallbacks: Optional[List] = initial_kwargs.get( + "context_window_fallbacks", self.context_window_fallbacks + ) + content_policy_fallbacks: Optional[List] = initial_kwargs.get( + "content_policy_fallbacks", self.content_policy_fallbacks + ) + initial_kwargs["original_function"] = self._acompletion + initial_kwargs["messages"] = messages + [ + { + "role": "system", + "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", + }, + { + "role": "assistant", + "content": e.generated_content, + "prefix": True, + }, + ] + self._update_kwargs_before_fallbacks( + model=model_group, kwargs=initial_kwargs + ) + fallback_response = ( + await self.async_function_with_fallbacks_common_utils( + e=e, + disable_fallbacks=False, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + ) + ) + + # If fallback returns a streaming response, iterate over it + if hasattr(fallback_response, "__aiter__"): + async for fallback_item in fallback_response: # type: ignore + yield fallback_item + else: + # If fallback returns a non-streaming response, yield None + yield None + + except Exception as fallback_error: + # If fallback also fails, log and re-raise original error + verbose_router_logger.error( + f"Fallback also failed: {fallback_error}" + ) + raise e # Re-raise the original error + + return FallbackStreamWrapper(stream_with_fallbacks()) + async def _acompletion( self, model: str, messages: List[Dict[str, str]], **kwargs - ) -> Union[ModelResponse, CustomStreamWrapper]: + ) -> Union[ + ModelResponse, + CustomStreamWrapper, + ]: """ - Get an available deployment - call it with a semaphore over the call @@ -1092,9 +1187,9 @@ class Router: {} ) # this is a temporary dict to debug timeout issues try: - verbose_router_logger.debug( - f"Inside _acompletion()- model: {model}; kwargs: {kwargs}" - ) + input_kwargs_for_streaming_fallback = kwargs.copy() + input_kwargs_for_streaming_fallback["model"] = model + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) start_time = time.time() deployment = await self.async_get_available_deployment( @@ -1134,15 +1229,15 @@ class Router: ) self.total_calls[model_name] += 1 - _response = litellm.acompletion( - **{ - **data, - "messages": messages, - "caching": self.cache_responses, - "client": model_client, - **kwargs, - } - ) + input_kwargs = { + **data, + "messages": messages, + "caching": self.cache_responses, + "client": model_client, + **kwargs, + } + + _response = litellm.acompletion(**input_kwargs) logging_obj: Optional[LiteLLMLogging] = kwargs.get( "litellm_logging_obj", None @@ -1199,6 +1294,13 @@ class Router: parent_otel_span=parent_otel_span, ) + if isinstance(response, CustomStreamWrapper): + return await self._acompletion_streaming_iterator( + model_response=response, + messages=messages, + initial_kwargs=input_kwargs_for_streaming_fallback, + ) + return response except litellm.Timeout as e: deployment_request_timeout_param = _timeout_debug_deployment_dict.get( @@ -1577,7 +1679,8 @@ class Router: Wrapper around self.acompletion that catches exceptions and returns them as a result """ try: - return await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) # type: ignore + result = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) # type: ignore + return result except asyncio.CancelledError: verbose_router_logger.debug( "Received 'task.cancel'. Cancelling call w/ model={}.".format(model) @@ -1625,6 +1728,7 @@ class Router: ) for completed_task in done: result = await check_response(completed_task) + if result is not None: # Return the first successful result result._hidden_params["fastest_response_batch_completion"] = True @@ -3517,8 +3621,199 @@ class Router: #### [END] ASSISTANTS API #### + async def async_function_with_fallbacks_common_utils( # noqa: PLR0915 + self, + e: Exception, + disable_fallbacks: Optional[bool], + fallbacks: Optional[List], + context_window_fallbacks: Optional[List], + content_policy_fallbacks: Optional[List], + model_group: Optional[str], + args: tuple, + kwargs: dict, + ): + """ + Common utilities for async_function_with_fallbacks + """ + verbose_router_logger.debug(f"Traceback{traceback.format_exc()}") + original_exception = e + fallback_model_group = None + original_model_group: Optional[str] = kwargs.get("model") # type: ignore + fallback_failure_exception_str = "" + + if disable_fallbacks is True or original_model_group is None: + raise e + + input_kwargs = { + "litellm_router": self, + "original_exception": original_exception, + **kwargs, + } + + if "max_fallbacks" not in input_kwargs: + input_kwargs["max_fallbacks"] = self.max_fallbacks + if "fallback_depth" not in input_kwargs: + input_kwargs["fallback_depth"] = 0 + + try: + verbose_router_logger.info("Trying to fallback b/w models") + + # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] + is_non_standard_fallback_format = _check_non_standard_fallback_format( + fallbacks=fallbacks + ) + + if is_non_standard_fallback_format: + input_kwargs.update( + { + "fallback_model_group": fallbacks, + "original_model_group": original_model_group, + } + ) + + response = await run_async_fallback( + *args, + **input_kwargs, + ) + + return response + + if isinstance(e, litellm.ContextWindowExceededError): + if context_window_fallbacks is not None: + context_window_fallback_model_group: Optional[List[str]] = ( + self._get_fallback_model_group_from_fallbacks( + fallbacks=context_window_fallbacks, + model_group=model_group, + ) + ) + if context_window_fallback_model_group is None: + raise original_exception + + input_kwargs.update( + { + "fallback_model_group": context_window_fallback_model_group, + "original_model_group": original_model_group, + } + ) + + response = await run_async_fallback( + *args, + **input_kwargs, + ) + return response + + else: + error_message = "model={}. context_window_fallbacks={}. fallbacks={}.\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks".format( + model_group, context_window_fallbacks, fallbacks + ) + verbose_router_logger.info( + msg="Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ + to fallbacks, if available.{}".format( + error_message + ) + ) + + e.message += "\n{}".format(error_message) + elif isinstance(e, litellm.ContentPolicyViolationError): + if content_policy_fallbacks is not None: + content_policy_fallback_model_group: Optional[List[str]] = ( + self._get_fallback_model_group_from_fallbacks( + fallbacks=content_policy_fallbacks, + model_group=model_group, + ) + ) + if content_policy_fallback_model_group is None: + raise original_exception + + input_kwargs.update( + { + "fallback_model_group": content_policy_fallback_model_group, + "original_model_group": original_model_group, + } + ) + + response = await run_async_fallback( + *args, + **input_kwargs, + ) + return response + else: + error_message = "model={}. content_policy_fallback={}. fallbacks={}.\n\nSet 'content_policy_fallback' - https://docs.litellm.ai/docs/routing#fallbacks".format( + model_group, content_policy_fallbacks, fallbacks + ) + verbose_router_logger.info( + msg="Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ + to fallbacks, if available.{}".format( + error_message + ) + ) + + e.message += "\n{}".format(error_message) + if fallbacks is not None and model_group is not None: + verbose_router_logger.debug(f"inside model fallbacks: {fallbacks}") + ( + fallback_model_group, + generic_fallback_idx, + ) = get_fallback_model_group( + fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] + model_group=cast(str, model_group), + ) + ## if none, check for generic fallback + if fallback_model_group is None and generic_fallback_idx is not None: + fallback_model_group = fallbacks[generic_fallback_idx]["*"] + + if fallback_model_group is None: + verbose_router_logger.info( + f"No fallback model group found for original model_group={model_group}. Fallbacks={fallbacks}" + ) + if hasattr(original_exception, "message"): + original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={fallbacks}" # type: ignore + raise original_exception + + input_kwargs.update( + { + "fallback_model_group": fallback_model_group, + "original_model_group": original_model_group, + } + ) + + response = await run_async_fallback( + *args, + **input_kwargs, + ) + + return response + except Exception as new_exception: + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + verbose_router_logger.error( + "litellm.router.py::async_function_with_fallbacks() - Error occurred while trying to do fallbacks - {}\n{}\n\nDebug Information:\nCooldown Deployments={}".format( + str(new_exception), + traceback.format_exc(), + await _async_get_cooldown_deployments_with_debug_info( + litellm_router_instance=self, + parent_otel_span=parent_otel_span, + ), + ) + ) + fallback_failure_exception_str = str(new_exception) + + if hasattr(original_exception, "message"): + # add the available fallbacks to the exception + original_exception.message += ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore + model_group, + fallback_model_group, + ) + if len(fallback_failure_exception_str) > 0: + original_exception.message += ( # type: ignore + "\nError doing the fallback: {}".format( + fallback_failure_exception_str + ) + ) + + raise original_exception + @tracer.wrap() - async def async_function_with_fallbacks(self, *args, **kwargs): # noqa: PLR0915 + async def async_function_with_fallbacks(self, *args, **kwargs): """ Try calling the function_with_retries If it fails after num_retries, fall back to another model group @@ -3557,185 +3852,16 @@ class Router: ) return response except Exception as e: - verbose_router_logger.debug(f"Traceback{traceback.format_exc()}") - original_exception = e - fallback_model_group = None - original_model_group: Optional[str] = kwargs.get("model") # type: ignore - fallback_failure_exception_str = "" - - if disable_fallbacks is True or original_model_group is None: - raise e - - input_kwargs = { - "litellm_router": self, - "original_exception": original_exception, - **kwargs, - } - - if "max_fallbacks" not in input_kwargs: - input_kwargs["max_fallbacks"] = self.max_fallbacks - if "fallback_depth" not in input_kwargs: - input_kwargs["fallback_depth"] = 0 - - try: - verbose_router_logger.info("Trying to fallback b/w models") - - # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] - is_non_standard_fallback_format = _check_non_standard_fallback_format( - fallbacks=fallbacks - ) - - if is_non_standard_fallback_format: - input_kwargs.update( - { - "fallback_model_group": fallbacks, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - - return response - - if isinstance(e, litellm.ContextWindowExceededError): - if context_window_fallbacks is not None: - fallback_model_group: Optional[List[str]] = ( - self._get_fallback_model_group_from_fallbacks( - fallbacks=context_window_fallbacks, - model_group=model_group, - ) - ) - if fallback_model_group is None: - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response - - else: - error_message = "model={}. context_window_fallbacks={}. fallbacks={}.\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks".format( - model_group, context_window_fallbacks, fallbacks - ) - verbose_router_logger.info( - msg="Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ - to fallbacks, if available.{}".format( - error_message - ) - ) - - e.message += "\n{}".format(error_message) - elif isinstance(e, litellm.ContentPolicyViolationError): - if content_policy_fallbacks is not None: - fallback_model_group: Optional[List[str]] = ( - self._get_fallback_model_group_from_fallbacks( - fallbacks=content_policy_fallbacks, - model_group=model_group, - ) - ) - if fallback_model_group is None: - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response - else: - error_message = "model={}. content_policy_fallback={}. fallbacks={}.\n\nSet 'content_policy_fallback' - https://docs.litellm.ai/docs/routing#fallbacks".format( - model_group, content_policy_fallbacks, fallbacks - ) - verbose_router_logger.info( - msg="Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ - to fallbacks, if available.{}".format( - error_message - ) - ) - - e.message += "\n{}".format(error_message) - if fallbacks is not None and model_group is not None: - verbose_router_logger.debug(f"inside model fallbacks: {fallbacks}") - ( - fallback_model_group, - generic_fallback_idx, - ) = get_fallback_model_group( - fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] - model_group=cast(str, model_group), - ) - ## if none, check for generic fallback - if ( - fallback_model_group is None - and generic_fallback_idx is not None - ): - fallback_model_group = fallbacks[generic_fallback_idx]["*"] - - if fallback_model_group is None: - verbose_router_logger.info( - f"No fallback model group found for original model_group={model_group}. Fallbacks={fallbacks}" - ) - if hasattr(original_exception, "message"): - original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={fallbacks}" # type: ignore - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - - return response - except Exception as new_exception: - parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) - verbose_router_logger.error( - "litellm.router.py::async_function_with_fallbacks() - Error occurred while trying to do fallbacks - {}\n{}\n\nDebug Information:\nCooldown Deployments={}".format( - str(new_exception), - traceback.format_exc(), - await _async_get_cooldown_deployments_with_debug_info( - litellm_router_instance=self, - parent_otel_span=parent_otel_span, - ), - ) - ) - fallback_failure_exception_str = str(new_exception) - - if hasattr(original_exception, "message"): - # add the available fallbacks to the exception - original_exception.message += ". Received Model Group={}\nAvailable Model Group Fallbacks={}".format( # type: ignore - model_group, - fallback_model_group, - ) - if len(fallback_failure_exception_str) > 0: - original_exception.message += ( # type: ignore - "\nError doing the fallback: {}".format( - fallback_failure_exception_str - ) - ) - - raise original_exception + return await self.async_function_with_fallbacks_common_utils( + e, + disable_fallbacks, + fallbacks, + context_window_fallbacks, + content_policy_fallbacks, + model_group, + args, + kwargs, + ) def _handle_mock_testing_fallbacks( self, diff --git a/tests/local_testing/test_router_batch_completion.py b/tests/local_testing/test_router_batch_completion.py index 6fedb82a55..534db7ed01 100644 --- a/tests/local_testing/test_router_batch_completion.py +++ b/tests/local_testing/test_router_batch_completion.py @@ -130,6 +130,7 @@ async def test_batch_completion_fastest_response_unit_test(): @pytest.mark.asyncio async def test_batch_completion_fastest_response_streaming(): litellm.set_verbose = True + litellm._turn_on_debug() router = litellm.Router( model_list=[ diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 1d349a44e5..980693aa73 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -125,3 +125,21 @@ def test_handle_any_messages_to_chat_completion_str_messages_conversion_complex( result = handle_any_messages_to_chat_completion_str_messages_conversion(message) assert len(result) == 1 assert result[0]["input"] == json.dumps(message) + + +def test_convert_prefix_message_to_non_prefix_messages(): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_prefix_message_to_non_prefix_messages, + ) + + messages = [ + {"role": "assistant", "content": "value", "prefix": True}, + ] + result = convert_prefix_message_to_non_prefix_messages(messages) + assert result == [ + { + "role": "system", + "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", + }, + {"role": "assistant", "content": "value"}, + ] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cacdd51ee9..4795fdc372 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2,7 +2,7 @@ import copy import json import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient @@ -13,6 +13,7 @@ sys.path.insert( import litellm +from litellm.router_utils.fallback_event_handlers import run_async_fallback def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -1064,3 +1065,321 @@ def test_router_get_model_access_groups_team_only_models(): model_name="gpt-3.5-turbo", team_id="team_1" ) assert list(access_groups.keys()) == ["default-models"] + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator(): + """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.exceptions import MidStreamFallbackError + from litellm.types.utils import ModelResponseStream + + # Helper class for creating async iterators + class AsyncIterator: + def __init__(self, items, error_after=None): + self.items = items + self.index = 0 + self.error_after = error_after + + def __aiter__(self): + return self + + async def __anext__(self): + if self.error_after is not None and self.index >= self.error_after: + raise self.error_after + if self.index >= len(self.items): + raise StopAsyncIteration + item = self.items[self.index] + self.index += 1 + return item + + # Set up router with fallback configuration + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, + }, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key-2"}, + }, + ], + fallbacks=[{"gpt-4": ["gpt-3.5-turbo"]}], + set_verbose=True, + ) + + # Test data + messages = [{"role": "user", "content": "Hello"}] + initial_kwargs = {"model": "gpt-4", "stream": True, "temperature": 0.7} + + # Test 1: Successful streaming (no errors) + print("\n=== Test 1: Successful streaming ===") + + # Mock successful streaming response + mock_chunks = [ + MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), + MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]), + MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]), + ] + + mock_response = AsyncIterator(mock_chunks) + + setattr(mock_response, "model", "gpt-4") + setattr(mock_response, "custom_llm_provider", "openai") + setattr(mock_response, "logging_obj", MagicMock()) + + result = await router._acompletion_streaming_iterator( + model_response=mock_response, messages=messages, initial_kwargs=initial_kwargs + ) + + # Collect streamed chunks + collected_chunks = [] + async for chunk in result: + collected_chunks.append(chunk) + + assert len(collected_chunks) == 3 + assert all(chunk in mock_chunks for chunk in collected_chunks) + print("✓ Successfully streamed all chunks") + + # Test 2: MidStreamFallbackError with fallback + print("\n=== Test 2: MidStreamFallbackError with fallback ===") + + # Create error that should trigger after first chunk + error = MidStreamFallbackError( + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + ) + + class AsyncIteratorWithError: + def __init__(self, items, error_after_index): + self.items = items + self.index = 0 + self.error_after_index = error_after_index + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index >= len(self.items): + raise StopAsyncIteration + if self.index == self.error_after_index: + raise error + item = self.items[self.index] + self.index += 1 + return item + + mock_error_response = AsyncIteratorWithError( + mock_chunks, 1 + ) # Error after first chunk + + setattr(mock_error_response, "model", "gpt-4") + setattr(mock_error_response, "custom_llm_provider", "openai") + setattr(mock_error_response, "logging_obj", MagicMock()) + + # Mock the fallback response + fallback_chunks = [ + MagicMock(choices=[MagicMock(delta=MagicMock(content=" world"))]), + MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]), + ] + + mock_fallback_response = AsyncIterator(fallback_chunks) + + # Mock the fallback function + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=mock_fallback_response, + ) as mock_fallback_utils: + + collected_chunks = [] + result = await router._acompletion_streaming_iterator( + model_response=mock_error_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + async for chunk in result: + collected_chunks.append(chunk) + + # Verify fallback was called + assert mock_fallback_utils.called + call_args = mock_fallback_utils.call_args + + # Check that generated content was added to messages + fallback_kwargs = call_args.kwargs["kwargs"] + modified_messages = fallback_kwargs["messages"] + + # Should have original message + system message + assistant message with prefix + assert len(modified_messages) == 3 + assert modified_messages[0] == {"role": "user", "content": "Hello"} + assert modified_messages[1]["role"] == "system" + assert "continuation" in modified_messages[1]["content"] + assert modified_messages[2]["role"] == "assistant" + assert modified_messages[2]["content"] == "Hello" + assert modified_messages[2]["prefix"] == True + + # Verify fallback parameters + assert call_args.kwargs["disable_fallbacks"] == False + assert call_args.kwargs["model_group"] == "gpt-4" + + # Should get original chunk + fallback chunks + assert len(collected_chunks) == 3 # 1 original + 2 fallback + print("✓ Fallback system called correctly with proper message modification") + + # Test 3: Fallback failure + print("\n=== Test 3: Fallback failure ===") + + mock_error_response_2 = AsyncIteratorWithError(mock_chunks, 1) # Same error pattern + + # Mock fallback failure + fallback_error = Exception("Fallback also failed") + with patch.object( + router, "async_function_with_fallbacks_common_utils", side_effect=fallback_error + ): + + collected_chunks = [] + original_error = None + setattr(mock_error_response_2, "model", "gpt-4") + setattr(mock_error_response_2, "custom_llm_provider", "openai") + setattr(mock_error_response_2, "logging_obj", MagicMock()) + + try: + result = await router._acompletion_streaming_iterator( + model_response=mock_error_response_2, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + async for chunk in result: + collected_chunks.append(chunk) + except MidStreamFallbackError as e: + original_error = e + + # Should re-raise original MidStreamFallbackError, not fallback error + assert original_error is not None + assert isinstance(original_error, MidStreamFallbackError) + assert original_error.generated_content == "Hello" + print("✓ Original error re-raised when fallback fails") + + print("\n=== All tests passed! ===") + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_edge_cases(): + """Test edge cases for _acompletion_streaming_iterator.""" + from unittest.mock import MagicMock + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + set_verbose=True, + ) + + messages = [{"role": "user", "content": "Test"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + # Test: Empty generated content + empty_error = MidStreamFallbackError( + message="Error", + model="gpt-4", + llm_provider="openai", + generated_content="", # Empty content + ) + + class AsyncIteratorImmediateError: + def __aiter__(self): + return self + + async def __anext__(self): + raise empty_error + + mock_response = AsyncIteratorImmediateError() + + # Mock empty fallback response using AsyncIterator + class EmptyAsyncIterator: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + mock_fallback_response = EmptyAsyncIterator() + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + return_value=mock_fallback_response, + ) as mock_fallback_utils: + + collected_chunks = [] + async for chunk in router._acompletion_streaming_iterator( + model_response=mock_response, + messages=messages, + initial_kwargs=initial_kwargs, + ): + collected_chunks.append(chunk) + + # Should still call fallback even with empty content + assert mock_fallback_utils.called + fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"] + modified_messages = fallback_kwargs["messages"] + + # Should have assistant message with empty content + assert modified_messages[2]["content"] == "" + print("✓ Handles empty generated content correctly") + + print("✓ Edge case tests passed!") + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_common_utils(): + """Test the async_function_with_fallbacks_common_utils method""" + # Create a basic router for testing + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + }, + } + ], + max_fallbacks=5, + ) + + # Test case 1: disable_fallbacks=True should raise original exception + test_exception = Exception("Test error") + with pytest.raises(Exception, match="Test error"): + await router.async_function_with_fallbacks_common_utils( + e=test_exception, + disable_fallbacks=True, + fallbacks=None, + context_window_fallbacks=None, + content_policy_fallbacks=None, + model_group="gpt-3.5-turbo", + args=(), + kwargs=MagicMock(), + ) + + # Test case 2: original_model_group=None should raise original exception + with pytest.raises(Exception, match="Test error"): + await router.async_function_with_fallbacks_common_utils( + e=test_exception, + disable_fallbacks=False, + fallbacks=None, + context_window_fallbacks=None, + content_policy_fallbacks=None, + model_group="gpt-3.5-turbo", + args=(), + kwargs={}, # No model key + )