fix(azure_ai): count prompt tokens for streaming relays that carry no usage chunk

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

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:
         """

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),
         )

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,

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
 

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

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(),

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