fix(proxy): request stream usage upstream by default and strip it from client streams

litellm/proxy/common_request_processing.py

@@ -5,13 +5,15 @@
 import time
 import traceback
 from datetime import datetime
+from functools import lru_cache
 from typing import (
     TYPE_CHECKING,
     Any,
     AsyncGenerator,
     Callable,
     Dict,
     Literal,
+    Mapping,
     Optional,
     Tuple,
     Union,
@@ -38,6 +40,9 @@
 )
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
+from litellm.litellm_core_utils.get_supported_openai_params import (
+    get_supported_openai_params,
+)
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
@@ -244,6 +249,71 @@ async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None
             pass
 
 
+@lru_cache(maxsize=512)
+def _litellm_model_supports_stream_options(litellm_model: str) -> bool:
+    try:
+        supported_params = get_supported_openai_params(model=litellm_model)
+    except Exception:  # noqa: BLE001  # unmapped or malformed model strings must disable injection, not fail the request
+        return False
+    return supported_params is not None and "stream_options" in supported_params
+
+
+def _deployment_litellm_model(deployment: Mapping[str, object]) -> str | None:
+    litellm_params = deployment.get("litellm_params")
+    if isinstance(litellm_params, Mapping):
+        litellm_model = litellm_params.get("model")
+    else:
+        litellm_model = getattr(litellm_params, "model", None)
+    return litellm_model if isinstance(litellm_model, str) else None
+
+
+def _model_deployments_support_stream_options(
+    model: object,
+    llm_router: Router | None,
+    team_id: str | None,
+) -> bool:
+    if not isinstance(model, str):
+        return False
+    deployments = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router is not None else None
+    deployment_models = tuple(
+        litellm_model
+        for deployment in deployments or ()
+        if (litellm_model := _deployment_litellm_model(deployment)) is not None
+    )
+    candidate_models = deployment_models if deployment_models else (model,)
+    return all(_litellm_model_supports_stream_options(m) for m in candidate_models)
+
+
+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 = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {}
+    if data.get("stream", False) is not True:
+        return scrub
+    always_include = general_settings.get("always_include_stream_usage")
+    stream_options = data.get("stream_options")
+    if always_include is True:
+        if "stream_options" not in data:
+            return {**scrub, "stream_options": {"include_usage": True}}
+        if isinstance(stream_options, dict) and "include_usage" not in stream_options:
+            return {**scrub, "stream_options": {**stream_options, "include_usage": True}}
+        return scrub
+    if always_include is False or route_type != "acompletion":
+        return scrub
+    if isinstance(stream_options, dict) and stream_options.get("include_usage") is True:
+        return scrub
+    if not supports_stream_options():
+        return scrub
+    merged_stream_options = {**stream_options} if isinstance(stream_options, dict) else {}
+    return {
+        "stream_options": {**merged_stream_options, "include_usage": True},
+        "_litellm_strip_stream_usage": True,
+    }
+
+
 def _serialize_http_exception_detail(
     detail: Any,
 ) -> Tuple[str, Optional[dict]]:
@@ -1232,17 +1302,18 @@ async def common_processing_pre_call_logic(
         )
 
         ### AUTO STREAM USAGE TRACKING ###
-        # If always_include_stream_usage is enabled and this is a streaming request
-        # automatically add stream_options={'include_usage': True} if not already set
-        if (
-            general_settings.get("always_include_stream_usage", False) is True
-            and self.data.get("stream", False) is True
-        ):
-            # Only set if stream_options is not already provided by the client
-            if "stream_options" not in self.data:
-                self.data["stream_options"] = {"include_usage": True}
-            elif isinstance(self.data["stream_options"], dict) and "include_usage" not in self.data["stream_options"]:
-                self.data["stream_options"]["include_usage"] = True
+        self.data.update(
+            _stream_usage_tracking_updates(
+                data=self.data,
+                general_settings=general_settings,
+                route_type=route_type,
+                supports_stream_options=lambda: _model_deployments_support_stream_options(
+                    model=self.data.get("model"),
+                    llm_router=llm_router,
+                    team_id=user_api_key_dict.team_id,
+                ),
+            )
+        )
         ### CALL HOOKS ### - modify/reject incoming data before calling the model
 
         ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call

litellm/proxy/proxy_server.py

@@ -119,6 +119,7 @@
 from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
+    StreamingChoices,
     TextCompletionResponse,
     TokenCountResponse,
 )
@@ -7368,6 +7369,25 @@ def _serialize_streaming_chunk(chunk: BaseModel) -> Union[str, bytes]:
     return chunk.model_dump_json(exclude_none=True, exclude_unset=True)
 
 
+def _is_injected_stream_usage_artifact(chunk: object) -> bool:
+    if not isinstance(chunk, ModelResponseStream):
+        return False
+    if chunk.provider_specific_fields is not None:
+        return False
+    return all(_is_empty_streaming_choice(choice) for choice in chunk.choices or [])
+
+
+def _is_empty_streaming_choice(choice: StreamingChoices) -> bool:
+    if choice.finish_reason is not None:
+        return False
+    if getattr(choice, "logprobs", None) is not None:
+        return False
+    delta = getattr(choice, "delta", None)
+    if delta is None:
+        return True
+    return all(value is None for value in delta.model_dump().values())
+
+
 async def _apply_streaming_chunk_hooks(
     *,
     chunk: Any,
@@ -7447,6 +7467,7 @@ async def async_data_generator(
         needs_iterator_wrap = proxy_logging_obj.needs_iterator_wrap()
         needs_per_chunk_hook = proxy_logging_obj.needs_per_chunk_streaming_hook()
         is_raw_sse_stream = bool(request_data.get("_litellm_raw_sse_stream"))
+        strip_stream_usage = bool(request_data.get("_litellm_strip_stream_usage"))
         raw_sse_buffer = ""
 
         if needs_iterator_wrap:
@@ -7498,6 +7519,15 @@ async def async_data_generator(
                 fallback_model_from_metadata=fallback_model_from_metadata,
             )
 
+            if strip_stream_usage and _is_injected_stream_usage_artifact(chunk):
+                if pending_fallback_event:
+                    yield _format_fallback_metadata_sse_event(
+                        fallback_model=fallback_model_from_metadata,
+                        fallback_errors=fallback_errors,
+                    )
+                    fallback_metadata_event_sent = True
+                continue
+
             raw_passthrough = False
             if isinstance(chunk, BaseModel):
                 chunk = _serialize_streaming_chunk(chunk)
@@ -13470,6 +13500,7 @@ async def async_queue_request(
     data = {}
     try:
         data = await request.json()  # type: ignore
+        data.pop("_litellm_strip_stream_usage", None)
 
         # Include original request and headers in the data
         data["proxy_server_request"] = {

litellm/types/utils.py

@@ -3296,6 +3296,7 @@ def shared_backend_model_info(model_info: Dict[str, Any]) -> Dict[str, Any]:
         "model_file_id_mapping",
         "litellm_logging_obj",
         "litellm_call_id",
+        "_litellm_strip_stream_usage",
         "use_client",
         "id",
         "fallbacks",