feat(rust): stream native Responses calls
Some checks failed
LiteLLM Rust / rust-wheel (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled

Native Responses calls with stream=True used to decline to Python. The
bridge now opens the provider stream in Rust and hands its bytes to
Python's own Responses stream iterator, so callers get the same typed
events, ids, hooks and logging as the Python route.

Streams decline to Python when a cache is configured or when the model
needs a faked stream.
This commit is contained in:
narmstrong-berry 2026-10-15 16:30:19 +00:00 • committed by yujonglee
parent 0423ac1164
commit d6a4b71cd1
17 changed files with 782 additions and 66 deletions

View file

@ -32,6 +32,15 @@ struct PassThroughStream {
endpoint_type: &'static str,
}
/// Who logs a stream once it reaches the caller
#[derive(Clone, Copy, Debug)]
enum StreamLogging {
/// The adapter reproduces the pass-through iterator's billing over the delivered chunks
PassThrough(PassThroughStream),
/// The Python stream iterator the route hands its caller logs the stream itself
PythonIterator,
}
/// What the Messages stream iterator keeps for its end-of-stream billing.
struct DeliveredStream {
chunks: Py<PyList>,
@ -111,15 +120,21 @@ impl LegacyLogging {
}
}
fn stream_billing(&self) -> Option<PassThroughStream> {
fn stream_logging(&self) -> Option<StreamLogging> {
match self.operation {
LoggingOperation::Messages => Some(PassThroughStream {
LoggingOperation::Messages => Some(StreamLogging::PassThrough(PassThroughStream {
url_route: "/v1/messages",
endpoint_type: "anthropic",
}),
LoggingOperation::Completion | LoggingOperation::Responses | LoggingOperation::Ocr => {
None
}
})),
LoggingOperation::Responses => Some(StreamLogging::PythonIterator),
LoggingOperation::Completion | LoggingOperation::Ocr => None,
}
}
fn stream_billing(&self) -> Option<PassThroughStream> {
match self.stream_logging()? {
StreamLogging::PassThrough(billing) => Some(billing),
StreamLogging::PythonIterator => None,
}
}
@ -209,7 +224,10 @@ impl LegacyLogging {
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billing = match self.stream_logging().ok_or(HostError::MissingState)? {
StreamLogging::PassThrough(billing) => billing,
StreamLogging::PythonIterator => return Ok(()),
};
let billed = Streaming::Success.call(
py,
(
@ -535,7 +553,7 @@ impl LegacyLogging {
}
pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
if self.stream_billing().is_none() {
if self.stream_logging().is_none() {
return Err(HostError::MissingState.into());
}
if let Some(key) = &self.cache_key {
@ -551,6 +569,9 @@ impl LegacyLogging {
}
pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
if self.stream_billing().is_none() {
return Ok(());
}
let stream = self.stream.as_mut().ok_or(HostError::MissingState)?;
if stream.first_chunk.is_none() {
stream.first_chunk = Some(datetime(py, epoch_seconds())?);

View file

@ -99,6 +99,10 @@ impl InferenceHost {
self.view.argument(py, arguments, name)
}
pub fn request<'py>(&self, py: Python<'py>) -> PyResult<&Bound<'py, PyDict>> {
self.view.request(py)
}
pub fn close(&mut self) {
self.view.close();
}

View file

@ -1,21 +1,24 @@
use std::convert::Infallible;
use crate::{marshal::project_optional_fields, routes::inference::InferenceHost};
use bytes::Bytes;
use litellm_auth::AwsParams;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host_python::{InvokeError, PythonBinding, PythonHostCalls, PythonOwned, from_py};
use litellm_host_python::{
InvokeError, PythonBinding, PythonHostCalls, PythonOwned, from_py, to_py,
};
use litellm_inference_responses::{
Error,
ids::{self, IdScope},
route::Responses,
types::ResponsesCall,
types::{ResponsesCall, ResponsesStreamHead},
};
use litellm_router_types::LitellmParams;
use pyo3::{
exceptions::PyValueError,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
types::{PyBytes, PyDict},
};
pub(super) struct ResponsesPythonHost {
@ -184,15 +187,6 @@ impl PythonBinding for ResponsesPythonHost {
) -> Result<ResponsesCall, InvokeError<Error>> {
let call = project(&mut self.host, py, arguments).map_err(InvokeError::Python)?;
self.ids = id_scope(&self.host, py, arguments, &call).map_err(InvokeError::Python)?;
if call
.optional_params
.get("stream")
.is_some_and(|value| value == &serde_json::Value::Bool(true))
{
return Err(InvokeError::Native(Error::Unsupported(
"native Python responses streaming",
)));
}
Ok(call)
}
@ -210,24 +204,17 @@ impl PythonBinding for ResponsesPythonHost {
fn encode_stream_head(
&mut self,
_py: Python<'_>,
head: <Responses as litellm_host::protocol::Protocol>::StreamHead,
py: Python<'_>,
head: ResponsesStreamHead,
) -> PyResult<Py<PyAny>> {
let _ = head;
Err(pyo3::exceptions::PyRuntimeError::new_err(
"native Python Responses streaming is not supported",
))
py.import(super::ROUTE_HOST_MODULE)?
.getattr("stream_head")?
.call1((to_py(py, &head.headers)?, self.host.request(py)?))
.map(Bound::unbind)
}
fn encode_chunk(
&mut self,
_py: Python<'_>,
chunk: <Responses as litellm_host::protocol::Protocol>::Chunk,
) -> PyResult<Py<PyAny>> {
let _ = chunk;
Err(pyo3::exceptions::PyRuntimeError::new_err(
"native Python Responses streaming is not supported",
))
fn encode_chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult<Py<PyAny>> {
Ok(PyBytes::new(py, &chunk).into_any().unbind())
}
fn map_error(&self, py: Python<'_>, error: Error) -> PyResult<PyErr> {

View file

@ -12,7 +12,6 @@ use litellm_host_python::present;
use litellm_inference_responses::{ResponsesRoute, route::Responses};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use serde_json::Value;
pub(crate) use websocket::ResponsesWebSocketConnection;
use super::{
@ -58,15 +57,6 @@ fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> Py
"native HTTP responses provider",
));
}
if argument("stream")?
.map(|value| litellm_host_python::from_py::<Value>(&value))
.transpose()?
.is_some_and(|value| value == Value::Bool(true))
{
return Err(RustBridgeDeclined::new_err(
"native Python responses streaming",
));
}
let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE);
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost::new(host))
}

View file

@ -97,14 +97,6 @@ def run(execution: Execution, open_stream: Callable[[Execution, object], object]
return _hand_off(execution, _settled(execution.send(None)), open_stream, cleanup)
def synchronous(native: Callable[[_Call], object]) -> Callable[[_Call], object]:
return lambda call: run(_execution(native(call)), SyncStream)
def asynchronous(native: Callable[[_Call], object]) -> Callable[[_Call], Awaitable[object]]:
return lambda call: drive(_execution(native(call)), AsyncStream)
def _execution(value: object) -> Execution:
return cast("Execution", value) # cast-ok: native entrypoints return Execution per _native.pyi
@ -191,3 +183,15 @@ class SyncStream(Iterator[object]):
def _finish(self) -> None:
self._done = True
self._execution.close()
def synchronous(
native: Callable[[_Call], object], open_stream: Callable[[Execution, object], object] = SyncStream
) -> Callable[[_Call], object]:
return lambda call: run(_execution(native(call)), open_stream)
def asynchronous(
native: Callable[[_Call], object], open_stream: Callable[[Execution, object], object] = AsyncStream
) -> Callable[[_Call], Awaitable[object]]:
return lambda call: drive(_execution(native(call)), open_stream)

View file

@ -1,12 +1,16 @@
from __future__ import annotations
import os
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, cast # noqa: TID251 # narrows caller json the checker cannot see into
from typing import Final, Protocol, cast # noqa: TID251 # narrows caller json the checker cannot see into
import httpx
import litellm
from litellm import get_llm_provider
from litellm.litellm_core_utils import core_helpers
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig
from litellm.llms.databricks.responses.transformation import DatabricksResponsesAPIConfig
@ -17,6 +21,7 @@ from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfi
from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig
from litellm.rust_bridge import failures
from litellm.rust_bridge.public_call import inference_decline_reason, known_params, optional_str
from litellm.rust_bridge.responses.streams import StreamSource
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams, ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
@ -24,6 +29,7 @@ from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
PARAMETERS: Final = tuple(ResponsesAPIOptionalRequestParams.__annotations__)
_REQUEST_FIELDS: Final = frozenset((*PARAMETERS, "model", "input", "litellm_logging_obj"))
_NATIVE_PARAMETERS: Final = (*PARAMETERS, "api_version")
_MANTLE_PARAMETERS: Final = (
*_NATIVE_PARAMETERS,
@ -56,6 +62,15 @@ def body_params(request: Mapping[str, object]) -> Mapping[str, object]:
return known_params(PARAMETERS, request)
class _HeaderProcessor(Protocol):
def __call__(self, response_headers: httpx.Headers) -> Mapping[str, object]: ...
_PROCESS_HEADERS: Final = cast( # cast-ok: the shared header projection is untyped
_HeaderProcessor, core_helpers.process_response_headers
)
def connection_defaults(provider: str) -> tuple[str | None, str | None]:
if provider in ("databricks", "bedrock_mantle"):
return None, None
@ -92,6 +107,45 @@ def response(value: Mapping[str, object]) -> ResponsesAPIResponse:
return ResponsesAPIResponse.model_validate(value)
def stream_head(headers: Sequence[tuple[str, str]], request: Mapping[str, object]) -> Mapping[str, object]:
"""The opened stream's headers for call logging, and what Python's Responses stream iterator is built from"""
model, provider = _resolved_model(request)
config: Final = _responses_config(provider, model)
logger: Final = request.get("litellm_logging_obj")
if config is None or not isinstance(logger, Logging):
raise RuntimeError("a native Responses stream opened without its provider config or call logger")
return {
"additional_headers": dict(_PROCESS_HEADERS(httpx.Headers(list(headers)))),
"source": StreamSource(
headers=tuple(headers),
model=model,
config=config,
logging_obj=logger,
litellm_metadata=_litellm_metadata(request),
custom_llm_provider=provider,
request_data={
"input": request.get("input"),
**{name: request[name] for name in PARAMETERS if request.get(name) is not None},
"litellm_params": dict(
GenericLiteLLMParams(
**{name: value for name, value in request.items() if name not in _REQUEST_FIELDS}
)
),
},
call_kwargs=request,
),
}
def _litellm_metadata(
request: Mapping[str, object],
) -> dict[str, object] | None: # mutable-ok: the caller's own dict, which the stream iterator writes into
metadata: Final = request.get("litellm_metadata", {})
if not isinstance(metadata, dict):
return None
return cast(dict[str, object], metadata) # cast-ok: caller json
def map_failure(error: Exception, request: Mapping[str, object]) -> Exception:
_, provider = _resolved_model(request)
return failures.map_native_failure(
@ -123,9 +177,22 @@ def decline_reason(request: Mapping[str, object]) -> str | None:
inference_decline_reason(_MANTLE_PARAMETERS if provider == "bedrock_mantle" else _NATIVE_PARAMETERS, request)
or _provider_decline_reason(provider, model, request)
or _model_rules_decline_reason(config, model, request)
or _stream_decline_reason(config, provider, model, request)
)
def _stream_decline_reason(
config: OpenAIResponsesAPIConfig, provider: str, model: str, request: Mapping[str, object]
) -> str | None:
if request.get("stream") is not True:
return None
if litellm.cache is not None:
return "native Responses streaming does not implement caching"
if config.should_fake_stream(model=model, stream=True, custom_llm_provider=provider):
return "native Responses does not implement faked streams"
return None
def _responses_config(provider: str, model: str) -> OpenAIResponsesAPIConfig | None:
if provider != "bedrock_mantle":
return _CONFIGS.get(provider)

View file

@ -3,9 +3,11 @@ from __future__ import annotations
from collections.abc import Awaitable
from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator
from litellm.rust_bridge import execution
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.public_call import NativeCall
from litellm.rust_bridge.responses import streams
from litellm.types.llms.openai import ResponsesAPIResponse
@ -13,27 +15,27 @@ class NativeResponses(Protocol):
def __call__(
self,
request: NativeCall,
) -> ResponsesAPIResponse: ...
) -> ResponsesAPIResponse | SyncResponsesAPIStreamingIterator: ...
class NativeAresponses(Protocol):
def __call__(
self,
request: NativeCall,
) -> Awaitable[ResponsesAPIResponse]: ...
) -> Awaitable[ResponsesAPIResponse | ResponsesAPIStreamingIterator]: ...
def _responses_binding(value: object) -> NativeResponses | None:
if not callable(value):
return None
native: Final = execution.synchronous(value)
native: Final = execution.synchronous(value, streams.open_sync)
return cast("NativeResponses", native) # cast-ok: the execution adapter returns the route's result
def _aresponses_binding(value: object) -> NativeAresponses | None:
if not callable(value):
return None
native: Final = execution.asynchronous(value)
native: Final = execution.asynchronous(value, streams.open_async)
return cast("NativeAresponses", native) # cast-ok: the execution adapter returns the route's result

View file

@ -0,0 +1,123 @@
from __future__ import annotations
import datetime
from collections.abc import AsyncIterator, Iterator, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, cast # noqa: TID251 # Logging's litellm params are untyped
import httpx
from pydantic import TypeAdapter
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
ResponsesAPIStreamingIterator,
SyncResponsesAPIStreamingIterator,
)
from litellm.rust_bridge.execution import AsyncStream, Execution, SyncStream
from litellm.rust_bridge.host.callbacks_legacy_python import finalize
from litellm.types.utils import CallTypes
_HEAD: Final = TypeAdapter(Mapping[str, object])
class _CallLogger(Protocol):
litellm_params: Mapping[str, object]
model_call_details: MutableMapping[str, object] # mutable-ok: the Logging object's own call details
@dataclass(frozen=True, slots=True)
class StreamSource:
"""What Python's Responses handler hands its stream iterator, for a stream the native route opened"""
headers: Sequence[tuple[str, str]]
model: str
config: BaseResponsesAPIConfig
logging_obj: Logging
litellm_metadata: dict[str, object] | None # mutable-ok: the iterator writes stream timing into it
custom_llm_provider: str
request_data: Mapping[str, object]
call_kwargs: Mapping[str, object]
def _source(head: object) -> StreamSource:
source: Final = _HEAD.validate_python(head).get("source")
if not isinstance(source, StreamSource):
raise TypeError("the native Responses stream head carries no stream source")
return source
def _chunk(value: object) -> bytes:
if not isinstance(value, bytes):
raise TypeError(f"a native Responses stream chunk is bytes, got {type(value).__name__}")
return value
class _AsyncNativeBytes(httpx.AsyncByteStream):
def __init__(self, stream: AsyncStream) -> None:
self._stream: Final = stream
async def __aiter__(self) -> AsyncIterator[bytes]:
async for chunk in self._stream:
yield _chunk(chunk)
async def aclose(self) -> None:
await self._stream.aclose()
class _SyncNativeBytes(httpx.SyncByteStream):
def __init__(self, stream: SyncStream) -> None:
self._stream: Final = stream
def __iter__(self) -> Iterator[bytes]:
return map(_chunk, self._stream)
def close(self) -> None:
self._stream.close()
def _response(source: StreamSource, stream: httpx.AsyncByteStream | httpx.SyncByteStream) -> httpx.Response:
return httpx.Response(200, headers=list(source.headers), stream=stream)
def _open(source: StreamSource, stream: BaseResponsesAPIStreamingIterator, asynchronous: bool) -> None:
"""Python's route marks its logger async or sync, which picks the handlers the iterator's logging runs, and
`@client` sets the stream's hidden params"""
logger: Final = cast(_CallLogger, source.logging_obj) # cast-ok: Logging's litellm params are untyped
logger.litellm_params = {**logger.litellm_params, "aresponses": asynchronous}
logger.model_call_details["litellm_params"] = logger.litellm_params
opened_at: Final = datetime.datetime.now() # noqa: DTZ005 # Logging start times are naive local datetimes
finalize(stream, source.logging_obj, dict(source.call_kwargs), source.logging_obj.start_time, opened_at)
def open_async(execution: Execution, head: object) -> ResponsesAPIStreamingIterator:
source: Final = _source(head)
stream: Final = ResponsesAPIStreamingIterator(
response=_response(source, _AsyncNativeBytes(AsyncStream(execution, head))),
model=source.model,
responses_api_provider_config=source.config,
logging_obj=source.logging_obj,
litellm_metadata=source.litellm_metadata,
custom_llm_provider=source.custom_llm_provider,
request_data=dict(source.request_data),
call_type=CallTypes.responses.value,
)
_open(source, stream, asynchronous=True)
return stream
def open_sync(execution: Execution, head: object) -> SyncResponsesAPIStreamingIterator:
source: Final = _source(head)
stream: Final = SyncResponsesAPIStreamingIterator(
response=_response(source, _SyncNativeBytes(SyncStream(execution, head))),
model=source.model,
responses_api_provider_config=source.config,
logging_obj=source.logging_obj,
litellm_metadata=source.litellm_metadata,
custom_llm_provider=source.custom_llm_provider,
request_data=dict(source.request_data),
call_type=CallTypes.responses.value,
)
_open(source, stream, asynchronous=False)
return stream

View file

@ -34,14 +34,20 @@ def recording_server() -> Generator[RecordingServer]:
def native_responses() -> Iterator[FallbackProbe]:
probe: Final = FallbackProbe()
python_aresponses: Final = dispatch._PYTHON_ARESPONSES # pyright: ignore[reportPrivateUsage] # the fallback under observation
python_responses: Final = dispatch._PYTHON_RESPONSES # pyright: ignore[reportPrivateUsage] # the fallback under observation
async def python_fallback(*args: object, **kwargs: object) -> object:
probe.calls += 1
return await python_aresponses(*args, **kwargs)
def python_sync_fallback(*args: object, **kwargs: object) -> object:
probe.calls += 1
return python_responses(*args, **kwargs)
with (
rebound(catalog, "RULES", (RouteRule(Route.RESPONSES, Rollout.RUST_OPT_OUT), *catalog.RULES)),
rebound(dispatch, "_PYTHON_ARESPONSES", python_fallback),
rebound(dispatch, "_PYTHON_RESPONSES", python_sync_fallback),
):
yield probe

View file

@ -7,6 +7,8 @@ from litellm.responses import main as python_responses
PYTHON_ARESPONSES: Final = cast(Callable[..., Awaitable[object]], python_responses.aresponses) # noqa: TID251 # the Python baseline under comparison
PUBLIC_ARESPONSES: Final = cast(Callable[..., Awaitable[object]], litellm.aresponses)
PYTHON_RESPONSES: Final = cast(Callable[..., object], python_responses.responses) # noqa: TID251 # the Python baseline under comparison
PUBLIC_RESPONSES: Final = cast(Callable[..., object], litellm.responses)
@dataclass(slots=True)

View file

@ -0,0 +1,303 @@
"""Native Responses streams must be indistinguishable from the Python implementation.
Each case streams one public Responses call through Python and through the native route against the same recorded
OpenAI event stream, then compares what reached the provider, the events the caller read, and what the success
callback logged. The native leg fails instead of silently falling back to Python
"""
import asyncio
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
import pytest
from pydantic import BaseModel, JsonValue, TypeAdapter
import litellm
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator
from .support import (
PUBLIC_ARESPONSES,
PUBLIC_RESPONSES,
PYTHON_ARESPONSES,
PYTHON_RESPONSES,
FallbackProbe,
)
from .test_native_parity import (
DEPLOYMENT_METADATA,
FUNCTION_CALL_RESPONSE,
MANTLE_MODEL,
TEXT_MESSAGE,
TEXT_RESPONSE,
WEATHER_TOOL,
)
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.isolation import isolated_callback_registries
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
pytestmark = pytest.mark.requires_rust_extension
MODEL: Final = "openai/gpt-6.1-sol"
_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_LOGGED: Final = TypeAdapter(dict[str, object])
_COMPARED_LOG_FIELDS: Final = (
"call_type",
"stream",
"custom_llm_provider",
"model",
"model_id",
"prompt_tokens",
"completion_tokens",
"total_tokens",
"response_cost",
"status",
"api_base",
"response",
)
_COMPARED_HIDDEN_PARAMS: Final = ("custom_llm_provider", "model_id", "api_base")
_COMPARED_RESPONSE_HEADERS: Final = (
"llm_provider-content-type",
"llm_provider-x-request-id",
"llm_provider-openai-processing-ms",
"request-id",
)
_SYNC_HANDLER_GRACE: Final = 0.2
_UPSTREAM_HEADERS: Final = MappingProxyType({"x-request-id": "req_stream", "openai-processing-ms": "12"})
def _opened(response: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
return {**response, "status": "in_progress", "output": [], "usage": None}
def _events(response: Mapping[str, JsonValue]) -> tuple[tuple[str, object], ...]:
"""The event stream OpenAI sends for a response with one output item"""
output: Final = response["output"]
assert isinstance(output, list)
item: Final = output[-1]
assert isinstance(item, dict)
opened: Final = (
("response.created", {"type": "response.created", "sequence_number": 0, "response": _opened(response)}),
(
"response.in_progress",
{"type": "response.in_progress", "sequence_number": 1, "response": _opened(response)},
),
(
"response.output_item.added",
{
"type": "response.output_item.added",
"sequence_number": 2,
"output_index": 0,
"item": {**item, "status": "in_progress"},
},
),
)
body: Final = (
(
("response.output_text.delta", _text_delta(item, 3, "Par")),
("response.output_text.delta", _text_delta(item, 4, "is")),
)
if item["type"] == "message"
else (
(
"response.function_call_arguments.delta",
{
"type": "response.function_call_arguments.delta",
"sequence_number": 3,
"item_id": item["id"],
"output_index": 0,
"delta": item["arguments"],
},
),
)
)
closed: Final = (
(
"response.output_item.done",
{"type": "response.output_item.done", "sequence_number": 5, "output_index": 0, "item": item},
),
("response.completed", {"type": "response.completed", "sequence_number": 6, "response": dict(response)}),
)
return (*opened, *body, *closed)
def _text_delta(item: Mapping[str, JsonValue], sequence: int, delta: str) -> dict[str, JsonValue]:
return {
"type": "response.output_text.delta",
"sequence_number": sequence,
"item_id": item["id"],
"output_index": 0,
"content_index": 0,
"delta": delta,
"logprobs": [],
}
@dataclass(frozen=True, slots=True)
class Leg:
request_path: str
request_body: Mapping[str, JsonValue]
events: tuple[Mapping[str, JsonValue], ...]
logged: Mapping[str, object]
hidden_params: Mapping[str, object]
hooks: tuple[str, ...]
fell_back: bool
ASPECTS: Final = ("dispatch", "request", "events", "logged", "hooks")
@dataclass(frozen=True, slots=True)
class Case:
name: str
kwargs: Mapping[str, object]
upstream: Mapping[str, JsonValue] = TEXT_RESPONSE
synchronous: bool = False
def params(self) -> Iterator[object]:
for aspect in ASPECTS:
yield pytest.param(self, aspect, id=f"{self.name}-{aspect}")
CASES: Final = (
Case("plain_text", {"input": "What is the capital of France?"}),
Case(
"function_tool",
{"input": "Weather in Paris?", "tools": [dict(WEATHER_TOOL)], "tool_choice": "auto"},
upstream=FUNCTION_CALL_RESPONSE,
),
Case(
"deployment_ids_and_encrypted_content",
{"input": "hi", "litellm_metadata": dict(DEPLOYMENT_METADATA)},
upstream={**TEXT_RESPONSE, "output": [{**TEXT_MESSAGE, "encrypted_content": "enc"}]},
),
Case("bedrock_mantle", {"model": MANTLE_MODEL, "input": "hi"}),
Case("sync_plain_text", {"input": "What is the capital of France?"}, synchronous=True),
)
def _event(value: object) -> Mapping[str, JsonValue]:
assert isinstance(value, BaseModel)
return _OBJECT.validate_python(value.model_dump(mode="json", exclude_none=True))
def _hidden_params(stream: object) -> Mapping[str, object]:
hidden: Final = _LOGGED.validate_python(getattr(stream, "_hidden_params"))
headers: Final = _LOGGED.validate_python(hidden["additional_headers"])
return MappingProxyType(
{
**{name: hidden.get(name) for name in _COMPARED_HIDDEN_PARAMS},
**{name: headers.get(name) for name in _COMPARED_RESPONSE_HEADERS},
}
)
def _sync_read(native: bool, kwargs: Mapping[str, object]) -> tuple[object, tuple[object, ...]]:
stream: Final = (PUBLIC_RESPONSES if native else PYTHON_RESPONSES)(**kwargs)
assert isinstance(stream, SyncResponsesAPIStreamingIterator)
return stream, tuple(stream)
async def _async_stream(native: bool, kwargs: Mapping[str, object]) -> ResponsesAPIStreamingIterator:
stream: Final = await (PUBLIC_ARESPONSES if native else PYTHON_ARESPONSES)(**kwargs)
assert isinstance(stream, ResponsesAPIStreamingIterator)
return stream
async def _read(case: Case, kwargs: Mapping[str, object], native: bool) -> tuple[object, tuple[object, ...]]:
if case.synchronous:
return await asyncio.to_thread(_sync_read, native, kwargs)
stream: Final = await _async_stream(native, kwargs)
return stream, tuple([event async for event in stream])
async def _leg(case: Case, server: RecordingServer, probe: FallbackProbe | None) -> Leg:
with isolated_callback_registries():
recorder: Final = RecordingLogger()
server.enqueue(ResponseSpec(body=None, events=_events(case.upstream), headers=dict(_UPSTREAM_HEADERS)))
kwargs: Final = {
"model": MODEL,
"api_key": "sk-test",
"api_base": server.base_url,
"callbacks": [recorder],
"stream": True,
**case.kwargs,
}
fallbacks_before: Final = 0 if probe is None else probe.calls
stream, events = await _read(case, kwargs, native=probe is not None)
await drain_logging()
success: Final = (
await asyncio.to_thread(recorder.wait_for, "log_success_event")
if case.synchronous
else await recorder.wait_for_async("async_log_success_event")
)
sent: Final = server.requests[-1]
standard: Final = _LOGGED.validate_python(_LOGGED.validate_python(success[0].kwargs)["standard_logging_object"])
return Leg(
request_path=sent.path,
request_body=_OBJECT.validate_python(sent.body),
events=tuple(map(_event, events)),
logged=MappingProxyType({name: standard.get(name) for name in _COMPARED_LOG_FIELDS}),
hidden_params=_hidden_params(stream),
hooks=recorder.names,
fell_back=probe is not None and probe.calls > fallbacks_before,
)
@pytest.fixture
def parity_server(recording_server: RecordingServer) -> RecordingServer:
recording_server.expected_requests = 2
return recording_server
@pytest.mark.asyncio
@pytest.mark.parametrize(("case", "aspect"), tuple(param for case in CASES for param in case.params()))
async def test_native_responses_streams_match_python(
case: Case, aspect: str, parity_server: RecordingServer, native_responses: FallbackProbe
) -> None:
python: Final = await _leg(case, parity_server, probe=None)
native: Final = await _leg(case, parity_server, probe=native_responses)
match aspect:
case "dispatch":
assert not native.fell_back
case "request":
assert (native.request_path, native.request_body) == (python.request_path, python.request_body)
case "events":
assert native.events == python.events
assert len(native.events) == len(_events(case.upstream))
case "logged":
assert (native.logged, native.hidden_params) == (python.logged, python.hidden_params)
case _:
assert native.hooks == python.hooks
async def _truncated_leg(server: RecordingServer, probe: FallbackProbe | None) -> tuple[str, str, tuple[str, ...]]:
with isolated_callback_registries():
recorder: Final = RecordingLogger()
server.enqueue(ResponseSpec(body=None, events=_events(TEXT_RESPONSE)[:-1]))
stream: Final = await _async_stream(
probe is not None,
{
"model": MODEL,
"api_key": "sk-test",
"api_base": server.base_url,
"callbacks": [recorder],
"stream": True,
"input": "hi",
},
)
with pytest.raises(litellm.APIConnectionError) as raised:
_ = [event async for event in stream]
await drain_logging()
await recorder.wait_for_async("async_log_failure_event")
await asyncio.sleep(_SYNC_HANDLER_GRACE)
return type(raised.value).__name__, str(raised.value), recorder.names
@pytest.mark.asyncio
async def test_a_native_stream_that_ends_early_fails_like_python(
parity_server: RecordingServer, native_responses: FallbackProbe
) -> None:
python: Final = await _truncated_leg(parity_server, probe=None)
native: Final = await _truncated_leg(parity_server, probe=native_responses)
assert native_responses.calls == 0
assert native == python

View file

@ -69,5 +69,4 @@ async def call(server: RecordingServer, *, asynchronous: bool = True, **kwargs:
SYNC_DEPLOYMENT_HOOK: Final = gap("native sync calls skip the async pre-call deployment hook")
STREAMING: Final = gap("native Responses streaming declines to Python", raises=RuntimeError)
BODY_HOOKS: Final = gap("native Responses does not run the handler body hooks")

View file

@ -13,7 +13,6 @@ from tests.test_litellm_rust.responses.data import (
REQUEST,
RESPONSES_EVENTS,
RESPONSES_RESPONSE,
STREAMING,
call,
)
from tests.test_litellm_rust.support.body_hooks import LegacyAgenticLoop, PlannedAgenticLoop, agentic_hooks
@ -34,7 +33,6 @@ def server(server: RecordingServer) -> Iterator[RecordingServer]:
@pytest.mark.asyncio
@STREAMING
async def test_body_streaming_deployment_hook_rewrites_chunks_before_delivery(server: RecordingServer) -> None:
seen: Final[list[str]] = []

View file

@ -62,7 +62,6 @@ async def test_provider_failure_is_terminal_and_shared_with_callbacks(
@pytest.mark.parametrize(
"options",
(
{"stream": True},
{"extra_body": {"provider_option": True}},
{"mock_response": "mock"},
{"num_retries": 1},

View file

@ -3,7 +3,7 @@ from typing import Final
import pytest
from tests.test_litellm_rust.responses.data import RESPONSES_EVENTS, RESPONSES_RESPONSE, STREAMING, call
from tests.test_litellm_rust.responses.data import RESPONSES_EVENTS, RESPONSES_RESPONSE, call
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed
from tests.test_litellm_rust.support.results import consume, response_id
@ -20,7 +20,6 @@ def server(server: RecordingServer) -> Iterator[RecordingServer]:
@pytest.mark.asyncio
@STREAMING
async def test_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk(
server: RecordingServer,
) -> None:
@ -45,7 +44,6 @@ async def test_stream_relays_provider_events_and_logs_success_once_after_the_las
@pytest.mark.asyncio
@STREAMING
async def test_sync_stream_relays_provider_events_and_logs_success_once(server: RecordingServer) -> None:
server.enqueue(STREAM)
recorder: Final = RecordingLogger()

View file

@ -395,3 +395,47 @@ def test_bedrock_mantle_connection_defaults_ignore_the_global_settings(monkeypat
monkeypatch.setattr(litellm, "api_key", "global")
monkeypatch.setattr(litellm, "api_base", "https://global.invalid")
assert connection_defaults("bedrock_mantle") == (None, None)
_STREAMING_MODEL_COST: Final = MappingProxyType(
{
"gpt-batch-only": {
"litellm_provider": "openai",
"mode": "responses",
"input_cost_per_token": 1e-6,
"output_cost_per_token": 1e-6,
"supports_native_streaming": False,
},
"gpt-streams": {
"litellm_provider": "openai",
"mode": "responses",
"input_cost_per_token": 1e-6,
"output_cost_per_token": 1e-6,
},
}
)
@pytest.mark.parametrize(
("model", "stream", "expected"),
[
pytest.param("openai/gpt-streams", True, None, id="native_stream"),
pytest.param(
"openai/gpt-batch-only", True, "native Responses does not implement faked streams", id="faked_stream"
),
pytest.param("openai/gpt-batch-only", False, None, id="faked_stream_model_without_stream"),
],
)
def test_streams_python_fakes_stay_on_python(
monkeypatch: pytest.MonkeyPatch, model: str, stream: bool, expected: str | None
) -> None:
monkeypatch.setattr(litellm, "model_cost", dict(_STREAMING_MODEL_COST))
assert decline_reason({"model": model, "input": "hi", "api_key": "sk-test", "stream": stream}) == expected
@pytest.mark.parametrize(
("stream", "expected"), [(True, "native Responses streaming does not implement caching"), (False, None)]
)
def test_cached_streams_stay_on_python(monkeypatch: pytest.MonkeyPatch, stream: bool, expected: str | None) -> None:
monkeypatch.setattr(litellm, "cache", litellm.Cache())
assert decline_reason({"model": "gpt-4.1", "input": "hi", "api_key": "sk-test", "stream": stream}) == expected

View file

@ -0,0 +1,169 @@
import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import datetime
from typing import Final
import pytest
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, SyncResponsesAPIStreamingIterator
from litellm.rust_bridge.execution import Complete, Step, StreamChunk, StreamOpen, drive, run
from litellm.rust_bridge.host.routes.responses import stream_head
from litellm.rust_bridge.responses.streams import StreamSource, open_async, open_sync
from litellm.types.llms.openai import ResponsesAPIStreamEvents
_HEADERS: Final = (("content-type", "text/event-stream"), ("x-request-id", "req_native"))
_RESPONSE: Final = {
"id": "resp_native",
"object": "response",
"created_at": 1,
"model": "gpt-4o",
"status": "completed",
"output": [
{
"type": "message",
"id": "msg_native",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Paris", "annotations": []}],
}
],
}
_EVENTS: Final = (
{
"type": "response.output_text.delta",
"sequence_number": 0,
"item_id": "msg_native",
"output_index": 0,
"content_index": 0,
"delta": "Paris",
},
{"type": "response.completed", "sequence_number": 1, "response": _RESPONSE},
)
def _sse(event: Mapping[str, object]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
@dataclass(slots=True)
class _Execution:
"""Opens with `head`, then yields each chunk before completing"""
head: object
chunks: tuple[object, ...]
sent: list[object] = field(default_factory=list) # mutable-ok: records what the driver fed back
closed: bool = False
def send(self, value: object) -> Step:
self.sent.append(value)
index: Final = len(self.sent) - 2
if index < 0:
return StreamOpen(self.head)
return StreamChunk(self.chunks[index]) if index < len(self.chunks) else Complete(None)
def throw(self, error: BaseException) -> Step:
raise error
def close(self) -> None:
self.closed = True
def _logger() -> Logging:
logger: Final = Logging(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="responses",
start_time=datetime.now(), # noqa: DTZ005 # Logging start times are naive local datetimes
litellm_call_id="native-stream",
function_id="native-stream",
)
logger.update_environment_variables(litellm_params={}, optional_params={}, model="gpt-4o")
return logger
def _request(logger: object, **fields: object) -> dict[str, object]:
return {"model": "openai/gpt-4o", "input": "hi", "stream": True, "litellm_logging_obj": logger, **fields}
def _execution(logger: Logging, chunks: tuple[object, ...] = tuple(map(_sse, _EVENTS))) -> _Execution:
return _Execution(head=stream_head(_HEADERS, _request(logger)), chunks=chunks)
@pytest.mark.asyncio
async def test_drive_streams_native_bytes_as_typed_events() -> None:
logger: Final = _logger()
execution: Final = _execution(logger)
stream: Final = await drive(execution, open_async)
assert isinstance(stream, ResponsesAPIStreamingIterator)
events: Final = [event async for event in stream]
assert [event.type for event in events] == [
ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
]
assert events[0].delta == "Paris"
assert events[-1].response.output[0].content[0].text == "Paris"
assert execution.closed
assert logger.model_call_details["litellm_params"]["aresponses"] is True
hidden: Final = stream._hidden_params
assert hidden["additional_headers"]["llm_provider-x-request-id"] == "req_native"
assert (hidden["litellm_call_id"], hidden["litellm_model_name"]) == ("native-stream", "openai/gpt-4o")
def test_sync_stream_streams_native_bytes_as_typed_events() -> None:
logger: Final = _logger()
execution: Final = _execution(logger)
stream: Final = run(execution, open_sync)
assert isinstance(stream, SyncResponsesAPIStreamingIterator)
assert [event.type for event in stream] == [
ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA,
ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
]
assert execution.closed
assert logger.model_call_details["litellm_params"]["aresponses"] is False
@pytest.mark.asyncio
async def test_a_head_without_a_stream_source_closes_the_execution() -> None:
execution: Final = _Execution(head={"additional_headers": {}}, chunks=())
with pytest.raises(TypeError, match="no stream source"):
await drive(execution, open_async)
assert execution.closed
@pytest.mark.asyncio
async def test_a_chunk_that_is_not_bytes_is_rejected() -> None:
stream: Final = await drive(_execution(_logger(), chunks=("event: text\n\n",)), open_async)
assert isinstance(stream, ResponsesAPIStreamingIterator)
with pytest.raises(TypeError, match="chunk is bytes, got str"):
_ = [event async for event in stream]
def test_stream_head_projects_headers_and_builds_the_iterator_request() -> None:
metadata: Final = {"model_info": {"id": "deployment"}}
request: Final = _request(_logger(), litellm_metadata=metadata, temperature=None, max_output_tokens=12)
head: Final = stream_head(_HEADERS, request)
assert head["additional_headers"]["llm_provider-x-request-id"] == "req_native"
source: Final = head["source"]
assert isinstance(source, StreamSource)
assert (source.model, source.custom_llm_provider) == ("gpt-4o", "openai")
assert source.litellm_metadata is metadata
assert source.request_data["max_output_tokens"] == 12
assert "temperature" not in source.request_data
assert "litellm_logging_obj" not in source.request_data["litellm_params"]
def test_stream_head_without_a_call_logger_raises() -> None:
with pytest.raises(RuntimeError, match="call logger"):
stream_head(_HEADERS, _request(None))