From ee6e76e1f9f0bbf450f647cca5e1e7c82ff2ccf4 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 27 Jun 2025 20:01:12 -0700 Subject: [PATCH] Bedrock Passthrough cost tracking (`/invoke` + `/converse` routes - streaming + non-streaming) (#12123) * refactor(passthrough_endpoints-success-handler): refactor llm passthrough logging logic isolate the llm translation work to enable cost tracking on sdk * feat: initial implementation of passthrough SDK cost calculation enables bedrock passthrough cost tracking to work * feat(cost_calculator.py): working cost calculation for bedrock passthrough * feat(litellm_logging.py): consider allm_passthrough in cost tracking allows async calls (e.g. via proxy) to work * feat(bedrock/passthrough): working event stream decoding for bedrock passthrough calls + logging instrumentation for passthrough sdk calls (log on stream completion) Enables bedrock streaming cost calculation * feat(litellm_logging.py): support streaming passthrough cost tracking * feat(passthrough/main.py): working async streaming cost calculation Closes https://github.com/BerriAI/litellm/issues/11359 * feat(proxy_server.py): fix passthrough routing when llm router enabled * feat: further fixes * feat(bedrock/): working bedrock passthrough cost tracking (non-streaming) * feat(litellm_logging.py): working usage tracking for bedrock passthrough calls ensures tokens are logged * feat(bedrock/passthrough): add converse passthrough cost tracking support * feat(base_llm/passthrough): remove redundant function * refactor(litellm_logging.py): refactor function to be below 50 LOC * test: update test * test: remove redundant test --- litellm/cost_calculator.py | 8 +- litellm/litellm_core_utils/litellm_logging.py | 124 +++++++++++--- .../litellm_core_utils/streaming_handler.py | 26 +++ .../base_llm/passthrough/transformation.py | 46 ++++- litellm/llms/bedrock/chat/__init__.py | 30 +++- litellm/llms/bedrock/common_utils.py | 66 ++++++++ .../bedrock/passthrough/transformation.py | 146 +++++++++++++++- litellm/passthrough/main.py | 146 ++++++++++++++-- litellm/proxy/common_request_processing.py | 67 +++++--- .../pass_through_endpoints/success_handler.py | 84 ++++++--- litellm/proxy/route_llm_request.py | 6 + litellm/router.py | 32 +++- litellm/types/utils.py | 11 ++ .../test_bedrock_completion.py | 160 ++++++++++++++++-- .../llms/azure/test_azure_common_utils.py | 8 +- 15 files changed, 841 insertions(+), 119 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cd37592d40..d65411e0f9 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2,8 +2,9 @@ ## File for 'response_cost' calculation in Logging import time from functools import lru_cache -from typing import Any, List, Literal, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union, cast +from httpx import Response from pydantic import BaseModel import litellm @@ -93,6 +94,9 @@ from litellm.utils import ( token_counter, ) +if TYPE_CHECKING: + pass + def _cost_per_token_custom_pricing_helper( prompt_tokens: float = 0, @@ -655,6 +659,7 @@ def completion_cost( # noqa: PLR0915 potential_model_names = [selected_model] if model is not None: potential_model_names.append(model) + for idx, model in enumerate(potential_model_names): try: verbose_logger.debug( @@ -964,6 +969,7 @@ def response_cost_calculator( ResponsesAPIResponse, LiteLLMRealtimeStreamLoggingObject, OpenAIModerationResponse, + Response, ], model: str, custom_llm_provider: Optional[str], diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 713c3358b5..a8b9ac8bfa 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -14,6 +14,7 @@ import uuid from datetime import datetime as dt_object from functools import lru_cache from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -26,6 +27,7 @@ from typing import ( cast, ) +from httpx import Response from pydantic import BaseModel import litellm @@ -81,6 +83,7 @@ from litellm.types.rerank import RerankResponse from litellm.types.router import CustomPricingLiteLLMParams from litellm.types.utils import ( CallTypes, + CostResponseTypes, DynamicPromptManagementParamLiteral, EmbeddingResponse, ImageResponse, @@ -148,6 +151,8 @@ from .initialize_dynamic_callback_params import ( ) from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache +if TYPE_CHECKING: + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, @@ -1254,6 +1259,46 @@ class Logging(LiteLLMLoggingBaseClass): self.completion_start_time = completion_start_time self.model_call_details["completion_start_time"] = self.completion_start_time + def normalize_logging_result(self, result: Any) -> Any: + """ + Some endpoints return a different type of result than what is expected by the logging system. + This function is used to normalize the result to the expected type. + """ + logging_result = result + if self.call_type == CallTypes.arealtime.value and isinstance(result, list): + combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( + results=result + ) + logging_result = ( + RealtimeAPITokenUsageProcessor.create_logging_realtime_object( + usage=combined_usage_object, + results=result, + ) + ) + + elif ( + self.call_type == CallTypes.llm_passthrough_route.value + or self.call_type == CallTypes.allm_passthrough_route.value + ) and isinstance(result, Response): + from litellm.utils import ProviderConfigManager + + provider_config = ProviderConfigManager.get_provider_passthrough_config( + provider=self.model_call_details.get("custom_llm_provider", ""), + model=self.model, + ) + if provider_config is not None: + logging_result = provider_config.logging_non_streaming_response( + model=self.model, + custom_llm_provider=self.model_call_details.get( + "custom_llm_provider", "" + ), + httpx_response=result, + request_data=self.model_call_details.get("request_data", {}), + logging_obj=self, + endpoint=self.model_call_details.get("endpoint", ""), + ) + return logging_result + def _success_handler_helper_fn( self, result=None, @@ -1287,28 +1332,8 @@ class Logging(LiteLLMLoggingBaseClass): ## if model in model cost map - log the response cost ## else set cost to None - logging_result = result + logging_result = self.normalize_logging_result(result=result) - if self.call_type == CallTypes.arealtime.value and isinstance(result, list): - combined_usage_object = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( - results=result - ) - logging_result = ( - RealtimeAPITokenUsageProcessor.create_logging_realtime_object( - usage=combined_usage_object, - results=result, - ) - ) - - # self.model_call_details[ - # "response_cost" - # ] = handle_realtime_stream_cost_calculation( - # results=result, - # combined_usage_object=combined_usage_object, - # custom_llm_provider=self.custom_llm_provider, - # litellm_model_name=self.model, - # ) - # self.model_call_details["combined_usage_object"] = combined_usage_object if ( standard_logging_object is None and result is not None @@ -1426,6 +1451,59 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {str(e)}") + def _flush_passthrough_collected_chunks_helper( + self, + raw_bytes: List[bytes], + provider_config: "BasePassthroughConfig", + ) -> Optional["CostResponseTypes"]: + all_chunks = provider_config._convert_raw_bytes_to_str_lines(raw_bytes) + complete_streaming_response = provider_config.handle_logging_collected_chunks( + all_chunks=all_chunks, + litellm_logging_obj=self, + model=self.model, + custom_llm_provider=self.model_call_details.get("custom_llm_provider", ""), + endpoint=self.model_call_details.get("endpoint", ""), + ) + return complete_streaming_response + + def flush_passthrough_collected_chunks( + self, + raw_bytes: List[bytes], + provider_config: "BasePassthroughConfig", + ): + """ + Flush collected chunks from the logging object + This is used to log the collected chunks once streaming is done on passthrough endpoints + + 1. Decode the raw bytes to string lines + 2. Get the complete streaming response from the provider config + 3. Log the complete streaming response (trigger success handler) + This is used for passthrough endpoints + """ + complete_streaming_response = self._flush_passthrough_collected_chunks_helper( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + + if complete_streaming_response is not None: + + self.success_handler(result=complete_streaming_response) + return + + async def async_flush_passthrough_collected_chunks( + self, + raw_bytes: List[bytes], + provider_config: "BasePassthroughConfig", + ): + complete_streaming_response = self._flush_passthrough_collected_chunks_helper( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + + if complete_streaming_response is not None: + await self.async_success_handler(result=complete_streaming_response) + return + def success_handler( # noqa: PLR0915 self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs ): @@ -1968,6 +2046,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["async_complete_streaming_response"] = ( complete_streaming_response ) + try: if self.model_call_details.get("cache_hit", False) is True: self.model_call_details["response_cost"] = 0.0 @@ -2047,6 +2126,7 @@ class Logging(LiteLLMLoggingBaseClass): ) self.has_run_logging(event_type="async_success") + for callback in callbacks: # check if callback can run for this request litellm_params = self.model_call_details.get("litellm_params", {}) @@ -2709,6 +2789,8 @@ class Logging(LiteLLMLoggingBaseClass): return result elif isinstance(result, ResponseCompletedEvent): return result.response + else: + return None return None def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 529e7c8399..98fb94922f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1891,3 +1891,29 @@ def generic_chunk_has_all_required_fields(chunk: dict) -> bool: decision = all(key in _all_fields for key in chunk) return decision + + +def convert_generic_chunk_to_model_response_stream( + chunk: GChunk, +) -> ModelResponseStream: + from litellm.types.utils import Delta + + model_response_stream = ModelResponseStream( + id=str(uuid.uuid4()), + model="", + choices=[ + StreamingChoices( + index=chunk.get("index", 0), + delta=Delta( + content=chunk["text"], + tool_calls=chunk.get("tool_use", None), + ), + ) + ], + finish_reason=chunk["finish_reason"] if chunk["is_finished"] else None, + ) + + if "usage" in chunk and chunk["usage"] is not None: + setattr(model_response_stream, "usage", chunk["usage"]) + + return model_response_stream diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index e157a8ffde..60d89c1610 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -1,10 +1,13 @@ from abc import abstractmethod -from typing import TYPE_CHECKING, Optional, Tuple, Union +from typing import TYPE_CHECKING, List, Optional, Tuple, Union from ..base_utils import BaseLLMModelInfo if TYPE_CHECKING: - from httpx import URL, Headers + from httpx import URL, Headers, Response + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import CostResponseTypes from ..chat.transformation import BaseLLMException @@ -101,3 +104,42 @@ class BasePassthroughConfig(BaseLLMModelInfo): return BaseLLMException( status_code=status_code, message=error_message, headers=headers ) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: "Response", + request_data: dict, + logging_obj: "LiteLLMLoggingObj", + endpoint: str, + ) -> Optional["CostResponseTypes"]: + pass + + def handle_logging_collected_chunks( + self, + all_chunks: List[str], + litellm_logging_obj: "LiteLLMLoggingObj", + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + return None + + def _convert_raw_bytes_to_str_lines(self, raw_bytes: List[bytes]) -> List[str]: + """ + Converts a list of raw bytes into a list of string lines, similar to aiter_lines() + + Args: + raw_bytes: List of bytes chunks from aiter.bytes() + + Returns: + List of string lines, with each line being a complete data: {} chunk + """ + # Combine all bytes and decode to string + combined_str = b"".join(raw_bytes).decode("utf-8") + + # Split by newlines and filter out empty lines + lines = [line.strip() for line in combined_str.split("\n") if line.strip()] + + return lines diff --git a/litellm/llms/bedrock/chat/__init__.py b/litellm/llms/bedrock/chat/__init__.py index c3f6aef6d2..8cd0e94e68 100644 --- a/litellm/llms/bedrock/chat/__init__.py +++ b/litellm/llms/bedrock/chat/__init__.py @@ -1,2 +1,30 @@ +from typing import Optional + from .converse_handler import BedrockConverseLLM -from .invoke_handler import BedrockLLM +from .invoke_handler import ( + AmazonAnthropicClaudeStreamDecoder, + AmazonDeepSeekR1StreamDecoder, + AWSEventStreamDecoder, + BedrockLLM, +) + + +def get_bedrock_event_stream_decoder( + invoke_provider: Optional[str], model: str, sync_stream: bool, json_mode: bool +): + if invoke_provider and invoke_provider == "anthropic": + decoder: AWSEventStreamDecoder = AmazonAnthropicClaudeStreamDecoder( + model=model, + sync_stream=sync_stream, + json_mode=json_mode, + ) + return decoder + elif invoke_provider and invoke_provider == "deepseek_r1": + decoder = AmazonDeepSeekR1StreamDecoder( + model=model, + sync_stream=sync_stream, + ) + return decoder + else: + decoder = AWSEventStreamDecoder(model=model) + return decoder diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index ee2a6cda8d..2a8fdc148b 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -2,6 +2,7 @@ Common utilities used across bedrock chat/embedding/image generation """ +import json import os from typing import TYPE_CHECKING, List, Literal, Optional, Union @@ -458,3 +459,68 @@ class BedrockModelInfo(BaseLLMModelInfo): ): return "converse" return "invoke" + + +class BedrockEventStreamDecoderBase: + """ + Base class for event stream decoding for Bedrock + """ + + _response_stream_shape_cache = None + + def __init__(self): + from botocore.parsers import EventStreamJSONParser + + self.parser = EventStreamJSONParser() + + def get_response_stream_shape(self): + if self._response_stream_shape_cache is None: + from botocore.loaders import Loader + from botocore.model import ServiceModel + + loader = Loader() + bedrock_service_dict = loader.load_service_model( + "bedrock-runtime", "service-2" + ) + bedrock_service_model = ServiceModel(bedrock_service_dict) + self._response_stream_shape_cache = bedrock_service_model.shape_for( + "ResponseStream" + ) + + return self._response_stream_shape_cache + + def _parse_message_from_event(self, event) -> Optional[str]: + response_dict = event.to_response_dict() + parsed_response = self.parser.parse( + response_dict, self.get_response_stream_shape() + ) + + if response_dict["status_code"] != 200: + decoded_body = response_dict["body"].decode() + if isinstance(decoded_body, dict): + error_message = decoded_body.get("message") + elif isinstance(decoded_body, str): + error_message = decoded_body + else: + error_message = "" + exception_status = response_dict["headers"].get(":exception-type") + error_message = exception_status + " " + error_message + raise BedrockError( + status_code=response_dict["status_code"], + message=( + json.dumps(error_message) + if isinstance(error_message, dict) + else error_message + ), + ) + if "chunk" in parsed_response: + chunk = parsed_response.get("chunk") + if not chunk: + return None + return chunk.get("bytes").decode() # type: ignore[no-any-return] + else: + chunk = response_dict.get("body") + if not chunk: + return None + + return chunk.decode() # type: ignore[no-any-return] diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 6ed6d7829e..d7221ff4b7 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -1,15 +1,26 @@ -from typing import TYPE_CHECKING, Optional, Tuple +import json +from typing import TYPE_CHECKING, List, Optional, Tuple, cast +from httpx import Response + +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockModelInfo +from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import CostResponseTypes + if TYPE_CHECKING: from httpx import URL -class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BasePassthroughConfig): +class BedrockPassthroughConfig( + BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig +): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in endpoint @@ -51,3 +62,132 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BasePassthroughConf api_base=api_base, model=model, ) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: dict, + logging_obj: Logging, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + from litellm import encoding + from litellm.types.utils import LlmProviders, ModelResponse + from litellm.utils import ProviderConfigManager + + if "invoke" in endpoint: + chat_config_model = "invoke/" + model + elif "converse" in endpoint: + chat_config_model = "converse/" + model + else: + return None + + provider_chat_config = ProviderConfigManager.get_provider_chat_config( + provider=LlmProviders(custom_llm_provider), + model=chat_config_model, + ) + + if provider_chat_config is None: + raise ValueError(f"No provider config found for model: {model}") + + litellm_model_response: ModelResponse = provider_chat_config.transform_response( + model=model, + messages=[{"role": "user", "content": "no-message-pass-through-endpoint"}], + raw_response=httpx_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + optional_params={}, + litellm_params={}, + api_key="", + request_data=request_data, + encoding=encoding, + ) + + return litellm_model_response + + def _convert_raw_bytes_to_str_lines(self, raw_bytes: List[bytes]) -> List[str]: + from botocore.eventstream import EventStreamBuffer + + all_chunks = [] + event_stream_buffer = EventStreamBuffer() + for chunk in raw_bytes: + event_stream_buffer.add_data(chunk) + for event in event_stream_buffer: + message = self._parse_message_from_event(event) + if message is not None: + all_chunks.append(message) + + return all_chunks + + def handle_logging_collected_chunks( + self, + all_chunks: List[str], + litellm_logging_obj: "LiteLLMLoggingObj", + model: str, + custom_llm_provider: str, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + """ + 1. Convert all_chunks to a ModelResponseStream + 2. combine model_response_stream to model_response + 3. Return the model_response + """ + + from litellm.litellm_core_utils.streaming_handler import ( + convert_generic_chunk_to_model_response_stream, + generic_chunk_has_all_required_fields, + ) + from litellm.llms.bedrock.chat import get_bedrock_event_stream_decoder + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, + ) + from litellm.main import stream_chunk_builder + from litellm.types.utils import GenericStreamingChunk, ModelResponseStream + + all_translated_chunks = [] + if "invoke" in endpoint: + invoke_provider = AmazonInvokeConfig.get_bedrock_invoke_provider(model) + if invoke_provider is None: + raise ValueError( + f"Invalid invoke provider: {invoke_provider}, for model: {model}" + ) + obj = get_bedrock_event_stream_decoder( + invoke_provider=invoke_provider, + model=model, + sync_stream=True, + json_mode=False, + ) + elif "converse" in endpoint: + obj = get_bedrock_event_stream_decoder( + invoke_provider=None, + model=model, + sync_stream=True, + json_mode=False, + ) + else: + return None + + for chunk in all_chunks: + message = json.loads(chunk) + translated_chunk = obj._chunk_parser(chunk_data=message) + + if isinstance( + translated_chunk, dict + ) and generic_chunk_has_all_required_fields(cast(dict, translated_chunk)): + chunk_obj = convert_generic_chunk_to_model_response_stream( + cast(GenericStreamingChunk, translated_chunk) + ) + elif isinstance(translated_chunk, ModelResponseStream): + chunk_obj = translated_chunk + else: + continue + + all_translated_chunks.append(chunk_obj) + + if len(all_translated_chunks) > 0: + model_response = stream_chunk_builder( + chunks=all_translated_chunks, + ) + return model_response + return None diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 5e95b0a81b..59fab1b336 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -5,7 +5,17 @@ This module is used to pass through requests to the LLM APIs. import asyncio import contextvars from functools import partial -from typing import TYPE_CHECKING, Any, Coroutine, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Coroutine, + Generator, + List, + Optional, + Union, + cast, +) import httpx from httpx._types import CookieTypes, QueryParamTypes, RequestFiles @@ -20,6 +30,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig @@ -42,7 +53,12 @@ async def allm_passthrough_route( cookies: Optional[CookieTypes] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, -) -> Union[httpx.Response, Coroutine[Any, Any, httpx.Response]]: +) -> Union[ + httpx.Response, + Coroutine[Any, Any, httpx.Response], + Generator[Any, Any, Any], + AsyncGenerator[Any, Any], +]: """ Async: Reranks a list of documents based on their relevance to the query """ @@ -50,6 +66,26 @@ async def allm_passthrough_route( loop = asyncio.get_event_loop() kwargs["allm_passthrough_route"] = True + model, custom_llm_provider, api_key, api_base = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + ) + + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + provider_config = cast( + Optional["BasePassthroughConfig"], kwargs.get("provider_config") + ) or ProviderConfigManager.get_provider_passthrough_config( + provider=LlmProviders(custom_llm_provider), + model=model, + ) + + if provider_config is None: + raise Exception(f"Provider {custom_llm_provider} not found") + func = partial( llm_passthrough_route, method=method, @@ -76,20 +112,24 @@ async def allm_passthrough_route( if asyncio.iscoroutine(init_response): response = await init_response + try: response.raise_for_status() except httpx.HTTPStatusError as e: error_text = await e.response.aread() error_text_str = error_text.decode("utf-8") raise Exception(error_text_str) + else: response = init_response + return response + except Exception as e: # For passthrough routes, we need to get the provider config to properly handle errors from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager - + # Get the provider using the same logic as llm_passthrough_route _, resolved_custom_llm_provider, _, _ = get_llm_provider( model=model, @@ -97,7 +137,7 @@ async def allm_passthrough_route( api_base=api_base, api_key=api_key, ) - + # Get provider config if available provider_config = None if resolved_custom_llm_provider: @@ -111,11 +151,11 @@ async def allm_passthrough_route( except Exception: # If we can't get provider config, pass None pass - + if provider_config is None: # If no provider config available, raise the original exception raise e - + raise base_llm_http_handler._handle_error( e=e, provider_config=provider_config, @@ -142,22 +182,31 @@ def llm_passthrough_route( cookies: Optional[CookieTypes] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, **kwargs, -) -> Union[httpx.Response, Coroutine[Any, Any, httpx.Response]]: +) -> Union[ + httpx.Response, + Coroutine[Any, Any, httpx.Response], + Generator[Any, Any, Any], + AsyncGenerator[Any, Any], +]: """ Pass through requests to the LLM APIs. Step 1. Build the request Step 2. Send the request Step 3. Return the response - - [TODO] Refactor this into a provider-config pattern, once we expand this to non-vllm providers. """ + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + if client is None: if allm_passthrough_route: client = litellm.module_level_aclient else: client = litellm.module_level_client + litellm_logging_obj = cast("LiteLLMLoggingObj", kwargs.get("litellm_logging_obj")) + model, custom_llm_provider, api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, @@ -165,11 +214,15 @@ def llm_passthrough_route( api_key=api_key, ) - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.types.utils import LlmProviders - from litellm.utils import ProviderConfigManager - litellm_params_dict = get_litellm_params(**kwargs) + litellm_logging_obj.update_environment_variables( + model=model, + litellm_params=litellm_params_dict, + optional_params={}, + endpoint=endpoint, + custom_llm_provider=custom_llm_provider, + request_data=data if data else json, + ) provider_config = cast( Optional["BasePassthroughConfig"], kwargs.get("provider_config") @@ -215,7 +268,7 @@ def llm_passthrough_route( model=model, ) - ## SWAP MODEL IN JSON BODY + ## SWAP MODEL IN JSON BODY [TODO: REFACTOR TO A provider_config.transform_request method] if json and isinstance(json, dict) and "model" in json: json["model"] = model @@ -237,12 +290,27 @@ def llm_passthrough_route( request_data=data or json or {}, ) + # Update logging object with streaming status + litellm_logging_obj.stream = is_streaming_request + try: response = client.client.send(request=request, stream=is_streaming_request) if asyncio.iscoroutine(response): - return response + if is_streaming_request: + return _async_streaming(response, litellm_logging_obj, provider_config) + else: + return response response.raise_for_status() - return response + + if ( + hasattr(response, "iter_bytes") and is_streaming_request + ): # yield the chunk, so we can store it in the logging object + + return _sync_streaming(response, litellm_logging_obj, provider_config) + else: + + # For non-streaming responses, yield the entire response + return response except Exception as e: if provider_config is None: raise e @@ -250,3 +318,49 @@ def llm_passthrough_route( e=e, provider_config=provider_config, ) + + +def _sync_streaming( + response: httpx.Response, + litellm_logging_obj: "LiteLLMLoggingObj", + provider_config: "BasePassthroughConfig", +): + from litellm.utils import executor + + try: + raw_bytes: List[bytes] = [] + for chunk in response.iter_bytes(): # type: ignore + raw_bytes.append(chunk) + yield chunk + + executor.submit( + litellm_logging_obj.flush_passthrough_collected_chunks, + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + except Exception as e: + raise e + + +async def _async_streaming( + response: Coroutine[Any, Any, httpx.Response], + litellm_logging_obj: "LiteLLMLoggingObj", + provider_config: "BasePassthroughConfig", +): + try: + iter_response = await response + raw_bytes: List[bytes] = [] + + async for chunk in iter_response.aiter_bytes(): # type: ignore + + raw_bytes.append(chunk) + yield chunk + + asyncio.create_task( + litellm_logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + ) + except Exception as e: + raise e diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f943549a41..3a33b3c5f9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -108,14 +108,26 @@ async def create_streaming_response( final_status_code = default_status_code try: + + # Handle coroutine that returns a generator + if asyncio.iscoroutine(generator): + generator = await generator + + # Now get the first chunk from the actual generator first_chunk_value = await generator.__anext__() + if first_chunk_value is not None: - error_code_from_chunk = await _parse_event_data_for_error(first_chunk_value) - if error_code_from_chunk is not None: - final_status_code = error_code_from_chunk - verbose_proxy_logger.debug( - f"Error detected in first stream chunk. Status code set to: {final_status_code}" + try: + error_code_from_chunk = await _parse_event_data_for_error( + first_chunk_value ) + if error_code_from_chunk is not None: + final_status_code = error_code_from_chunk + verbose_proxy_logger.debug( + f"Error detected in first stream chunk. Status code set to: {final_status_code}" + ) + except Exception as e: + verbose_proxy_logger.debug(f"Error parsing first chunk value: {e}") except StopAsyncIteration: # Generator was empty. Default status @@ -152,6 +164,7 @@ async def create_streaming_response( with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): yield first_chunk_value async for chunk in generator: + with tracer.trace(DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE): yield chunk @@ -432,6 +445,8 @@ class ProxyBaseLLMRequestProcessing: ) if self._is_streaming_request( data=self.data, is_streaming_request=is_streaming_request + ) or self._is_streaming_response( + response ): # use generate_responses to stream responses custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, @@ -448,11 +463,28 @@ class ProxyBaseLLMRequestProcessing: **additional_headers, ) if route_type == "allm_passthrough_route": - return StreamingResponse( - content=response.aiter_bytes(), - status_code=response.status_code, - headers=custom_headers, - ) + # Check if response is an async generator + if self._is_streaming_response(response): + + if asyncio.iscoroutine(response): + generator = await response + else: + generator = response + + # For passthrough routes, stream directly without error parsing + # since we're dealing with raw binary data (e.g., AWS event streams) + return StreamingResponse( + content=generator, + status_code=status.HTTP_200_OK, + headers=custom_headers, + ) + else: + # Traditional HTTP response with aiter_bytes + return StreamingResponse( + content=response.aiter_bytes(), + status_code=response.status_code, + headers=custom_headers, + ) else: selected_data_generator = select_data_generator( response=response, @@ -557,7 +589,6 @@ class ProxyBaseLLMRequestProcessing: This uses standard Python inspection to detect streaming/async iterator objects rather than relying on specific wrapper classes. """ - import asyncio import inspect from collections.abc import AsyncGenerator, AsyncIterator @@ -569,20 +600,6 @@ class ProxyBaseLLMRequestProcessing: if isinstance(response, (AsyncIterator, AsyncGenerator)): return True - # Check for __aiter__ method (async iterator protocol) - if hasattr(response, "__aiter__") and callable(getattr(response, "__aiter__")): - return True - - # Check if it's a coroutine that might yield an async generator - if asyncio.iscoroutine(response): - return True - - # Check for common streaming HTTP response patterns - if hasattr(response, "aiter_bytes") and callable( - getattr(response, "aiter_bytes") - ): - return True - return False def _is_streaming_request( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 4075f4097d..ce576d5ac7 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -1,6 +1,6 @@ import json from datetime import datetime -from typing import Optional, Union +from typing import Any, Optional, Union from urllib.parse import urlparse import httpx @@ -88,29 +88,25 @@ class PassThroughEndpointLogging: cache_hit=False, **kwargs, ) - - - async def pass_through_async_success_handler( + def normalize_llm_passthrough_logging_payload( self, httpx_response: httpx.Response, response_body: Optional[dict], + request_body: dict, logging_obj: LiteLLMLoggingObj, url_route: str, result: str, start_time: datetime, end_time: datetime, cache_hit: bool, - request_body: dict, - passthrough_logging_payload: PassthroughStandardLoggingPayload, **kwargs, ): - standard_logging_response_object: Optional[ - PassThroughEndpointLoggingResultValues - ] = None - logging_obj.model_call_details["passthrough_logging_payload"] = ( - passthrough_logging_payload - ) + return_dict = { + "standard_logging_response_object": None, + "kwargs": kwargs, + } + standard_logging_response_object: Optional[Any] = None if self.is_vertex_route(url_route): vertex_passthrough_logging_handler_result = ( VertexPassthroughLoggingHandler.vertex_passthrough_handler( @@ -166,7 +162,33 @@ class PassThroughEndpointLogging: cohere_passthrough_logging_handler_result["result"] ) kwargs = cohere_passthrough_logging_handler_result["kwargs"] - elif self.is_assemblyai_route(url_route): + return_dict["standard_logging_response_object"] = ( + standard_logging_response_object + ) + return_dict["kwargs"] = kwargs + return return_dict + + async def pass_through_async_success_handler( + self, + httpx_response: httpx.Response, + response_body: Optional[dict], + logging_obj: LiteLLMLoggingObj, + url_route: str, + result: str, + start_time: datetime, + end_time: datetime, + cache_hit: bool, + request_body: dict, + passthrough_logging_payload: PassthroughStandardLoggingPayload, + **kwargs, + ): + standard_logging_response_object: Optional[ + PassThroughEndpointLoggingResultValues + ] = None + logging_obj.model_call_details["passthrough_logging_payload"] = ( + passthrough_logging_payload + ) + if self.is_assemblyai_route(url_route): if ( AssemblyAIPassthroughLoggingHandler._should_log_request( httpx_response.request.method @@ -189,18 +211,37 @@ class PassThroughEndpointLogging: elif self.is_langfuse_route(url_route): # Don't log langfuse pass-through requests return - + else: + normalized_llm_passthrough_logging_payload = ( + self.normalize_llm_passthrough_logging_payload( + httpx_response=httpx_response, + response_body=response_body, + request_body=request_body, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + **kwargs, + ) + ) + standard_logging_response_object = ( + normalized_llm_passthrough_logging_payload[ + "standard_logging_response_object" + ] + ) + kwargs = normalized_llm_passthrough_logging_payload["kwargs"] if standard_logging_response_object is None: standard_logging_response_object = StandardPassThroughResponseObject( response=httpx_response.text ) - + kwargs = self._set_cost_per_request( logging_obj=logging_obj, passthrough_logging_payload=passthrough_logging_payload, kwargs=kwargs, ) - await self._handle_logging( logging_obj=logging_obj, @@ -244,11 +285,10 @@ class PassThroughEndpointLogging: if route in parsed_url.path: return True return False - def _set_cost_per_request( - self, - logging_obj: LiteLLMLoggingObj, + self, + logging_obj: LiteLLMLoggingObj, passthrough_logging_payload: PassthroughStandardLoggingPayload, kwargs: dict, ): @@ -265,8 +305,8 @@ class PassThroughEndpointLogging: kwargs["response_cost"] = passthrough_logging_payload.get( "cost_per_request" ) - logging_obj.model_call_details["response_cost"] = passthrough_logging_payload.get( - "cost_per_request" + logging_obj.model_call_details["response_cost"] = ( + passthrough_logging_payload.get("cost_per_request") ) - + return kwargs diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 3d7882740b..06d50ccdc5 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -121,20 +121,24 @@ async def route_request( data["model"] in router_model_names or data["model"] in llm_router.get_model_ids() ): + return getattr(llm_router, f"{route_type}")(**data) elif ( llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias ): + return getattr(llm_router, f"{route_type}")(**data) elif data["model"] in llm_router.deployment_names: + return getattr(llm_router, f"{route_type}")( **data, specific_deployment=True ) elif data["model"] not in router_model_names: + if llm_router.router_general_settings.pass_through_all_models: return getattr(litellm, f"{route_type}")(**data) elif ( @@ -154,7 +158,9 @@ async def route_request( elif user_model is not None: return getattr(litellm, f"{route_type}")(**data) elif route_type == "allm_passthrough_route": + return getattr(litellm, f"{route_type}")(**data) + # if no route found then it's a bad request route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError( diff --git a/litellm/router.py b/litellm/router.py index 3af1167dd9..77f99ea84b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -792,6 +792,7 @@ class Router: generate_content, generate_content_stream, ) + self.agenerate_content = self.factory_function( agenerate_content, call_type="agenerate_content" ) @@ -2479,6 +2480,7 @@ class Router: """ handler_name = original_function.__name__ function_name = "_ageneric_api_call_with_fallbacks" + passthrough_on_no_deployment = kwargs.pop("passthrough_on_no_deployment", False) self._update_kwargs_before_fallbacks( model=model, kwargs=kwargs, @@ -2491,12 +2493,17 @@ class Router: f"Inside _ageneric_api_call() - handler: {handler_name}, model: {model}; kwargs: {kwargs}" ) parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) - deployment = await self.async_get_available_deployment( - model=model, - request_kwargs=kwargs, - messages=kwargs.get("messages", None), - specific_deployment=kwargs.pop("specific_deployment", None), - ) + try: + deployment = await self.async_get_available_deployment( + model=model, + request_kwargs=kwargs, + messages=kwargs.get("messages", None), + specific_deployment=kwargs.pop("specific_deployment", None), + ) + except Exception as e: + if passthrough_on_no_deployment: + return await original_function(model=model, **kwargs) + raise e self._update_kwargs_with_deployment( deployment=deployment, kwargs=kwargs, function_name=function_name @@ -3323,7 +3330,6 @@ class Router: "aretrieve_fine_tuning_job", "alist_files", "aimage_edit", - "allm_passthrough_route", "agenerate_content", "agenerate_content_stream", ): @@ -3331,6 +3337,12 @@ class Router: original_function=original_function, **kwargs, ) + elif call_type == "allm_passthrough_route": + return await self._ageneric_api_call_with_fallbacks( + original_function=original_function, + passthrough_on_no_deployment=True, + **kwargs, + ) elif call_type in ( "aget_responses", "adelete_responses", @@ -4143,8 +4155,10 @@ class Router: ) # Determine cooldown time with priority: deployment config > response header > router default - deployment_cooldown = kwargs.get("litellm_params", {}).get("cooldown_time", None) - + deployment_cooldown = kwargs.get("litellm_params", {}).get( + "cooldown_time", None + ) + header_cooldown = None if exception_headers is not None: header_cooldown = litellm.utils._get_retry_after_from_exception_header( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5e38402031..9d76f04d97 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -276,6 +276,8 @@ class CallTypes(Enum): responses = "responses" aresponses = "aresponses" alist_input_items = "alist_input_items" + llm_passthrough_route = "llm_passthrough_route" + allm_passthrough_route = "allm_passthrough_route" ######################################################### # Google GenAI Native Call Types @@ -2512,3 +2514,12 @@ class CallbacksByType(TypedDict): success: List[str] failure: List[str] success_and_failure: List[str] + + +CostResponseTypes = Union[ + ModelResponse, + TextCompletionResponse, + EmbeddingResponse, + ImageResponse, + TranscriptionResponse, +] diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 7b7ddf0a3d..4bab551d9f 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -3216,7 +3216,9 @@ def test_bedrock_meta_llama_function_calling(): print(response) -def test_bedrock_passthrough(): +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_bedrock_passthrough(sync_mode: bool): import litellm litellm._turn_on_debug() @@ -3238,20 +3240,86 @@ def test_bedrock_passthrough(): "anthropic_beta": ["claude-code-20250219"], } - response = litellm.llm_passthrough_route( - model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", - method="POST", - endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke", - data=data, - ) + if sync_mode: + response = litellm.llm_passthrough_route( + model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + method="POST", + endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke", + data=data, + ) + else: + response = await litellm.allm_passthrough_route( + model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + method="POST", + endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke", + data=data, + ) print(response.text) assert response.status_code == 200 -def test_bedrock_streaming_passthrough(): +@pytest.mark.asyncio +async def test_bedrock_converse__streaming_passthrough(monkeypatch): import litellm + from litellm.integrations.custom_logger import CustomLogger + import asyncio + + class MockCustomLogger(CustomLogger): + pass + + mock_custom_logger = MockCustomLogger() + monkeypatch.setattr(litellm, "callbacks", [mock_custom_logger]) + + litellm._turn_on_debug() + + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "text": "Write an article about impact of high inflation to GDP of a country" + } + ], + } + ], + "system": [{"text": "You are an economist with access to lots of data"}], + "inferenceConfig": {"maxTokens": 100, "temperature": 0.5}, + } + with patch.object(mock_custom_logger, "async_log_success_event") as mock_callback: + response = await litellm.allm_passthrough_route( + model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + method="POST", + endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/converse-stream", + data=data, + ) + async for chunk in response: + print(chunk) + + await asyncio.sleep(1) + + mock_callback.assert_called_once() + print(mock_callback.call_args.kwargs.keys()) + assert "response_cost" in mock_callback.call_args.kwargs["kwargs"] + assert mock_callback.call_args.kwargs["kwargs"]["response_cost"] > 0 + assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"] + + +@pytest.mark.asyncio +async def test_bedrock_streaming_passthrough(monkeypatch): + import litellm + import time + import asyncio + from unittest.mock import MagicMock + from litellm.integrations.custom_logger import CustomLogger + + class MockCustomLogger(CustomLogger): + pass + + mock_custom_logger = MockCustomLogger() + monkeypatch.setattr(litellm, "callbacks", [mock_custom_logger]) litellm._turn_on_debug() @@ -3272,12 +3340,72 @@ def test_bedrock_streaming_passthrough(): "anthropic_beta": ["claude-code-20250219"], } - response = litellm.llm_passthrough_route( - model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", - method="POST", - endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke-with-response-stream", - data=data, - stream=True, - ) + with patch.object(mock_custom_logger, "async_log_success_event") as mock_callback: + response = await litellm.allm_passthrough_route( + model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + method="POST", + endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke-with-response-stream", + data=data, + ) + async for chunk in response: + print(chunk) - assert response.status_code == 200 + await asyncio.sleep(1) + + mock_callback.assert_called_once() + # check standard logging payload created + print(mock_callback.call_args.kwargs.keys()) + assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"] + assert "response_cost" in mock_callback.call_args.kwargs["kwargs"] + + +@pytest.mark.asyncio +async def test_bedrock_streaming_passthrough(monkeypatch): + import litellm + import time + import asyncio + from unittest.mock import MagicMock + from litellm.integrations.custom_logger import CustomLogger + + class MockCustomLogger(CustomLogger): + pass + + mock_custom_logger = MockCustomLogger() + monkeypatch.setattr(litellm, "callbacks", [mock_custom_logger]) + + litellm._turn_on_debug() + + data = { + "max_tokens": 512, + "messages": [{"role": "user", "content": "Hey"}], + "system": [ + { + "type": "text", + "text": "Analyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title that captures the new topic. Format your response as a JSON object with two fields: 'isNewTopic' (boolean) and 'title' (string, or null if isNewTopic is false). Only include these fields, no other text.", + } + ], + "temperature": 0, + "metadata": { + "user_id": "5dd07c33da27e6d2968d94ea20bf47a7b090b6b158b82328d54da2909a108e84" + }, + "anthropic_version": "bedrock-2023-05-31", + "anthropic_beta": ["claude-code-20250219"], + } + + with patch.object(mock_custom_logger, "async_log_success_event") as mock_callback: + response = await litellm.allm_passthrough_route( + model="bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0", + method="POST", + endpoint="/model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke-with-response-stream", + data=data, + ) + async for chunk in response: + print(chunk) + + await asyncio.sleep(1) + + mock_callback.assert_called_once() + # check standard logging payload created + print(mock_callback.call_args.kwargs.keys()) + assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"] + assert "response_cost" in mock_callback.call_args.kwargs["kwargs"] diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 477549bd79..e076b66b7d 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -27,7 +27,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_TENANT_ID", raising=False) monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) - + with patch( "litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id" ) as mock_entra_token, patch( @@ -429,6 +429,8 @@ def test_select_azure_base_url_called(setup_mocks): "image_edit", "agenerate_content_stream", "agenerate_content", + "allm_passthrough_route", + "llm_passthrough_route", ] ], ) @@ -1033,8 +1035,8 @@ def test_with_existing_azure_ad_token_from_env(setup_mocks): # mock get_secret_str("AZURE_AD_TOKEN") to "test-token" with patch("litellm.llms.azure.common_utils.get_secret_str") as mock_get_secret_str: # Configure the mock to return "test-token" when called with "AZURE_AD_TOKEN" - mock_get_secret_str.side_effect = ( - lambda key: "test-token" if key == "AZURE_AD_TOKEN" else None + mock_get_secret_str.side_effect = lambda key: ( + "test-token" if key == "AZURE_AD_TOKEN" else None ) litellm_params = GenericLiteLLMParams()