mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(proxy): gate default stream usage injection on provider support and neutralize client-sent strip marker
Bytez and OCI param maps raise on stream_options when drop_params is unset, so the default injection would have broken every streamed chat completion routed to them. Injection now only happens when every router deployment behind the requested model (wildcards and aliases included) declares stream_options in its supported OpenAI params; providers that do not declare it either reject the param or already stream usage natively, so skipping them keeps old behavior instead of erroring. _litellm_strip_stream_usage arriving in the client request body is now overwritten at ingress (and popped in the experimental queue endpoint), so a client can no longer suppress the usage chunk it explicitly requested by planting the internal marker.
This commit is contained in:
parent
43efd02d35
commit
a770b437d5
3 changed files with 191 additions and 9 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"] = {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue