diff --git a/litellm/main.py b/litellm/main.py index af4cdbcb49a..c2c450967a5 100644 --- a/litellm/main.py +++ b/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: diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3eea47bcd5a..9408898ad89 100644 --- a/tests/test_litellm/test_main.py +++ b/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 diff --git a/tests/test_litellm/test_stream_chunk_builder_citations.py b/tests/test_litellm/test_stream_chunk_builder_citations.py new file mode 100644 index 00000000000..55fce727ab2 --- /dev/null +++ b/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"}]