fix(main): report response_cost and Anthropic citations from stream_chunk_builder

litellm/main.py

@@ -8590,6 +8590,29 @@ def stream_chunk_builder_text_completion(chunks: list, messages: list | None = N
     return TextCompletionResponse(**response)
 
 
+def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None:
+    usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None)
+    if isinstance(usage_cost, (int, float)):
+        return float(usage_cost)
+    if logging_obj is not None:
+        return None
+    provider_hint: Final = response._hidden_params.get(  # pyright: ignore[reportPrivateUsage]  # no public accessor
+        "custom_llm_provider"
+    )
+    try:
+        return litellm.completion_cost(completion_response=response, custom_llm_provider=provider_hint)
+    except Exception:
+        return None
+
+
+def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
+    response_cost: Final = _stream_builder_response_cost(response, logging_obj)
+    if response_cost is None:
+        return
+    hidden_params: Final = response._hidden_params  # pyright: ignore[reportPrivateUsage]  # no public accessor
+    hidden_params["response_cost"] = response_cost
+
+
 def stream_chunk_builder(
     chunks: list,
     messages: list | None = None,
@@ -8690,6 +8713,8 @@ def stream_chunk_builder(
                     "cost",
                     logging_obj._response_cost_calculator(result=response),
                 )
+            _set_stream_builder_response_cost(response, logging_obj)
+
             processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
             return response
 
@@ -8814,18 +8839,24 @@ def stream_chunk_builder(
         ]
 
         if len(provider_specific_chunks) > 0:
-            combined_provider_fields: Final[dict[str, object]] = {}
-            for chunk in provider_specific_chunks:
-                fields = chunk["choices"][0]["delta"]["provider_specific_fields"]
-                if isinstance(fields, dict):
-                    for key, value in fields.items():
-                        if key not in combined_provider_fields:
-                            combined_provider_fields[key] = value
-                        elif isinstance(value, list) and isinstance(combined_provider_fields[key], list):
-                            # For lists like web_search_results, take the last (most complete) one
-                            combined_provider_fields[key] = value
-                        else:
-                            combined_provider_fields[key] = value
+            provider_field_dicts: Final = tuple(
+                fields
+                for chunk in provider_specific_chunks
+                for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
+                if isinstance(fields, dict)
+            )
+            streamed_citations: Final = tuple(
+                fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
+            )
+            citation_fields: Final = (
+                {"citations": [list(streamed_citations)]} if streamed_citations else {}  # mutable-ok: JSON dict field
+            )
+            combined_provider_fields: Final = {  # mutable-ok: Message.provider_specific_fields is a plain dict field
+                key: value
+                for fields in (citation_fields, *provider_field_dicts)
+                for key, value in fields.items()
+                if key != "citation"
+            }
 
             if combined_provider_fields:
                 _choice = cast(Choices, response.choices[0])
@@ -8862,6 +8893,8 @@ def stream_chunk_builder(
         if litellm.include_cost_in_streaming_usage and logging_obj is not None:
             setattr(usage, "cost", logging_obj._response_cost_calculator(result=response))
 
+        _set_stream_builder_response_cost(response, logging_obj)
+
         processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
         return response
     except Exception as e:

tests/test_litellm/test_main.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+from datetime import datetime
 import contextlib
 import copy
 import json
@@ -21,7 +22,8 @@ import litellm
 from litellm import main as litellm_main
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
-from litellm.types.utils import Usage
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
+from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
 
 
 async def _async_fake_bedrock_image_details(image_url):
@@ -3071,3 +3073,93 @@ async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking(
     assert expected_cost > 0
     assert speech_event.response_cost == pytest.approx(expected_cost)
     assert speech_event.logged_response_cost == pytest.approx(expected_cost)
+
+
+def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-cost",
+        created=1724900000,
+        model=model,
+        object="chat.completion.chunk",
+        choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content=content, role="assistant"))],
+    )
+
+
+def test_stream_chunk_builder_sets_hidden_response_cost_for_known_model():
+    chunks: Final = [
+        _stream_builder_text_chunk("gpt-4o", "Hello "),
+        _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"),
+    ]
+
+    response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}])
+
+    assert response is not None
+    prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage)
+    expected_cost: Final = prompt_cost + completion_cost
+    assert expected_cost > 0
+    assert response._hidden_params["response_cost"] == pytest.approx(expected_cost)
+
+
+def test_stream_chunk_builder_unknown_model_leaves_response_cost_unset():
+    chunks: Final = [
+        _stream_builder_text_chunk("totally-unknown-model-xyz", "Hello "),
+        _stream_builder_text_chunk("totally-unknown-model-xyz", "world.", finish_reason="stop"),
+    ]
+
+    response: Final = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}])
+
+    assert response is not None
+    assert response._hidden_params.get("response_cost") is None
+    assert response.choices[0].message.content == "Hello world."
+
+
+def _stream_builder_logging_obj() -> LiteLLMLogging:
+    logging_obj: Final = LiteLLMLogging(
+        model="gpt-4o",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=True,
+        call_type="completion",
+        start_time=datetime.now(),
+        litellm_call_id="test-call-id",
+        function_id="test-function-id",
+    )
+    logging_obj.update_environment_variables(
+        model="gpt-4o",
+        user=None,
+        optional_params={},
+        litellm_params={"custom_llm_provider": "openai"},
+    )
+    return logging_obj
+
+
+def test_stream_chunk_builder_reports_streaming_usage_cost_when_enabled(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
+    chunks: Final = [
+        _stream_builder_text_chunk("gpt-4o", "Hello "),
+        _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"),
+    ]
+
+    response: Final = litellm.stream_chunk_builder(
+        chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj()
+    )
+
+    assert response is not None
+    usage_cost: Final = getattr(response.usage, "cost", None)
+    assert usage_cost is not None
+    assert usage_cost > 0
+    assert response._hidden_params["response_cost"] == pytest.approx(usage_cost)
+
+
+def test_stream_chunk_builder_defers_cost_to_logging_obj_when_usage_cost_absent(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
+    chunks: Final = [
+        _stream_builder_text_chunk("gpt-4o", "Hello "),
+        _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"),
+    ]
+
+    response: Final = litellm.stream_chunk_builder(
+        chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj()
+    )
+
+    assert response is not None
+    assert response._hidden_params.get("response_cost") is None

tests/test_litellm/test_stream_chunk_builder_citations.py

@@ -0,0 +1,85 @@
+from typing import Final
+
+from litellm import stream_chunk_builder
+from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
+
+_CITATION_ONE: Final = {
+    "type": "char_location",
+    "cited_text": "The grass is green.",
+    "document_index": 0,
+    "document_title": "My Document",
+    "start_char_index": 0,
+    "end_char_index": 20,
+}
+_CITATION_TWO: Final = {
+    "type": "char_location",
+    "cited_text": "The sky is blue.",
+    "document_index": 0,
+    "document_title": "My Document",
+    "start_char_index": 20,
+    "end_char_index": 36,
+}
+
+
+def _chunk(delta: Delta, finish_reason: str | None = None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-citations",
+        created=1724900000,
+        model="claude-opus-5",
+        object="chat.completion.chunk",
+        choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=delta)],
+    )
+
+
+def test_stream_chunk_builder_collects_every_streamed_citation():
+    chunks: Final = [
+        _chunk(Delta(content="The grass is green", role="assistant")),
+        _chunk(Delta(content="", provider_specific_fields={"citation": _CITATION_ONE})),
+        _chunk(Delta(content=" and the sky is blue.")),
+        _chunk(Delta(content="", provider_specific_fields={"citation": _CITATION_TWO})),
+        _chunk(Delta(content=""), finish_reason="stop"),
+    ]
+
+    response: Final = stream_chunk_builder(chunks=chunks)
+
+    assert response is not None
+    fields: Final = response.choices[0].message.provider_specific_fields
+    assert fields is not None
+    assert fields["citations"] == [[_CITATION_ONE, _CITATION_TWO]]
+    assert "citation" not in fields
+    assert response.choices[0].message.content == "The grass is green and the sky is blue."
+
+
+def test_stream_chunk_builder_keeps_other_provider_fields_alongside_citations():
+    thinking_blocks: Final = [{"type": "thinking", "thinking": "checking the document", "signature": "sig"}]
+    chunks: Final = [
+        _chunk(Delta(content="Green.", role="assistant")),
+        _chunk(Delta(content="", provider_specific_fields={"citation": _CITATION_ONE})),
+        _chunk(Delta(content="", provider_specific_fields={"thinking_blocks": thinking_blocks})),
+        _chunk(Delta(content=""), finish_reason="stop"),
+    ]
+
+    response: Final = stream_chunk_builder(chunks=chunks)
+
+    assert response is not None
+    fields: Final = response.choices[0].message.provider_specific_fields
+    assert fields is not None
+    assert fields["citations"] == [[_CITATION_ONE]]
+    assert fields["thinking_blocks"] == thinking_blocks
+    assert "citation" not in fields
+
+
+def test_stream_chunk_builder_without_citation_deltas_sets_no_citations_key():
+    chunks: Final = [
+        _chunk(Delta(content="Hello", role="assistant")),
+        _chunk(Delta(content="", provider_specific_fields={"web_search_results": [{"url": "https://example.com"}]})),
+        _chunk(Delta(content=""), finish_reason="stop"),
+    ]
+
+    response: Final = stream_chunk_builder(chunks=chunks)
+
+    assert response is not None
+    fields: Final = response.choices[0].message.provider_specific_fields
+    assert fields is not None
+    assert "citations" not in fields
+    assert fields["web_search_results"] == [{"url": "https://example.com"}]