From d887a3d1027091bc2104ee910ed12f9dba6c268f Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Thu, 31 Jul 2025 21:28:06 -0700 Subject: [PATCH] Anthropic - mid stream fallbacks p2 (add token usage across both calls) (#13170) * 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 * fix(router.py): return combined usage ensures accurate usage tracking on clientside for stream w/ fallbacks --- litellm/cost_calculator.py | 12 +++++++--- litellm/router.py | 49 ++++++++++++++++++++++++++++++++++---- 2 files changed, 54 insertions(+), 7 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index c8892cd26a5..6fecc7fa976 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -789,6 +789,7 @@ def completion_cost( # noqa: PLR0915 from litellm.llms.recraft.cost_calculator import ( cost_calculator as recraft_image_cost_calculator, ) + return recraft_image_cost_calculator( model=model, image_response=completion_response, @@ -797,6 +798,7 @@ def completion_cost( # noqa: PLR0915 from litellm.llms.gemini.image_generation.cost_calculator import ( cost_calculator as gemini_image_cost_calculator, ) + return gemini_image_cost_calculator( model=model, image_response=completion_response, @@ -867,7 +869,10 @@ def completion_cost( # noqa: PLR0915 from litellm.proxy._experimental.mcp_server.cost_calculator import ( MCPCostCalculator, ) - return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj) + + return MCPCostCalculator.calculate_mcp_tool_call_cost( + litellm_logging_obj=litellm_logging_obj + ) # Calculate cost based on prompt_tokens, completion_tokens if ( "togethercomputer" in model @@ -1318,7 +1323,7 @@ class BaseTokenUsageProcessor: combined.completion_tokens_details = CompletionTokensDetails() # Check what keys exist in the model's completion_tokens_details - for attr in dir(usage.completion_tokens_details): + for attr in usage.completion_tokens_details.model_fields: if not attr.startswith("_") and not callable( getattr(usage.completion_tokens_details, attr) ): @@ -1326,7 +1331,8 @@ class BaseTokenUsageProcessor: combined.completion_tokens_details, attr, 0 ) new_val = getattr(usage.completion_tokens_details, attr, 0) - if new_val is not None: + + if new_val is not None and current_val is not None: setattr( combined.completion_tokens_details, attr, diff --git a/litellm/router.py b/litellm/router.py index 2673df1504b..ca724006fdf 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -23,6 +23,7 @@ from functools import lru_cache from typing import ( TYPE_CHECKING, Any, + AsyncGenerator, Callable, Dict, List, @@ -146,7 +147,7 @@ from litellm.types.services import ServiceTypes from litellm.types.utils import GenericBudgetConfigType, LiteLLMBatch from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -1092,7 +1093,7 @@ class Router: from litellm.exceptions import MidStreamFallbackError class FallbackStreamWrapper(CustomStreamWrapper): - def __init__(self, async_generator): + def __init__(self, async_generator: AsyncGenerator): # Copy attributes from the original model_response super().__init__( completion_stream=async_generator, @@ -1103,7 +1104,7 @@ class Router: self._async_generator = async_generator def __aiter__(self): - return self._async_generator + return self async def __anext__(self): return await self._async_generator.__anext__() @@ -1113,6 +1114,15 @@ class Router: async for item in model_response: yield item except MidStreamFallbackError as e: + from litellm.main import stream_chunk_builder + + complete_response_object = stream_chunk_builder( + chunks=model_response.chunks + ) + complete_response_object_usage = cast( + Optional[Usage], + getattr(complete_response_object, "usage", None), + ) try: # Use the router's fallback system model_group = cast(str, initial_kwargs.get("model")) @@ -1156,6 +1166,37 @@ class Router: # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): async for fallback_item in fallback_response: # type: ignore + if ( + fallback_item + and isinstance(fallback_item, ModelResponseStream) + and hasattr(fallback_item, "usage") + ): + from litellm.cost_calculator import ( + BaseTokenUsageProcessor, + ) + + usage = cast( + Optional[Usage], + getattr(fallback_item, "usage", None), + ) + if usage is not None: + usage_objects = [usage] + else: + usage_objects = [] + + if ( + complete_response_object_usage is not None + and hasattr(complete_response_object_usage, "usage") + and complete_response_object_usage.usage is not None # type: ignore + ): + usage_objects.append(complete_response_object_usage) + + combined_usage = ( + BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=usage_objects + ) + ) + setattr(fallback_item, "usage", combined_usage) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -1166,7 +1207,7 @@ class Router: verbose_router_logger.error( f"Fallback also failed: {fallback_error}" ) - raise e # Re-raise the original error + raise fallback_error return FallbackStreamWrapper(stream_with_fallbacks())