mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(responses): keep deployment stream off retrieves and reject oversized cursors
This commit is contained in:
parent
9d6f615392
commit
76bc3c39f3
3 changed files with 66 additions and 3 deletions
|
|
@ -7,6 +7,7 @@ from pydantic import TypeAdapter
|
|||
RETRIEVE_STREAM_QUERY_KEYS: Final = frozenset({"stream", "starting_after"})
|
||||
RETRIEVE_INCLUDE_QUERY_KEYS: Final = frozenset({"include", "include[]", "include_obfuscation"})
|
||||
_QUERY_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_MAX_STARTING_AFTER_DIGITS: Final = len(str(2**63 - 1))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -43,7 +44,7 @@ def _parsed_starting_after(value: object) -> int | None | InvalidRetrieveStreamQ
|
|||
return None
|
||||
case int() if not isinstance(value, bool) and value >= 0:
|
||||
return value
|
||||
case str() if value.isascii() and value.isdigit():
|
||||
case str() if value.isascii() and value.isdigit() and len(value) <= _MAX_STARTING_AFTER_DIGITS:
|
||||
return int(value)
|
||||
case _:
|
||||
return InvalidRetrieveStreamQuery(
|
||||
|
|
@ -72,3 +73,11 @@ def requests_streamed_retrieval(request: Mapping[str, object]) -> bool:
|
|||
and isinstance(extra_query, Mapping)
|
||||
and _parsed_stream(_QUERY_ADAPTER.validate_python(extra_query).get("stream")) is True
|
||||
)
|
||||
|
||||
|
||||
def without_retrieve_stream_defaults(
|
||||
deployment_params: Mapping[str, object], request: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
if "response_id" not in request:
|
||||
return deployment_params
|
||||
return {key: value for key, value in deployment_params.items() if key not in RETRIEVE_STREAM_QUERY_KEYS}
|
||||
|
|
|
|||
|
|
@ -148,7 +148,7 @@ from litellm.llms.openai_like.model_info import (
|
|||
get_openai_compatible_model_info,
|
||||
)
|
||||
from litellm.responses.bridge_check import keep_registered_mode
|
||||
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval
|
||||
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults
|
||||
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.router_strategy.complexity_router.context_compaction import (
|
||||
|
|
@ -6372,7 +6372,7 @@ class Router:
|
|||
generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "")
|
||||
|
||||
response_kwargs: Final = {
|
||||
**data,
|
||||
**without_retrieve_stream_defaults(data, kwargs),
|
||||
"caching": self.cache_responses,
|
||||
**kwargs,
|
||||
"model": model_name,
|
||||
|
|
|
|||
|
|
@ -412,6 +412,7 @@ async def test_a_resume_that_times_out_raises_a_timeout(sync_mode: bool) -> None
|
|||
{"stream": True, "starting_after": -1},
|
||||
{"extra_query": {"stream": "true", "starting_after": "abc"}},
|
||||
{"extra_query": {"stream": "1"}},
|
||||
{"extra_query": {"stream": "true", "starting_after": "9" * 5000}},
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
|
|
@ -435,6 +436,7 @@ async def test_a_malformed_resume_query_is_rejected_before_any_upstream_call(
|
|||
({"extra_query": {"stream": "1"}}, {}),
|
||||
({"stream": True, "starting_after": -1}, {"stream": "true"}),
|
||||
({"extra_query": {"stream": "true", "starting_after": "abc"}}, {"stream": "true"}),
|
||||
({"extra_query": {"stream": "true", "starting_after": "9" * 5000}}, {"stream": "true"}),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
|
|
@ -538,6 +540,58 @@ async def test_a_routed_resume_requested_in_extra_query_gets_the_deployment_stre
|
|||
assert call.request.extensions["timeout"]["read"] == 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_routed_plain_retrieve_on_a_deployment_that_forces_stream_answers_the_response_object() -> None:
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "background-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.5",
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
"stream": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
with respx.mock() as mock_router:
|
||||
route: Final = _retrieve_route(mock_router, "openai").mock(
|
||||
return_value=httpx.Response(200, json=_stored_response("completed"))
|
||||
)
|
||||
retrieved: Final = await router.aget_responses(model="background-model", response_id=_STORED_RESPONSE_ID)
|
||||
(call,) = route.calls
|
||||
|
||||
assert isinstance(retrieved, ResponsesAPIResponse)
|
||||
assert dict(call.request.url.params) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_routed_create_on_a_deployment_that_forces_stream_still_streams() -> None:
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "background-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.5",
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
"stream": True,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
with respx.mock() as mock_router:
|
||||
route: Final = mock_router.post("https://api.openai.com/v1/responses").mock(
|
||||
return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3)
|
||||
)
|
||||
stream: Final = await router.aresponses(model="background-model", input="hi", background=True)
|
||||
event_types: Final = await _event_types(stream)
|
||||
(call,) = route.calls
|
||||
|
||||
assert event_types == ("response.output_text.delta", "response.completed")
|
||||
assert json.loads(call.request.content)["stream"] is True
|
||||
|
||||
class _SuccessLog(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue