diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py index 0abf7c11280..71e2742fa2e 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -37,6 +37,7 @@ from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -690,18 +691,21 @@ class ModelEndpoint: async def _invoke_sdk( self, route: str, - kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs + kwargs: dict[str, Any], # mutable-ok: SDK call kwargs built by _sdk_kwargs ) -> object: if route == ROUTE_MESSAGES: return await litellm.anthropic.messages.acreate(**kwargs) if route == ROUTE_CHAT: - if kwargs.get("stream"): - stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field - kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion - "include_usage": True, - **stream_options, + if not kwargs.get("stream"): + return await litellm.acompletion(**kwargs) + stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field + return await litellm.acompletion( + **{ + **kwargs, + "stream_options": {"include_usage": True, **stream_options}, + STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options, } - return await litellm.acompletion(**kwargs) + ) return await litellm.aresponses(**kwargs) async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 473666e5ba1..45349bb0e5a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -138,6 +138,7 @@ from litellm.types.interactions import ( InteractionsAPIResponse, InteractionsAPIStreamingResponse, ) +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.openai import ( AllMessageValues, Batch, @@ -698,6 +699,7 @@ class Logging(LiteLLMLoggingBaseClass): self._pre_call_session_id: str = session_id_var.get() _sid: Final = kwargs.get("litellm_session_id") if kwargs else None self.litellm_session_id: str = str(_sid) if _sid else "" + self.stream_usage_injected: Final = kwargs is not None and kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True # supports_correlation_logging is False for calls originating from the # sync client entry point (wrapper() in utils.py): a plain OS thread # has no per-call context isolation the way an asyncio Task does, and diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ec426d6da5..b210085c5cc 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -342,6 +342,10 @@ class CustomStreamWrapper: self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False + self.caller_requested_stream_usage: Final = ( + self.check_send_stream_usage(getattr(logging_obj, "stream_options", None)) + and getattr(logging_obj, "stream_usage_injected", False) is not True + ) self.tool_call = False # keep track of the returned chunks - used for calculating the input/output tokens for stream options self.chunks: list[ModelResponseStream] = [] @@ -2070,7 +2074,7 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: complete_streaming_response: Final = self._build_complete_streaming_response() - response = self.model_response_creator() + response = self._new_usage_chunk() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2222,27 +2226,10 @@ class CustomStreamWrapper: ) if _has_usage: - # Store a copy ONLY when usage stripping below will mutate - # the chunk. For non-usage chunks (vast majority), store - # directly to avoid expensive model_copy() per chunk. - self.chunks.append(processed_chunk.model_copy()) - - # Strip usage from the outgoing chunk so it's not sent twice - # (once in the chunk, once in _hidden_params). - obj_dict = processed_chunk.model_dump() - if "usage" in obj_dict: - del obj_dict["usage"] - processed_chunk = self.model_response_creator( - chunk=obj_dict, - hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage - Mapping[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ), - ) - is_empty = is_model_response_stream_empty( - model_response=cast(ModelResponseStream, processed_chunk) - ) - if is_empty: + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: continue + processed_chunk = usage_free_chunk else: # No usage data — safe to store directly without copying self.chunks.append(processed_chunk) @@ -2273,13 +2260,18 @@ class CustomStreamWrapper: if processed_chunk is None: continue - choice = processed_chunk.choices[0] - if isinstance(choice, StreamingChoices): - self._accumulate_streamed_delta(choice.delta) + if processed_chunk.choices: + choice = processed_chunk.choices[0] + if isinstance(choice, StreamingChoices): + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) - # RETURN RESULT - self.chunks.append(processed_chunk) - return processed_chunk + if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None: + self.chunks.append(processed_chunk) + return processed_chunk + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: + continue + return usage_free_chunk except (StopAsyncIteration, StopIteration): return await self._finalize_completed_stream(cache_hit=cache_hit) except httpx.TimeoutException as e: # if httpx read timeout error occues @@ -2306,7 +2298,7 @@ class CustomStreamWrapper: async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: complete_streaming_response: Final = await self._async_build_complete_streaming_response() - response: Final = self.model_response_creator() + response: Final = self._new_usage_chunk() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2463,6 +2455,22 @@ class CustomStreamWrapper: return False return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + def _new_usage_chunk(self) -> ModelResponseStream: + usage_chunk: Final = self.model_response_creator() + if self.caller_requested_stream_usage: + usage_chunk.choices = [] + return usage_chunk + + def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + self.chunks.append(chunk.model_copy()) + usage_free_chunk: Final = self.model_response_creator( + chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"}, + hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage + Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) + ), + ) + return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk + def _accumulate_streamed_delta(self, delta: object) -> None: # Shared by the sync, async, and non-aiohttp iteration sites so answer # text and the disqualifying-content latch stay in step across all three. diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 59db23820fc..cf616873a2e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -265,6 +265,7 @@ from litellm.proxy.litellm_pre_call_utils import ( sanitize_for_log, ) from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -551,22 +552,29 @@ def model_deployments_support_stream_options( return all(_litellm_model_supports_stream_options(m) for m in candidate_models) +_STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD) + + def stream_usage_tracking_updates( data: Mapping[str, object], general_settings: Mapping[str, object], route_type: str, supports_stream_options: Callable[[], bool], ) -> Mapping[str, object]: - scrub: Final = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {} + scrub: Final = {flag: False for flag in _STREAM_USAGE_FLAGS if flag in data} if data.get("stream", False) is not True: return scrub always_include: Final = general_settings.get("always_include_stream_usage") stream_options: Final = data.get("stream_options") if always_include is True: if "stream_options" not in data: - return {**scrub, "stream_options": {"include_usage": True}} + return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True} if isinstance(stream_options, dict) and "include_usage" not in stream_options: - return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return { + **scrub, + "stream_options": {**stream_options, "include_usage": True}, + STREAM_USAGE_INJECTED_FIELD: True, + } return scrub if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data: return {**scrub, "stream_options": {"include_usage": False}} @@ -580,6 +588,7 @@ def stream_usage_tracking_updates( return { "stream_options": {**merged_stream_options, "include_usage": True}, "_litellm_strip_stream_usage": True, + STREAM_USAGE_INJECTED_FIELD: True, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index 06ea3ce4b20..7fb87080cd1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -494,7 +494,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): sse: Final = AnthropicSSERestorer(step) events: Final = ResponsesStreamRestorer(step, rehydrate) carries: Final[CarryWindows] = {} - last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. + last_chunk = None # rebind-ok: tracks the most recent chunk with choices for the final flush. async for chunk in response: if isinstance(chunk, (bytes, str)): @@ -506,7 +506,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail): yield event continue restored_chunk = detached(chunk) - last_chunk = restored_chunk + if getattr(restored_chunk, "choices", None): + last_chunk = restored_chunk for choice in attribute_items(restored_chunk, "choices"): await self._restore_choice(choice, carries, session_id) yield restored_chunk diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 09f53cecb72..d34787cd5e9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -268,7 +268,7 @@ def sanitize_for_log(value: object) -> str: from litellm.router import Router from litellm.secret_managers.main import get_secret_bool -from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES +from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES, STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.router import BaggageMappedTagsStamp, holds_api_key_secret_reference from litellm.types.services import ServiceTypes @@ -2370,7 +2370,15 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) + frozenset( + { + "secret_fields", + "proxy_server_request", + "litellm_logging_obj", + VECTOR_STORE_CALLER_METADATA_KWARG, + STREAM_USAGE_INJECTED_FIELD, + } + ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8d6c9dd925b..b48b5ab1c1e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -976,6 +976,7 @@ from litellm.secret_managers.main import ( ) from litellm.tracing.config import is_lens_tracing_enabled from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ( AnthropicMessagesRequest, AnthropicResponse, @@ -18219,6 +18220,7 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) + data.pop(STREAM_USAGE_INJECTED_FIELD, None) # Include original request and headers in the data data["proxy_server_request"] = { diff --git a/litellm/router.py b/litellm/router.py index a18ad7c3a97..1ce1e0a5a70 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -361,10 +361,11 @@ from litellm.router_utils.routing_groups import ( validate_routing_strategy, ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch +from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD, RoutingStrategyName from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolParam, @@ -4673,7 +4674,8 @@ class Router: deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) ): kwargs.pop(key, None) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=self.default_litellm_params) + defaults: Final = _LITELLM_PARAMS_ADAPTER.validate_python(self.default_litellm_params) + self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=defaults) model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] @@ -4817,6 +4819,19 @@ class Router: data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), ) + deployment_litellm_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + if router_config_requests_stream_usage( + _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + ( + _LITELLM_PARAMS_ADAPTER.validate_python( + deployment_litellm_params + if isinstance(deployment_litellm_params, Mapping) + else deployment_litellm_params.model_dump(exclude_none=True) + ), + defaults, + ), + ): + kwargs[STREAM_USAGE_INJECTED_FIELD] = True self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): diff --git a/litellm/router_utils/stream_usage_injection.py b/litellm/router_utils/stream_usage_injection.py new file mode 100644 index 00000000000..946543d9b83 --- /dev/null +++ b/litellm/router_utils/stream_usage_injection.py @@ -0,0 +1,17 @@ +from collections.abc import Mapping, Sequence + + +def router_config_requests_stream_usage( + request_kwargs: Mapping[str, object], configured_params: Sequence[Mapping[str, object]] +) -> bool: + if request_kwargs.get("stream") is not True or "stream_options" in request_kwargs: + return False + return any(_includes_usage(params.get("stream_options")) for params in configured_params) + + +def _includes_usage(stream_options: object) -> bool: + match stream_options: + case {"include_usage": True}: + return True + case _: + return False diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 5f318cc3c69..0ed156524be 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -54,6 +54,7 @@ RoutingStrategyName: TypeAlias = Literal[ TRUSTED_CALLBACK_VARS_FIELD: Final = "litellm_trusted_callback_vars" ADDRESSED_RESPONSE_ID_FIELD: Final = "_litellm_addressed_response_id" +STREAM_USAGE_INJECTED_FIELD: Final = "_litellm_stream_usage_injected" WIRE_NAME: Final = "wire_name" @@ -397,6 +398,7 @@ class ProxyRequestState: trusted_callback_vars: Mapping[str, str] | None = field(default=None, metadata=wire(TRUSTED_CALLBACK_VARS_FIELD)) addressed_response_id: str | None = field(default=None, metadata=wire(ADDRESSED_RESPONSE_ID_FIELD)) strip_stream_usage: bool | None = field(default=None, metadata=wire("_litellm_strip_stream_usage")) + stream_usage_injected: bool | None = field(default=None, metadata=wire(STREAM_USAGE_INJECTED_FIELD)) client_side_timeout: bool | None = None model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 379eddb3155..1aad9baf06a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7962,12 +7962,16 @@ class TextCompletionStreamWrapper: text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception - delta: Final = chunk["choices"][0]["delta"] - text_choices["text"] = delta["content"] - text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] - response["choices"] = [text_choices] + + if chunk["choices"]: + delta: Final = chunk["choices"][0]["delta"] + text_choices["text"] = delta["content"] + text_choices["reasoning_content"] = delta.get("reasoning_content") + text_choices["index"] = chunk["choices"][0]["index"] + text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + response["choices"] = [text_choices] + else: + response["choices"] = [] # only pass usage when stream_options["include_usage"] is True if self.stream_options and self.stream_options.get("include_usage", False) is True: diff --git a/tests/integration/compatibility/test_openai_consumer.py b/tests/integration/compatibility/test_openai_consumer.py index a49f20f4711..6150995723b 100644 --- a/tests/integration/compatibility/test_openai_consumer.py +++ b/tests/integration/compatibility/test_openai_consumer.py @@ -459,7 +459,6 @@ async def test_openai_clients_parse_text_completion_streams_and_choices(gateway: def test_text_completion_usage_chunk_has_empty_choices(gateway: Gateway) -> None: - pytest.skip("BUG: the /v1/completions include_usage chunk carries choices [{index: 0}] instead of []") identity: Final = f"cmpl-usage-{uuid.uuid4().hex}" prompt: Final = f"raw-stream {identity}" with wire_server(_text_provider(identity, {prompt})) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index b46503703c7..58785c89033 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -312,7 +312,7 @@ def _chat_caller_stream_usage_event(event: dict[str, JsonValue], model: str) -> "object": event["object"], "created": 1, "model": model, - "choices": [{"index": 0, "delta": {}}], + "choices": [], "usage": {**_json_object_value(event["usage"]), "cost": 4.05e-6}, } diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index 2badf810b6a..d4f9533fcdc 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -1,6 +1,7 @@ import json import re from collections.abc import Callable, Mapping +from itertools import chain from typing import Final from uuid import uuid4 @@ -263,13 +264,17 @@ def _assert_chunk_envelope(chunks: tuple[dict[str, JsonValue], ...], model: str, return next(iter(identities)) +def _content_chunks(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(chunk for chunk in chunks if chunk.get("usage") is None) + + def _finish_reasons(chunks: tuple[dict[str, JsonValue], ...]) -> list[JsonValue]: - reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in chunks) + reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in _content_chunks(chunks)) return [reason for reason in reasons if reason is not None] def _deltas(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: - return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in chunks) + return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in _content_chunks(chunks)) def _sdk_base(gateway: Gateway) -> str: @@ -301,12 +306,12 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ @pytest.mark.parametrize( ("stream_options", "usage_chunks"), [ - pytest.param({"stream_options": {"include_usage": True}}, [_USAGE], id="include_usage"), + pytest.param({"stream_options": {"include_usage": True}}, [(True, *_USAGE)], id="include_usage"), pytest.param({}, [], id="no_usage"), ], ) def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_stream( - gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[int, int, int]] + gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) @@ -322,9 +327,9 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ identity: Final = _assert_chunk_envelope(chunks, model, response.text) assert "".join(str(delta.get("content") or "") for delta in _deltas(chunks)) == _answer(marker), response.text assert _finish_reasons(chunks) == ["stop"], response.text - assert [_token_usage(chunk) for chunk in chunks if chunk.get("usage") is not None] == usage_chunks, ( - response.text - ) + assert [ + (chunk["choices"] == [], *_token_usage(chunk)) for chunk in chunks if chunk.get("usage") is not None + ] == usage_chunks, response.text _assert_success_row( _spend_row(_BY_REQUEST_ID, identity), model, _HAIKU, "acompletion", _mantle_cost(gateway, _HAIKU) ) @@ -620,15 +625,14 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: ) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, @@ -688,15 +692,14 @@ async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gatew chunks: Final = tuple([chunk async for chunk in stream]) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, diff --git a/tests/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 9055f6f06c4..226accdbf54 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -15,7 +15,7 @@ import openai import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -723,6 +723,33 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm assert len(post_requests) == 1 +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_chat_stream_usage_the_proxy_adds_keeps_its_marker_out_of_the_stored_request( + gateway: Gateway, tmp_path: Path +) -> None: + response_id: Final = f"chatcmpl-proxy-usage-{uuid4()}" + with ( + wire_server( + _answering_model_listing(lambda request: _chat_stream(response_id, "sort_key", include_usage=True)) + ) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + response: Final = isolated.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "stream marker"}], "stream": True}, + ) + assert response.status_code == 200, response.text + assert {string_value(event["id"]) for event in _sse_events(response.text)} == {response_id}, response.text + stored_request: Final = object_value(_stored_row(response_id)["proxy_server_request"]) + assert stored_request["messages"] == [{"role": "user", "content": "stream marker"}], stored_request + assert "_litellm_stream_usage_injected" not in stored_request, stored_request + observed: Final = _provider_calls(wire.drain()) + assert [_json_object(request.body)["stream_options"] for request in observed] == [{"include_usage": True}] + + def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gateway, tmp_path: Path) -> None: arguments: Final = '{"sort_key":"created_at"}' tool_content: Final = "tool-result-created_at" diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py index d7537a78db8..1d2a8b25792 100644 --- a/tests/integration/streaming/test_stream_contracts.py +++ b/tests/integration/streaming/test_stream_contracts.py @@ -3,6 +3,7 @@ import json import threading import uuid from collections.abc import Callable +from itertools import chain from pathlib import Path from queue import SimpleQueue from typing import Final @@ -13,9 +14,10 @@ from hypothesis import Phase, example, given, settings from hypothesis import strategies as st from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server -from openai import OpenAI +from openai import OpenAI, omit +from openai.types.chat import ChatCompletionChunk from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state @@ -197,6 +199,7 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert {chunk.id for chunk in chunks} == {identity} usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) assert len(usages) == (1 if include else 0) + assert [len(chunk.choices) for chunk in chunks if chunk.usage is not None] == ([0] if include else []) if include: assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4 requests: Final = wire.drain() @@ -214,6 +217,119 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert float(rows[0]["spend"]) == pytest.approx(0.019) +_STREAM_USAGE_INJECTED_FLAG: Final = "_litellm_stream_usage_injected" + + +def _registered_model_names(gateway: Gateway) -> tuple[JsonValue, ...]: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + return tuple(object_value(entry).get("model_name") for entry in entries) + + +def _registered(gateway: Gateway, model: str) -> None: + eventually(lambda: _registered_model_names(gateway), lambda model_names: model in model_names, seconds=60) + + +def _sdk_stream( + gateway: Gateway, + route_prefix: str, + model: str, + identity: str, + *, + include_usage: bool | None, + extra_body: dict[str, JsonValue], +) -> tuple[ChatCompletionChunk, ...]: + with OpenAI( + api_key=gateway.key, + base_url=str(gateway.client.base_url).rstrip("/") + route_prefix, + timeout=30, + max_retries=0, + ) as client: + stream: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": identity}], + stream=True, + stream_options=omit if include_usage is None else {"include_usage": include_usage}, + extra_body=extra_body, + ) + with stream: + return tuple(stream) + + +def _streamed_text(chunks: tuple[ChatCompletionChunk, ...]) -> str: + choices: Final = chain.from_iterable(chunk.choices for chunk in chunks) + return "".join(choice.delta.content or "" for choice in choices) + + +def _usage_chunk_shapes(chunks: tuple[ChatCompletionChunk, ...]) -> list[tuple[int, int]]: + return [(len(chunk.choices), chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None] + + +@pytest.mark.parametrize( + ("route_prefix", "sends", "extra_body"), + [pytest.param("/v1", 2, {}, id="chat_miss_then_cache_hit"), pytest.param("/queue", 1, {"priority": 0}, id="queue")], +) +def test_caller_requested_stream_usage_has_no_choices_even_with_a_client_sent_injection_flag( + gateway: Gateway, route_prefix: str, sends: int, extra_body: dict[str, JsonValue] +) -> None: + identity: Final = "stream-usage-flag-" + uuid.uuid4().hex + with ( + wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=wire.url + "/v1") + _registered(gateway, model) + streams: Final = tuple( + _sdk_stream( + gateway, + route_prefix, + model, + identity, + include_usage=True, + extra_body={**extra_body, _STREAM_USAGE_INJECTED_FLAG: True}, + ) + for _ in range(sends) + ) + assert [_streamed_text(chunks) for chunks in streams] == ["Hello 雪 café"] * sends + assert [_usage_chunk_shapes(chunks) for chunks in streams] == [[(0, 15)]] * sends + requests: Final = wire.drain() + assert len(requests) == 1 + assert _STREAM_USAGE_INJECTED_FLAG not in json.loads(requests[0].body) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_stream_usage_the_deployment_adds_keeps_one_empty_choice(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "deployment-stream-usage-" + uuid.uuid4().hex + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire: + shared_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **shared_config, + "model_list": [ + { + "model_name": "deployment-usage", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-deployment-usage-key", + "api_base": wire.url + "/v1", + "stream_options": {"include_usage": True}, + }, + } + ], + "general_settings": {**shared_config["general_settings"], "always_include_stream_usage": False}, + } + path: Final = tmp_path / "deployment-stream-usage.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + chunks: Final = _sdk_stream( + candidate, "/v1", "deployment-usage", identity, include_usage=None, extra_body={} + ) + assert _streamed_text(chunks) == "Hello 雪 café" + assert _usage_chunk_shapes(chunks) == [(1, 15)] + requests: Final = wire.drain() + assert len(requests) == 1 + assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True + + @pytest.mark.covers("other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream") def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None: identity: Final = "messages-empty-choices-" + uuid.uuid4().hex diff --git a/tests/unit/harness/test_endpoint.py b/tests/unit/harness/test_endpoint.py index 19bd6d0450e..d1447fd40cc 100644 --- a/tests/unit/harness/test_endpoint.py +++ b/tests/unit/harness/test_endpoint.py @@ -272,6 +272,7 @@ async def test_sdk_chat_stream(monkeypatch: pytest.MonkeyPatch) -> None: assert lines[-1] == "data: [DONE]" assert json.loads(lines[0][6:])["choices"][0]["delta"]["content"] == "he" assert calls[0]["stream_options"] == {"include_usage": True} + assert calls[0]["_litellm_stream_usage_injected"] is True assert (ep.usage.input_tokens, ep.usage.output_tokens) == (6, 2) assert ep.usage.cost == pytest.approx(6 * 0.001 + 2 * 0.002) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 68aabc3c01e..550ecd3c0f7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -603,6 +603,91 @@ async def test_streaming_handler_with_usage( assert chunk_has_usage +def _usage_test_logging( + caller_stream_options: dict[str, bool] | None, stream_usage_injected: bool = False +) -> Logging: + logging_obj: Final = Logging( + model="hosted_vllm/m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="12345", + function_id="1245", + kwargs={"_litellm_stream_usage_injected": True} if stream_usage_injected else {}, + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="hosted_vllm/m", stream_options=caller_stream_options + ) + return logging_obj + + +@pytest.mark.parametrize( + ("caller_stream_options", "stream_usage_injected", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, False, 0), + (None, False, 1), + ({"include_usage": True}, True, 1), + ], + ids=["caller-asked", "provider-handler-injected", "proxy-injected"], +) +@pytest.mark.parametrize("plain_iterator", [False, True]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_usage_chunk_has_empty_choices_only_for_callers_that_asked( + sync_mode: bool, + plain_iterator: bool, + caller_stream_options: dict[str, bool] | None, + stream_usage_injected: bool, + usage_chunk_choice_count: int, +): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ModelResponseStream(id="c1", choices=[], provider_specific_fields={"citations": ["a"]}), + ModelResponseStream(id="c1", choices=[], usage=usage), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=( + iter(provider_chunks) if plain_iterator else ModelResponseListIterator(model_responses=provider_chunks) + ), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(caller_stream_options, stream_usage_injected), + stream_options={"include_usage": True}, + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + assert [len(chunk.choices) for chunk in chunks] == [1, 1, 0, usage_chunk_choice_count] + assert chunks[2].provider_specific_fields == {"citations": ["a"]} + assert [getattr(chunk, "usage", None) is not None for chunk in chunks] == [False, False, False, True] + assert chunks[-1].usage.total_tokens == 6 + + +@pytest.mark.asyncio +async def test_plain_iterator_stream_keeps_usage_on_its_chunk_when_usage_was_not_requested(): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream( + id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], usage=usage + ), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(provider_chunks), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(None), + ) + + chunks: Final = [chunk async for chunk in wrapper] + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] + assert chunks[-1].usage.total_tokens == 6 + + @pytest.mark.parametrize("sync_mode", [False]) @pytest.mark.asyncio @pytest.mark.flaky(reruns=3) diff --git a/tests/unit/llms/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index f465bbea663..9e2754eef68 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -364,7 +364,8 @@ def test_text_completion_stream_forwards_stream_options_and_usage(monkeypatch): ) ) - assert tuple(chunk.choices[0].text for chunk in chunks) == ("hello", None, None) + assert tuple(chunk.choices[0].text for chunk in chunks[:-1]) == ("hello", None) + assert chunks[-1].choices == [] assert tuple(chunk.usage for chunk in chunks[:-1]) == (None, None) assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, diff --git a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py index 12690d5822b..56c0aee3efd 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py @@ -543,8 +543,9 @@ def test_nova_streaming_sends_stream_flag_and_parses_chunks_and_usage(respx_mock "stream": True, "stream_options": {"include_usage": True}, } - assert tuple(chunk.choices[0].delta.content for chunk in chunks) == ("Nova ", "stream", None, None) - assert tuple(chunk.choices[0].finish_reason for chunk in chunks) == (None, None, "stop", None) + assert tuple(chunk.choices[0].delta.content for chunk in chunks[:-1]) == ("Nova ", "stream", None) + assert chunks[-1].choices == [] + assert tuple(chunk.choices[0].finish_reason for chunk in chunks[:-1]) == (None, None, "stop") assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, 1, diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index bb624add673..6ea3039edd8 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1894,7 +1894,9 @@ def test_gemini_tool_use(respx_mock: respx.MockRouter) -> None: assert [(call.function.name, json.loads(call.function.arguments)) for call in tool_calls] == [ ("get_weather", {"location": "Lima, Peru"}) ] - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["tool_calls"] + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] @pytest.mark.respx(assert_all_called=True) diff --git a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index fb1cf4ced2d..0726ebb5ee1 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -632,7 +632,7 @@ class TestVertexGemmaCompletion: assert chunks[0].choices[0].delta.content == "READY" assert chunks[1].choices[0].finish_reason == "stop" if stream_options and stream_options["include_usage"]: - assert chunks[-1].choices[0].delta.content is None + assert chunks[-1].choices == [] assert chunks[-1].usage.total_tokens == 114 assert span.get_metric("_dd.llmobs.total_tokens") == 114 else: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 63fcdfbb0c5..df2fc8e1da2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1,6 +1,7 @@ import asyncio import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock import pytest @@ -1975,6 +1976,20 @@ class TestProxyWiring: class TestStreamUsage: + @pytest.mark.asyncio + async def test_trailing_flush_survives_a_usage_chunk_with_no_choices(self): + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + held: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Mail [EMAI"))]) + usage_chunk: Final = ModelResponseStream( + choices=[], usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + ) + + out: Final = await _restore_stream(guardrail, [held, usage_chunk]) + + assert out[1] is not None and out[1].choices == [] + assert [chunk.choices[0].delta.content for chunk in out[2:]] == ["[EMAI"] + assert all(getattr(chunk, "usage", None) is None for chunk in out[2:]) + @pytest.mark.asyncio async def test_trailing_flush_does_not_repeat_usage(self): """With n>=2 and include_usage, the last chunk carries usage; a flush copied from it must not. diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f0baceecdaa..db0a74d360a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8536,6 +8536,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", ) diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index 07a61c8587b..a86314e3063 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -6,10 +6,13 @@ Pins (PR2): from __future__ import annotations +from collections.abc import Callable +from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi.testclient import TestClient from litellm import Router from litellm.proxy import proxy_server @@ -85,6 +88,25 @@ def test_queue_chat_completions_happy(client, auth_as, patched_queue): } +def test_queue_chat_completions_drops_client_sent_stream_usage_flags( + client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], patched_queue: MagicMock +) -> None: + payload: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "priority": 0, + "_litellm_strip_stream_usage": True, + "_litellm_stream_usage_injected": True, + } + with auth_as(): + response: Final = client.post("/queue/chat/completions", json=payload) + assert response.status_code == 200 + forwarded: Final = patched_queue.schedule_acompletion.call_args.kwargs + assert "_litellm_strip_stream_usage" not in forwarded + assert "_litellm_stream_usage_injected" not in forwarded + assert forwarded["messages"] == [{"role": "user", "content": "hi"}] + + def test_queue_chat_completions_no_router_error(client, auth_as, queue_no_router): """Pins ``POST /queue/chat/completions`` (error: no llm_router).""" payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d742956bb59..a226ce8a4a9 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -9181,6 +9181,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_default_preserves_other_client_stream_options_keys(self): data = {"stream": True, "stream_options": {"include_obfuscation": True}} @@ -9189,6 +9190,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_obfuscation": True, "include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_client_requested_usage_is_left_untouched_and_not_stripped(self): data = {"stream": True, "stream_options": {"include_usage": True}} @@ -9205,6 +9207,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"]["include_usage"] is True assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_explicit_false_flag_disables_injection_entirely(self): data = {"stream": True} @@ -9230,6 +9233,36 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert "_litellm_strip_stream_usage" not in data + @pytest.mark.parametrize("client_stream_options", [None, {"include_obfuscation": True}]) + def test_flag_true_marks_usage_injected_for_clients_that_did_not_ask( + self, client_stream_options: dict[str, bool] | None + ): + data: Final = { + "stream": True, + **({} if client_stream_options is None else {"stream_options": client_stream_options}), + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"]["include_usage"] is True + assert data["_litellm_stream_usage_injected"] is True + + def test_flag_true_leaves_usage_unmarked_when_client_asked_and_neutralizes_client_marker(self): + data: Final = {"stream": True, "stream_options": {"include_usage": True}, "_litellm_stream_usage_injected": True} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"] == {"include_usage": True} + assert data["_litellm_stream_usage_injected"] is False + def test_flag_true_respects_client_explicit_include_usage_false(self): data = {"stream": True, "stream_options": {"include_usage": False}} diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c2a5b01fe4e..3f8e3bc93d1 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -936,6 +936,26 @@ def test_body_snapshot_excludes_team_callback_credentials() -> None: }, proxy_request +def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request() -> None: + from litellm.proxy.common_request_processing import stream_usage_tracking_updates + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + + request: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True} + proxy_request: Final = {"body": dict(request)} + data: Final = { + **request, + **stream_usage_tracking_updates(request, {}, "acompletion", lambda: True), + "proxy_server_request": proxy_request, + } + + refresh_proxy_server_request_body_snapshot(data) + + assert data[STREAM_USAGE_INJECTED_FIELD] is True + assert proxy_request["body"]["stream_options"] == {"include_usage": True} + assert STREAM_USAGE_INJECTED_FIELD not in proxy_request["body"] + + @pytest.mark.asyncio @pytest.mark.parametrize("pre_call_ran", [False, True]) async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_logs( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f5d13aecb6f..c3d3575b8b3 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -3868,6 +3868,9 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): model="gpt-3.5-turbo-instruct", choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], ), + TextCompletionResponse( + id="cmpl-1", object="text_completion", created=1, model="gpt-3.5-turbo-instruct", choices=[] + ), ] response = stream_chunk_builder_text_completion(chunks=chunks, messages=[{"role": "user", "content": "say hello"}]) @@ -8853,7 +8856,8 @@ async def test_text_completion_stream_include_usage_returns_usage_on_final_chunk "stream": True, "stream_options": {"include_usage": True}, } - assert "".join(chunk.choices[0].text or "" for chunk in chunks) == "reply" + assert "".join(chunk.choices[0].text or "" for chunk in chunks if chunk.choices) == "reply" + assert chunks[-1].choices == [] final_usage: Final = chunks[-1].usage assert (final_usage.prompt_tokens, final_usage.completion_tokens, final_usage.total_tokens) == (3, 2, 5) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 13ac55fa937..b119d216542 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -29670,6 +29670,73 @@ def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: py assert [d["model_info"]["id"] for d in filtered] == ["only"] +@pytest.mark.parametrize( + ("deployment_stream_options", "default_stream_options", "request_stream_options", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, None, None, 1), + (None, {"include_usage": True}, None, 1), + ({"include_usage": True}, None, {"include_usage": True}, 0), + (None, None, {"include_usage": True}, 0), + ], + ids=["deployment-config", "router-default", "request-over-deployment", "request-only"], +) +@pytest.mark.asyncio +async def test_router_configured_stream_usage_keeps_a_choice_on_the_usage_chunk( + deployment_stream_options: dict[str, bool] | None, + default_stream_options: dict[str, bool] | None, + request_stream_options: dict[str, bool] | None, + usage_chunk_choice_count: int, +): + deployment_params: Final = {"model": "openai/gpt-5-mini", "api_key": "fake", "mock_response": "hi there"} + router: Final = Router( + model_list=[ + { + "model_name": "m", + "litellm_params": ( + deployment_params + if deployment_stream_options is None + else {**deployment_params, "stream_options": deployment_stream_options} + ), + } + ], + default_litellm_params=None if default_stream_options is None else {"stream_options": default_stream_options}, + ) + request_kwargs: Final = {} if request_stream_options is None else {"stream_options": request_stream_options} + + response: Final = await router.acompletion( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True, **request_kwargs + ) + chunks: Final = [chunk async for chunk in response] + + assert getattr(chunks[-1], "usage", None) is not None + assert len(chunks[-1].choices) == usage_chunk_choice_count + assert all(len(chunk.choices) == 1 for chunk in chunks[:-1]) + + +@pytest.mark.parametrize( + ("stream_options", "marked_injected"), + [({"include_usage": True}, True), (None, False)], + ids=["include-usage", "no-stream-options"], +) +def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litellm_params( + stream_options: dict[str, bool] | None, marked_injected: bool +): + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + from litellm.types.router import LiteLLM_Params + + kwargs: Final[dict[str, object]] = {"stream": True, "metadata": {}} + Router()._update_kwargs_with_deployment( + deployment={ + "model_name": "gpt-5-mini", + "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), + "model_info": {"id": "pydantic-params"}, + }, + kwargs=kwargs, + ) + + assert (kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True) is marked_injected + + @pytest.mark.asyncio async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 788d405f99f..5a2cefd08b1 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -263,6 +263,7 @@ INTERNAL_STATE_NAMES: Final = ( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", "acompletion",