diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 38874768ca8..52b50167190 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -420,6 +420,21 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk + _response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) + if ( + _chunk_type + in ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + ) + 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 + ) _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj) if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED: diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 41a3ded7022..e523b3771af 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,4 +1,5 @@ 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 @@ -1239,3 +1240,18 @@ 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 5e0e794d93e..c593946ac4a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -31,8 +31,15 @@ def _sse_event(payload: dict) -> bytes: def _mock_config() -> Mock: mock_config = Mock(spec=BaseResponsesAPIConfig) - mock_responses_api_response = Mock(spec=ResponsesAPIResponse) - mock_responses_api_response.id = "resp_ttft" + mock_responses_api_response = ResponsesAPIResponse( + id="resp_ttft", + created_at=0, + status="completed", + model="gpt-4o-mini", + object="response", + output=[], + usage=ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2), + ) def _transform(model, parsed_chunk, logging_obj): evt_type = parsed_chunk.get("type") @@ -54,6 +61,8 @@ def _make_iterator( sse_events: list[bytes], logging_obj: LiteLLMLoggingObj, trailing_error: Optional[Exception] = None, + config: Mock | None = None, + request_data: dict | None = None, ) -> ResponsesAPIStreamingIterator: async def aiter_bytes(): for evt in sse_events: @@ -68,10 +77,11 @@ def _make_iterator( return ResponsesAPIStreamingIterator( response=mock_response, model="gpt-4o-mini", - responses_api_provider_config=_mock_config(), + responses_api_provider_config=config or _mock_config(), logging_obj=logging_obj, litellm_metadata={}, custom_llm_provider="openai", + request_data=request_data, ) @@ -329,6 +339,86 @@ def test_run_post_success_hooks_does_not_report_generation_time_as_overhead(): assert "litellm_overhead_time_ms" not in iterator.completed_response._hidden_params +def _mock_config_with_completed_response(response: ResponsesAPIResponse) -> Mock: + mock_config = Mock(spec=BaseResponsesAPIConfig) + + def _transform(model, parsed_chunk, logging_obj): + evt_type = parsed_chunk.get("type") + if evt_type == "response.completed": + completed = Mock(spec=ResponseCompletedEvent) + completed.type = ResponsesAPIStreamEvents.RESPONSE_COMPLETED + completed.response = response + return completed + stub = Mock() + stub.type = evt_type + if evt_type == "response.output_text.delta": + stub.delta = parsed_chunk.get("delta") + return stub + + mock_config.transform_streaming_response.side_effect = _transform + return mock_config + + +def _responses_api_response_without_usage() -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_no_usage", + created_at=int(datetime(2025, 1, 1).timestamp()), + status="completed", + model="gpt-4o-mini", + object="response", + output=[], + usage=None, + ) + + +@pytest.mark.asyncio +async def test_completed_event_without_usage_gets_text_estimate(): + """A response.completed event carrying usage: null still bills: the + iterator estimates usage from the request input and generated text.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "count these input tokens please"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage is not None + assert usage.input_tokens > 0 + assert usage.output_tokens > 0 + assert usage.total_tokens == usage.input_tokens + usage.output_tokens + + +@pytest.mark.asyncio +async def test_completed_event_with_usage_is_left_untouched(): + """Provider-reported usage on response.completed wins over the estimate.""" + response = _responses_api_response_with_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "count these input tokens please"}, + ) + + async for _ in iterator: + pass + + usage = iterator.completed_response.response.usage + assert usage.input_tokens == 20 + assert usage.output_tokens == 60 + assert usage.total_tokens == 80 + + def _responses_api_response_with_usage() -> ResponsesAPIResponse: return ResponsesAPIResponse( id="resp_lit6427",