fix(responses): recount tokens when a streamed response completes without usage

litellm/responses/streaming_iterator.py

@@ -32,6 +32,9 @@
 )
 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 @@ def __init__(
         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 @@ def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None:
                     _delta: Final = getattr(openai_responses_api_chunk, "delta", None)
                     if isinstance(_delta, str):
                         self._generated_content += _delta
+                elif _event_type in _TOOL_ARGUMENTS_DELTA_EVENTS:
+                    _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,
@@ -420,7 +428,32 @@ def _process_chunk(self, chunk: str) -> ResponsesAPIStreamingResponse | None:
                     openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
                 ):
                     self.completed_response = openai_responses_api_chunk
-                    _stamp_responses_usage_cost(getattr(openai_responses_api_chunk, "response", None), self.logging_obj)
+                    _response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
+                    _typed_response: Final[ResponsesAPIResponse | None] = (
+                        ResponsesAPIResponse.model_construct(**_response_obj)  # pyright: ignore[reportUnknownArgumentType]  # the model_constructed terminal event leaves response as an untyped dict
+                        if isinstance(_response_obj, dict)
+                        else _response_obj
+                        if isinstance(_response_obj, ResponsesAPIResponse)
+                        else None
+                    )
+                    if (
+                        _typed_response is not None
+                        and _chunk_type
+                        in (
+                            openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
+                            openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
+                        )
+                        and _typed_response.usage is None
+                    ):
+                        _typed_response.usage = _estimate_usage_safely(
+                            self.model or "",
+                            self.request_data.get("input"),
+                            self.request_data,
+                            self._generated_content + self._generated_tool_arguments,
+                        )
+                    if _typed_response is not None and _typed_response is not _response_obj:
+                        openai_responses_api_chunk.response = _typed_response  # pyright: ignore[reportAttributeAccessIssue]  # reached only on the dict path, which only response-carrying terminal events produce
+                    _stamp_responses_usage_cost(_typed_response, self.logging_obj)
 
                     if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
                         self._handle_logging_failed_response()
@@ -1332,6 +1365,56 @@ def _usage_as_model(usage: object) -> ResponseAPIUsage | None:
         return None
 
 
+_TOOL_ARGUMENTS_DELTA_EVENTS: Final = frozenset(
+    {
+        ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA,
+        ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA,
+        ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
+    }
+)
+
+
+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 _estimate_usage_safely(
+    model: str,
+    request_input: object,
+    responses_api_request: Mapping[str, object],
+    generated_text: str,
+) -> ResponseAPIUsage | None:
+    try:
+        return _estimate_usage_from_text(
+            model=model,
+            request_input=request_input,
+            responses_api_request=responses_api_request,
+            generated_text=generated_text,
+        )
+    except Exception as e:
+        verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e)
+        return None
+
+
 def _stamp_responses_usage_cost(
     response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
 ) -> None: