From bbbdccb82d88f3489a1a311109f86d87bb1782e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:30:13 -0700 Subject: [PATCH] fix(azure_ai): count relayed image prompt tokens without fetching the image --- basedpyright-code-budget.json | 8 +++--- .../streaming_chunk_builder_utils.py | 5 +++- litellm/main.py | 3 ++ .../openai_passthrough_logging_handler.py | 2 +- ruff-strict-budget.json | 4 +-- .../test_azure_passthrough_transformation.py | 28 +++++++++++++++++++ type-discipline-budget.json | 4 +-- 7 files changed, 44 insertions(+), 10 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 4dc7a0d8671..f63fb2a497a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15277 + "limit": 15276 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38266 }, "reportUnknownParameterType": { - "limit": 19580 + "limit": 19579 }, "reportUnknownVariableType": { - "limit": 29821 + "limit": 29817 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1beb06f52e7..1dca64f1a80 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -994,6 +994,7 @@ class ChunkProcessor: completion_output: str, messages: Sequence | None = None, reasoning_tokens: int | None = None, + use_default_image_token_count: bool = False, ) -> Usage: """ Calculate usage for the given chunks. @@ -1018,7 +1019,9 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages) + returned_usage.prompt_tokens = prompt_tokens or token_counter( + model=model, messages=messages, use_default_image_token_count=use_default_image_token_count + ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 diff --git a/litellm/main.py b/litellm/main.py index a7fa5560d30..d125f1c1d7f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8674,6 +8674,7 @@ def stream_chunk_builder( start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, + use_default_image_token_count: bool = False, ) -> ModelResponse | TextCompletionResponse | None: try: if chunks is None: @@ -8747,6 +8748,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=0, + use_default_image_token_count=use_default_image_token_count, ) setattr(response, "usage", usage) @@ -8924,6 +8926,7 @@ def stream_chunk_builder( completion_output=completion_output, messages=messages, reasoning_tokens=reasoning_tokens, + use_default_image_token_count=use_default_image_token_count, ) setattr(response, "usage", usage) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 9805c7d7fae..e21105f760c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -561,7 +561,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): # Build complete response from chunks complete_streaming_response: Final = litellm.stream_chunk_builder( - chunks=all_openai_chunks, messages=messages + chunks=all_openai_chunks, messages=messages, use_default_image_token_count=True ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index cfe737867f1..907317e81b1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2912 + "limit": 2910 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 114 + "limit": 113 }, "SIM113": { "limit": 3 diff --git a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py index d0657679983..b07b7915277 100644 --- a/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py +++ b/tests/test_litellm/llms/azure/passthrough/test_azure_passthrough_transformation.py @@ -153,6 +153,34 @@ def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_fr assert response.usage.completion_tokens > 0 +def test_azure_passthrough_streaming_chunks_count_remote_image_prompt_tokens_without_fetching_the_image(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": "http://127.0.0.1:9/doc.png"}}, + ], + } + ] + logging_obj = MagicMock() + logging_obj.model_call_details = {"request_data": {"messages": messages, "stream": True}} + + response = AzurePassthroughConfig().handle_logging_collected_chunks( + all_chunks=[chunk for chunk in _azure_chat_completion_chunks() if '"usage"' not in chunk], + litellm_logging_obj=logging_obj, + model="gpt-4.1-mini", + custom_llm_provider="azure", + endpoint="openai/deployments/gpt-4.1-mini/chat/completions", + ) + + assert isinstance(response, ModelResponse) + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter( + model="gpt-4.1-mini", messages=messages, use_default_image_token_count=True + ) + + def test_azure_passthrough_streaming_chunks_for_unknown_endpoint_return_none(): response = AzurePassthroughConfig().handle_logging_collected_chunks( all_chunks=_azure_chat_completion_chunks(), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 74b54148418..78cd84f5b35 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22172 + "limit": 22166 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16452 + "limit": 16446 }, "LIT011": { "limit": 5506