diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index d891b5a1f0c..2fb5e4670cf 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -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, @@ -111,15 +120,21 @@ impl LegacyLogging { } } - fn stream_billing(&self) -> Option { + fn stream_logging(&self) -> Option { 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 { + 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) -> 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) -> 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())?); diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index a5be7f51828..92cbd651bee 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -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(); } diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs index f660530713c..9c477533977 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/host.rs @@ -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> { 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: ::StreamHead, + py: Python<'_>, + head: ResponsesStreamHead, ) -> PyResult> { - 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: ::Chunk, - ) -> PyResult> { - 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> { + Ok(PyBytes::new(py, &chunk).into_any().unbind()) } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index fe7e4ef88bb..54ffc07bdbe 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -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)) - .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::(py, call, asynchronous, ResponsesPythonHost::new(host)) } diff --git a/litellm/rust_bridge/execution.py b/litellm/rust_bridge/execution.py index 430de8d1ad6..0f8511cefed 100644 --- a/litellm/rust_bridge/execution.py +++ b/litellm/rust_bridge/execution.py @@ -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) diff --git a/litellm/rust_bridge/host/routes/responses.py b/litellm/rust_bridge/host/routes/responses.py index 1d68a4ec230..7ab879a01a6 100644 --- a/litellm/rust_bridge/host/routes/responses.py +++ b/litellm/rust_bridge/host/routes/responses.py @@ -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) diff --git a/litellm/rust_bridge/responses/entrypoints.py b/litellm/rust_bridge/responses/entrypoints.py index 909d00d8131..c4234c5800c 100644 --- a/litellm/rust_bridge/responses/entrypoints.py +++ b/litellm/rust_bridge/responses/entrypoints.py @@ -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 diff --git a/litellm/rust_bridge/responses/streams.py b/litellm/rust_bridge/responses/streams.py new file mode 100644 index 00000000000..ab6b493e3ca --- /dev/null +++ b/litellm/rust_bridge/responses/streams.py @@ -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 diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/conftest.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/conftest.py index ea800c5b2b0..bc7b49f8c89 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/conftest.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/conftest.py @@ -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 diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/support.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/support.py index c182fd176c6..f8f06f7a7e5 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/support.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/support.py @@ -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) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_stream_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_stream_parity.py new file mode 100644 index 00000000000..3d42ad66401 --- /dev/null +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_stream_parity.py @@ -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 diff --git a/tests/test_litellm_rust/responses/data.py b/tests/test_litellm_rust/responses/data.py index 604a7de7442..cd6ed4905ff 100644 --- a/tests/test_litellm_rust/responses/data.py +++ b/tests/test_litellm_rust/responses/data.py @@ -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") diff --git a/tests/test_litellm_rust/responses/test_body_hooks.py b/tests/test_litellm_rust/responses/test_body_hooks.py index 59ba3413da9..81e685b48aa 100644 --- a/tests/test_litellm_rust/responses/test_body_hooks.py +++ b/tests/test_litellm_rust/responses/test_body_hooks.py @@ -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]] = [] diff --git a/tests/test_litellm_rust/responses/test_failures.py b/tests/test_litellm_rust/responses/test_failures.py index 6dad1bc7278..2aaad3e863b 100644 --- a/tests/test_litellm_rust/responses/test_failures.py +++ b/tests/test_litellm_rust/responses/test_failures.py @@ -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}, diff --git a/tests/test_litellm_rust/responses/test_streaming.py b/tests/test_litellm_rust/responses/test_streaming.py index fe284bfa3fc..f2d7bf7610e 100644 --- a/tests/test_litellm_rust/responses/test_streaming.py +++ b/tests/test_litellm_rust/responses/test_streaming.py @@ -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() diff --git a/tests/unit/rust_bridge/host/routes/test_responses.py b/tests/unit/rust_bridge/host/routes/test_responses.py index a62e5b6a172..9ab548a25ff 100644 --- a/tests/unit/rust_bridge/host/routes/test_responses.py +++ b/tests/unit/rust_bridge/host/routes/test_responses.py @@ -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 diff --git a/tests/unit/rust_bridge/host/routes/test_streams.py b/tests/unit/rust_bridge/host/routes/test_streams.py new file mode 100644 index 00000000000..b42d63e9909 --- /dev/null +++ b/tests/unit/rust_bridge/host/routes/test_streams.py @@ -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))