diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 750954b69f7..d25ef8a3039 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -5,6 +5,7 @@ import math import time import traceback from datetime import datetime +from functools import lru_cache from typing import ( TYPE_CHECKING, Any, @@ -39,6 +40,9 @@ from litellm.constants import ( ) 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, @@ -245,25 +249,64 @@ 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]) -> Optional[str]: + 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: Optional[Router], +) -> bool: + if not isinstance(model, str): + return False + deployments = llm_router.get_model_list(model_name=model) if llm_router is not None else None + deployment_models = tuple( + litellm_model + for deployment in deployments or () + for litellm_model in (_deployment_litellm_model(deployment),) + if litellm_model 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 {} + 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 {"stream_options": {"include_usage": True}} + return {**scrub, "stream_options": {"include_usage": True}} if isinstance(stream_options, dict) and "include_usage" not in stream_options: - return {"stream_options": {**stream_options, "include_usage": True}} - return {} + return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return scrub if always_include is False or route_type != "acompletion": - return {} + return scrub if isinstance(stream_options, dict) and stream_options.get("include_usage") is True: - return {} + 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}, @@ -1264,6 +1307,10 @@ class ProxyBaseLLMRequestProcessing: 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, + ), ) ) ### CALL HOOKS ### - modify/reject incoming data before calling the model diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ab55e885cda..a60ea2da019 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13500,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"] = { diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 24b2007217e..28fb97006d8 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,7 +1,7 @@ import asyncio import copy import datetime -from typing import AsyncGenerator, Optional +from typing import AsyncGenerator, Callable, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -5113,10 +5113,22 @@ class TestStreamingClientDisconnectBilling: proxy_logging_obj._arelease_max_parallel_requests_on_disconnect.assert_awaited_once() -def _apply_stream_usage_tracking(data: dict, general_settings: dict, route_type: str) -> None: +def _apply_stream_usage_tracking( + data: dict, + general_settings: dict, + route_type: str, + supports_stream_options: Callable[[], bool] = lambda: True, +) -> None: from litellm.proxy.common_request_processing import _stream_usage_tracking_updates - data.update(_stream_usage_tracking_updates(data=data, general_settings=general_settings, route_type=route_type)) + data.update( + _stream_usage_tracking_updates( + data=data, + general_settings=general_settings, + route_type=route_type, + supports_stream_options=supports_stream_options, + ) + ) class TestApplyStreamUsageTracking: @@ -5203,3 +5215,125 @@ class TestApplyStreamUsageTracking: assert "stream_options" not in data assert "_litellm_strip_stream_usage" not in data + + def test_default_skips_injection_when_provider_lacks_stream_options_support(self): + data = {"stream": True, "model": "bytez-model"} + + _apply_stream_usage_tracking( + data=data, + general_settings={}, + route_type="acompletion", + supports_stream_options=lambda: False, + ) + + assert "stream_options" not in data + assert "_litellm_strip_stream_usage" not in data + + def test_client_supplied_strip_marker_is_neutralized(self): + data = { + "stream": True, + "stream_options": {"include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["_litellm_strip_stream_usage"] is False + assert data["stream_options"] == {"include_usage": True} + + def test_client_supplied_strip_marker_is_neutralized_with_flag_true(self): + data = { + "stream": True, + "stream_options": {"include_usage": True}, + "_litellm_strip_stream_usage": True, + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["_litellm_strip_stream_usage"] is False + + def test_client_supplied_strip_marker_is_neutralized_on_non_streaming_request(self): + data = {"_litellm_strip_stream_usage": True} + + _apply_stream_usage_tracking(data=data, general_settings={}, route_type="acompletion") + + assert data["_litellm_strip_stream_usage"] is False + + +class TestModelDeploymentsSupportStreamOptions: + def _support(self, model, llm_router=None) -> bool: + from litellm.proxy.common_request_processing import ( + _model_deployments_support_stream_options, + ) + + return _model_deployments_support_stream_options(model=model, llm_router=llm_router) + + def test_openai_compatible_deployment_supports_stream_options(self): + router = litellm.Router( + model_list=[ + { + "model_name": "azure-nano", + "litellm_params": { + "model": "azure/gpt-5.4-nano", + "api_key": "fake", + "api_base": "https://example.openai.azure.com", + }, + } + ] + ) + + assert self._support("azure-nano", router) is True + + def test_deployment_on_provider_rejecting_stream_options_is_not_injected(self): + router = litellm.Router( + model_list=[ + { + "model_name": "tiny", + "litellm_params": {"model": "bytez/openai-community/gpt2", "api_key": "fake"}, + } + ] + ) + + assert self._support("tiny", router) is False + + def test_mixed_provider_model_group_is_not_injected(self): + router = litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "oci/cohere.command-r-plus", "api_key": "fake"}, + }, + ] + ) + + assert self._support("mixed", router) is False + + def test_wildcard_route_resolves_provider_support(self): + router = litellm.Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "fake"}, + } + ] + ) + + assert self._support("openai/gpt-4o", router) is True + + def test_provider_prefixed_model_without_router_is_resolved_directly(self): + assert self._support("openai/gpt-4o", None) is True + assert self._support("bytez/openai-community/gpt2", None) is False + + def test_unmapped_model_name_is_not_injected(self): + assert self._support("some-unmapped-public-alias", None) is False + + def test_non_string_model_is_not_injected(self): + assert self._support(None, None) is False