From 9ea9aa2e7b1ef97f80e2141f0bc299498e351143 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 00:00:23 -0700 Subject: [PATCH] fix(azure_ai): count prompt tokens for streaming relays that carry no usage chunk --- basedpyright-code-budget.json | 8 ++--- .../streaming_chunk_builder_utils.py | 4 +-- .../llms/azure/passthrough/transformation.py | 20 ++++++++++- litellm/main.py | 4 +-- .../openai_passthrough_logging_handler.py | 7 ++-- ruff-strict-budget.json | 4 +-- .../test_azure_passthrough_transformation.py | 33 +++++++++++++++++-- type-discipline-budget.json | 4 +-- 8 files changed, 67 insertions(+), 17 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef1bd675c5e..4dc7a0d8671 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5570 }, "reportMissingTypeArgument": { - "limit": 15278 + "limit": 15277 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38276 + "limit": 38271 }, "reportUnknownParameterType": { - "limit": 19581 + "limit": 19580 }, "reportUnknownVariableType": { - "limit": 29825 + "limit": 29821 }, "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 0e01577b20e..1beb06f52e7 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -209,7 +209,7 @@ def apply_grounding_request_counts( class ChunkProcessor: - def __init__(self, chunks: list, messages: list | None = None): + def __init__(self, chunks: list, messages: Sequence | None = None): self.chunks = self._sort_chunks(chunks) self.messages = messages self.first_chunk = chunks[0] @@ -992,7 +992,7 @@ class ChunkProcessor: chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, - messages: list | None = None, + messages: Sequence | None = None, reasoning_tokens: int | None = None, ) -> Usage: """ diff --git a/litellm/llms/azure/passthrough/transformation.py b/litellm/llms/azure/passthrough/transformation.py index 7654958fc0b..0b8b216faae 100644 --- a/litellm/llms/azure/passthrough/transformation.py +++ b/litellm/llms/azure/passthrough/transformation.py @@ -1,8 +1,9 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional import httpx from httpx import Response +from pydantic import BaseModel, ValidationError from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.azure.common_utils import BaseAzureLLM @@ -17,6 +18,22 @@ if TYPE_CHECKING: from litellm.types.utils import CostResponseTypes +class RelayedChatRequest(BaseModel): + messages: Sequence[Mapping[str, object]] | None = None + + +class RelayedCallDetails(BaseModel): + request_data: RelayedChatRequest | None = None + + +def _relayed_messages(litellm_logging_obj: Logging) -> Sequence[Mapping[str, object]] | None: + try: + details: Final = RelayedCallDetails.model_validate(litellm_logging_obj.model_call_details) + except ValidationError: + return None + return details.request_data.messages if details.request_data else None + + class AzurePassthroughConfig(BasePassthroughConfig): def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in request_data @@ -137,4 +154,5 @@ class AzurePassthroughConfig(BasePassthroughConfig): all_chunks=all_chunks, litellm_logging_obj=litellm_logging_obj, model=model, + messages=_relayed_messages(litellm_logging_obj), ) diff --git a/litellm/main.py b/litellm/main.py index 2929790f2bd..a7fa5560d30 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8553,7 +8553,7 @@ def config_completion(**kwargs): ) -def stream_chunk_builder_text_completion(chunks: list, messages: list | None = None) -> TextCompletionResponse: +def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse: id: Final = chunks[0]["id"] object: Final = chunks[0]["object"] created: Final = chunks[0]["created"] @@ -8670,7 +8670,7 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o def stream_chunk_builder( chunks: list, - messages: list | None = None, + messages: Sequence | None = None, start_time=None, end_time=None, logging_obj: Optional["Logging"] = None, 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 a39ea342f0a..9805c7d7fae 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 @@ -4,7 +4,7 @@ OpenAI Passthrough Logging Handler Handles cost tracking and logging for OpenAI passthrough endpoints, specifically /chat/completions. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final from urllib.parse import urlparse @@ -516,6 +516,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): all_chunks: Sequence[str], litellm_logging_obj: LiteLLMLoggingObj, model: str, + messages: Sequence[Mapping[str, object]] | None = None, ) -> ModelResponse | TextCompletionResponse | None: """ Builds complete response from raw chunks for OpenAI streaming responses. @@ -559,7 +560,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): return None # Build complete response from chunks - complete_streaming_response: Final = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, messages=messages + ) return complete_streaming_response diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 3354610bfcf..cfe737867f1 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2914 + "limit": 2912 }, "C401": { "limit": 8 @@ -201,7 +201,7 @@ "limit": 310 }, "SIM103": { - "limit": 115 + "limit": 114 }, "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 68fbc99b5d5..d0657679983 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 @@ -3,6 +3,8 @@ from unittest.mock import MagicMock import httpx +import litellm + from litellm.llms.azure.passthrough.transformation import AzurePassthroughConfig from litellm.types.utils import ModelResponse @@ -101,8 +103,15 @@ def _sse_line(payload: dict) -> str: def _azure_chat_completion_chunks() -> list[str]: head = {"id": "chatcmpl-abc123", "object": "chat.completion.chunk", "created": 1700000000, "model": "gpt-4.1-mini"} return [ - _sse_line({**head, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}]}), - _sse_line({**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]}), + _sse_line( + { + **head, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hello!"}, "finish_reason": None}], + } + ), + _sse_line( + {**head, "choices": [{"index": 0, "delta": {"content": " How can I assist?"}, "finish_reason": None}]} + ), _sse_line({**head, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}), _sse_line({**head, "choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}}), "data: [DONE]", @@ -124,6 +133,26 @@ def test_azure_passthrough_streaming_chat_chunks_build_the_complete_response(): assert response.usage.completion_tokens == 8 +def test_azure_passthrough_streaming_chunks_without_usage_count_prompt_tokens_from_the_relayed_request(): + messages = [{"role": "user", "content": "Say hi in three words"}] + 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.choices[0].message.content == "Hello! How can I assist?" + assert response.usage.prompt_tokens > 0 + assert response.usage.prompt_tokens == litellm.token_counter(model="gpt-4.1-mini", messages=messages) + assert response.usage.completion_tokens > 0 + + 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 30bd580c4e2..74b54148418 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22178 + "limit": 22172 }, "LIT002": { "limit": 26745 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16458 + "limit": 16452 }, "LIT011": { "limit": 5506