diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 2dfca3d308f..c23678c51ae 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -164,6 +164,7 @@ jobs: tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/rag_endpoints + tests/test_litellm/proxy/rerank_endpoints tests/test_litellm/proxy/realtime_endpoints tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/config_resolvers diff --git a/litellm/litellm_core_utils/json_fragment_accumulator.py b/litellm/litellm_core_utils/json_fragment_accumulator.py new file mode 100644 index 00000000000..81d18dd0119 --- /dev/null +++ b/litellm/litellm_core_utils/json_fragment_accumulator.py @@ -0,0 +1,97 @@ +import json +from typing import Final, cast # noqa: TID251 # raw_decode returns tuple[Any, int]; no cast-free unpack + + +class JSONFragmentAccumulator: + """ + Buffers a JSON value that arrives piecemeal over a stream (SSE data split + across TCP packets, one shard per network read, etc) without the O(n^2) + cost of repeated `buffer += fragment` string concatenation, and without + the O(n^2) cost of re-copying the unconsumed remainder on every peeled + value when one payload holds many concatenated JSON values. + + Fragments are appended to a list in O(1). The buffer is only rebuilt into + a single string, and only decoded, when a caller asks for a value via + `pop_next_value`, and `could_close_json` lets callers skip that rebuild + entirely for fragments that plainly cannot close a JSON value yet. Once + rebuilt, consumed values are dropped by advancing a cursor rather than + slicing a new string, so draining N concatenated values already sitting + in the buffer costs O(n) total, not O(n^2). + """ + + def __init__(self) -> None: + self._chunks: list[str] = [] # mutable-ok: O(1) append; string concat would copy the buffer each time + self._buffer: str = ( + "" # mutable-ok: lazily materialized join of _chunks, rebuilt only when _chunks is non-empty + ) + self._offset: int = 0 # mutable-ok: cursor past already-consumed values; avoids re-slicing on every pop + self._could_close: bool = False # mutable-ok: cached heuristic; rescanning past fragments was itself O(n^2) + + def __bool__(self) -> bool: + return bool(self._chunks) or self._offset < len(self._buffer) + + def append(self, fragment: str) -> None: + self._chunks.append(fragment) # mutable-ok: see __init__ + stripped: Final = fragment.rstrip() + if stripped: + self._could_close = stripped[-1] in ("}", "]") # mutable-ok: see __init__ + + def could_close_json(self) -> bool: + """ + Whether the buffer's logical last non-whitespace byte is "}" or "]", + i.e. whether a JSON value could plausibly be complete. Tracked + incrementally in `append` rather than rescanned here, so a run of + blank keepalive fragments (e.g. from a malformed upstream stream) + can't make this, or the join+parse it gates, cost O(n^2). + """ + return self._could_close + + def _materialize(self) -> None: + if not self._chunks: + return + unconsumed: Final = self._buffer[self._offset :] + self._buffer = unconsumed + "".join(self._chunks) # mutable-ok: merge pending fragments, once per append batch + self._offset = 0 # mutable-ok: see __init__ + self._chunks = [] # mutable-ok: see __init__ + + def pop_next_value(self) -> tuple[bool, object]: + """ + Attempt to decode one complete JSON value from the front of the + buffer. On success, advances a cursor past that value (keeping any + unconsumed tail, e.g. a second concatenated value, in place rather + than copying it) and returns (True, value). If the buffer is empty + or holds no complete value yet, it is left untouched and this + returns (False, None). + """ + self._materialize() + length: Final = len(self._buffer) + start = self._offset + while start < length and self._buffer[start].isspace(): + start += 1 + if start >= length: + self._offset = start # mutable-ok: see __init__ + return False, None + decoder: Final = json.JSONDecoder() + try: + raw_value: Final = decoder.raw_decode(self._buffer, start) + except json.JSONDecodeError: + return False, None + decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode returns tuple[Any, int] + self._offset = end_index # mutable-ok: see __init__ + if self._offset >= len(self._buffer): + self._buffer = "" # mutable-ok: see __init__ + self._offset = 0 # mutable-ok: see __init__ + self._could_close = False # mutable-ok: buffer is empty, nothing can close + return True, decoded + + def snapshot(self) -> str: + self._materialize() + return self._buffer[self._offset :] + + def set(self, value: str) -> None: + """Replace the buffer's contents with a single fragment.""" + self._chunks = [] # mutable-ok: see __init__ + self._buffer = value # mutable-ok: see __init__ + self._offset = 0 # mutable-ok: see __init__ + stripped: Final = value.rstrip() + self._could_close = bool(stripped) and stripped[-1] in ("}", "]") # mutable-ok: see __init__ diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index d9bb0d7abff..5842cc2679d 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -18,6 +18,7 @@ from litellm.anthropic_beta_headers_manager import ( ) from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -654,7 +655,7 @@ class ModelResponseIterator: # For handling partial JSON chunks from fragmentation # See: https://github.com/BerriAI/litellm/issues/17473 - self.accumulated_json: str = "" + self._json_buffer = JSONFragmentAccumulator() self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json" # Track current content block type to avoid emitting tool calls for non-tool blocks @@ -678,6 +679,14 @@ class ModelResponseIterator: self._current_server_tool_id: str | None = None self._container_id: str | None = None + @property + def accumulated_json(self) -> str: + return self._json_buffer.snapshot() + + @accumulated_json.setter + def accumulated_json(self, value: str) -> None: + self._json_buffer.set(value) + def check_empty_tool_call_args(self) -> bool: """ Check if the tool call block so far has been an empty string @@ -1149,31 +1158,39 @@ class ModelResponseIterator: container: Final = message_delta["delta"].get("container") return finish_reason, usage, container - def _handle_accumulated_json_chunk(self, data_str: str) -> ModelResponseStream | None: + def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None: """ Handle partial JSON chunks by accumulating them until valid JSON is received. This fixes network fragmentation issues where SSE data chunks may be split across TCP packets. See: https://github.com/BerriAI/litellm/issues/17473 + Mid-stream, defer parsing until the buffer's last byte can close a value: + attempting a parse after every fragment of one large object is O(n^2) and + holds the GIL, freezing the event loop. At end of stream (is_final) no more + data is coming, so drain whatever complete values remain regardless of the + trailing byte. + Args: data_str: The JSON string to parse (without "data:" prefix) + is_final: True when called from the end-of-stream drain, where the + trailing-byte heuristic no longer applies Returns: ModelResponseStream if JSON is complete, None if still accumulating """ - # Accumulate JSON data - self.accumulated_json += data_str + self._json_buffer.append(data_str) - # Try to parse the accumulated JSON - try: - data_json: Final = json.loads(self.accumulated_json) - self.accumulated_json = "" # Reset after successful parsing - return self.chunk_parser(chunk=data_json) - except json.JSONDecodeError: - # If it's not valid JSON yet, continue to the next chunk + if not is_final and not self._json_buffer.could_close_json(): return None + while True: + found, decoded = self._json_buffer.pop_next_value() + if not found: + return None + if isinstance(decoded, dict): + return self.chunk_parser(chunk=decoded) + def _parse_sse_data(self, str_line: str) -> ModelResponseStream | None: """ Parse SSE data line, handling both complete and partial JSON chunks. @@ -1209,13 +1226,10 @@ class ModelResponseIterator: chunk = self.response_iterator.__next__() except StopIteration: # If we have accumulated JSON when stream ends, try to parse it - if self.accumulated_json: - try: - data_json = json.loads(self.accumulated_json) - self.accumulated_json = "" - return self.chunk_parser(chunk=data_json) - except json.JSONDecodeError: - pass + if self._json_buffer: + result = self._handle_accumulated_json_chunk(data_str="", is_final=True) + if result is not None: + return result raise StopIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") @@ -1258,13 +1272,10 @@ class ModelResponseIterator: chunk = await self.async_response_iterator.__anext__() except StopAsyncIteration: # If we have accumulated JSON when stream ends, try to parse it - if self.accumulated_json: - try: - data_json = json.loads(self.accumulated_json) - self.accumulated_json = "" - return self.chunk_parser(chunk=data_json) - except json.JSONDecodeError: - pass + if self._json_buffer: + result = self._handle_accumulated_json_chunk(data_str="", is_final=True) + if result is not None: + return result raise StopAsyncIteration except ValueError as e: raise RuntimeError(f"Error receiving chunk from stream: {e}") diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index d9d62e0719f..492e0050cfe 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -214,6 +214,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: system text -> message(role=system, input_text) user text -> message(role=user, input_text) user image -> message(role=user, input_image) + user document -> message(role=user, input_file) user tool_result -> function_call_output assistant text -> message(role=assistant, output_text) assistant thinking -> reasoning @@ -268,6 +269,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: {"type": "input_image", "image_url": url}, block.get("prompt_cache_breakpoint") ) ) + elif btype == "document": + file_part = self._translate_anthropic_document_block_to_file_part(block) + if file_part: + user_parts.append( + with_prompt_cache_breakpoint(file_part, block.get("prompt_cache_breakpoint")) + ) elif btype == "tool_result": tool_use_id = block.get("tool_use_id", "") inner = block.get("content") diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 1cc72f265eb..a2a1acec80c 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -29,6 +29,7 @@ class BedrockRerankHandler(BaseAWSLLM): async def arerank( self, prepared_request: BedrockPreparedRequest, + logging_obj: LitellmLogging, timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ): @@ -40,6 +41,7 @@ class BedrockRerankHandler(BaseAWSLLM): headers=dict(prepared_request["prepped"].headers), data=prepared_request["body"], timeout=timeout, + logging_obj=logging_obj, ) response.raise_for_status() except httpx.HTTPStatusError as err: @@ -98,6 +100,7 @@ class BedrockRerankHandler(BaseAWSLLM): if _is_async: return self.arerank( prepared_request, + logging_obj=logging_obj, timeout=timeout, client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None, ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b7ddb55ae89..bb582f677be 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1203,6 +1203,7 @@ class BaseLLMHTTPHandler: headers=headers, data=json.dumps(request_data), timeout=timeout, + logging_obj=logging_obj, ) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py index 5f0ab5e56af..d06bcb22119 100644 --- a/litellm/llms/together_ai/chat/transformation.py +++ b/litellm/llms/together_ai/chat/transformation.py @@ -4,17 +4,24 @@ Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/compl Docs: https://docs.together.ai/docs/chat-overview """ -from collections.abc import Callable, Container -from typing import Final +from collections.abc import Callable, Container, Coroutine +from typing import ( + Final, + Literal, + cast, # noqa: TID251 # rebuilding a TypedDict minus keys has no checked spelling + overload, +) import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError +from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_function_calling, supports_response_schema from ...openai.chat.gpt_transformation import OpenAIGPTConfig TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call") +LITELLM_INTERNAL_ASSISTANT_FIELDS: Final = frozenset({"thinking_blocks", "provider_specific_fields"}) FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling" STRUCTURED_OUTPUTS_DOCS_URL: Final = "https://docs.together.ai/docs/inference/chat/structured-outputs" @@ -102,7 +109,50 @@ def _drop_response_format(passed_params: Container[str], model: str, drop_params ) +def _without_litellm_internal_fields(message: AllMessageValues) -> AllMessageValues: + if message["role"] != "assistant" or LITELLM_INTERNAL_ASSISTANT_FIELDS.isdisjoint(message): + return message + return cast( # cast-ok: rebuilding the same TypedDict minus internal keys loses the narrowed type + "AllMessageValues", + { # mutable-ok: TypedDict rebuild minus internal keys + key: value for key, value in message.items() if key not in LITELLM_INTERNAL_ASSISTANT_FIELDS + }, + ) + + class TogetherAIChatConfig(OpenAIGPTConfig): + @overload + def _transform_messages( + self, + messages: list[AllMessageValues], # mutable-ok: inherited contract + model: str, + is_async: Literal[True], + ) -> Coroutine[object, object, list[AllMessageValues]]: ... # mutable-ok: inherited contract + + @overload + def _transform_messages( + self, + messages: list[AllMessageValues], # mutable-ok: inherited contract + model: str, + is_async: Literal[False] = False, + ) -> list[AllMessageValues]: ... # mutable-ok: inherited contract + + def _transform_messages( + self, + messages: list[AllMessageValues], # mutable-ok: inherited contract + model: str, + is_async: bool = False, + ) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: # mutable-ok: inherited contract + """Together consumes replayed assistant `reasoning_content` (preserved thinking via + `chat_template_kwargs: {"clear_thinking": false}`), so it must stay in the payload; + only litellm-internal fields are stripped before sending.""" + stripped: Final = [ # mutable-ok: super() requires a list + _without_litellm_internal_fields(message) for message in messages + ] + if is_async: + return super()._transform_messages(stripped, model, is_async=True) + return super()._transform_messages(stripped, model, is_async=False) + def map_openai_params( self, non_default_params: dict, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index d298670aa7a..d12ba24eda4 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -23,6 +23,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE, DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, ) +from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( _encode_tool_call_id_with_signature, ) @@ -3087,7 +3088,7 @@ class ModelResponseIterator: self.streaming_response = streaming_response self.response = response self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json" - self.accumulated_json = "" + self._json_buffer = JSONFragmentAccumulator() self.sent_first_chunk = False self.logging_obj = logging_obj self.response_headers = response_headers or {} @@ -3095,6 +3096,14 @@ class ModelResponseIterator: self.cumulative_tool_call_index: int = 0 self.has_seen_tool_calls: bool = False + @property + def accumulated_json(self) -> str: + return self._json_buffer.snapshot() + + @accumulated_json.setter + def accumulated_json(self, value: str) -> None: + self._json_buffer.set(value) + @staticmethod def _check_streaming_error(chunk: dict) -> None: """Detect embedded errors (e.g. 429 RESOURCE_EXHAUSTED) in streaming chunks and raise VertexAIError.""" @@ -3298,8 +3307,8 @@ class ModelResponseIterator: return self.chunk_parser(chunk=json_chunk) def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]: - message: Final = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "" - self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip() + message: Final = (litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "").replace("\n\n", "") + self._json_buffer.append(message) # Mid-stream, defer parsing until the buffer's last byte can close a value: # attempting a parse after every fragment of one large object is O(n^2) and @@ -3307,27 +3316,23 @@ class ModelResponseIterator: # data is coming, so drain whatever complete values remain regardless of the # trailing byte, otherwise a complete leading value sitting behind a truncated # trailing one would be silently dropped. - if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"): + if not is_final and not self._json_buffer.could_close_json(): return None # Peel one complete JSON value from the front of the buffer and keep the # unconsumed tail. Running json.loads over the whole buffer would fail # forever once it held more than one concatenated value ("Extra data") while # never resetting the buffer, so the buffer grew without bound and pinned the - # core. raw_decode reports where the value ended, so concatenated values drain - # one call at a time. A leading non-dict value (never emitted by Gemini in - # practice) is consumed and skipped so it cannot block the dict values behind it. - decoder: Final = json.JSONDecoder() - while self.accumulated_json: - try: - raw_value = decoder.raw_decode(self.accumulated_json) - except json.JSONDecodeError: + # core. pop_next_value reports where the value ended, so concatenated values + # drain one call at a time. A leading non-dict value (never emitted by Gemini + # in practice) is consumed and skipped so it cannot block the dict values + # behind it. + while True: + found, decoded = self._json_buffer.pop_next_value() + if not found: return None - decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int] - self.accumulated_json = self.accumulated_json[end_index:].strip() if isinstance(decoded, dict): return self.chunk_parser(chunk=decoded) - return None def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]: try: @@ -3351,7 +3356,7 @@ class ModelResponseIterator: try: chunk: Final = self.response_iterator.__next__() except StopIteration: - if self.chunk_type == "accumulated_json" and self.accumulated_json: + if self.chunk_type == "accumulated_json" and self._json_buffer: result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True) if result is not None: return result @@ -3375,7 +3380,7 @@ class ModelResponseIterator: try: chunk: Final = await self.async_response_iterator.__anext__() except StopAsyncIteration: - if self.chunk_type == "accumulated_json" and self.accumulated_json: + if self.chunk_type == "accumulated_json" and self._json_buffer: result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True) if result is not None: return result diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 628e569e1b8..ed35691c6ec 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3546,6 +3546,8 @@ class SpendLogsMetadata(TypedDict): litellm_overhead_time_ms: float | None # LiteLLM overhead time in milliseconds attempted_retries: int | None # Number of retries attempted (0 = first attempt succeeded) max_retries: int | None # Max retries configured for this request + attempted_fallbacks: ReadOnly[int | None] # Number of fallbacks attempted (0 = primary model group served) + original_model_group: ReadOnly[str | None] # Model group requested before any fallbacks cost_breakdown: CostBreakdown | None # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) compression_savings: CompressionSavingsMetadata | None autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index 45b190c1f9d..dd5803796b7 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -90,12 +90,15 @@ async def rerank( fastapi_response.headers.update( ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, + call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None), model_id=model_id, cache_key=cache_key, api_base=api_base, version=version, + response_cost=hidden_params.get("response_cost", None), model_region=getattr(user_api_key_dict, "allowed_model_region", ""), request_data=data, + hidden_params=hidden_params, **additional_headers, ) ) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index d931d712a92..38da38ead2b 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -131,6 +131,8 @@ def _get_spend_logs_metadata( litellm_overhead_time_ms=None, attempted_retries=None, max_retries=None, + attempted_fallbacks=None, + original_model_group=None, cost_breakdown=None, compression_savings=None, autorouter_savings=autorouter_savings, diff --git a/litellm/router.py b/litellm/router.py index d07effd0d90..3d6230e0428 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -141,6 +141,7 @@ from litellm.router_utils.cooldown_handlers import ( is_advisor_orchestration_failure, ) from litellm.router_utils.fallback_event_handlers import ( + AttemptedFallbackTargets, _check_non_standard_fallback_format, get_fallback_model_group, run_async_fallback, @@ -6930,6 +6931,20 @@ class Router: If it fails after num_retries, fall back to another model group """ model_group: Final[str | None] = kwargs.get("model") + if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets): + _fallback_metadata_key: Final = _get_router_metadata_variable_name( + function_name=getattr(kwargs.get("original_function"), "__name__", None) + ) + _sibling_metadata_key: Final = ( + "metadata" if _fallback_metadata_key == "litellm_metadata" else "litellm_metadata" + ) + if isinstance(_sibling_metadata := kwargs.get(_sibling_metadata_key), dict): + _sibling_metadata.pop("attempted_fallbacks", None) + _sibling_metadata.pop("original_model_group", None) + if isinstance(_fallback_metadata := kwargs.get(_fallback_metadata_key), dict): + _fallback_metadata["attempted_fallbacks"] = 0 + if model_group is not None: + _fallback_metadata["original_model_group"] = model_group include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False) fallbacks: Final[list | None] = kwargs.get("fallbacks", self.fallbacks) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3c9a4097321..acdc7df5bd1 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -374,11 +374,13 @@ async def run_async_fallback( kwargs["model"] = mg elif isinstance(mg, dict): kwargs.update(mg) + fallback_depth = fallback_depth + 1 kwargs[metadata_variable_name] = { + "original_model_group": original_model_group, **(kwargs.get(metadata_variable_name) or {}), "model_group": kwargs.get("model", None), + "attempted_fallbacks": fallback_depth, } - fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks kwargs["attempted_targets"] = attempted diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 2a4f3acc896..b55756b94ea 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py b/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py new file mode 100644 index 00000000000..d16230afe3a --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py @@ -0,0 +1,236 @@ +import json +import time +from unittest.mock import patch + +from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator + + +def test_initial_state_is_empty(): + accumulator = JSONFragmentAccumulator() + assert not accumulator + assert accumulator.could_close_json() is False + assert accumulator.snapshot() == "" + + +def test_could_close_json_true_only_when_last_fragment_closes_a_value(): + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": ') + assert accumulator.could_close_json() is False + + accumulator.append("1}") + assert accumulator.could_close_json() is True + + +def test_could_close_json_looks_past_trailing_blank_fragments(): + """A whitespace-only or empty fragment (e.g. the flush call at end of + stream) must not mask a real closing byte in an earlier fragment.""" + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": 1}') + accumulator.append("") + accumulator.append(" \n") + assert accumulator.could_close_json() is True + + +def test_pop_next_value_on_empty_buffer_returns_false_without_touching_state(): + accumulator = JSONFragmentAccumulator() + found, value = accumulator.pop_next_value() + assert found is False + assert value is None + + +def test_pop_next_value_on_incomplete_buffer_leaves_buffer_untouched(): + accumulator = JSONFragmentAccumulator() + accumulator.append('{"candidates": [{"content":') + + found, value = accumulator.pop_next_value() + + assert found is False + assert value is None + assert accumulator.snapshot() == '{"candidates": [{"content":' + + +def test_pop_next_value_decodes_single_complete_object_and_clears_buffer(): + accumulator = JSONFragmentAccumulator() + accumulator.append('{"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}') + + found, value = accumulator.pop_next_value() + + assert found is True + assert value == {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]} + assert accumulator.snapshot() == "" + assert not accumulator + + +def test_pop_next_value_reassembles_a_value_split_across_many_fragments(): + obj = {"candidates": [{"content": {"parts": [{"text": "x" * 5000}]}}]} + blob = json.dumps(obj) + fragments = [blob[i : i + 37] for i in range(0, len(blob), 37)] + assert len(fragments) > 10, "need a genuinely multi-fragment payload" + + accumulator = JSONFragmentAccumulator() + found = False + value = None + for fragment in fragments: + accumulator.append(fragment) + if accumulator.could_close_json(): + found, value = accumulator.pop_next_value() + + assert found is True + assert value == obj + + +def test_pop_next_value_peels_one_value_and_keeps_remainder(): + """Two concatenated envelopes in the buffer must both surface, one per + call, instead of json.loads's "Extra data" failure wedging the buffer.""" + obj = '{"a": 1}' + accumulator = JSONFragmentAccumulator() + accumulator.append(obj + obj) + + first_found, first_value = accumulator.pop_next_value() + assert first_found is True + assert first_value == {"a": 1} + assert accumulator.snapshot() == obj, "second value must remain buffered" + + second_found, second_value = accumulator.pop_next_value() + assert second_found is True + assert second_value == {"a": 1} + assert not accumulator + + +def test_pop_next_value_skips_non_ascii_whitespace_between_concatenated_values(): + """A separator like U+00A0 (non-breaking space) between two concatenated + values must not strand the second value forever. `raw_decode` only skips + the narrow `json.decoder.WHITESPACE` set, so the accumulator's own + whitespace skip must be as tolerant as `str.strip()` was before this + class replaced it, not merely match `raw_decode`'s narrower set.""" + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": 1}' + "\xa0" + '{"a": 2}') + + first_found, first_value = accumulator.pop_next_value() + assert first_found is True + assert first_value == {"a": 1} + + second_found, second_value = accumulator.pop_next_value() + assert second_found is True, "the second value must not be permanently stranded" + assert second_value == {"a": 2} + assert not accumulator + + +def test_pop_next_value_advances_past_a_non_dict_leading_value(): + accumulator = JSONFragmentAccumulator() + accumulator.append("[1, 2]" + '{"a": 1}') + + first_found, first_value = accumulator.pop_next_value() + assert first_found is True + assert first_value == [1, 2] + + second_found, second_value = accumulator.pop_next_value() + assert second_found is True + assert second_value == {"a": 1} + + +def test_set_and_snapshot_roundtrip(): + accumulator = JSONFragmentAccumulator() + accumulator.set('{"a": 1}') + assert accumulator.snapshot() == '{"a": 1}' + assert accumulator + + accumulator.set("") + assert accumulator.snapshot() == "" + assert not accumulator + + +def test_append_never_calls_raw_decode(): # test-quality-ok: TQ002 - laziness contract has no caller-observable proxy other than spying on the stdlib call it must defer + """Appending must be O(1) bookkeeping only; the O(n) join+decode is + deferred entirely to pop_next_value.""" + accumulator = JSONFragmentAccumulator() + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: + for fragment in ['{"a":', " 1", "}"]: + accumulator.append(fragment) + assert spy.call_count == 0 + + +def test_pop_next_value_calls_raw_decode_at_most_once_per_value(): + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": 1}' * 3) + + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: + for _ in range(3): + found, _ = accumulator.pop_next_value() + assert found is True + assert spy.call_count == 3 + + +def test_accumulation_of_many_fragments_is_not_quadratic(): + """Regression guard: appending 1000 shards must stay O(n) total, not the + O(n^2) cost of repeated `buffer += fragment` string concatenation.""" + accumulator = JSONFragmentAccumulator() + shard = "x" * 2048 + + start = time.perf_counter() + for _ in range(1000): + accumulator.append(shard) + elapsed_ms = (time.perf_counter() - start) * 1000 + + assert elapsed_ms < 50, f"1000-fragment append took {elapsed_ms:.1f} ms (expected < 50 ms); O(n^2) regression?" + + +def test_draining_many_concatenated_values_is_not_quadratic(): + """ + Regression guard: peeling N JSON values already sitting in one buffer, + one pop_next_value() call per value with no new fragments in between, + must be O(n) total. Re-copying the shrinking remainder on every pop + (slicing a new string instead of advancing a cursor) makes total drain + time scale with the square of the buffer size. + + Uses a doubling ratio rather than an absolute ms budget so it isn't + flaky on a slower or busier CI runner: doubling the input should + roughly double an O(n) drain's time but roughly quadruple an O(n^2) + drain's time, and that ratio holds regardless of machine speed. + """ + + def drain_time_ms(n: int) -> float: + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": 1}' * n) + start = time.perf_counter() + drained = 0 + while True: + found, _ = accumulator.pop_next_value() + if not found: + break + drained += 1 + assert drained == n + return (time.perf_counter() - start) * 1000 + + small_ms = drain_time_ms(40_000) + large_ms = drain_time_ms(80_000) + + ratio = large_ms / max(small_ms, 0.001) + assert ratio < 3.0, ( + f"doubling drained values scaled time by {ratio:.2f}x ({small_ms:.1f} ms -> {large_ms:.1f} ms); " + "expected roughly 2x for O(n); O(n^2) regression?" + ) + + +def test_could_close_json_after_many_blank_fragments_is_not_quadratic(): + """ + Regression test: a hostile upstream can send malformed JSON that never + closes, followed by thousands of blank keepalive fragments. Rescanning + every blank fragment on each could_close_json() call would make N calls + cost O(n^2) total; it must be O(1) regardless of how many blank + fragments preceded it. + """ + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": ') # never closes + + start = time.perf_counter() + for _ in range(20_000): + accumulator.append("") + accumulator.could_close_json() + elapsed_ms = (time.perf_counter() - start) * 1000 + + assert accumulator.could_close_json() is False + assert elapsed_ms < 300, ( + f"20000 blank-fragment could_close_json() calls took {elapsed_ms:.1f} ms " + "(expected < 300 ms); quadratic rescan regression?" + ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index f6cd6ac6734..bd750a47f63 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -1008,6 +1008,143 @@ def test_multiple_partial_chunks_accumulation(): assert result3.choices[0].delta.content == "Hello" +def test_accumulated_json_partial_fragment_returns_none_without_parsing(): + """ + Regression test: before the shared JSONFragmentAccumulator, every partial + fragment triggered a `json.loads` attempt over the whole growing buffer, + unlike Vertex which already deferred parsing until the buffer could close. + A fragment that can't close a JSON value must not trigger a decode attempt. + """ + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) + iterator.chunk_type = "accumulated_json" + + with patch.object( + json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode + ) as spy: + result = iterator._handle_accumulated_json_chunk( + '{"type":"content_block_delta","index":0,"delta":' + ) + assert result is None + assert spy.call_count == 0, "incomplete buffer should not be parsed" + + +def test_accumulated_json_does_not_reparse_every_fragment(): + """ + Regression test for the O(n^2) json.loads-per-fragment anti-pattern: a + payload split across many fragments must be parsed ~once, not once per + fragment. + """ + text = "x" * 200_000 + blob = json.dumps( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}} + ) + fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] + assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" + + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) + iterator.chunk_type = "accumulated_json" + + parsed = None + with patch.object( + json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode + ) as spy: + for fragment in fragments: + out = iterator._handle_accumulated_json_chunk(fragment) + if out is not None: + parsed = out + parse_calls = spy.call_count + + assert parsed is not None, "the reassembled chunk must still parse" + assert parsed.choices[0].delta.content == text + assert parse_calls <= 2, ( + f"raw_decode was called {parse_calls} times for {len(fragments)} fragments; " + "the O(n^2) per-fragment re-parse has regressed" + ) + + +def test_accumulated_json_concatenated_envelopes_do_not_wedge(): + """ + Regression test: Anthropic's single `json.loads(self.accumulated_json)` + call raised "Extra data" on two concatenated envelopes and, since the + buffer was never reset on that failure, returned None forever while + growing without bound. The shared accumulator peels one value at a time + and keeps the remainder, so both values surface across two calls. + """ + obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) + iterator.chunk_type = "accumulated_json" + + first = iterator._handle_accumulated_json_chunk(obj + obj) + assert first is not None + assert first.choices[0].delta.content == "a" + + second = iterator._handle_accumulated_json_chunk("") + assert second is not None + assert second.choices[0].delta.content == "a" + + assert iterator.accumulated_json == "" + + +def test_accumulated_json_heuristic_passes_but_value_still_incomplete(): + """ + A buffer whose newest fragment ends in '}' can still be genuinely + incomplete (an inner object closed, the outer one didn't). The + heuristic must let the parse attempt through, and pop_next_value + finding nothing must propagate as None rather than raising. + """ + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True, json_mode=False + ) + iterator.chunk_type = "accumulated_json" + + result = iterator._handle_accumulated_json_chunk('{"type": {"nested": 1}') + assert result is None + + +def test_accumulated_json_setter_and_sync_end_of_stream_drain(): + """ + The accumulated_json setter and __next__'s StopIteration drain branch: + a buffered partial JSON must still parse and return when the + underlying stream ends, instead of being silently dropped. + """ + obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' + iterator = ModelResponseIterator( + streaming_response=iter([]), sync_stream=True, json_mode=False + ) + iterator.chunk_type = "accumulated_json" + iterator.accumulated_json = obj # exercises the setter + + result = iterator.__next__() + assert result is not None + assert result.choices[0].delta.content == "a" + + +def test_accumulated_json_async_end_of_stream_drain(): + """Async twin of the sync end-of-stream drain test: __anext__'s + StopAsyncIteration branch must also parse a buffered value.""" + import asyncio + + obj = '{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"a"}}' + iterator = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=False, json_mode=False + ) + iterator.chunk_type = "accumulated_json" + iterator.accumulated_json = obj + mock_async_iterator = MagicMock() + mock_async_iterator.__anext__ = AsyncMock(side_effect=StopAsyncIteration) + iterator.async_response_iterator = mock_async_iterator + + result = asyncio.run(iterator.__anext__()) + assert result is not None + assert result.choices[0].delta.content == "a" + + def test_web_search_tool_result_no_extra_tool_calls(): """ Test that web_search_tool_result blocks don't emit tool call chunks. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 44c956c8ee8..56f106e407c 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -1688,6 +1688,87 @@ class TestToolResultDocuments: ] +class TestUserContentDocuments: + """Documents in plain user content must survive translation (LIT-6144): each + document block becomes an input_file part of the user message, in block order, + exactly like image blocks become input_image parts. Untranslatable documents + are dropped without disturbing the surrounding parts.""" + + PDF_B64 = "JVBERi0xLjQKJSBQT05H" + PDF_DATA_URI = "data:application/pdf;base64,JVBERi0xLjQKJSBQT05H" + PDF_URL = "https://example.com/report.pdf" + EXPLICIT = {"mode": "explicit"} + + def _translate(self, user_content): + return _ADAPTER.translate_messages_to_responses_input([{"role": "user", "content": user_content}]) + + @staticmethod + def _user_content(items): + return next(item for item in items if item.get("type") == "message" and item.get("role") == "user")["content"] + + def _base64_document(self, **extra): + return { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": self.PDF_B64}, + **extra, + } + + def test_document_then_text_keeps_block_order(self): + content = self._user_content( + self._translate([self._base64_document(), {"type": "text", "text": "what does the pdf say?"}]) + ) + assert content == [ + {"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}, + {"type": "input_text", "text": "what does the pdf say?"}, + ] + + def test_document_title_becomes_filename(self): + content = self._user_content(self._translate([self._base64_document(title="quarterly-report.pdf")])) + assert content == [ + {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} + ] + + def test_url_document_becomes_file_url_part(self): + content = self._user_content( + self._translate([{"type": "document", "source": {"type": "url", "url": self.PDF_URL}}]) + ) + assert content == [{"type": "input_file", "file_url": self.PDF_URL}] + + def test_document_only_content_still_produces_user_message(self): + content = self._user_content(self._translate([self._base64_document()])) + assert content == [{"type": "input_file", "filename": "document.pdf", "file_data": self.PDF_DATA_URI}] + + def test_empty_base64_data_drops_only_the_document_part(self): + content = self._user_content( + self._translate( + [ + {"type": "text", "text": "still here"}, + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": ""}}, + ] + ) + ) + assert content == [{"type": "input_text", "text": "still here"}] + + def test_non_dict_source_drops_only_the_document_part(self): + content = self._user_content( + self._translate([{"type": "text", "text": "still here"}, {"type": "document", "source": self.PDF_URL}]) + ) + assert content == [{"type": "input_text", "text": "still here"}] + + def test_document_breakpoint_rides_on_the_file_part(self): + content = self._user_content( + self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)]) + ) + assert content == [ + { + "type": "input_file", + "filename": "document.pdf", + "file_data": self.PDF_DATA_URI, + "prompt_cache_breakpoint": self.EXPLICIT, + } + ] + + def _contains_key(value, key) -> bool: if isinstance(value, dict): return key in value or any(_contains_key(v, key) for v in value.values()) diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index b2a2046b131..253edad57e9 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -77,7 +77,7 @@ def test_bedrock_rerank_header_forwarding_sync(model): with ( patch.object(client, "post") as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -170,7 +170,7 @@ async def test_bedrock_rerank_header_forwarding_async(model): with ( patch.object(client, "post", new_callable=AsyncMock) as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -241,7 +241,7 @@ def test_bedrock_rerank_timeout_sync(): with ( patch.object(client, "post") as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -285,7 +285,7 @@ async def test_bedrock_rerank_timeout_async(): with ( patch.object(client, "post", new_callable=AsyncMock) as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -340,7 +340,7 @@ def test_bedrock_rerank_extra_headers_and_headers_merge(): with ( patch.object(client, "post") as mock_post, - patch( + patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info, ), @@ -400,3 +400,32 @@ def test_bedrock_rerank_extra_headers_and_headers_merge(): except Exception as e: pytest.fail(f"Failed to merge and forward headers: {str(e)}") + + +@pytest.mark.asyncio +async def test_bedrock_rerank_records_llm_api_duration(): + """The bedrock rerank handler must feed httpx timing into the logging obj, so the + proxy can emit x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank.""" + import httpx + + def handle(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=bedrock_rerank_response) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + with patch( # test-quality-ok: boto credential lookup needs live AWS; the HTTP boundary is already a MockTransport + "litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", + return_value=create_mock_credentials(), + ): + response = await litellm.arerank( + model="bedrock/arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", + query=test_query, + documents=test_documents, + top_n=3, + client=client, + aws_region_name="us-east-1", + ) + + assert response._hidden_params["litellm_overhead_time_ms"] is not None + assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"] diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index a93b14d45f3..694a01cda5f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2528,6 +2528,37 @@ def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monke assert _collect_ws_project_quota_callbacks() == (quota,) +@pytest.mark.asyncio +async def test_async_rerank_records_llm_api_duration(): + """arerank must feed the httpx timing into the logging obj, so the proxy can emit + x-litellm-overhead-duration-ms / x-litellm-timing-* on /rerank.""" + + def handle(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "rerank-1", + "results": [{"index": 0, "relevance_score": 0.9}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}}, + }, + ) + + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + + response = await litellm.arerank( + model="cohere/rerank-v3.5", + query="what is the capital of france", + documents=["paris", "berlin"], + top_n=1, + api_key="fake-key", + client=client, + ) + + assert response._hidden_params["litellm_overhead_time_ms"] is not None + assert response._hidden_params["_response_ms"] >= response._hidden_params["litellm_overhead_time_ms"] + + class _JSONBodyVideoConfig(OpenAIVideoConfig): def use_multipart_form_data(self) -> bool: return False diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py index 3848a9c7e6c..4573b1012a1 100644 --- a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -1,5 +1,6 @@ import json import logging +from collections.abc import Mapping, Sequence from unittest.mock import MagicMock import httpx @@ -326,6 +327,96 @@ def test_streaming_chunk_preserves_tool_call_index_and_id(): assert continuation["function"]["arguments"] == '{"city": "San' +REPLAYED_ASSISTANT_MESSAGE = { + "role": "assistant", + "content": "The digit sum is 11.", + "reasoning_content": "The secret number is 47. 4 + 7 = 11.", + "thinking_blocks": [{"type": "thinking", "thinking": "The secret number is 47.", "signature": ""}], + "provider_specific_fields": {"thinking_blocks": [{"type": "thinking", "thinking": "The secret number is 47."}]}, +} + +PRESERVED_THINKING_MESSAGES = [ + {"role": "user", "content": "Pick a secret two-digit number and tell me only its digit sum."}, + REPLAYED_ASSISTANT_MESSAGE, + {"role": "user", "content": "What was the secret number?"}, +] + + +def _assert_internal_fields_stripped_reasoning_kept(transformed_messages: Sequence[Mapping[str, object]]): + assistant_message = transformed_messages[1] + assert assistant_message["reasoning_content"] == REPLAYED_ASSISTANT_MESSAGE["reasoning_content"] + assert "thinking_blocks" not in assistant_message + assert "provider_specific_fields" not in assistant_message + assert assistant_message["content"] == REPLAYED_ASSISTANT_MESSAGE["content"] + assert transformed_messages[0] == PRESERVED_THINKING_MESSAGES[0] + assert transformed_messages[2] == PRESERVED_THINKING_MESSAGES[2] + + +def test_transform_request_keeps_reasoning_content_strips_internal_fields(): + request = TogetherAIChatConfig().transform_request( + model=REASONING_MODEL, + messages=[dict(message) for message in PRESERVED_THINKING_MESSAGES], + optional_params={}, + litellm_params={"custom_llm_provider": "together_ai"}, + headers={}, + ) + + _assert_internal_fields_stripped_reasoning_kept(request["messages"]) + + +async def test_async_transform_request_keeps_reasoning_content_strips_internal_fields(): + request = await TogetherAIChatConfig().async_transform_request( + model=REASONING_MODEL, + messages=[dict(message) for message in PRESERVED_THINKING_MESSAGES], + optional_params={}, + litellm_params={"custom_llm_provider": "together_ai"}, + headers={}, + ) + + _assert_internal_fields_stripped_reasoning_kept(request["messages"]) + + +def test_completion_sends_chat_template_kwargs_and_preserved_reasoning(): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + captured_requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together-preserved", + "object": "chat.completion", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "47"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + litellm.completion( + model=f"together_ai/{REASONING_MODEL}", + messages=[dict(message) for message in PRESERVED_THINKING_MESSAGES], + chat_template_kwargs={"clear_thinking": False}, + api_key="fake-key", + client=client, + ) + + request_body = json.loads(captured_requests[0].content) + assert request_body["chat_template_kwargs"] == {"clear_thinking": False} + assert "extra_body" not in request_body + _assert_internal_fields_stripped_reasoning_kept(request_body["messages"]) + + def test_together_ai_config_alias_points_at_chat_config(): assert litellm.TogetherAIConfig is litellm.TogetherAIChatConfig config = litellm.TogetherAIConfig(max_tokens=10) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 3d882deeb52..790a32506f2 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3002,8 +3002,11 @@ def test_accumulated_json_does_not_reparse_every_fragment(): The buffer only becomes a complete JSON object on the final fragment, so a correct implementation parses it ~once, not once per fragment. We assert the - full chunk still parses correctly AND that json.loads is not called on every - fragment (which is what made it quadratic). + full chunk still parses correctly AND that the buffer is not decoded on + every fragment (which is what made it quadratic). Post-migration to the + shared JSONFragmentAccumulator, decoding goes through + `json.JSONDecoder.raw_decode`, not `json.loads` (see the equivalent + Anthropic tests) so the spy targets that call, not `json.loads`. """ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, @@ -3024,7 +3027,9 @@ def test_accumulated_json_does_not_reparse_every_fragment(): assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" parsed = None - with patch("json.loads", wraps=json.loads) as spy: + with patch.object( + json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode + ) as spy: for fragment in fragments: out = iterator.handle_accumulated_json_chunk(chunk=fragment) if out is not None: @@ -3035,14 +3040,16 @@ def test_accumulated_json_does_not_reparse_every_fragment(): assert parsed.choices[0].delta.content == text, "content must be preserved intact" assert parse_calls <= 2, ( - f"json.loads was called {parse_calls} times for {len(fragments)} " + f"raw_decode was called {parse_calls} times for {len(fragments)} " "fragments; the O(n^2) per-fragment re-parse has regressed" ) def test_accumulated_json_partial_fragment_returns_none_without_parsing(): - """A fragment that cannot complete the JSON must not trigger a json.loads - parse of the whole growing buffer (issue #26181).""" + """A fragment that cannot complete the JSON must not trigger a decode + attempt over the whole growing buffer (issue #26181). Decoding goes + through `json.JSONDecoder.raw_decode` post-JSONFragmentAccumulator + migration, not `json.loads`.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator, ) @@ -3054,7 +3061,9 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing(): ) iterator.chunk_type = "accumulated_json" - with patch("json.loads", wraps=json.loads) as spy: + with patch.object( + json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode + ) as spy: result = iterator.handle_accumulated_json_chunk( chunk='{"candidates": [{"content": {"parts": [{"text": "partial' ) @@ -5552,3 +5561,21 @@ def test_accumulated_json_skips_non_dict_leading_value(): assert len(out) == 1 assert out[0].choices[0].delta.content == "a" + + +def test_accumulated_json_async_end_of_stream_drains_buffered_value(): + """Async twin of test_accumulated_json_end_of_stream_drains_all_buffered_values: + __anext__'s StopAsyncIteration branch must also parse a buffered value.""" + import asyncio + from unittest.mock import AsyncMock, MagicMock + + obj = '{"candidates":[{"content":{"parts":[{"text":"a"}]}}],"usageMetadata":{}}' + iterator = _accumulating_gemini_iterator() + iterator.accumulated_json = obj + mock_async_iterator = MagicMock() + mock_async_iterator.__anext__ = AsyncMock(side_effect=StopAsyncIteration) + iterator.async_response_iterator = mock_async_iterator + + result = asyncio.run(iterator.__anext__()) + assert result is not None + assert result.choices[0].delta.content == "a" diff --git a/tests/test_litellm/proxy/rerank_endpoints/__init__.py b/tests/test_litellm/proxy/rerank_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py new file mode 100644 index 00000000000..9f11ff6f20d --- /dev/null +++ b/tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py @@ -0,0 +1,120 @@ +""" +Tests for rerank_endpoints/endpoints.py response headers. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request, Response + +import litellm.proxy.common_request_processing as common_request_processing_mod +import litellm.proxy.proxy_server as proxy_server_mod +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.rerank_endpoints.endpoints import rerank +from litellm.types.utils import RerankResponse + +HIDDEN_PARAMS = { + "model_id": "deployment-1", + "api_base": "https://bedrock-agent-runtime.us-east-1.amazonaws.com", + "response_cost": 0.002, + "_response_ms": 1500.5, + "litellm_overhead_time_ms": 12.5, + "callback_duration_ms": 1.25, + "timing_llm_api_ms": 1488.0, + "timing_pre_processing_ms": 10.0, + "timing_post_processing_ms": 2.5, + "timing_message_copy_ms": 0.01, +} + + +def _build_request() -> Request: + body = json.dumps({"model": "rerank-model", "query": "q", "documents": ["a", "b"]}).encode() + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + scope={ + "type": "http", + "method": "POST", + "path": "/rerank", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + }, + receive=receive, + ) + + +async def _call_rerank(hidden_params: dict = HIDDEN_PARAMS) -> Response: + response = RerankResponse(id="rerank-1", results=[{"index": 0, "relevance_score": 0.9}]) + response._hidden_params = dict(hidden_params) + + fastapi_response = Response() + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging_obj.update_request_status = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs): + return {**kwargs["data"], "litellm_call_id": "call-123"} + + async def fake_route_request(**kwargs): + async def _call(): + return response + + return _call() + + with ( + patch.object(proxy_server_mod, "add_litellm_data_to_request", fake_add_litellm_data_to_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "route_request", fake_route_request), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "proxy_logging_obj", proxy_logging_obj), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "llm_router", MagicMock()), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + patch.object(proxy_server_mod, "version", "1.2.3"), # test-quality-ok: the rerank route reads these proxy_server module globals; no injection seam on the FastAPI handler + ): + await rerank( + request=_build_request(), + fastapi_response=fastapi_response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + return fastapi_response + + +@pytest.mark.asyncio +async def test_rerank_emits_latency_and_cost_headers(): + """/rerank must surface the same hidden_params-derived headers as /chat/completions.""" + fastapi_response = await _call_rerank() + + assert fastapi_response.headers["x-litellm-call-id"] == "call-123" + assert fastapi_response.headers["x-litellm-response-duration-ms"] == "1500.5" + assert fastapi_response.headers["x-litellm-overhead-duration-ms"] == "12.5" + assert fastapi_response.headers["x-litellm-callback-duration-ms"] == "1.25" + assert fastapi_response.headers["x-litellm-response-cost"] == "0.002" + + +@pytest.mark.asyncio +async def test_rerank_emits_detailed_timing_headers_when_enabled(): + """LITELLM_DETAILED_TIMING must also work on /rerank, not just /chat/completions.""" + with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test + fastapi_response = await _call_rerank() + + assert fastapi_response.headers["x-litellm-timing-llm-api-ms"] == "1488.0" + assert fastapi_response.headers["x-litellm-timing-pre-processing-ms"] == "10.0" + assert fastapi_response.headers["x-litellm-timing-post-processing-ms"] == "2.5" + assert fastapi_response.headers["x-litellm-timing-message-copy-ms"] == "0.01" + + +@pytest.mark.asyncio +async def test_rerank_emits_zero_response_cost_header(): + """A free deployment costs 0.0, which is a real cost and must not be dropped.""" + fastapi_response = await _call_rerank({**HIDDEN_PARAMS, "response_cost": 0.0}) + + assert fastapi_response.headers["x-litellm-response-cost"] == "0.0" + + +@pytest.mark.asyncio +async def test_rerank_omits_detailed_timing_headers_when_disabled(): + with patch.object(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False): # test-quality-ok: LITELLM_DETAILED_TIMING is a module constant; toggling it is the behavior under test + fastapi_response = await _call_rerank() + + assert "x-litellm-timing-llm-api-ms" not in fastapi_response.headers diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 2b062d9020d..8c15ead8983 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2629,7 +2629,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -2725,7 +2725,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -2819,7 +2819,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 843e1d1296f..5b9d591ea56 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3642,3 +3642,200 @@ def test_caller_forged_autorouter_savings_is_discarded(bucket): ) metadata = json.loads(payload["metadata"]) assert metadata["autorouter_savings"] is None + + +def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata(): + """ + Test that fallback info (attempted_fallbacks, original_model_group) from metadata + is included in the spend logs metadata JSON. + """ + kwargs = { + "model": "gpt-3.5-turbo", + "litellm_params": { + "metadata": { + "user_api_key": "sk-test-key", + "attempted_fallbacks": 2, + "original_model_group": "azure-gpt-fallback", + } + }, + "standard_logging_object": StandardLoggingPayload( + id="test-fallback-123", + call_type="completion", + stream=False, + response_cost=0.001, + status="success", + total_tokens=100, + prompt_tokens=50, + completion_tokens=50, + startTime=1234567890.0, + endTime=1234567891.0, + completionStartTime=None, + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ), + model="gpt-3.5-turbo", + model_id="model-123", + model_group="openai", + custom_llm_provider="openai", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_alias=None, + user_api_key_team_id=None, + user_api_key_org_id=None, + user_api_key_user_id=None, + user_api_key_team_alias=None, + spend_logs_metadata=None, + requester_ip_address=None, + requester_metadata=None, + user_api_key_end_user_id=None, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address=None, + messages=[], + response={}, + error_str=None, + model_parameters={}, + hidden_params=StandardLoggingHiddenParams( + model_id="model-123", + cache_key=None, + api_base="https://api.openai.com", + response_cost="0.001", + litellm_overhead_time_ms=None, + additional_headers=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ), + ), + } + + response_obj = { + "id": "test-response-retry", + "choices": [{"message": {"content": "Hello!"}}], + "usage": { + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + }, + } + + start_time = datetime.datetime.now(timezone.utc) + end_time = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + metadata = json.loads(payload["metadata"]) + + assert ( + metadata.get("attempted_fallbacks") == 2 + ), f"Expected attempted_fallbacks=2, got {metadata.get('attempted_fallbacks')}" + assert ( + metadata.get("original_model_group") == "azure-gpt-fallback" + ), f"Expected original_model_group=azure-gpt-fallback, got {metadata.get('original_model_group')}" + + +def test_get_logging_payload_handles_missing_fallback_info_gracefully(): + """ + Test that fallback fields are None when not present in metadata (backward compatibility). + """ + kwargs = { + "model": "gpt-3.5-turbo", + "litellm_params": { + "metadata": { + "user_api_key": "sk-test-key", + } + }, + "standard_logging_object": StandardLoggingPayload( + id="test-no-fallback-456", + call_type="completion", + stream=False, + response_cost=0.001, + status="success", + total_tokens=100, + prompt_tokens=50, + completion_tokens=50, + startTime=1234567890.0, + endTime=1234567891.0, + completionStartTime=None, + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-3.5-turbo", model_map_value=None + ), + model="gpt-3.5-turbo", + model_id="model-123", + model_group="openai", + custom_llm_provider="openai", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_alias=None, + user_api_key_team_id=None, + user_api_key_org_id=None, + user_api_key_user_id=None, + user_api_key_team_alias=None, + spend_logs_metadata=None, + requester_ip_address=None, + requester_metadata=None, + user_api_key_end_user_id=None, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address=None, + messages=[], + response={}, + error_str=None, + model_parameters={}, + hidden_params=StandardLoggingHiddenParams( + model_id="model-123", + cache_key=None, + api_base="https://api.openai.com", + response_cost="0.001", + litellm_overhead_time_ms=None, + additional_headers=None, + batch_models=None, + litellm_model_name=None, + usage_object=None, + ), + ), + } + + response_obj = { + "id": "test-response-no-fallback", + "choices": [{"message": {"content": "Hello!"}}], + "usage": { + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + }, + } + + start_time = datetime.datetime.now(timezone.utc) + end_time = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + metadata = json.loads(payload["metadata"]) + + assert ( + metadata.get("attempted_fallbacks") is None + ), "attempted_fallbacks should be None when not provided" + assert ( + metadata.get("original_model_group") is None + ), "original_model_group should be None when not provided" diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 24477248a8a..3e02838b88f 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -308,7 +308,11 @@ async def test_run_async_fallback_handles_explicitly_none_metadata(): metadata=None, ) - assert router.received_kwargs["metadata"] == {"model_group": "azure-group"} + assert router.received_kwargs["metadata"] == { + "model_group": "azure-group", + "attempted_fallbacks": 1, + "original_model_group": "openai-group", + } @pytest.mark.asyncio @@ -843,3 +847,45 @@ class TestRunAsyncFallbackTriggersCooldown: ) mock_trigger.assert_not_called() + + +@pytest.mark.asyncio +async def test_run_async_fallback_stamps_fallback_info_into_metadata(): + """Spend logs are built from the request metadata of the nested call, so the + fallback signal has to be stamped there before recursing.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + ) + + metadata = router.received_kwargs["metadata"] + assert metadata["attempted_fallbacks"] == 1 + assert metadata["original_model_group"] == "primary-model" + assert metadata["model_group"] == "fallback-model" + + +@pytest.mark.asyncio +async def test_run_async_fallback_preserves_original_model_group_on_nested_fallback(): + """A second-level fallback receives the first fallback target as its + original_model_group argument, so the first-stamped value must survive the hop.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["second-fallback"], + original_model_group="first-fallback", + original_exception=RuntimeError("first fallback failed"), + max_fallbacks=3, + fallback_depth=1, + metadata={"attempted_fallbacks": 1, "original_model_group": "primary-model"}, + ) + + metadata = router.received_kwargs["metadata"] + assert metadata["attempted_fallbacks"] == 2 + assert metadata["original_model_group"] == "primary-model" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 56fd6df446f..e59b00f53ff 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10325,3 +10325,168 @@ async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispa result = await wrapped(model="primary") assert result == "ok" mock_anthropic.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_stamps_zero_attempted_fallbacks(): + """A request served by the primary model group records attempted_fallbacks=0 and + the requested model group in metadata, mirroring the x-litellm-attempted-fallbacks header.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + metadata = {} + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + ) + + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "gpt-3.5-turbo" + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_stamps_route_bucket_not_litellm_metadata(): + """A chat completion carrying both metadata buckets gets stamped in the route's bucket + (metadata), matching where run_async_fallback rewrites, so the two never diverge.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + metadata = {} + litellm_metadata = {"client_key": "client_value"} + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + litellm_metadata=litellm_metadata, + ) + + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "gpt-3.5-turbo" + assert litellm_metadata["client_key"] == "client_value" + assert "attempted_fallbacks" not in litellm_metadata + assert "original_model_group" not in litellm_metadata + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_overrides_client_supplied_stamp_values(): + """Client-supplied attempted_fallbacks and original_model_group are replaced on entry, + so a reused metadata dict or a spoofed value cannot leak stale attribution into logs.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + metadata = {"attempted_fallbacks": 99, "original_model_group": "stale-group"} + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + ) + + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "gpt-3.5-turbo" + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_stamps_despite_forged_reentry_params(): + """A client injecting fallback_depth or a JSON-shaped attempted_targets via request + litellm params cannot skip the entry stamp; only the router's own in-process + AttemptedFallbackTargets instance marks a genuine re-entrant hop.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + metadata = {"attempted_fallbacks": 99, "original_model_group": "spoofed-group"} + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + fallback_depth=3, + attempted_targets={"keys": ["spoofed-group"]}, + ) + + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "gpt-3.5-turbo" + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_skips_stamp_on_genuine_reentrant_hop(): + """A re-entrant hop carrying the router's own AttemptedFallbackTargets instance keeps + the per-hop metadata that run_async_fallback wrote instead of resetting it to zero.""" + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + metadata = {"attempted_fallbacks": 1, "original_model_group": "prod-chat"} + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + attempted_targets=AttemptedFallbackTargets(keys=frozenset(("prod-chat",))), + ) + + assert metadata["attempted_fallbacks"] == 1 + assert metadata["original_model_group"] == "prod-chat" + + +@pytest.mark.asyncio +async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_bucket(): + """Spend logs read a truthy litellm_metadata dict in preference to metadata, so spoofed + stamp keys planted in the bucket the route does not own are removed on entry instead of + flowing into the spend log row.""" + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, + } + ] + ) + metadata = {} + litellm_metadata = { + "attempted_fallbacks": 99, + "original_model_group": "spoofed-group", + "client_key": "client_value", + } + + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hey"}], + metadata=metadata, + litellm_metadata=litellm_metadata, + ) + + assert "attempted_fallbacks" not in litellm_metadata + assert "original_model_group" not in litellm_metadata + assert litellm_metadata["client_key"] == "client_value" + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "gpt-3.5-turbo"