From de9aa48cd62e73b5b291d6c5a2a952dd35f09f00 Mon Sep 17 00:00:00 2001 From: kerry Date: Wed, 16 Sep 2026 01:02:29 +0000 Subject: [PATCH] fix(responses): count multimodal input and tool-call output in the streamed usage fallback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/responses/streaming_iterator.py | 38 ++++++++- litellm/responses/utils.py | 16 ---- .../responses/test_streaming_iterator.py | 78 ++++++++++++++++++- 3 files changed, 112 insertions(+), 20 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 52b50167190..6cb0331620f 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -32,6 +32,9 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ) from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.integrations.custom_logger import converted_stream_requested from litellm.types.llms.openai import ( @@ -257,6 +260,7 @@ class BaseResponsesAPIStreamingIterator: self._failure_handled = False # Track if failure handler has been called self._yielded_first_chunk = False self._generated_content = "" + self._generated_tool_arguments = "" self._completed_response_cached = False self._completed_response_logged = False self._completed_response_cache_hit: bool | None = None @@ -352,6 +356,10 @@ class BaseResponsesAPIStreamingIterator: _delta: Final = getattr(openai_responses_api_chunk, "delta", None) if isinstance(_delta, str): self._generated_content += _delta + elif _event_type == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA: + _args_delta: Final = getattr(openai_responses_api_chunk, "delta", None) + if isinstance(_args_delta, str): + self._generated_tool_arguments += _args_delta _stream_model_id: Final = _model_id_from_metadata(self.litellm_metadata) if _event_type in ( ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -432,8 +440,11 @@ class BaseResponsesAPIStreamingIterator: and _response_obj is not None and _response_obj.usage is None ): - _response_obj.usage = ResponseAPILoggingUtils.estimate_usage_from_text( - self.model or "", self.request_data.get("input"), self._generated_content + _response_obj.usage = _estimate_usage_from_text( + self.model or "", + self.request_data.get("input"), + self.request_data, + self._generated_content + self._generated_tool_arguments, ) _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) @@ -1347,6 +1358,29 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None: return None +def _estimate_usage_from_text( + model: str, + request_input: object, + responses_api_request: Mapping[str, object], + generated_text: str, +) -> ResponseAPIUsage: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped + input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union + responses_api_request=dict(responses_api_request), + ) + input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped + model=model, messages=messages + ) + output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped + model=model, text=generated_text, count_response_tokens=True + ) + return ResponseAPIUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ) + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index e523b3771af..41a3ded7022 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,5 +1,4 @@ import base64 -import json import re from collections.abc import Iterable, Mapping, Sequence from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload @@ -1240,18 +1239,3 @@ class ResponseAPILoggingUtils: setattr(chat_usage, "cost", response_api_usage.cost) return chat_usage - - @staticmethod - def estimate_usage_from_text(model: str, request_input: object, generated_text: str) -> ResponseAPIUsage: - input_text: Final = request_input if isinstance(request_input, str) else json.dumps(request_input, default=str) - input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped - model=model, text=input_text - ) - output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped - model=model, text=generated_text, count_response_tokens=True - ) - return ResponseAPIUsage( - input_tokens=input_tokens, - output_tokens=output_tokens, - total_tokens=input_tokens + output_tokens, - ) diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c593946ac4a..c3721353d42 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -5,12 +5,13 @@ completion_start_time = end_time.""" import json from datetime import datetime -from typing import Optional +from typing import Final, Optional from unittest.mock import Mock, patch import httpx import pytest +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( @@ -351,8 +352,10 @@ def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock return completed stub = Mock() stub.type = evt_type - if evt_type == "response.output_text.delta": + if "delta" in parsed_chunk: stub.delta = parsed_chunk.get("delta") + if "item" in parsed_chunk: + stub.item = parsed_chunk.get("item") return stub mock_config.transform_streaming_response.side_effect = _transform @@ -718,3 +721,74 @@ async def test_streaming_logging_copy_keeps_client_usage_when_response_fails_val assert isinstance(client_usage, ResponseAPIUsage) assert client_usage.input_tokens == 29 assert client_usage.cost == pytest.approx(0.0001) + + +@pytest.mark.asyncio +async def test_completed_event_without_usage_counts_tool_call_arguments(): + """A function-call-only stream still bills output tokens: streamed + function_call_arguments deltas feed the text estimate.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event( + { + "type": "response.output_item.added", + "item": {"type": "function_call", "name": "get_weather", "call_id": "call_1"}, + } + ), + _sse_event( + { + "type": "response.function_call_arguments.delta", + "delta": '{"location": "San Francisco", "unit": "celsius"}', + } + ), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "what is the weather in san francisco"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens + + +@pytest.mark.asyncio +async def test_completed_event_without_usage_counts_multimodal_input_as_messages(): + """Multimodal request input is counted as chat messages, not as a JSON blob: + a huge base64 image must not inflate the estimated input tokens.""" + image_input: Final = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "what is in this image"}, + { + "type": "input_image", + "image_url": "data:image/png;base64," + "A" * 4000, + }, + ], + } + ] + json_count: Final = litellm.token_counter(model="gpt-4o-mini", text=json.dumps(image_input)) + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "it is a cat"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": image_input}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.input_tokens < json_count / 2