From 61ee2f5ecbb602fb42a8f8b9d2f4316bdc0678a4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:00:57 -0700 Subject: [PATCH 1/6] fix(anthropic): carry user-content document blocks through the /v1/messages responses bridge --- .../responses_adapters/transformation.py | 7 ++ .../test_responses_adapters_transformation.py | 81 +++++++++++++++++++ 2 files changed, 88 insertions(+) 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/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()) From 1ba66b1555fdbe088e8cc6ce55d67ad7dae4cd63 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:48:58 -0700 Subject: [PATCH 2/6] fix(together_ai): strip internal thinking fields from outbound messages, pin chat_template_kwargs passthrough --- .../llms/together_ai/chat/transformation.py | 54 ++++++++++- .../test_together_ai_chat_transformation.py | 90 +++++++++++++++++++ 2 files changed, 142 insertions(+), 2 deletions(-) diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py index 3162a34f1b9..6bf53ee348b 100644 --- a/litellm/llms/together_ai/chat/transformation.py +++ b/litellm/llms/together_ai/chat/transformation.py @@ -4,18 +4,25 @@ 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 Container +from collections.abc import Container, Coroutine from types import MappingProxyType -from typing import Final +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 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"}) PLAIN_TEXT_RESPONSE_FORMAT: Final = MappingProxyType({"type": "text"}) FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling" @@ -61,7 +68,50 @@ def _tool_params_to_drop(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 get_supported_openai_params(self, model: str) -> list: supports_fc: Final = _function_calling_verdict(model) supported_params: Final = super().get_supported_openai_params(model) 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 0b9fd5364f9..110d93cbc0d 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 @@ -268,6 +268,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: list): + 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 = [] + + 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) From bb22742025065c10baba94b7d587ab000c6174ff Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:54:25 -0700 Subject: [PATCH 3/6] fix(rerank): emit latency and cost headers on /rerank (#35419) * fix(rerank): emit latency and cost headers on /rerank Thread the logging object into the rerank httpx calls and pass hidden_params through to get_custom_headers, so x-litellm-overhead-duration-ms, x-litellm-response-duration-ms, x-litellm-response-cost, x-litellm-call-id and the LITELLM_DETAILED_TIMING x-litellm-timing-* headers show up on rerank like they do on chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(rerank): keep zero response cost in the /rerank cost header Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: assign the new rerank endpoint tests to the proxy-endpoints shard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: suppress TQ008 on the rerank header tests with reasons Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: milan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin --- .github/workflows/test-unit.yml | 1 + litellm/llms/bedrock/rerank/handler.py | 3 + litellm/llms/custom_httpx/llm_http_handler.py | 1 + litellm/proxy/rerank_endpoints/endpoints.py | 3 + .../test_bedrock_rerank_header_forwarding.py | 39 +++++- .../custom_httpx/test_llm_http_handler.py | 31 +++++ .../proxy/rerank_endpoints/__init__.py | 0 .../proxy/rerank_endpoints/test_endpoints.py | 120 ++++++++++++++++++ 8 files changed, 193 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/rerank_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/rerank_endpoints/test_endpoints.py 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/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/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/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/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 From 8a61b286d1c28792c4b1574a7acda246da156c1e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:03:43 -0700 Subject: [PATCH 4/6] test(together_ai): annotate preserved-thinking test helper and capture list --- .../together_ai/chat/test_together_ai_chat_transformation.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 110d93cbc0d..f6d4b73f27a 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 @@ -283,7 +284,7 @@ PRESERVED_THINKING_MESSAGES = [ ] -def _assert_internal_fields_stripped_reasoning_kept(transformed_messages: list): +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 @@ -320,7 +321,7 @@ async def test_async_transform_request_keeps_reasoning_content_strips_internal_f def test_completion_sends_chat_template_kwargs_and_preserved_reasoning(): from litellm.llms.custom_httpx.http_handler import HTTPHandler - captured_requests = [] + captured_requests: list[httpx.Request] = [] def respond(request: httpx.Request) -> httpx.Response: captured_requests.append(request) From 6684bb343b0b10a78835cb592a980b9bd61d18ee Mon Sep 17 00:00:00 2001 From: Deepanshu Lulla Date: Tue, 25 Aug 2026 19:15:47 -0400 Subject: [PATCH 5/6] perf(streaming): add shared JSONFragmentAccumulator for Vertex and Anthropic (#36610) * perf(streaming): add shared JSONFragmentAccumulator for Vertex and Anthropic Vertex's handle_accumulated_json_chunk and Anthropic's _handle_accumulated_json_chunk each independently accumulated SSE fragments into a JSON envelope with self.accumulated_json += fragment. Because the attribute holds a live reference, CPython copies the whole prior buffer on every fragment, making buffer assembly O(n^2) in total payload size. Anthropic additionally had no completeness heuristic at all and called json.loads on the whole buffer after every fragment, and could wedge forever on two concatenated envelopes. Add JSONFragmentAccumulator in litellm_core_utils/: fragments append to a list in O(1), a could_close_json heuristic lets callers skip the join+parse entirely until a value could plausibly be complete, and pop_next_value peels one JSON value off the front of the buffer at a time using json.JSONDecoder().raw_decode, keeping any unconsumed remainder instead of failing on concatenated values. Migrate both providers onto it; Anthropic's __next__/__anext__ end-of-stream handlers now delegate to _handle_accumulated_json_chunk(is_final=True) instead of duplicating the parse-and-reset logic inline. Fixes #31861. * test(streaming): close diff-coverage gaps in JSONFragmentAccumulator migration Codecov flagged 11 uncovered lines in the migration: the accumulated_json setter, __next__/__anext__'s end-of-stream drain branches in both providers, and the pop_next_value "not found" path when a buffer's newest fragment ends in "}" but is genuinely incomplete (an inner object closed, the outer one didn't). Add targeted tests for each. * fix(streaming): make JSONFragmentAccumulator's completeness heuristic O(1) could_close_json rescanned every preceding blank fragment on each call, so a hostile upstream that sends malformed JSON (never closing) followed by many blank keepalive fragments could drive that scan, and the join+parse it gates, to O(n^2) total. Track the last non-blank fragment's trailing byte incrementally in append/pop_next_value/set instead of rescanning the buffer. Reported by automated review on PR #36610. * fix(streaming): make JSONFragmentAccumulator.pop_next_value O(1) per call pop_next_value previously rebuilt the full remaining string and sliced a new remainder on every call, so draining N concatenated JSON values already sitting in one buffer cost O(n^2) total. Replace the rebuild-and- slice with a materialize-once cursor: pending fragments are joined into the buffer only when new ones have arrived since the last pop, and consumed values are dropped by advancing an offset instead of copying the remaining string. * test(streaming): make JSONFragmentAccumulator drain regression test CI-stable The 80k-value drain test used an absolute ms budget that flaked on a busier CI runner (233.5ms vs a 150ms budget calibrated on a quiet machine). Replace it with a doubling-ratio check: draining twice as many concatenated values should take roughly 2x as long for O(n), not ~4x for O(n^2), and that ratio holds regardless of machine speed. * test(streaming): suppress TQ002 on the append-laziness spy test A test-quality gate (TQ002: don't assert only that a mock was called) landed upstream since this branch's last rebase and now flags test_append_never_calls_raw_decode. The test verifies append() defers all decoding to pop_next_value, which has no caller-observable proxy other than spying on the stdlib call it must avoid making. * test(vertex): spy on raw_decode instead of json.loads in accumulator regression tests Post-migration to the shared JSONFragmentAccumulator, Vertex's decode path goes through json.JSONDecoder.raw_decode, not json.loads. The two O(n^2)/partial-fragment regression tests still patched json.loads, which that path never calls, so both passed unconditionally regardless of whether the underlying implementation regressed. Verified by simulating an eager-reparse regression: both tests now fail against it and pass against the correct implementation. * fix(streaming): widen JSONFragmentAccumulator's whitespace skip to match str.strip() pop_next_value's whitespace skip only matched json.decoder.WHITESPACE's ASCII set, narrower than str.strip() (Unicode-aware) which the O(1)-cursor rewrite replaced. A non-ASCII separator like U+00A0 between two concatenated JSON values on one SSE line made raw_decode fail on it, and the buffer never advanced past that byte again, permanently stranding everything after it for the rest of the stream. Use str.isspace() to match str.strip()'s tolerance. --------- Co-authored-by: Deepanshu --- .../json_fragment_accumulator.py | 97 +++++++ litellm/llms/anthropic/chat/handler.py | 61 +++-- .../vertex_and_google_ai_studio_gemini.py | 39 +-- .../test_json_fragment_accumulator.py | 236 ++++++++++++++++++ .../chat/test_anthropic_chat_handler.py | 137 ++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 41 ++- 6 files changed, 562 insertions(+), 49 deletions(-) create mode 100644 litellm/litellm_core_utils/json_fragment_accumulator.py create mode 100644 tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py 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/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/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/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" From 75bf9f945260d5bb152d93fbd118ef47caf7d1c2 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 25 Aug 2026 16:22:01 -0700 Subject: [PATCH 6/6] fix(router): persist attempted_fallbacks and original_model_group into spend logs metadata (#38107) --- litellm/proxy/_types.py | 2 + .../spend_tracking/spend_tracking_utils.py | 2 + litellm/router.py | 15 ++ .../router_utils/fallback_event_handlers.py | 4 +- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- .../test_spend_management_endpoints.py | 6 +- .../test_spend_tracking_utils.py | 197 ++++++++++++++++++ .../test_fallback_event_handlers.py | 48 ++++- tests/test_litellm/test_router.py | 165 +++++++++++++++ 9 files changed, 435 insertions(+), 6 deletions(-) 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/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/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"