mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(responses): resume background response streams with stream and starting_after
Retrieving a response with stream=true and starting_after=N now returns the
replayed event stream from OpenAI and Azure instead of one JSON object, on the
SDK (get_responses / aget_responses) and on the proxy's GET /v1/responses/{id}
Providers that cannot replay a stored stream drop the params under drop_params
and raise UnsupportedParamsError without it. Malformed values answer 400 with
the param named. A replay that reaches its terminal event settles the
background response's bill once and passes the provider's usage through
without estimating one. An empty replay (cursor at the end) ends normally
This commit is contained in:
parent
1176e608b1
commit
577980d4b6
23 changed files with 1047 additions and 67 deletions
|
|
@ -30,7 +30,7 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
||||
ManagedObjectCreatorAttribution,
|
||||
|
|
@ -290,12 +290,14 @@ class CheckResponsesCost:
|
|||
sees provider env vars, so it fails for every deployment whose credentials
|
||||
live in the config; the row then never leaves ``queued``.
|
||||
"""
|
||||
if not via_router:
|
||||
return await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
|
||||
router_response = await self.llm_router.aget_responses(
|
||||
response_id=response_id, litellm_metadata=litellm_metadata
|
||||
retrieved: Final[object] = (
|
||||
await self.llm_router.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
|
||||
if via_router
|
||||
else await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
|
||||
)
|
||||
return cast(ResponsesAPIResponse, router_response)
|
||||
if not isinstance(retrieved, ResponsesAPIResponse):
|
||||
raise TypeError(f"Retrieving {response_id} did not return a response object")
|
||||
return retrieved
|
||||
|
||||
async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
|
||||
return await self._get_response(
|
||||
|
|
|
|||
|
|
@ -120,7 +120,7 @@ def is_unbilled_non_inference_call(
|
|||
origin: Final = metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) if metadata is not None else None
|
||||
if origin == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN:
|
||||
return False
|
||||
if not is_background_response(response):
|
||||
if not is_background_response(_finished_response(response)):
|
||||
return True
|
||||
return origin == BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN or _is_managed_background_call(metadata)
|
||||
|
||||
|
|
|
|||
|
|
@ -323,7 +323,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[str, dict]:
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
"""
|
||||
Transform the get response API request into a URL and data
|
||||
|
||||
|
|
@ -331,7 +331,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
- GET /v1/responses/{response_id}
|
||||
"""
|
||||
get_url: Final = self._construct_url_for_response_id_in_path(api_base=api_base, response_id=response_id)
|
||||
data: Final[dict] = {}
|
||||
data: Final[dict[str, str]] = {}
|
||||
verbose_logger.debug("get response url=%s", get_url)
|
||||
return get_url, data
|
||||
|
||||
|
|
|
|||
|
|
@ -69,6 +69,9 @@ class BaseResponsesAPIConfig(ABC):
|
|||
def supports_encrypted_agent_messages(self) -> bool:
|
||||
return False
|
||||
|
||||
def supports_streamed_retrieval(self) -> bool:
|
||||
return self.custom_llm_provider in (LlmProviders.OPENAI, LlmProviders.AZURE)
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
|
|
@ -209,7 +212,7 @@ class BaseResponsesAPIConfig(ABC):
|
|||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[str, dict]:
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
|
|
|
|||
|
|
@ -764,18 +764,23 @@ class AsyncHTTPHandler:
|
|||
follow_redirects: bool | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
max_response_bytes: int | None = None,
|
||||
stream: bool = False,
|
||||
):
|
||||
# Set follow_redirects to UseClientDefault if None
|
||||
_follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT
|
||||
query: Final = httpx.QueryParams(params).merge(HTTPHandler.extract_query_params(url))
|
||||
request_headers: Final = httpx.Headers(headers)
|
||||
|
||||
params = params or {}
|
||||
params.update(HTTPHandler.extract_query_params(url))
|
||||
if stream:
|
||||
return await self._send_streaming_get(
|
||||
url, params=query, headers=request_headers, follow_redirects=follow_redirects, timeout=timeout
|
||||
)
|
||||
|
||||
if max_response_bytes is not None:
|
||||
return await self._get_with_response_limit(
|
||||
url,
|
||||
params=httpx.QueryParams(params),
|
||||
headers=httpx.Headers(headers),
|
||||
params=query,
|
||||
headers=request_headers,
|
||||
max_bytes=max_response_bytes,
|
||||
follow_redirects=self.client.follow_redirects if follow_redirects is None else follow_redirects,
|
||||
timeout=self.client.timeout if timeout is None else httpx.Timeout(timeout),
|
||||
|
|
@ -783,13 +788,41 @@ class AsyncHTTPHandler:
|
|||
|
||||
response: Final = await self.client.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=headers,
|
||||
params=query,
|
||||
headers=request_headers,
|
||||
follow_redirects=_follow_redirects,
|
||||
timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT,
|
||||
)
|
||||
return response
|
||||
|
||||
async def _send_streaming_get(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
params: httpx.QueryParams,
|
||||
headers: httpx.Headers,
|
||||
follow_redirects: bool | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> httpx.Response:
|
||||
start_time: Final = time.time()
|
||||
request: Final = self.client.build_request(
|
||||
"GET", url, params=params, headers=headers, timeout=USE_CLIENT_DEFAULT if timeout is None else timeout
|
||||
)
|
||||
try:
|
||||
response: Final = await self.client.send(
|
||||
request,
|
||||
stream=True,
|
||||
follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects,
|
||||
)
|
||||
_anchor_handler_to(response, self)
|
||||
response.raise_for_status()
|
||||
except httpx.TimeoutException as e:
|
||||
_raise_timeout(e, start_time)
|
||||
except httpx.HTTPStatusError as e:
|
||||
await _raise_masked_async_error(e, True)
|
||||
else:
|
||||
return response
|
||||
|
||||
async def _get_with_response_limit(
|
||||
self,
|
||||
url: str,
|
||||
|
|
@ -1468,15 +1501,20 @@ class HTTPHandler:
|
|||
headers: Mapping[str, Any] | None = None,
|
||||
follow_redirects: bool | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
stream: bool = False,
|
||||
):
|
||||
# Set follow_redirects to UseClientDefault if None
|
||||
_follow_redirects: Final = follow_redirects if follow_redirects is not None else USE_CLIENT_DEFAULT
|
||||
params = params or {}
|
||||
params.update(self.extract_query_params(url))
|
||||
query: Final = httpx.QueryParams(params).merge(self.extract_query_params(url))
|
||||
|
||||
if stream:
|
||||
return self._send_streaming_get(
|
||||
url, params=query, headers=headers, follow_redirects=follow_redirects, timeout=timeout
|
||||
)
|
||||
|
||||
response: Final = self.client.get(
|
||||
url,
|
||||
params=params,
|
||||
params=query,
|
||||
headers=headers,
|
||||
follow_redirects=_follow_redirects,
|
||||
timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT,
|
||||
|
|
@ -1484,6 +1522,34 @@ class HTTPHandler:
|
|||
|
||||
return response
|
||||
|
||||
def _send_streaming_get(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
params: httpx.QueryParams,
|
||||
headers: Mapping[str, str] | None,
|
||||
follow_redirects: bool | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
) -> httpx.Response:
|
||||
start_time: Final = time.time()
|
||||
request: Final = self.client.build_request(
|
||||
"GET", url, params=params, headers=headers, timeout=USE_CLIENT_DEFAULT if timeout is None else timeout
|
||||
)
|
||||
try:
|
||||
response: Final = self.client.send(
|
||||
request,
|
||||
stream=True,
|
||||
follow_redirects=USE_CLIENT_DEFAULT if follow_redirects is None else follow_redirects,
|
||||
)
|
||||
_anchor_handler_to(response, self)
|
||||
response.raise_for_status()
|
||||
except httpx.TimeoutException as e:
|
||||
_raise_timeout(e, start_time)
|
||||
except httpx.HTTPStatusError as e:
|
||||
_raise_masked_sync_error(e, True)
|
||||
else:
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def extract_query_params(url: str) -> dict[str, str]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -371,6 +371,14 @@ def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) -
|
|||
return {"json": data}
|
||||
|
||||
|
||||
def _stream_query_params(stream: bool, starting_after: int | None) -> Mapping[str, object]:
|
||||
if not stream:
|
||||
return {}
|
||||
if starting_after is None:
|
||||
return {"stream": "true"}
|
||||
return {"stream": "true", "starting_after": starting_after}
|
||||
|
||||
|
||||
def _google_genai_streaming_hidden_params(
|
||||
*,
|
||||
api_base: str,
|
||||
|
|
@ -3480,12 +3488,20 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: Mapping[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
stream: bool = False,
|
||||
starting_after: int | None = None,
|
||||
litellm_metadata: Mapping[str, object] = MappingProxyType({}),
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None = None,
|
||||
_is_async: bool = False,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]:
|
||||
) -> (
|
||||
ResponsesAPIResponse
|
||||
| SyncResponsesAPIStreamingIterator
|
||||
| Coroutine[object, object, ResponsesAPIResponse | ResponsesAPIStreamingIterator]
|
||||
):
|
||||
"""
|
||||
Get a response by ID
|
||||
Uses GET /v1/responses/{response_id} endpoint in the responses API
|
||||
|
|
@ -3498,7 +3514,11 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
extra_headers=extra_headers,
|
||||
extra_query=extra_query,
|
||||
extra_body=extra_body,
|
||||
stream=stream,
|
||||
starting_after=starting_after,
|
||||
litellm_metadata=litellm_metadata,
|
||||
timeout=timeout,
|
||||
client=client,
|
||||
shared_session=shared_session,
|
||||
|
|
@ -3527,20 +3547,23 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
query: Final = {**data, **(extra_query or {}), **_stream_query_params(stream, starting_after)}
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"complete_input_dict": query,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = sync_httpx_client.get(url=url, headers=headers, params=data)
|
||||
response: Final = sync_httpx_client.get(
|
||||
url=url, headers=headers, params=query, stream=stream, timeout=timeout if stream else None
|
||||
)
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
raise self._handle_error(
|
||||
|
|
@ -3548,6 +3571,18 @@ class BaseLLMHTTPHandler:
|
|||
provider_config=responses_api_provider_config,
|
||||
)
|
||||
|
||||
if stream:
|
||||
return SyncResponsesAPIStreamingIterator(
|
||||
response=response,
|
||||
model="",
|
||||
logging_obj=logging_obj,
|
||||
responses_api_provider_config=responses_api_provider_config,
|
||||
litellm_metadata=dict(litellm_metadata),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_data={"litellm_params": dict(litellm_params)},
|
||||
replays_stored_stream=True,
|
||||
)
|
||||
|
||||
return responses_api_provider_config.transform_get_response_api_response(
|
||||
raw_response=response,
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -3561,11 +3596,15 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
custom_llm_provider: str | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: Mapping[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
stream: bool = False,
|
||||
starting_after: int | None = None,
|
||||
litellm_metadata: Mapping[str, object] = MappingProxyType({}),
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
client: HTTPHandler | AsyncHTTPHandler | None = None,
|
||||
shared_session: Optional["ClientSession"] = None,
|
||||
) -> ResponsesAPIResponse:
|
||||
) -> ResponsesAPIResponse | ResponsesAPIStreamingIterator:
|
||||
"""
|
||||
Async version of get_responses
|
||||
"""
|
||||
|
|
@ -3603,20 +3642,23 @@ class BaseLLMHTTPHandler:
|
|||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
query: Final = {**data, **(extra_query or {}), **_stream_query_params(stream, starting_after)}
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": data,
|
||||
"complete_input_dict": query,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = await async_httpx_client.get(url=url, headers=headers, params=data)
|
||||
response: Final = await async_httpx_client.get(
|
||||
url=url, headers=headers, params=query, stream=stream, timeout=timeout if stream else None
|
||||
)
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error retrieving response: %s", e)
|
||||
|
|
@ -3625,6 +3667,18 @@ class BaseLLMHTTPHandler:
|
|||
provider_config=responses_api_provider_config,
|
||||
)
|
||||
|
||||
if stream:
|
||||
return ResponsesAPIStreamingIterator(
|
||||
response=response,
|
||||
model="",
|
||||
logging_obj=logging_obj,
|
||||
responses_api_provider_config=responses_api_provider_config,
|
||||
litellm_metadata=dict(litellm_metadata),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_data={"litellm_params": dict(litellm_params)},
|
||||
replays_stored_stream=True,
|
||||
)
|
||||
|
||||
return responses_api_provider_config.transform_get_response_api_response(
|
||||
raw_response=response,
|
||||
logging_obj=logging_obj,
|
||||
|
|
|
|||
|
|
@ -241,7 +241,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[str, dict]:
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
"""
|
||||
Transform the get response API request into a URL and data.
|
||||
|
||||
|
|
@ -252,7 +252,7 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
"""
|
||||
encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id")
|
||||
url: Final = f"{api_base}/{encoded_response_id}"
|
||||
data: Final[dict] = {}
|
||||
data: Final[dict[str, str]] = {}
|
||||
return url, data
|
||||
|
||||
def transform_get_response_api_response(
|
||||
|
|
|
|||
|
|
@ -1060,7 +1060,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[str, dict]:
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
"""
|
||||
Transform the get response API request into a URL and data
|
||||
|
||||
|
|
@ -1069,7 +1069,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
"""
|
||||
encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id")
|
||||
url: Final = f"{api_base}/{encoded_response_id}"
|
||||
data: Final[dict] = {}
|
||||
data: Final[dict[str, str]] = {}
|
||||
return url, data
|
||||
|
||||
def transform_get_response_api_response(
|
||||
|
|
|
|||
|
|
@ -310,10 +310,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[str, dict]:
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id")
|
||||
url: Final = f"{api_base}/{encoded_response_id}"
|
||||
data: Final[dict] = {}
|
||||
data: Final[dict[str, str]] = {}
|
||||
return url, data
|
||||
|
||||
def transform_get_response_api_response(
|
||||
|
|
|
|||
|
|
@ -33584,6 +33584,12 @@
|
|||
"schema": {
|
||||
"$ref": "#/components/schemas/ResponsesAPIResponse"
|
||||
}
|
||||
},
|
||||
"text/event-stream": {
|
||||
"schema": {
|
||||
"description": "Server sent events when stream=true",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
|
|
|
|||
|
|
@ -3155,7 +3155,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data=self.data,
|
||||
request=request,
|
||||
)
|
||||
if route_type == "aresponses":
|
||||
if route_type in ("aresponses", "aget_responses"):
|
||||
selected_data_generator = (
|
||||
ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
|
||||
original_stream_response=response,
|
||||
|
|
|
|||
|
|
@ -102,11 +102,11 @@ router: Final = APIRouter()
|
|||
_RESPONSES_WS_CONFIG_VALUE_ADAPTER: Final[TypeAdapter[object | None]] = TypeAdapter(object | None)
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_PROXY_SERVER_REQUEST_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
_RETRIEVE_QUERY_PARAMS: Final = ("include", "include[]", "include_obfuscation", "starting_after", "stream")
|
||||
|
||||
_ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg
|
||||
|
||||
RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {200: {"model": ResponsesAPIResponse}}
|
||||
RESPONSES_API_CREATE_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {
|
||||
RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {
|
||||
200: {
|
||||
"model": ResponsesAPIResponse,
|
||||
"content": {
|
||||
|
|
@ -250,19 +250,19 @@ async def _resolve_cursor_model_variant_before_auth(request: Request) -> None:
|
|||
"/v1/responses",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
tags=["responses"],
|
||||
responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS,
|
||||
responses=RESPONSES_API_RESPONSE_SCHEMAS,
|
||||
)
|
||||
@router.post(
|
||||
"/responses",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
tags=["responses"],
|
||||
responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS,
|
||||
responses=RESPONSES_API_RESPONSE_SCHEMAS,
|
||||
)
|
||||
@router.post(
|
||||
"/openai/v1/responses",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
tags=["responses"],
|
||||
responses=RESPONSES_API_CREATE_RESPONSE_SCHEMAS,
|
||||
responses=RESPONSES_API_RESPONSE_SCHEMAS,
|
||||
)
|
||||
async def responses_api(
|
||||
request: Request,
|
||||
|
|
@ -825,6 +825,14 @@ async def get_response(
|
|||
data: Final = await read_request_body(request=request)
|
||||
read_settles_billing: Final = not carries_client_side_overrides(data)
|
||||
data["response_id"] = response_id
|
||||
retrieve_query: Final = {
|
||||
key: values[0] if len(values) == 1 else values
|
||||
for key in _RETRIEVE_QUERY_PARAMS
|
||||
if (values := request.query_params.getlist(key))
|
||||
}
|
||||
if retrieve_query:
|
||||
body_extra_query: Final[object] = data.get("extra_query")
|
||||
data["extra_query"] = {**(body_extra_query if isinstance(body_extra_query, dict) else {}), **retrieve_query}
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
response: Final[object] = await processor.base_process_llm_request(
|
||||
|
|
@ -836,7 +844,12 @@ async def get_response(
|
|||
llm_router=llm_router,
|
||||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
select_data_generator=partial(
|
||||
_streamed_retrieval_data_generator,
|
||||
partial(select_data_generator, responses_stream_errors=True),
|
||||
response_id,
|
||||
read_settles_billing,
|
||||
),
|
||||
model=None,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
|
|
@ -910,12 +923,24 @@ async def _bill_finished_background_response(
|
|||
fastapi_response: Response,
|
||||
read_settles_billing: bool,
|
||||
) -> None:
|
||||
_report_charge(
|
||||
await _settle_finished_background_read(response, response_id, user_api_key_dict, read_settles_billing),
|
||||
fastapi_response,
|
||||
)
|
||||
|
||||
|
||||
async def _settle_finished_background_read(
|
||||
response: object,
|
||||
response_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
read_settles_billing: bool,
|
||||
) -> "_BilledCharge | None":
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
if not isinstance(response, ResponsesAPIResponse) or not is_background_response(response):
|
||||
return
|
||||
return None
|
||||
if response.status not in TERMINAL_RESPONSE_STATUSES:
|
||||
return
|
||||
return None
|
||||
managed_files_hook: Final = _managed_files_hook(proxy_logging_obj)
|
||||
track_untracked: Final = (
|
||||
partial(
|
||||
|
|
@ -926,12 +951,11 @@ async def _bill_finished_background_response(
|
|||
if managed_files_hook is not None
|
||||
else None
|
||||
)
|
||||
charge: Final = await _settle_background_billing(
|
||||
return await _settle_background_billing(
|
||||
(lambda cost_check: cost_check.bill_if_finished(response_id, response, track_untracked))
|
||||
if read_settles_billing
|
||||
else (lambda cost_check: cost_check.bill_from_deployment(response_id, track_untracked))
|
||||
)
|
||||
_report_charge(charge, fastapi_response)
|
||||
|
||||
|
||||
class _BilledCharge(Protocol):
|
||||
|
|
@ -1081,27 +1105,53 @@ def _streamed_background_data_generator(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
request_tags=_background_response_spend_tags(request_data),
|
||||
)
|
||||
return _bill_streamed_background_response(response, chunks, track_untracked)
|
||||
return _bill_streamed_background_response(
|
||||
response, chunks, partial(_bill_finished_stream, track_untracked=track_untracked)
|
||||
)
|
||||
|
||||
|
||||
def _streamed_retrieval_data_generator(
|
||||
select_data_generator: Callable[..., AsyncIterator[str | bytes]],
|
||||
response_id: str,
|
||||
read_settles_billing: bool,
|
||||
response: object,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_data: Mapping[str, object],
|
||||
request: Request | None = None,
|
||||
) -> AsyncIterator[str | bytes]:
|
||||
chunks: Final = select_data_generator(
|
||||
response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, request=request
|
||||
)
|
||||
return _bill_streamed_background_response(
|
||||
response,
|
||||
chunks,
|
||||
partial(
|
||||
_settle_finished_background_read,
|
||||
response_id=response_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
read_settles_billing=read_settles_billing,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _bill_streamed_background_response(
|
||||
stream: object,
|
||||
chunks: AsyncIterator[str | bytes],
|
||||
track_untracked: Callable[[str, ResponsesAPIResponse], Awaitable[bool]],
|
||||
bill: Callable[[ResponsesAPIResponse], Awaitable[object]],
|
||||
) -> AsyncIterator[str | bytes]:
|
||||
async for chunk in chunks:
|
||||
completed: object = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream)
|
||||
if not isinstance(completed, ResponsesAPIResponse):
|
||||
yield chunk
|
||||
continue
|
||||
await _bill_finished_stream(completed, track_untracked)
|
||||
await bill(completed)
|
||||
yield chunk
|
||||
async for remaining_chunk in chunks:
|
||||
yield remaining_chunk
|
||||
return
|
||||
late_completed: Final = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream)
|
||||
if isinstance(late_completed, ResponsesAPIResponse):
|
||||
await _bill_finished_stream(late_completed, track_untracked)
|
||||
await bill(late_completed)
|
||||
|
||||
|
||||
async def _bill_finished_stream(
|
||||
|
|
|
|||
|
|
@ -106,7 +106,17 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
MCPTool = Any
|
||||
|
||||
from .streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
from .retrieve_stream_query import (
|
||||
RETRIEVE_STREAM_QUERY_KEYS,
|
||||
InvalidRetrieveStreamQuery,
|
||||
RetrieveStreamQuery,
|
||||
parse_retrieve_stream_query,
|
||||
)
|
||||
from .streaming_iterator import (
|
||||
BaseResponsesAPIStreamingIterator,
|
||||
ResponsesAPIStreamingIterator,
|
||||
SyncResponsesAPIStreamingIterator,
|
||||
)
|
||||
|
||||
__all__ = (
|
||||
"acancel_responses",
|
||||
|
|
@ -192,6 +202,7 @@ def mock_responses_api_response(
|
|||
|
||||
|
||||
_STR_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_STR_DICT: Final = TypeAdapter(dict[str, object])
|
||||
_FORWARDED_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
|
|
@ -2013,6 +2024,50 @@ def delete_responses(
|
|||
)
|
||||
|
||||
|
||||
def _requested_retrieve_stream_query(
|
||||
stream: bool | None,
|
||||
starting_after: int | None,
|
||||
extra_query: Mapping[str, object] | None,
|
||||
custom_llm_provider: str | None,
|
||||
) -> RetrieveStreamQuery:
|
||||
explicit: Final = {
|
||||
key: value for key, value in (("stream", stream), ("starting_after", starting_after)) if value is not None
|
||||
}
|
||||
requested: Final = {**explicit, **(extra_query or {})}
|
||||
parsed: Final = parse_retrieve_stream_query(requested.get("stream"), requested.get("starting_after"))
|
||||
if isinstance(parsed, InvalidRetrieveStreamQuery):
|
||||
raise litellm.BadRequestError(
|
||||
message=parsed.message,
|
||||
model=None,
|
||||
llm_provider=custom_llm_provider,
|
||||
body={"message": parsed.message, "type": "invalid_request_error", "param": parsed.param},
|
||||
)
|
||||
return parsed
|
||||
|
||||
|
||||
def _supported_retrieve_stream_query(
|
||||
requested: RetrieveStreamQuery,
|
||||
responses_api_provider_config: BaseResponsesAPIConfig,
|
||||
custom_llm_provider: str,
|
||||
drop_params: object,
|
||||
) -> RetrieveStreamQuery:
|
||||
if not requested.stream:
|
||||
return RetrieveStreamQuery(stream=False, starting_after=None)
|
||||
if responses_api_provider_config.supports_streamed_retrieval():
|
||||
return requested
|
||||
if drops_unsupported_params(drop_params):
|
||||
return RetrieveStreamQuery(stream=False, starting_after=None)
|
||||
raise UnsupportedParamsError(
|
||||
status_code=400,
|
||||
llm_provider=custom_llm_provider,
|
||||
message=(
|
||||
f"{custom_llm_provider} does not support streaming a stored response (stream=true on "
|
||||
"GET /v1/responses/{response_id}). To drop it, set `litellm.drop_params=True` or for proxy:\n\n"
|
||||
"`litellm_settings:\n drop_params: true`\n"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@client
|
||||
async def aget_responses(
|
||||
response_id: str,
|
||||
|
|
@ -2024,8 +2079,10 @@ async def aget_responses(
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
stream: bool | None = None,
|
||||
starting_after: int | None = None,
|
||||
**kwargs,
|
||||
) -> ResponsesAPIResponse:
|
||||
) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator:
|
||||
"""
|
||||
Async: Fetch a response by its ID.
|
||||
|
||||
|
|
@ -2058,6 +2115,8 @@ async def aget_responses(
|
|||
extra_query=extra_query,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
stream=stream,
|
||||
starting_after=starting_after,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
|
@ -2071,7 +2130,7 @@ async def aget_responses(
|
|||
response = init_response
|
||||
|
||||
# Update the responses_api_response_id with the model_id
|
||||
if not isinstance(response, ResponsesAPIResponse): # pyright: ignore[reportUnnecessaryIsInstance] # handlers can return non-ResponsesAPIResponse objects at runtime
|
||||
if not isinstance(response, ResponsesAPIResponse):
|
||||
return response
|
||||
return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id(
|
||||
responses_api_response=response,
|
||||
|
|
@ -2101,8 +2160,14 @@ def get_responses(
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
stream: bool | None = None,
|
||||
starting_after: int | None = None,
|
||||
**kwargs,
|
||||
) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]:
|
||||
) -> (
|
||||
ResponsesAPIResponse
|
||||
| SyncResponsesAPIStreamingIterator
|
||||
| Coroutine[object, object, ResponsesAPIResponse | ResponsesAPIStreamingIterator]
|
||||
):
|
||||
"""
|
||||
Fetch a response by its ID.
|
||||
|
||||
|
|
@ -2146,6 +2211,19 @@ def get_responses(
|
|||
if responses_api_provider_config is None:
|
||||
raise ValueError(f"GET responses is not supported for {custom_llm_provider}")
|
||||
|
||||
stream_query: Final = _supported_retrieve_stream_query(
|
||||
requested=_requested_retrieve_stream_query(
|
||||
stream=stream,
|
||||
starting_after=starting_after,
|
||||
extra_query=extra_query,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
),
|
||||
responses_api_provider_config=responses_api_provider_config,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
drop_params=litellm_params.drop_params,
|
||||
)
|
||||
litellm_metadata: Final = _STR_DICT.validate_python(kwargs.get("litellm_metadata") or {})
|
||||
|
||||
local_vars.update(kwargs)
|
||||
|
||||
# Pre Call logging
|
||||
|
|
@ -2169,7 +2247,13 @@ def get_responses(
|
|||
litellm_params=litellm_params,
|
||||
logging_obj=litellm_logging_obj,
|
||||
extra_headers=extra_headers,
|
||||
extra_query={
|
||||
key: value for key, value in (extra_query or {}).items() if key not in RETRIEVE_STREAM_QUERY_KEYS
|
||||
},
|
||||
extra_body=extra_body,
|
||||
stream=stream_query.stream,
|
||||
starting_after=stream_query.starting_after,
|
||||
litellm_metadata=litellm_metadata,
|
||||
timeout=timeout or request_timeout,
|
||||
_is_async=_is_async,
|
||||
client=kwargs.get("client"),
|
||||
|
|
@ -2180,7 +2264,7 @@ def get_responses(
|
|||
if isinstance(response, ResponsesAPIResponse):
|
||||
response = ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id(
|
||||
responses_api_response=response,
|
||||
litellm_metadata=kwargs.get("litellm_metadata", {}),
|
||||
litellm_metadata=litellm_metadata,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
|
|
|
|||
54
litellm/responses/retrieve_stream_query.py
Normal file
54
litellm/responses/retrieve_stream_query.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
RETRIEVE_STREAM_QUERY_KEYS: Final = frozenset({"stream", "starting_after"})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RetrieveStreamQuery:
|
||||
stream: bool
|
||||
starting_after: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InvalidRetrieveStreamQuery:
|
||||
param: Literal["stream", "starting_after"]
|
||||
message: str
|
||||
|
||||
|
||||
def _parsed_stream(value: object) -> bool | InvalidRetrieveStreamQuery:
|
||||
match value:
|
||||
case None:
|
||||
return False
|
||||
case bool():
|
||||
return value
|
||||
case str() if value.lower() in ("true", "false"):
|
||||
return value.lower() == "true"
|
||||
case _:
|
||||
return InvalidRetrieveStreamQuery("stream", f"Invalid 'stream': expected true or false, got {value!r}")
|
||||
|
||||
|
||||
def _parsed_starting_after(value: object) -> int | None | InvalidRetrieveStreamQuery:
|
||||
match value:
|
||||
case None:
|
||||
return None
|
||||
case int() if not isinstance(value, bool) and value >= 0:
|
||||
return value
|
||||
case str() if value.isascii() and value.isdigit():
|
||||
return int(value)
|
||||
case _:
|
||||
return InvalidRetrieveStreamQuery(
|
||||
"starting_after", f"Invalid 'starting_after': expected an integer >= 0, got {value!r}"
|
||||
)
|
||||
|
||||
|
||||
def parse_retrieve_stream_query(
|
||||
stream: object, starting_after: object
|
||||
) -> RetrieveStreamQuery | InvalidRetrieveStreamQuery:
|
||||
parsed_stream: Final = _parsed_stream(stream)
|
||||
if isinstance(parsed_stream, InvalidRetrieveStreamQuery):
|
||||
return parsed_stream
|
||||
parsed_starting_after: Final = _parsed_starting_after(starting_after)
|
||||
if isinstance(parsed_starting_after, InvalidRetrieveStreamQuery):
|
||||
return parsed_starting_after
|
||||
return RetrieveStreamQuery(stream=parsed_stream, starting_after=parsed_starting_after)
|
||||
|
|
@ -366,6 +366,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
|
||||
_pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
|
||||
_earlier_rounds_spend: BilledRoundsSpend | None = None
|
||||
_replays_stored_stream: bool = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -560,7 +561,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED,
|
||||
):
|
||||
_response_obj: Final[object] = getattr(openai_responses_api_chunk, "response", None)
|
||||
_estimate_wanted: Final[bool] = _chunk_type in (
|
||||
_estimate_wanted: Final[bool] = not self._replays_stored_stream and _chunk_type in (
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
)
|
||||
|
|
@ -1055,7 +1056,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
"""
|
||||
created: Final = self._created_response
|
||||
generated: Final = self._generated_content + self._generated_tool_arguments
|
||||
if self.completed_response is not None or created is None or not generated:
|
||||
if self._replays_stored_stream or self.completed_response is not None or created is None or not generated:
|
||||
return None
|
||||
usage: Final = _estimate_usage_safely(
|
||||
self.model or "", self.request_data.get("input"), self.request_data, generated
|
||||
|
|
@ -1100,6 +1101,8 @@ class BaseResponsesAPIStreamingIterator:
|
|||
def _raise_if_ended_without_terminal_event(self) -> None:
|
||||
if self.completed_response is not None:
|
||||
return
|
||||
if self._replays_stored_stream and not self._yielded_first_chunk:
|
||||
return
|
||||
error: Final = self._stream_ended_early_error()
|
||||
self._handle_failure(error)
|
||||
if self._output_started:
|
||||
|
|
@ -1147,6 +1150,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
custom_llm_provider: str | None = None,
|
||||
request_data: dict[str, object] | None = None,
|
||||
call_type: str | None = None,
|
||||
replays_stored_stream: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
response,
|
||||
|
|
@ -1158,6 +1162,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
request_data,
|
||||
call_type,
|
||||
)
|
||||
self._replays_stored_stream = replays_stored_stream
|
||||
self.stream_iterator = SSEDecoder().aiter_bytes(response.aiter_bytes())
|
||||
|
||||
def __aiter__(self):
|
||||
|
|
@ -1170,6 +1175,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
await self._await_pending_logging()
|
||||
raise
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self.response.aclose()
|
||||
|
||||
async def _next_event(self) -> ResponsesAPIStreamingResponse:
|
||||
try:
|
||||
self._check_max_streaming_duration()
|
||||
|
|
@ -1239,6 +1247,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
custom_llm_provider: str | None = None,
|
||||
request_data: dict[str, object] | None = None,
|
||||
call_type: str | None = None,
|
||||
replays_stored_stream: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
response,
|
||||
|
|
@ -1250,11 +1259,15 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
request_data,
|
||||
call_type,
|
||||
)
|
||||
self._replays_stored_stream = replays_stored_stream
|
||||
self.stream_iterator = SSEDecoder().iter_bytes(response.iter_bytes())
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def close(self) -> None:
|
||||
self.response.close()
|
||||
|
||||
def __next__(self):
|
||||
try:
|
||||
self._check_max_streaming_duration()
|
||||
|
|
|
|||
|
|
@ -2018,16 +2018,15 @@ def client(original_function):
|
|||
# MODEL CALL
|
||||
result = original_function(*args, **kwargs)
|
||||
end_time = datetime.datetime.now()
|
||||
if _is_streaming_request(
|
||||
kwargs=kwargs,
|
||||
call_type=call_type,
|
||||
):
|
||||
if _is_streaming_request(kwargs=kwargs, call_type=call_type) or _is_converted_stream_result(result):
|
||||
if "complete_response" in kwargs and kwargs["complete_response"] is True:
|
||||
chunks = []
|
||||
for idx, chunk in enumerate(result):
|
||||
chunks.append(chunk)
|
||||
return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None))
|
||||
else:
|
||||
logging_obj.stream = True
|
||||
logging_obj.model_call_details["stream"] = True
|
||||
# RETURN RESULT
|
||||
update_response_metadata: _ResponseMetadataUpdater = litellm_utils.update_response_metadata
|
||||
update_response_metadata(
|
||||
|
|
@ -2564,6 +2563,8 @@ _NEVER_STREAMING_CALL_TYPES: Final = frozenset(
|
|||
}
|
||||
)
|
||||
|
||||
_RESULT_DECIDES_STREAMING_CALL_TYPES: Final = frozenset({"get_responses", "aget_responses"})
|
||||
|
||||
|
||||
def _is_streaming_request(
|
||||
kwargs: dict[str, object],
|
||||
|
|
@ -2575,7 +2576,7 @@ def _is_streaming_request(
|
|||
- if "stream=True" in kwargs (litellm chat completion, litellm text completion, litellm messages)
|
||||
- if call_type is generate_content_stream or agenerate_content_stream (litellm google genai)
|
||||
"""
|
||||
if call_type in _NEVER_STREAMING_CALL_TYPES:
|
||||
if call_type in _NEVER_STREAMING_CALL_TYPES or call_type in _RESULT_DECIDES_STREAMING_CALL_TYPES:
|
||||
return False
|
||||
if "stream" in kwargs and kwargs["stream"] is True:
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import parse_qsl, urlsplit
|
||||
|
|
@ -297,9 +298,6 @@ def test_streamed_response_ids_are_stable_and_retrievable_on_each_alias(
|
|||
|
||||
@pytest.mark.parametrize("prefix", _ALIASES)
|
||||
def test_retrieve_query_reaches_the_provider(gateway: Gateway, prefix: str) -> None:
|
||||
pytest.skip(
|
||||
"BUG: GET /v1/responses/{id} drops include and starting_after; streamed GET returns 500 for provider SSE"
|
||||
)
|
||||
normal_raw_id: Final = f"resp_query_normal_{uuid.uuid4().hex}"
|
||||
stream_raw_id: Final = f"resp_query_stream_{uuid.uuid4().hex}"
|
||||
normal_include: Final = ["message.output_text.logprobs", "reasoning.encrypted_content"]
|
||||
|
|
@ -413,6 +411,83 @@ def test_retrieve_query_reaches_the_provider(gateway: Gateway, prefix: str) -> N
|
|||
assert actual_queries == expected_queries, requests
|
||||
|
||||
|
||||
def _resumed_delta() -> bytes:
|
||||
event: Final = {
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 4,
|
||||
"item_id": "msg_resumed",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "resumed",
|
||||
"logprobs": [],
|
||||
}
|
||||
return f"data: {json.dumps(event)}\n\n".encode()
|
||||
|
||||
|
||||
def _sorted_query(request: Request) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(sorted(parse_qsl(urlsplit(request.target).query, keep_blank_values=True)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _ALIASES)
|
||||
def test_resuming_past_the_last_event_ends_the_stream_without_events(gateway: Gateway, prefix: str) -> None:
|
||||
raw_id: Final = f"resp_resume_end_{uuid.uuid4().hex}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
if request.method == "POST":
|
||||
return Reply(body=_response(raw_id))
|
||||
return Reply(content_type="text/event-stream", chunks=())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire)
|
||||
client: Final = _client(gateway, gateway.key, prefix)
|
||||
created: Final = client.responses.create(model=model, input="resume past the end")
|
||||
events: Final = tuple(client.responses.retrieve(created.id, stream=True, starting_after=7))
|
||||
requests: Final = drain_contract_requests(wire)
|
||||
|
||||
assert events == ()
|
||||
assert [(request.method, urlsplit(request.target).path) for request in requests] == [
|
||||
("POST", "/v1/responses"),
|
||||
("GET", f"/v1/responses/{raw_id}"),
|
||||
], requests
|
||||
assert _sorted_query(requests[1]) == (("starting_after", "7"), ("stream", "true")), requests[1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _ALIASES)
|
||||
def test_a_client_leaving_a_resumed_stream_releases_the_provider_connection(gateway: Gateway, prefix: str) -> None:
|
||||
raw_id: Final = f"resp_resume_leave_{uuid.uuid4().hex}"
|
||||
release: Final = threading.Event()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
if request.method == "POST":
|
||||
return Reply(body=_response(raw_id))
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(_resumed_delta(), b"data: [DONE]\n\n"),
|
||||
gate_after_first=release,
|
||||
heartbeat=b": keepalive\n\n",
|
||||
pause_between_chunks=0.05,
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire)
|
||||
client: Final = _client(gateway, gateway.key, prefix)
|
||||
created: Final = client.responses.create(model=model, input="leave a resumed stream")
|
||||
stream: Final = client.responses.retrieve(created.id, stream=True, starting_after=3)
|
||||
try:
|
||||
first: Final = next(iter(stream))
|
||||
stream.close()
|
||||
left_from: Final = wire.disconnected.get(timeout=30)
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
assert first.type == "response.output_text.delta"
|
||||
assert urlsplit(left_from).path == f"/v1/responses/{raw_id}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix", _ALIASES)
|
||||
def test_delete_forwards_the_provider_id_once_and_reports_deletion(
|
||||
gateway: Gateway,
|
||||
|
|
|
|||
|
|
@ -181,6 +181,20 @@ class TestNonInferenceCallBilling:
|
|||
assert is_unbilled_non_inference_call_from_params("aresponses", {}, self.BACKGROUND) is False
|
||||
assert is_unbilled_non_inference_call_from_params("acompletion", {"metadata": self.MANAGED_READ}, None) is False
|
||||
|
||||
def test_a_streamed_read_prices_like_a_plain_read_of_the_same_response(self):
|
||||
from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse
|
||||
|
||||
for stored in (self.BACKGROUND, self.FOREGROUND):
|
||||
finished = ResponsesAPIResponse(**{**stored, "created_at": 1, "output": []})
|
||||
completed_event = ResponseCompletedEvent(type="response.completed", response=finished)
|
||||
for metadata in (None, self.MANAGED_READ, self.STATUS_POLL):
|
||||
plain_read = is_unbilled_non_inference_call("aget_responses", metadata, finished)
|
||||
assert is_unbilled_non_inference_call("aget_responses", metadata, completed_event) is plain_read
|
||||
assert (
|
||||
is_unbilled_non_inference_call("aget_responses", metadata, completed_event.model_dump())
|
||||
is plain_read
|
||||
)
|
||||
|
||||
def test_a_tracked_streamed_background_create_is_free_in_every_shape_its_stream_logs(self):
|
||||
from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse
|
||||
|
||||
|
|
|
|||
|
|
@ -1207,11 +1207,8 @@ class TestCheckResponsesCost:
|
|||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status = "completed"
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
mock_aget.return_value = _terminal_response()
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
assert [_is_billed(call[1]["litellm_metadata"]) for call in mock_aget.call_args_list] == [False]
|
||||
|
|
@ -1252,7 +1249,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = MagicMock(status="completed")
|
||||
mock_aget.return_value = _terminal_response()
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
(billed,) = billed_logger.billed
|
||||
|
|
|
|||
|
|
@ -950,7 +950,13 @@ def _request_for(method: str, path: str, body: bytes):
|
|||
return {"type": "http.request", "body": body, "more_body": False}
|
||||
|
||||
return Request(
|
||||
{"type": "http", "method": method, "path": path, "headers": [(b"content-type", b"application/json")]},
|
||||
{
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"path": path,
|
||||
"query_string": b"",
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
},
|
||||
receive,
|
||||
)
|
||||
|
||||
|
|
@ -1269,6 +1275,179 @@ async def test_a_streamed_background_create_bills_its_finished_response_once_und
|
|||
assert billed["request_tags"] == ["nightly"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_streamed_read_bills_the_finished_job_once_before_the_client_sees_it_finish(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Reading a finished background job as a stream settles its bill like a plain read does, before
|
||||
the terminal event reaches the client, and a later read of the same job bills nothing more."""
|
||||
from fastapi import Response as FastAPIResponse
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.response_api_endpoints.endpoints import (
|
||||
_bill_finished_background_response,
|
||||
_streamed_retrieval_data_generator,
|
||||
)
|
||||
|
||||
spend_events: Final = _capture_spend_events(monkeypatch)
|
||||
stored: Final = _one_pending_background_response()
|
||||
finished: Final = _background_response(stored.response_id, "completed", 100)
|
||||
stream: Final = MagicMock(completed_response=None)
|
||||
|
||||
async def replayed_events(**_: object):
|
||||
yield "event: response.output_text.delta\n\n"
|
||||
stream.completed_response = MagicMock(response=finished)
|
||||
yield "event: response.completed\n\n"
|
||||
|
||||
status_before_each_event: Final[list[str]] = []
|
||||
with _proxy_over(stored.prisma_client, stored.proxy_logging, _configured_deployment_router(finished)):
|
||||
async for _ in _streamed_retrieval_data_generator(
|
||||
replayed_events, stored.response_id, True, stream, UserAPIKeyAuth(api_key="sk-reader"), {}
|
||||
):
|
||||
status_before_each_event.extend(row.status for row in stored.table.rows)
|
||||
await _bill_finished_background_response(
|
||||
response=finished,
|
||||
response_id=stored.response_id,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-reader"),
|
||||
fastapi_response=FastAPIResponse(),
|
||||
read_settles_billing=True,
|
||||
)
|
||||
await _drain_spend_events()
|
||||
|
||||
assert status_before_each_event == ["queued", "completed"]
|
||||
(billed,) = spend_events.payloads
|
||||
assert billed["prompt_tokens"] == 100
|
||||
|
||||
|
||||
_RESUMED_UPSTREAM: Final = "https://resume.example/v1"
|
||||
|
||||
|
||||
def _resumed_events(response_id: str) -> tuple[Mapping[str, object], ...]:
|
||||
return (
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 4,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "hi",
|
||||
"logprobs": [],
|
||||
},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"sequence_number": 5,
|
||||
"response": _background_response(response_id, "completed", 100).model_dump(mode="json"),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sse_events(text: str) -> tuple[Mapping[str, object], ...]:
|
||||
frames: Final = (frame for frame in text.split("\n\n") if "data: " in frame and "data: [DONE]" not in frame)
|
||||
return tuple(
|
||||
json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) for frame in frames
|
||||
)
|
||||
|
||||
|
||||
async def _read_through_proxy(
|
||||
monkeypatch: pytest.MonkeyPatch, query: Mapping[str, str], upstream: httpx.Response
|
||||
) -> tuple[httpx.Response, tuple[httpx.Request, ...]]:
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5.6",
|
||||
"litellm_params": {"model": "openai/gpt-5.6", "api_base": _RESUMED_UPSTREAM, "api_key": "fixture-key"},
|
||||
"model_info": {"id": "deployment-read"},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override)
|
||||
with respx.mock(assert_all_called=False) as transport:
|
||||
route: Final = transport.get(_RESUMED_UPSTREAM + "/responses/resp_read").mock(return_value=upstream)
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client:
|
||||
result: Final = await client.get(f"/v1/responses/{_proxy_issued_id('resp_read')}", params=query)
|
||||
return result, tuple(call.request for call in route.calls)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_streamed_read_through_the_proxy_replays_the_events_after_the_cursor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
wire: Final = "".join(
|
||||
f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _resumed_events("resp_read")
|
||||
)
|
||||
result, upstream_requests = await _read_through_proxy(
|
||||
monkeypatch,
|
||||
{"stream": "true", "starting_after": "3"},
|
||||
httpx.Response(200, content=wire, headers={"Content-Type": "text/event-stream"}),
|
||||
)
|
||||
events: Final = _sse_events(result.text)
|
||||
|
||||
assert result.status_code == 200, result.text
|
||||
assert result.headers["content-type"].startswith("text/event-stream")
|
||||
assert [event["type"] for event in events] == ["response.output_text.delta", "response.completed"]
|
||||
assert events[-1]["response"]["id"] == _proxy_issued_id("resp_read")
|
||||
assert [dict(request.url.params) for request in upstream_requests] == [{"stream": "true", "starting_after": "3"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_streamed_read_cut_before_its_terminal_event_ends_in_a_responses_failure_event(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
delta: Final = _resumed_events("resp_read")[0]
|
||||
result, _ = await _read_through_proxy(
|
||||
monkeypatch,
|
||||
{"stream": "true", "starting_after": "3"},
|
||||
httpx.Response(
|
||||
200,
|
||||
content=f"event: {delta['type']}\ndata: {json.dumps(delta)}\n\n",
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
),
|
||||
)
|
||||
events: Final = _sse_events(result.text)
|
||||
|
||||
assert result.status_code == 200, result.text
|
||||
assert [event["type"] for event in events] == ["response.output_text.delta", "response.failed"]
|
||||
assert "before any terminal event" in events[-1]["response"]["error"]["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_plain_read_through_the_proxy_still_answers_json(monkeypatch: pytest.MonkeyPatch):
|
||||
result, upstream_requests = await _read_through_proxy(
|
||||
monkeypatch,
|
||||
{"starting_after": "3"},
|
||||
httpx.Response(200, json=_background_response("resp_read", "in_progress", 100).model_dump(mode="json")),
|
||||
)
|
||||
|
||||
assert result.status_code == 200, result.text
|
||||
assert result.json()["id"] == _proxy_issued_id("resp_read")
|
||||
assert [dict(request.url.params) for request in upstream_requests] == [{}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"query, bad_param",
|
||||
[
|
||||
({"stream": "maybe"}, "stream"),
|
||||
({"stream": "true", "starting_after": "-1"}, "starting_after"),
|
||||
({"stream": "true", "starting_after": "abc"}, "starting_after"),
|
||||
],
|
||||
)
|
||||
async def test_a_malformed_resume_query_through_the_proxy_is_a_400(
|
||||
monkeypatch: pytest.MonkeyPatch, query: Mapping[str, str], bad_param: str
|
||||
):
|
||||
result, upstream_requests = await _read_through_proxy(monkeypatch, query, httpx.Response(500))
|
||||
|
||||
assert result.status_code == 400, result.text
|
||||
assert bad_param in result.json()["error"]["message"]
|
||||
assert upstream_requests == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_read_that_bills_a_response_before_its_stream_ends_leaves_the_stream_nothing_to_bill(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,9 @@
|
|||
import json
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import httpx
|
||||
|
|
@ -6,6 +12,9 @@ import pytest
|
|||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from tests.unit.proxy.conftest import httpx_transport
|
||||
|
||||
pytestmark: Final = pytest.mark.usefixtures(httpx_transport.__name__)
|
||||
|
|
@ -107,3 +116,358 @@ async def test_cancel_responses_404_surfaces_openai_api_error_with_exact_url(
|
|||
assert exc_info.value.status_code == 404
|
||||
assert len(requests) == 1
|
||||
assert str(requests[0].request.url) == expected_url
|
||||
|
||||
|
||||
StreamProvider: TypeAlias = Literal["openai", "azure"]
|
||||
_STORED_RESPONSE_ID: Final = "resp_background"
|
||||
_STREAM_HEADERS: Final = {"content-type": "text/event-stream"}
|
||||
|
||||
|
||||
def _stored_response(status: str) -> dict[str, object]:
|
||||
return {
|
||||
"id": _STORED_RESPONSE_ID,
|
||||
"object": "response",
|
||||
"created_at": 1761600000,
|
||||
"status": status,
|
||||
"background": True,
|
||||
"model": "gpt-5.5",
|
||||
"output": [],
|
||||
"usage": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15},
|
||||
}
|
||||
|
||||
|
||||
def _sse(event: dict[str, object]) -> bytes:
|
||||
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
|
||||
|
||||
|
||||
def _delta_event(sequence_number: int) -> dict[str, object]:
|
||||
return {
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": sequence_number,
|
||||
"item_id": "msg_1",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "hi",
|
||||
"logprobs": [],
|
||||
}
|
||||
|
||||
|
||||
def _replay_body(*events: dict[str, object]) -> bytes:
|
||||
return b"".join(_sse(event) for event in events)
|
||||
|
||||
|
||||
_REPLAY_AFTER_3: Final = _replay_body(
|
||||
_delta_event(4),
|
||||
{"type": "response.completed", "sequence_number": 5, "response": _stored_response("completed")},
|
||||
)
|
||||
|
||||
|
||||
def _retrieve_route(mock_router: respx.MockRouter, provider: StreamProvider) -> respx.Route:
|
||||
match provider:
|
||||
case "openai":
|
||||
return mock_router.get(f"https://api.openai.com/v1/responses/{_STORED_RESPONSE_ID}")
|
||||
case "azure":
|
||||
return mock_router.get(f"https://example-resource.openai.azure.com/openai/responses/{_STORED_RESPONSE_ID}")
|
||||
|
||||
|
||||
def _retrieve_kwargs(provider: StreamProvider) -> dict[str, object]:
|
||||
match provider:
|
||||
case "openai":
|
||||
return {"custom_llm_provider": "openai", "api_key": "sk-test", "api_base": "https://api.openai.com/v1"}
|
||||
case "azure":
|
||||
return {
|
||||
"custom_llm_provider": "azure",
|
||||
"api_key": "sk-test",
|
||||
"api_base": "https://example-resource.openai.azure.com",
|
||||
"api_version": "2025-03-01-preview",
|
||||
}
|
||||
|
||||
|
||||
async def _retrieve(sync_mode: bool, **kwargs: object) -> object:
|
||||
if sync_mode:
|
||||
return litellm.get_responses(response_id=_STORED_RESPONSE_ID, **kwargs)
|
||||
return await litellm.aget_responses(response_id=_STORED_RESPONSE_ID, **kwargs)
|
||||
|
||||
|
||||
async def _event_types(stream: object) -> tuple[str, ...]:
|
||||
if isinstance(stream, SyncResponsesAPIStreamingIterator):
|
||||
return tuple(event.type for event in stream)
|
||||
assert isinstance(stream, ResponsesAPIStreamingIterator)
|
||||
return tuple([event.type async for event in stream])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.parametrize("provider", ("openai", "azure"))
|
||||
@pytest.mark.parametrize(
|
||||
("resume_kwargs", "expected_query"),
|
||||
(
|
||||
({"stream": True, "starting_after": 3}, {"stream": "true", "starting_after": "3"}),
|
||||
({"stream": True, "starting_after": 0}, {"stream": "true", "starting_after": "0"}),
|
||||
({"stream": True}, {"stream": "true"}),
|
||||
({"extra_query": {"stream": "true", "starting_after": "3"}}, {"stream": "true", "starting_after": "3"}),
|
||||
(
|
||||
{"stream": True, "starting_after": 1, "extra_query": {"starting_after": 3, "include[]": "usage"}},
|
||||
{"stream": "true", "starting_after": "3", "include[]": "usage"},
|
||||
),
|
||||
),
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_retrieve_replays_the_stored_events_after_the_cursor(
|
||||
provider: StreamProvider,
|
||||
sync_mode: bool,
|
||||
resume_kwargs: dict[str, object],
|
||||
expected_query: dict[str, str],
|
||||
) -> None:
|
||||
with respx.mock() as mock_router:
|
||||
route: Final = _retrieve_route(mock_router, provider).mock(
|
||||
return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3)
|
||||
)
|
||||
stream: Final = await _retrieve(sync_mode, **_retrieve_kwargs(provider), **resume_kwargs)
|
||||
event_types: Final = await _event_types(stream)
|
||||
(call,) = route.calls
|
||||
|
||||
sent_query: Final = dict(call.request.url.params)
|
||||
api_version: Final = {"api-version": "2025-03-01-preview"} if provider == "azure" else {}
|
||||
assert sent_query == {**expected_query, **api_version}
|
||||
assert event_types == ("response.output_text.delta", "response.completed")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_without_stream_still_returns_the_response_object(sync_mode: bool) -> None:
|
||||
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 _retrieve(sync_mode, starting_after=3, **_retrieve_kwargs("openai"))
|
||||
(call,) = route.calls
|
||||
|
||||
assert isinstance(retrieved, ResponsesAPIResponse)
|
||||
assert retrieved.status == "completed"
|
||||
assert dict(call.request.url.params) == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_resuming_at_the_end_of_a_finished_stream_ends_without_events(sync_mode: bool) -> None:
|
||||
with respx.mock() as mock_router:
|
||||
_retrieve_route(mock_router, "openai").mock(return_value=httpx.Response(200, headers=_STREAM_HEADERS))
|
||||
stream: Final = await _retrieve(sync_mode, stream=True, starting_after=5, **_retrieve_kwargs("openai"))
|
||||
event_types: Final = await _event_types(stream)
|
||||
|
||||
assert event_types == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_resumed_stream_cut_before_its_terminal_event_raises(sync_mode: bool) -> None:
|
||||
with respx.mock() as mock_router:
|
||||
_retrieve_route(mock_router, "openai").mock(
|
||||
return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_replay_body(_delta_event(4)))
|
||||
)
|
||||
stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai"))
|
||||
with pytest.raises(litellm.APIConnectionError, match="before any terminal event"):
|
||||
await _event_types(stream)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_resumed_stream_relays_the_stored_response_without_inventing_usage(sync_mode: bool) -> None:
|
||||
stored_without_usage: Final = {key: value for key, value in _stored_response("completed").items() if key != "usage"}
|
||||
with respx.mock() as mock_router:
|
||||
_retrieve_route(mock_router, "openai").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
headers=_STREAM_HEADERS,
|
||||
content=_replay_body(
|
||||
_delta_event(4),
|
||||
{"type": "response.completed", "sequence_number": 5, "response": stored_without_usage},
|
||||
),
|
||||
)
|
||||
)
|
||||
stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai"))
|
||||
if isinstance(stream, SyncResponsesAPIStreamingIterator):
|
||||
events: Final = tuple(stream)
|
||||
else:
|
||||
assert isinstance(stream, ResponsesAPIStreamingIterator)
|
||||
events = tuple([event async for event in stream])
|
||||
|
||||
completed: Final = events[-1]
|
||||
assert completed.type == "response.completed"
|
||||
assert completed.response.usage is None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ClosableReplay(httpx.SyncByteStream, httpx.AsyncByteStream):
|
||||
body: bytes
|
||||
closed: threading.Event = field(default_factory=threading.Event)
|
||||
|
||||
def __iter__(self) -> Iterator[bytes]:
|
||||
yield self.body
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
yield self.body
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed.set()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self.closed.set()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_closing_a_resumed_stream_early_releases_the_upstream_connection(sync_mode: bool) -> None:
|
||||
upstream: Final = _ClosableReplay(_REPLAY_AFTER_3)
|
||||
with respx.mock() as mock_router:
|
||||
_retrieve_route(mock_router, "openai").mock(
|
||||
return_value=httpx.Response(200, headers=_STREAM_HEADERS, stream=upstream)
|
||||
)
|
||||
stream: Final = await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai"))
|
||||
if isinstance(stream, SyncResponsesAPIStreamingIterator):
|
||||
first_sync: Final = next(stream)
|
||||
stream.close()
|
||||
first_type: Final = first_sync.type
|
||||
else:
|
||||
assert isinstance(stream, ResponsesAPIStreamingIterator)
|
||||
first_async: Final = await anext(stream)
|
||||
await stream.aclose()
|
||||
first_type = first_async.type
|
||||
|
||||
assert first_type == "response.output_text.delta"
|
||||
assert upstream.closed.is_set()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_rejected_resume_keeps_the_provider_error(sync_mode: bool) -> None:
|
||||
error_message: Final = "Response cannot be streamed, it is more than 5 minutes old."
|
||||
with respx.mock() as mock_router:
|
||||
route: Final = _retrieve_route(mock_router, "openai").mock(
|
||||
return_value=httpx.Response(
|
||||
400, json={"error": {"message": error_message, "type": "invalid_request_error", "param": None}}
|
||||
)
|
||||
)
|
||||
with pytest.raises(openai.BadRequestError, match=error_message) as exc_info:
|
||||
await _retrieve(sync_mode, stream=True, starting_after=3, **_retrieve_kwargs("openai"))
|
||||
(call,) = route.calls
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert dict(call.request.url.params) == {"stream": "true", "starting_after": "3"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_resume_that_times_out_raises_a_timeout(sync_mode: bool) -> None:
|
||||
with respx.mock() as mock_router:
|
||||
route: Final = _retrieve_route(mock_router, "openai").mock(side_effect=httpx.ReadTimeout("slow"))
|
||||
with pytest.raises(litellm.Timeout):
|
||||
await _retrieve(sync_mode, stream=True, timeout=1.0, **_retrieve_kwargs("openai"))
|
||||
(call,) = route.calls
|
||||
|
||||
assert dict(call.request.url.params) == {"stream": "true"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"resume_kwargs",
|
||||
(
|
||||
{"stream": "maybe"},
|
||||
{"stream": True, "starting_after": -1},
|
||||
{"extra_query": {"stream": "true", "starting_after": "abc"}},
|
||||
{"extra_query": {"stream": "1"}},
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_malformed_resume_query_is_rejected_before_any_upstream_call(
|
||||
sync_mode: bool, resume_kwargs: dict[str, object]
|
||||
) -> None:
|
||||
with respx.mock(assert_all_called=False) as mock_router:
|
||||
route: Final = _retrieve_route(mock_router, "openai")
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
await _retrieve(sync_mode, **_retrieve_kwargs("openai"), **resume_kwargs)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert route.call_count == 0
|
||||
|
||||
|
||||
_UNSUPPORTED_PROVIDERS: Final = ("xai", "manus", "volcengine")
|
||||
|
||||
|
||||
def _unsupported_provider_kwargs(provider: str) -> dict[str, object]:
|
||||
return {
|
||||
"custom_llm_provider": provider,
|
||||
"api_key": f"{provider}-test",
|
||||
"api_base": f"https://{provider}.example/v1",
|
||||
}
|
||||
|
||||
|
||||
def _unsupported_provider_route(mock_router: respx.MockRouter, provider: str) -> respx.Route:
|
||||
return mock_router.get(url__regex=rf"^https://{provider}\.example/.*/responses/{_STORED_RESPONSE_ID}$")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.parametrize("provider", _UNSUPPORTED_PROVIDERS)
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_retrieve_on_an_unsupported_provider_is_a_400_without_drop_params(
|
||||
sync_mode: bool, provider: str
|
||||
) -> None:
|
||||
with respx.mock(assert_all_called=False) as mock_router:
|
||||
route: Final = _unsupported_provider_route(mock_router, provider)
|
||||
with pytest.raises(litellm.UnsupportedParamsError) as exc_info:
|
||||
await _retrieve(sync_mode, stream=True, starting_after=3, **_unsupported_provider_kwargs(provider))
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert route.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", (True, False))
|
||||
@pytest.mark.parametrize("provider", _UNSUPPORTED_PROVIDERS)
|
||||
@pytest.mark.asyncio
|
||||
async def test_streamed_retrieve_on_an_unsupported_provider_drops_the_stream_under_drop_params(
|
||||
sync_mode: bool, provider: str
|
||||
) -> None:
|
||||
with respx.mock() as mock_router:
|
||||
route: Final = _unsupported_provider_route(mock_router, provider).mock(
|
||||
return_value=httpx.Response(200, json=_stored_response("completed"))
|
||||
)
|
||||
retrieved: Final = await _retrieve(
|
||||
sync_mode, stream=True, starting_after=3, drop_params=True, **_unsupported_provider_kwargs(provider)
|
||||
)
|
||||
(call,) = route.calls
|
||||
|
||||
assert isinstance(retrieved, ResponsesAPIResponse)
|
||||
assert dict(call.request.url.params) == {}
|
||||
|
||||
|
||||
class _SuccessLog(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.logged: queue.Queue[object] = queue.Queue()
|
||||
|
||||
def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
|
||||
self.logged.put(response_obj)
|
||||
|
||||
|
||||
def test_a_sync_streamed_retrieve_logs_success_once_with_the_finished_response(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
success_log: Final = _SuccessLog()
|
||||
monkeypatch.setattr(litellm, "callbacks", [success_log])
|
||||
with respx.mock() as mock_router:
|
||||
_retrieve_route(mock_router, "openai").mock(
|
||||
return_value=httpx.Response(200, headers=_STREAM_HEADERS, content=_REPLAY_AFTER_3)
|
||||
)
|
||||
stream: Final = litellm.get_responses(
|
||||
response_id=_STORED_RESPONSE_ID, stream=True, starting_after=3, **_retrieve_kwargs("openai")
|
||||
)
|
||||
assert isinstance(stream, SyncResponsesAPIStreamingIterator)
|
||||
consumed: Final = tuple(stream)
|
||||
|
||||
first_logged: Final = success_log.logged.get(timeout=10)
|
||||
time.sleep(0.5)
|
||||
|
||||
assert len(consumed) == 2
|
||||
assert isinstance(first_logged, ResponsesAPIResponse)
|
||||
assert first_logged.status == "completed"
|
||||
assert success_log.logged.empty()
|
||||
|
|
|
|||
|
|
@ -244,6 +244,21 @@ def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_ser
|
|||
assert partial.usage.total_tokens == partial.usage.input_tokens + partial.usage.output_tokens
|
||||
|
||||
|
||||
def test_disconnecting_from_a_resumed_background_stream_bills_no_partial_answer():
|
||||
iterator: Final = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=_make_logging_obj(),
|
||||
custom_llm_provider="openai",
|
||||
replays_stored_stream=True,
|
||||
)
|
||||
for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)):
|
||||
iterator._process_chunk(json.dumps(event))
|
||||
|
||||
assert iterator.partial_response_on_disconnect() is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"events",
|
||||
[
|
||||
|
|
|
|||
3
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
3
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -67086,6 +67086,7 @@ export interface operations {
|
|||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["ResponsesAPIResponse"];
|
||||
"text/event-stream": string;
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
|
|
@ -70550,6 +70551,7 @@ export interface operations {
|
|||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["ResponsesAPIResponse"];
|
||||
"text/event-stream": string;
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
|
|
@ -81348,6 +81350,7 @@ export interface operations {
|
|||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["ResponsesAPIResponse"];
|
||||
"text/event-stream": string;
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue