mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
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
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:
parent
0423ac1164
commit
d6a4b71cd1
17 changed files with 782 additions and 66 deletions
|
|
@ -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())?);
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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> {
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
123
litellm/rust_bridge/responses/streams.py
Normal file
123
litellm/rust_bridge/responses/streams.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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]] = []
|
||||
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
169
tests/unit/rust_bridge/host/routes/test_streams.py
Normal file
169
tests/unit/rust_bridge/host/routes/test_streams.py
Normal 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))
|
||||
Loading…
Add table
Reference in a new issue