diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index a52f379f98b..133769dc7ec 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1409,14 +1409,10 @@ name = "litellm-ai-gateway" version = "0.1.0" dependencies = [ "axum", - "base64", "futures-channel", "futures-util", "litellm-config", "litellm-core", - "reqwest", - "rustls 0.23.42", - "rustls-native-certs", "serde", "serde_json", "sha2 0.10.9", @@ -2306,6 +2302,7 @@ version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ + "indexmap", "itoa", "memchr", "serde", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b33fed58fca..0a69cb86276 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -32,7 +32,7 @@ rustls = { version = "0.23", default-features = false, features = ["ring", "std" rustls-native-certs = "0.8" serial_test = { version = "4.0.1", default-features = false } serde = { version = "1.0", features = ["derive"] } -serde_json = { version = "1.0", features = ["float_roundtrip"] } +serde_json = { version = "1.0", features = ["float_roundtrip", "preserve_order"] } sha2 = "0.10" subtle = "2" thiserror = "2.0" diff --git a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md index 6d090cf4c8e..82301795687 100644 --- a/litellm-rust/crates/ai-gateway/ARCHITECTURE.md +++ b/litellm-rust/crates/ai-gateway/ARCHITECTURE.md @@ -1,13 +1,23 @@ # ai-gateway architecture -The Rust ai-gateway does LLM inference (realtime WebSocket). Spend tracking is an -API callback: it POSTs each finished session to the LiteLLM proxy, which records -spend and runs the usual callbacks. +The Rust ai-gateway is the Axum transport host for core routes. It authenticates +clients, selects deployments, adapts HTTP or WebSocket traffic, and supplies +terminal logging services. Provider transformation, authentication, HTTP and +WebSocket I/O, stream instrumentation, and session completion live in +`litellm-core`. + +Spend tracking is an API callback: the gateway's terminal logger POSTs each +finished call or session to the LiteLLM proxy, which records spend and runs the +usual callbacks. Streaming and WebSocket routes retain their core completion +owner until the stream or session ends, so committed calls produce one terminal +record. Realtime pool warmup is only connection preparation and produces zero +terminal records on success or failure. ```mermaid flowchart LR - C[client] <--> G[Rust ai-gateway
LLM inference] - G <--> O[OpenAI realtime] + C[client] <--> G[Rust ai-gateway
Axum transport] + G <--> K[litellm-core
route and provider I/O] + K <--> O[provider] G -. spend tracking callback .-> P[litellm proxy] F[litellm-config
load-time only] --> G F -. Python backend .-> P diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index f1834922c84..2d317e65159 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -17,14 +17,10 @@ required-features = ["server"] tracing.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-config.workspace = true -# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the -# Python proxy callbacks API. -reqwest.workspace = true # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } futures-util.workspace = true serde_json.workspace = true -base64.workspace = true axum = { workspace = true, features = ["ws"], optional = true } serde.workspace = true subtle = { workspace = true, optional = true } @@ -43,4 +39,5 @@ trace-parity = ["server", "dep:tower", "litellm-core/observability"] [dev-dependencies] futures-channel = "0.3" +tokio-tungstenite.workspace = true tower = { version = "0.5.3", features = ["util"] } diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 28d94bf16e8..48659bc571e 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -1,8 +1,8 @@ # LiteLLM Rust AI Gateway -A minimal Axum service that fronts OpenAI's realtime API. Clients open a -WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a deployment, -dials OpenAI upstream, and splices the two sockets frame-by-frame. +A minimal Axum service that hosts LiteLLM core routes. For realtime, clients open +a WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a +deployment, and adapts frames while core dials OpenAI and splices the session. ## Crates @@ -16,7 +16,9 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame. | litellm-python-interop | Domain-neutral PyO3 foundation: typed Python/Serde conversion, retained callbacks, and sync/async Python↔Tokio execution. | | litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. | -Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop. +Dependency direction is acyclic: config depends on core, the gateway depends on +config and core, and the Python bridge depends on core and Python interop. Its +optional `trace-parity` diagnostics also depend on the gateway. - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) @@ -94,12 +96,15 @@ stand-in only for the leanest possible build. ## Request logging -The gateway runs no spend logic. When a session ends it builds one +The gateway runs no spend logic. When core completes a call or session, the +gateway terminal logger builds one `StandardLoggingPayload` and POSTs it to `{LITELLM_PROXY_BASE_URL}/v1/rust_control_plane/logs` (admin-only, bearer = `LITELLM_MASTER_KEY`), and the proxy replays it through its normal callbacks (spend logs, Langfuse, etc.). The POST is non-blocking: a bounded channel drained by a background worker, dropping with a counter if the proxy is -down. It sends one payload per session. Both env vars are in the table above. +down. It sends one payload per completed call or session. Realtime connection +warmup sends none; only a connection handed to a serving session can complete +and emit a terminal payload. Both env vars are in the table above. Worker tuning, rarely needed: `LITELLM_LOG_CHANNEL_CAPACITY` (4096), `LITELLM_LOG_BATCH_SIZE` (256), `LITELLM_LOG_FLUSH_INTERVAL_MS` (500). diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 1ed44cf7048..07f48f5fa80 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -5,10 +5,6 @@ //! modules. Env-overridable tunables keep their `DEFAULT_*` value here; the env //! read + fallback happens at the host/config layer. -/// Provider attributed to realtime sessions in the logging payload. -#[cfg(feature = "server")] -pub(crate) const DEFAULT_PROVIDER: &str = "openai"; - /// HTTP path for the non-streaming Anthropic Messages route. #[cfg(feature = "server")] pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages"; diff --git a/litellm-rust/crates/ai-gateway/src/integrations/README.md b/litellm-rust/crates/ai-gateway/src/integrations/README.md deleted file mode 100644 index 16a162dac57..00000000000 --- a/litellm-rust/crates/ai-gateway/src/integrations/README.md +++ /dev/null @@ -1,127 +0,0 @@ -# LiteLLM Rust integrations - -This directory contains Rust-native equivalents of LiteLLM integration hooks. -The first supported surfaces are terminal custom loggers and pre/during-call -custom guardrails. - -## File layout - -Every integration is a folder: - -- `mod.rs` contains the implementation, trait, runner, or adapter -- `types.rs` contains the integration-local request, response, error, and future - types - -Do not add new flat integration files such as `custom_logger.rs`. Shared wire -contracts that are used by multiple integrations can stay in -`integrations/types.rs`. - -Call ordering and lifecycle timing live in `litellm-core/src/call_lifecycle`. -Call-type modules, such as OCR, adapt their request and response shapes into -that generic lifecycle runner. - -## CustomLogger - -Implement `CustomLogger` when Rust code needs to observe terminal success or -failure events. Method names intentionally match Python `CustomLogger` names. - -```rust -use litellm_ai_gateway::integrations::custom_logger::{ - CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, -}; - -struct RecordingLogger; - -impl CustomLogger for RecordingLogger { - fn async_log_success_event<'a>( - &'a self, - model_call_details: &'a ModelCallDetails, - response_obj: &'a CallbackValue, - timing: CallbackTiming, - ) -> LogFuture<'a> { - Box::pin(async move { - let model = &model_call_details.model; - let provider = &model_call_details.custom_llm_provider; - let call_type = model_call_details.call_type.to_string(); - let request_id = model_call_details.request_id.as_deref(); - let response_object = &response_obj.object; - let duration = timing.end_time - timing.start_time; - let standard_payload = model_call_details.standard_logging_payload.as_ref(); - - Ok(()) - }) - } - - fn async_log_failure_event<'a>( - &'a self, - model_call_details: &'a ModelCallDetails, - response_obj: Option<&'a CallbackValue>, - timing: CallbackTiming, - ) -> LogFuture<'a> { - Box::pin(async move { - let error = model_call_details.failure_error.as_ref(); - let response_object = response_obj.map(|value| value.object.as_str()); - let duration = timing.end_time - timing.start_time; - - Ok(()) - }) - } -} -``` - -Use `CustomLoggerRunner` to fan out terminal events to configured loggers. The -runner is a no-op when no loggers are configured, which is the expected fast -path for requests without callbacks. - -## CustomGuardrail - -Implement `CustomGuardrail` when Rust code needs to run pre-call or native -during-call checks. Method names intentionally match Python `CustomGuardrail` -entrypoints inherited from Python `CustomLogger`. - -```rust -use litellm_ai_gateway::integrations::custom_guardrail::{ - CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailEventHook, - GuardrailFuture, GuardrailRequest, -}; - -struct BlocklistedPromptGuardrail; - -impl CustomGuardrail for BlocklistedPromptGuardrail { - fn guardrail_name(&self) -> &str { - "blocklisted-prompt" - } - - fn supported_event_hooks(&self) -> &[GuardrailEventHook] { - &[GuardrailEventHook::PreCall] - } - - fn async_pre_call_hook<'a>( - &'a self, - _context: &'a GuardrailContext, - request: GuardrailRequest, - ) -> GuardrailFuture<'a> { - Box::pin(async move { - if request.data.to_string().contains("blocked phrase") { - return Ok(GuardrailDecision::Block( - litellm_ai_gateway::integrations::custom_guardrail::GuardrailError::blocked( - "blocked phrase detected", - ), - )); - } - Ok(GuardrailDecision::Allow(request)) - }) - } -} -``` - -Use `CustomGuardrailRunner::run_pre_call` for `pre_call` guardrails and -`CustomGuardrailRunner::run_during_call` for `during_call` guardrails. A -`GuardrailDecision::Mask` continues with modified request data. -`GuardrailDecision::Block` short-circuits the provider call. - -## Current boundary - -These are Rust-only primitives. Python callback and guardrail adapters are a -separate layer that should implement these Rust traits instead of changing the -runner interfaces. diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index 90c7a65ddf8..b26180c495e 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,3 +1,2 @@ pub mod audio_transcription; pub mod realtime_pool; -pub(crate) mod tls; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs index e910c92bffc..29c0efac041 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs @@ -425,7 +425,7 @@ pub fn upstream_key( #[cfg(test)] mod tests { use super::*; - use futures_util::SinkExt; + use futures_util::{SinkExt, StreamExt}; use std::net::SocketAddr; use tokio::net::TcpListener; use tokio_tungstenite::tungstenite::Message; @@ -478,15 +478,11 @@ mod tests { } fn key_for(base: &str) -> UpstreamKey { - UpstreamKey { - model: "gpt-realtime".to_string(), - api_key: "sk-test".to_string(), - api_base: Some(base.to_string()), - } + UpstreamKey::new("gpt-realtime", Some("sk-test"), Some(base)).unwrap() } #[tokio::test] - async fn warm_handoff_relays_buffered_session_created() { + async fn taking_a_warm_connection_consumes_one_pool_entry() { let base = spawn_fake_openai().await; let pool = RealtimePool::new_unspawned(test_config()); let key = key_for(&base); @@ -494,17 +490,7 @@ mod tests { pool.warm_now(&key).await; assert_eq!(pool.warm_len(&key), 2); - let handoff = pool.take(&key).expect("a warm socket should be available"); - assert_eq!(handoff.session_created.event_type, "session.created"); - assert_eq!( - handoff - .session_created - .data - .get("session") - .and_then(|s| s.get("id")) - .and_then(|v| v.as_str()), - Some("sess_fake") - ); + let _handoff = pool.take(&key).expect("a warm socket should be available"); // Taking one leaves one. assert_eq!(pool.warm_len(&key), 1); } @@ -577,29 +563,7 @@ mod tests { test_config().target_size, "background replenisher should warm up to target_size" ); - let handoff = pool.take(&key).expect("a warm socket should be available"); - assert_eq!(handoff.session_created.event_type, "session.created"); - } - - #[tokio::test] - async fn closed_upstream_socket_is_detected_dead() { - // A genuinely dead socket: dial the fake, read session.created, then drop - // the server by closing from our side and waiting for the close to land. - let base = spawn_fake_openai().await; - let pool = RealtimePool::new_unspawned(test_config()); - let key = key_for(&base); - pool.register(key.clone()); - - let mut conn = warm_one(&key).await.expect("warm one"); - // Close the upstream from the client side; the server echoes a close. - let _ = conn.tx.send(Message::Close(None)).await; - // Give the close a moment to arrive on rx. - tokio::time::sleep(Duration::from_millis(50)).await; - pool.insert_warm(key.clone(), conn); - - // Liveness check at take() should detect the close and discard it. - assert!(pool.take(&key).is_none()); - assert_eq!(pool.warm_len(&key), 0); + let _handoff = pool.take(&key).expect("a warm socket should be available"); } #[tokio::test] diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs deleted file mode 100644 index a2562f60345..00000000000 --- a/litellm-rust/crates/ai-gateway/src/io/tls.rs +++ /dev/null @@ -1,80 +0,0 @@ -//! Outbound WebSocket dials over a TLS config this crate builds once and owns. -//! -//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth` -//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that -//! `tokio-tungstenite` uses when handed no connector panics rather than guess -//! between them. Naming ring on a connector of our own settles that for these -//! dials without touching the process-wide default, and building the config -//! once keeps the platform trust store, which `tokio-tungstenite` would -//! otherwise re-read on every dial, off the dial path. - -use std::io; -use std::sync::{Arc, OnceLock}; - -use rustls::{ClientConfig, RootCertStore}; -use tokio::net::TcpStream; -use tokio_tungstenite::tungstenite::Error; -use tokio_tungstenite::tungstenite::client::IntoClientRequest; -use tokio_tungstenite::tungstenite::error::TlsError; -use tokio_tungstenite::tungstenite::handshake::client::Response; -use tokio_tungstenite::{ - Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, -}; - -static TLS_CONFIG: OnceLock> = OnceLock::new(); - -fn build_config() -> Result> { - let native = rustls_native_certs::load_native_certs(); - let roots = { - let mut store = RootCertStore::empty(); - let (added, _ignored) = store.add_parsable_certificates(native.certs); - if added == 0 { - return Err(Box::new(Error::Io(io::Error::other(format!( - "no usable native root certificates: {:?}", - native.errors - ))))); - } - store - }; - - ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) - .with_safe_default_protocol_versions() - .map(|builder| builder.with_root_certificates(roots).with_no_client_auth()) - .map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error)))) -} - -fn tls_config() -> Result, Box> { - if let Some(config) = TLS_CONFIG.get() { - return Ok(Arc::clone(config)); - } - let built = Arc::new(build_config()?); - Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) -} - -pub(crate) async fn connect_upstream( - request: R, -) -> Result<(WebSocketStream>, Response), Box> -where - R: IntoClientRequest + Unpin, -{ - let request = request.into_client_request().map_err(Box::new)?; - let connector = match request.uri().scheme_str() { - Some("wss") => Some(Connector::Rustls(tls_config()?)), - _ => None, - }; - connect_async_tls_with_config(request, None, false, connector) - .await - .map_err(Box::new) -} - -#[cfg(test)] -mod tests { - use super::build_config; - - #[test] - fn builds_a_usable_config_with_both_provider_features_enabled() { - let config = build_config().expect("a client config"); - - assert!(!config.crypto_provider().cipher_suites.is_empty()); - } -} diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index 4cfd784dda4..a1b1aa1c278 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -1,8 +1,8 @@ //! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router. //! //! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment -//! (simple-shuffle) → `io::realtime::realtime()` invokes OpenAI. The -//! server owns transport + config; routing lives in the `router` crate. +//! (simple-shuffle) → `core::realtime()` invokes OpenAI. The server owns config +//! and Axum adaptation; routing and provider execution live in core. //! //! The binary requires the `server` feature (declared in `Cargo.toml` via //! `required-features`), so cargo skips it unless that feature is on. Everything @@ -135,12 +135,8 @@ where fn register_deployments(router: &Router, pool: &RealtimePool) { for deployment in router.deployments() { let params = &deployment.litellm_params; - let provider_model = params - .model - .strip_prefix("openai/") - .unwrap_or(¶ms.model); if let Some(key) = upstream_key( - provider_model, + ¶ms.model, params.api_key.as_deref(), params.api_base.as_deref(), ) { diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md b/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md index 3301576bb85..4e7e43ab9d4 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/README.md @@ -1,9 +1,11 @@ # Realtime route (`GET /v1/realtime`) -Proxies OpenAI's realtime WebSocket. `mod.rs` is the axum surface (handler + -socket↔events adapter); `service.rs` is the pure logic (select a deployment, then -splice client ↔ upstream). The pool itself lives in -`crates/providers/src/realtime_pool.rs`. +Proxies OpenAI's realtime WebSocket. `mod.rs` is the Axum surface and converts +socket frames to core events. `service.rs` selects a deployment and calls +`litellm_core::realtime::realtime`. Core owns provider resolution, upstream +dialing, event splicing, usage observation, and session completion. The pool +lives in `ai-gateway/src/io/realtime_pool.rs` and stores core `WarmConnection` +values as a transport optimization. ## Connection pooling @@ -46,6 +48,11 @@ unprompted on connect, we pre-read exactly that one frame and relay it on handof and we send nothing else on the socket before a client exists — so the client's first `session.update` behaves identically either way. +Warmup has no user-visible call lifecycle. Whether it succeeds or fails, it emits +zero terminal records. Ownership transfers only when a warm connection is handed +to the core realtime route; that serving session emits exactly one terminal record +when splicing completes or fails. + ### Sizing Each warm socket serves **exactly one** session (realtime isn't multiplexed), so the diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index bafe342dbe5..58370605954 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -7,19 +7,40 @@ //! socket we fresh-dial exactly as before — the pool is never on the critical path //! for correctness, only latency. -use std::time::Duration; use std::sync::Arc; +use std::time::Duration; use crate::io::realtime_pool::{RealtimePool, upstream_key}; use futures_util::{Sink, Stream}; use litellm_core::error::Error; -use litellm_core::integrations::custom_logger::{CustomLogger, CustomLoggerRunner}; +use litellm_core::integrations::custom_logger::{CustomLogger, CustomLoggerRunner, LogFuture}; use litellm_core::integrations::types::{RequestMetadata, StandardLoggingMetadata}; -use litellm_core::lifecycle::{CallLifecycleContext, ExecutedCall}; -use litellm_core::realtime::{RealtimeRequest, realtime}; +use litellm_core::lifecycle::{ + CallLifecycleContext, Clock, ExecutedCall, TerminalDispatcher, TerminalRecord, +}; use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::realtime::{RealtimeRequest, realtime}; use litellm_core::router::Router; +struct GatewayRealtimeServices { + runner: CustomLoggerRunner, +} + +impl Clock for GatewayRealtimeServices { + fn now(&self) -> f64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) + } +} + +impl TerminalDispatcher for GatewayRealtimeServices { + fn dispatch<'a>(&'a self, terminal: &'a TerminalRecord) -> LogFuture<'a> { + self.runner.dispatch(terminal) + } +} + /// Select a deployment for `model` and splice the client stream to the provider. /// /// `pool` supplies a pre-warmed upstream when one is available; otherwise we @@ -45,29 +66,28 @@ where .get_available_deployment(model) .ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?; let params = &deployment.litellm_params; - // Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model. - let provider_model = params - .model - .strip_prefix("openai/") - .unwrap_or(¶ms.model); - let connection = upstream_key( - provider_model, + ¶ms.model, params.api_key.as_deref(), params.api_base.as_deref(), - ).ok_or_else(|| Error::Auth("missing realtime provider API key".to_string()))?; - let warm = pool.take(&connection); - let context = CallLifecycleContext::new("realtime", model, "openai", call_id) - .with_metadata(StandardLoggingMetadata { + ); + let warm = connection.as_ref().and_then(|key| pool.take(key)); + let context = CallLifecycleContext::new("realtime", model, "openai", call_id).with_metadata( + StandardLoggingMetadata { user_api_key_hash: metadata.user_api_key_hash, user_api_key_user_id: metadata.user_api_key_user_id, user_api_key_team_id: metadata.user_api_key_team_id, ..Default::default() - }); + }, + ); Ok(realtime( - &CustomLoggerRunner::new(loggers.as_ref().clone()), + &GatewayRealtimeServices { + runner: CustomLoggerRunner::new(loggers.as_ref().clone()), + }, RealtimeRequest { - connection, + model: params.model.clone(), + api_key: params.api_key.clone(), + api_base: params.api_base.clone(), warm, idle_timeout, }, diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs index 1912a51ff17..b11fb7b411d 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs @@ -212,7 +212,11 @@ async fn bridge( }; let client_in = Box::pin(stream.filter_map(|message| async move { match message { - Ok(Message::Text(text)) => serde_json::from_str::(&text).ok(), + Ok(Message::Text(text)) => Some( + serde_json::from_str::(&text) + .map_err(|error| litellm_core::Error::InvalidRequest(error.to_string())), + ), + Err(error) => Some(Err(litellm_core::Error::Network(error.to_string()))), _ => None, } })); diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs index 9bf56d5cc02..36e64bee259 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/service.rs @@ -51,7 +51,7 @@ pub async fn run( client_out: Out, ) -> Result, Error> where - In: Stream + Unpin + Send, + In: Stream> + Unpin + Send, Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { @@ -76,7 +76,7 @@ where ..Default::default() }); responses_websocket( - &GatewayResponsesServices::new(loggers), + Arc::new(GatewayResponsesServices::new(loggers)), ResponsesWebSocketRequest { model: provider_model.to_string(), api_key: params.api_key.clone(), diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index aa0a94521d4..75daa5a6912 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -46,8 +46,9 @@ core/src/messages/ authoritative roots into a native request after callbacks, and sends it through `http_utils::buffered_post`. Reducto upload, Azure Document Intelligence polling, and HTTP document URL conversion are declined at admission until they have an -implementation on this settled-request path. `audio_transcription` and -`realtime` remain in flight. +implementation on this settled-request path. Audio transcription provider I/O, +realtime WebSocket dialing and splicing, and Responses WebSocket dialing and +splicing are also core-owned. The invariant is one function body owns the route lifecycle. The conceptual shape is: @@ -75,7 +76,7 @@ requirements traits (an `OcrServices` bundle) declare exactly what a route needs passed a catch-all gateway environment. The ordinary Rust client supplies native defaults; callers override implementations at construction. -`CallServices` opens one request-scoped sessions per call, owning per-call state +`CallServices` opens one request-scoped session per call, owning per-call state (timing, logging state, retained host objects, deferred completion, correlation state). No-op call services are the default, and their presence must not move provider behavior into a host or force callback payload materialization on an @@ -108,21 +109,20 @@ service construction dependencies do not leak into service interfaces ## Not the target -The current `CallLifecycleHooks` shape is not the final public service API: it -folds lifecycle sequencing into a stateless generic transformation interface, -needs `Send` futures, cannot express the full replacement and error contracts, -and does not model request-scoped retained ownership or Python caller-task -driving. It is a stepping stone, not the contract to build new routes against. - Do not introduce a dynamic `TypeId` service map, a shared gateway callback environment reused from core or the bridge, per-callback JSON serialization, or a full Effect layer API. Services traits plus constructors and a scoped call owner are the minimum design; add more machinery only when concrete consumers require it. -## Gateway migration +## Streaming ownership -Existing gateway-hosted OCR, transcription and WebSocket provider execution -predates this boundary and must move here as those routes migrate. Keeping a -gateway callback as a `CallServices` implementation does not justify keeping -provider orchestration beside it in the gateway. +Core owns provider sessions through completion. Streaming HTTP calls transfer a +`StreamingCall` whose completion registration keeps terminal dispatch alive +until the host finishes or drops the stream. Realtime and Responses WebSocket +entrypoints retain the provider connection while they splice events, then emit +exactly one terminal record after the committed session completes or fails. + +Realtime pool warmup is not a user call and must emit zero terminal records on +both success and failure. A warmed connection transfers into `realtime`; only +that serving session owns completion and terminal dispatch. diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 9a96b9d1140..bf36a6b9fb7 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -7,7 +7,7 @@ use super::client::http_client; use super::types::ProviderAudioTranscriptionRequest; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub async fn execute_audio_transcription_provider_call( +pub(super) async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, ) -> Result { let body = serde_json::to_vec(&request.body) diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index 811acd4a8e5..40af2e3c86d 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -8,9 +8,7 @@ pub mod types; use serde_json::Value; -pub use handler::execute_audio_transcription_provider_call; pub use lifecycle::{AudioRoute, AudioServices, DefaultAudioServices}; -pub use prepare::prepare_audio_transcription_provider_call; pub use types::{AudioRouteRequest, AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index bbef97341a9..25139da1d1c 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -18,7 +18,7 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv } #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] -pub fn prepare_audio_transcription_provider_call( +pub(super) fn prepare_audio_transcription_provider_call( request: AudioTranscriptionRequest<'_>, ) -> Result { let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) diff --git a/litellm-rust/crates/core/src/call_lifecycle/README.md b/litellm-rust/crates/core/src/call_lifecycle/README.md deleted file mode 100644 index 692e249ef27..00000000000 --- a/litellm-rust/crates/core/src/call_lifecycle/README.md +++ /dev/null @@ -1,167 +0,0 @@ -# Call lifecycle - -`litellm_core::call_lifecycle` is the shared execution wrapper for LiteLLM call -types migrated to Rust. It owns lifecycle ordering, phase timing, and trace -observer calls. It must not know about OCR, chat, messages, responses, -completions, provider auth, request transforms, or response normalization. - -Call-type modules own their domain behavior. For example, OCR owns document -payloads, OCR provider transforms, safe document fetch, guardrail payload shape, -callback payload shape, and provider HTTP execution. - -## Runtime order - -Every wrapped call runs in this order: - -1. `async_pre_call_hook` -2. `async_during_call_hook` -3. provider call -4. `async_log_success_event` or `async_log_failure_event` - -`async_pre_call_hook` receives the initial LiteLLM request shape. It is where -pre-call custom guardrails run. - -`async_during_call_hook` converts the initial request into the provider-ready -request. It is where provider config selection, parameter mapping, auth/header -resolution, request transforms, and during-call guardrails belong. - -The provider call receives only the provider-ready request. It should execute -I/O and call the provider response transform. - -Success and failure callbacks receive `CallLifecycleTiming`. Callback failures -must not replace the original provider or guardrail result. - -## Trace contract - -The lifecycle runner records: - -- full call start and end time -- `pre_call` phase timing -- `during_call` phase timing -- `provider_call` phase timing -- `success_callback` phase timing -- `failure_callback` phase timing - -`CallLifecycleObserver` receives phase start and end events. The default -observer is a no-op. Future OTEL support should implement this observer instead -of editing OCR, chat, messages, responses, completions, or provider modules. - -## Required shape - -Each migrated call type should use this folder shape: - -```text -litellm-rust/crates/ai-gateway/src// - mod.rs # thin public entrypoint - types.rs # public request, prepared request, provider request, response types - prepare.rs # model/provider/callback/guardrail setup - hooks.rs # CallLifecycleHooks implementation - handler.rs # provider I/O and response normalization - tests.rs # call-type lifecycle and handler tests -``` - -Provider transforms can live in `litellm-rust/crates/core/src/providers/...`. -Shared call-type helpers can live beside the call type, but generic lifecycle -code stays in this folder. - -## Core API - -The prepared request implements `CallLifecycleRequest`: - -```rust -impl CallLifecycleRequest for PreparedMessagesRequest { - fn lifecycle_context(&self) -> CallLifecycleContext { - CallLifecycleContext::new( - "messages", - self.model.clone(), - self.custom_llm_provider.clone(), - self.litellm_call_id.clone(), - ) - } -} -``` - -The call-type hooks implement `CallLifecycleHooks`: - -```rust -impl CallLifecycleHooks< - PreparedMessagesRequest, - ProviderMessagesRequest, - MessagesResponse, -> for MessagesLifecycleHooks { - fn async_pre_call_hook(...) { - // run pre-call custom guardrails against the LiteLLM request shape - } - - fn async_during_call_hook(...) { - // map params, validate env, transform request, run during-call guardrails - } - - fn async_log_success_event(...) { - // call async_log_success_event on configured custom loggers - } - - fn async_log_failure_event(...) { - // call async_log_failure_event without swallowing the original error - } -} -``` - -The public entrypoint stays thin: - -```rust -pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { - let PreparedMessagesCall { request, hooks } = prepare_messages_call(request)?; - - CallLifecycle::default() - .run_request(request, &hooks, execute_messages_provider_call) - .await -} -``` - -Use `run_request` for new call types. Keep `run` available only for specialized -tests or existing code that already has a `CallLifecycleContext`. - -## Adding a new call type - -1. Add `/types.rs` - -Define the public request accepted by the bridge, the prepared request used by -the lifecycle runner, and the provider request consumed by the handler. - -2. Implement `CallLifecycleRequest` - -Return `call_type`, `model`, `custom_llm_provider`, and `litellm_call_id`. -Do not put provider-specific logic here. - -3. Add `/prepare.rs` - -Resolve model/provider once, generate or preserve `litellm_call_id`, construct -callback and guardrail runners, and return `PreparedCall`. - -4. Add `/hooks.rs` - -Implement `CallLifecycleHooks`. Put pre-call guardrail payload construction, -provider config selection, param mapping, request transform, during-call -guardrail payload construction, and callback payload construction here. - -5. Add `/handler.rs` - -Execute the provider request and normalize the provider response. Do not repeat -provider-specific transforms here; call the provider config. - -6. Add tests - -Cover hook order, success callback payload, failure callback payload, pre-call -guardrail blocking before provider I/O, during-call body mutation, and provider -error mapping. - -## Review checklist - -- Core lifecycle has no call-type or provider-specific branches -- Public call-type entrypoint only prepares and calls `run_request` -- Provider behavior lives behind provider config/transformation code -- Hook method names map to the Python custom logger and guardrail concepts -- Phase timing is recorded once in lifecycle, not separately per call type -- Callback failures never hide the original provider or guardrail error -- Tests prove the provider socket is not touched when pre-call guardrails block diff --git a/litellm-rust/crates/core/src/chat_completions/lifecycle.rs b/litellm-rust/crates/core/src/chat_completions/lifecycle.rs index c2fba6c2b02..9abd7e26293 100644 --- a/litellm-rust/crates/core/src/chat_completions/lifecycle.rs +++ b/litellm-rust/crates/core/src/chat_completions/lifecycle.rs @@ -1,29 +1,12 @@ use crate::Error; -use crate::lifecycle::{ - ActionBinding, ActionKind, Delivery, ErrorDisposition, FailurePolicy, Lifecycle, - LifecycleRoute, Outcome, Owner, ResultPolicy, -}; +use crate::lifecycle::program::{CallProgram, ProgramOptions, actions_for}; +use crate::lifecycle::{ActionBinding, Lifecycle, LifecycleRoute, Outcome}; use super::chat_completions_decline_reason; use serde_json::{Map, Value}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Operation { - Setup, - DeploymentPre, - Prepare, - Send, - DeploymentSuccess, - DeploymentFailure, - SyncSuccess, - AsyncSuccess, - SyncSuccessIfNeeded, - SyncFailure, - AsyncFailure, - Restore, - Complete(Outcome), -} +pub use crate::lifecycle::program::{Observations, Operation, Transition}; #[derive(Clone, Debug)] pub struct Admission { @@ -39,17 +22,6 @@ pub struct Options { pub internal_call: bool, } -#[derive(Clone, Copy, Debug, Default)] -pub struct Observations { - pub logger_available: bool, - pub has_fallbacks: bool, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Transition { - pub error: ErrorDisposition, -} - #[derive(Debug, PartialEq, Eq)] pub struct Decline(&'static str); @@ -61,10 +33,7 @@ impl Decline { #[derive(Debug)] pub struct ChatCompletionsState { - operation: Operation, - outcome: Outcome, - asynchronous: bool, - internal_call: bool, + program: CallProgram, } #[derive(Debug)] @@ -95,15 +64,16 @@ impl LifecycleRoute for ChatCompletionsRoute { return Ok(Err(Decline(reason))); } Ok(Ok(ChatCompletionsState { - operation: Operation::Setup, - outcome: Outcome::Success, - asynchronous: options.asynchronous, - internal_call: options.internal_call, + program: CallProgram::new(ProgramOptions { + asynchronous: options.asynchronous, + internal_call: options.internal_call, + pre_call: false, + }), })) } fn operation(state: &Self::State) -> Operation { - state.operation + state.program.operation() } fn advance( @@ -111,95 +81,25 @@ impl LifecycleRoute for ChatCompletionsRoute { outcome: Outcome, observations: Observations, ) -> Result { - use Operation::*; - - if matches!(state.operation, Complete(_)) { - return Err(Error::InvalidRequest( - "chat completions lifecycle is already complete".into(), - )); - } - let failure = - if observations.logger_available && !(state.asynchronous && state.internal_call) { - SyncFailure - } else { - Restore - }; - let error = if outcome != Outcome::Success && state.operation != DeploymentFailure { - state.outcome = outcome; - ErrorDisposition::Replace - } else { - ErrorDisposition::Preserve - }; - state.operation = match (state.operation, outcome) { - (Restore, _) => Complete(state.outcome), - (DeploymentFailure, _) => failure, - (_, Outcome::Abort) => Restore, - (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, - (Prepare | Send, Outcome::Failure) if state.asynchronous => DeploymentFailure, - (_, Outcome::Failure) => failure, - (Setup, Outcome::Success) if state.asynchronous => DeploymentPre, - (Setup | DeploymentPre, Outcome::Success) => Prepare, - (Prepare, Outcome::Success) => Send, - (Send, Outcome::Success) if state.asynchronous => DeploymentSuccess, - (Send, Outcome::Success) => SyncSuccess, - (DeploymentSuccess, Outcome::Success) => { - if state.internal_call || observations.has_fallbacks { - SyncSuccessIfNeeded - } else { - AsyncSuccess - } - } - (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded, - (SyncFailure, Outcome::Success) if state.asynchronous => AsyncFailure, - (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => { - Restore - } - (Complete(_), _) => unreachable!(), - }; - Ok(Transition { error }) + state.program.advance(outcome, observations).ok_or_else(|| { + Error::InvalidRequest("chat completions lifecycle is already complete".into()) + }) } fn actions_for(operation: Operation, _: &Observations) -> &'static [ActionBinding] { - match operation { - Operation::Prepare | Operation::Send => &PROVIDER_ACTION, - Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => { - &FAILURE_ACTION - } - Operation::Restore => &RESTORE_ACTION, - Operation::Complete(_) => &[], - _ => &CALLBACK_ACTION, - } + actions_for(operation) } } -const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::ProviderCall, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Replace, - on_error: FailurePolicy::Propagate, - owner: Owner::Core, -}]; -const CALLBACK_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::TerminalSuccess, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::RecordAndContinue, - owner: Owner::Route, -}]; -const FAILURE_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::TerminalFailure, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::PreserveOriginalFailure, - owner: Owner::Route, -}]; -const RESTORE_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::Restore, - delivery: Delivery::InlineDirect, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::Propagate, - owner: Owner::Core, -}]; +impl Lifecycle { + pub fn commitment(&self) -> crate::lifecycle::Commitment { + self.state.program.commitment() + } + + pub fn failure_stage(&self) -> Option { + self.state.program.failure_stage() + } +} pub fn machine( admission: &Admission, diff --git a/litellm-rust/crates/core/src/lifecycle/execution.rs b/litellm-rust/crates/core/src/lifecycle/execution.rs index fa88b134536..53dbc6f292b 100644 --- a/litellm-rust/crates/core/src/lifecycle/execution.rs +++ b/litellm-rust/crates/core/src/lifecycle/execution.rs @@ -188,6 +188,47 @@ impl CallLifecycle { ClockImpl: Clock, ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, ProviderFuture: Future>, + { + self.run_with_usage( + context, + request, + policy, + dispatcher, + clock, + provider_call, + |_| None, + ) + .await + } + + pub(crate) async fn run_with_usage< + InitialReq, + ProviderReq, + Resp, + Policy, + Dispatcher, + ClockImpl, + ProviderCall, + ProviderFuture, + ResponseUsage, + >( + &self, + mut context: CallLifecycleContext, + request: InitialReq, + policy: &Policy, + dispatcher: &Dispatcher, + clock: &ClockImpl, + provider_call: ProviderCall, + response_usage: ResponseUsage, + ) -> ExecutedCall + where + Resp: Serialize, + Policy: RequestPolicy, + Dispatcher: TerminalDispatcher, + ClockImpl: Clock, + ProviderCall: FnOnce(ProviderReq) -> ProviderFuture, + ProviderFuture: Future>, + ResponseUsage: FnOnce(&Resp) -> Option, { let start_time = clock.now(); let request = match policy.async_pre_call_hook(&context, request).await { @@ -204,6 +245,9 @@ impl CallLifecycle { }; match provider_call(provider_request).await { Ok(response) => { + if let Some(usage) = response_usage(&response) { + context.usage = usage; + } let terminal = context.terminal( CallbackTiming::new(start_time, clock.now()), TerminalClassification::Success, diff --git a/litellm-rust/crates/core/src/lifecycle/mod.rs b/litellm-rust/crates/core/src/lifecycle/mod.rs index a042fc1ac8a..460e45e9393 100644 --- a/litellm-rust/crates/core/src/lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/lifecycle/mod.rs @@ -3,6 +3,7 @@ pub mod executed; pub mod execution; pub mod machine; pub mod ocr; +pub mod program; mod streaming; pub mod terminal; pub mod types; @@ -14,6 +15,7 @@ pub use execution::{ TerminalDispatcher, }; pub use machine::{Lifecycle, LifecycleRoute}; +pub use program::{Commitment, FailureStage}; pub use streaming::{ BytesStream, StreamingCall, StreamingCompletion, StreamingMetadata, StreamingObserver, StreamingSource, diff --git a/litellm-rust/crates/core/src/lifecycle/ocr.rs b/litellm-rust/crates/core/src/lifecycle/ocr.rs index ea412210e36..cd47d417f59 100644 --- a/litellm-rust/crates/core/src/lifecycle/ocr.rs +++ b/litellm-rust/crates/core/src/lifecycle/ocr.rs @@ -1,10 +1,10 @@ use crate::Error; use crate::ocr::{OcrAdmissionRequest, prepare}; -use super::{ - ActionBinding, ActionKind, Delivery, ErrorDisposition, FailurePolicy, LifecycleRoute, Outcome, - Owner, ResultPolicy, -}; +use super::program::{CallProgram, ProgramOptions, actions_for}; +use super::{ActionBinding, LifecycleRoute, Outcome}; + +pub use super::program::{Observations, Operation, Transition}; #[derive(Debug)] pub enum NativeOutcome { @@ -45,42 +45,9 @@ pub struct Identity { pub generated_call_id: bool, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Operation { - Setup, - DeploymentPre, - Prepare, - PreCall, - Send, - DeploymentSuccess, - DeploymentFailure, - SyncSuccess, - AsyncSuccess, - SyncSuccessIfNeeded, - SyncFailure, - AsyncFailure, - Restore, - Complete(Outcome), -} - -#[derive(Clone, Copy, Debug, Default)] -pub struct Observations { - pub logger_available: bool, - pub has_fallbacks: bool, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Transition { - pub operation: Operation, - pub error: ErrorDisposition, -} - #[derive(Debug)] pub struct OcrState { - operation: Operation, - outcome: Outcome, - asynchronous: bool, - internal_call: bool, + program: CallProgram, identity: Identity, } @@ -103,6 +70,14 @@ impl Lifecycle { pub fn identity(&self) -> &Identity { &self.state.identity } + + pub fn commitment(&self) -> super::Commitment { + self.state.program.commitment() + } + + pub fn failure_stage(&self) -> Option { + self.state.program.failure_stage() + } } impl LifecycleRoute for OcrRoute { @@ -132,10 +107,11 @@ impl LifecycleRoute for OcrRoute { let generated_call_id = options.call_id.is_none(); let call_id = options.call_id.unwrap_or_else(generate_call_id); Ok(Ok(OcrState { - operation: Operation::Setup, - outcome: Outcome::Success, - asynchronous: options.asynchronous, - internal_call: options.internal_call, + program: CallProgram::new(ProgramOptions { + asynchronous: options.asynchronous, + internal_call: options.internal_call, + pre_call: true, + }), identity: Identity { requested_model: admission.model.clone(), call_id, @@ -146,7 +122,7 @@ impl LifecycleRoute for OcrRoute { } fn operation(state: &Self::State) -> Self::Operation { - state.operation + state.program.operation() } fn advance( @@ -154,115 +130,20 @@ impl LifecycleRoute for OcrRoute { outcome: Self::Outcome, observations: Self::Observation, ) -> Result { - use Operation::*; - - if matches!(state.operation, Complete(_)) { - return Err(Error::InvalidRequest( - "OCR lifecycle is already complete".into(), - )); - } - let failure = - if observations.logger_available && !(state.asynchronous && state.internal_call) { - SyncFailure - } else { - Restore - }; - let error = if outcome != Outcome::Success && state.operation != DeploymentFailure { - state.outcome = outcome; - ErrorDisposition::Replace - } else { - ErrorDisposition::Preserve - }; - state.operation = match (state.operation, outcome) { - (Restore, _) => Complete(state.outcome), - (DeploymentFailure, _) => failure, - (_, Outcome::Abort) => Restore, - (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, - (Prepare | PreCall | Send, Outcome::Failure) if state.asynchronous => DeploymentFailure, - (_, Outcome::Failure) => failure, - (Setup, Outcome::Success) if state.asynchronous => DeploymentPre, - (Setup | DeploymentPre, Outcome::Success) => Prepare, - (Prepare, Outcome::Success) => PreCall, - (PreCall, Outcome::Success) => Send, - (Send, Outcome::Success) if state.asynchronous => DeploymentSuccess, - (Send, Outcome::Success) => SyncSuccess, - (DeploymentSuccess, Outcome::Success) => { - if state.internal_call || observations.has_fallbacks { - SyncSuccessIfNeeded - } else { - AsyncSuccess - } - } - (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded, - (SyncFailure, Outcome::Success) if state.asynchronous => AsyncFailure, - (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => { - Restore - } - (Complete(_), _) => unreachable!(), - }; - Ok(Transition { - operation: state.operation, - error, - }) + state + .program + .advance(outcome, observations) + .ok_or_else(|| Error::InvalidRequest("OCR lifecycle is already complete".into())) } fn actions_for( operation: Self::Operation, _context: &Self::Context, ) -> &'static [ActionBinding] { - match operation { - Operation::Prepare | Operation::Send => &PROVIDER_ACTION, - Operation::PreCall => &PRE_CALL_ACTION, - Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => { - &FAILURE_ACTION - } - Operation::Restore => &RESTORE_ACTION, - Operation::Complete(_) => &[], - _ => &CALLBACK_ACTION, - } + actions_for(operation) } } -const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::ProviderCall, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Replace, - on_error: FailurePolicy::Propagate, - owner: Owner::Core, -}]; - -const PRE_CALL_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::RequestPolicy, - delivery: Delivery::InlineDirect, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::Propagate, - owner: Owner::Route, -}]; - -const CALLBACK_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::TerminalSuccess, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::RecordAndContinue, - owner: Owner::Route, -}]; - -const FAILURE_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::TerminalFailure, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::PreserveOriginalFailure, - owner: Owner::Route, -}]; - -const RESTORE_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::Restore, - delivery: Delivery::InlineDirect, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::Propagate, - owner: Owner::Core, -}]; - fn generate_call_id() -> String { let id = (rand::random::() & !(0xf000_u128 << 64 | 0xc000_u128 << 48)) | (0x4000_u128 << 64 | 0x8000_u128 << 48); @@ -282,6 +163,7 @@ mod tests { use std::rc::Rc; use super::*; + use crate::lifecycle::ErrorDisposition; use crate::ocr::types::OcrDocument; fn request() -> OcrAdmissionRequest { @@ -461,6 +343,10 @@ mod tests { } assert!(Rc::ptr_eq(&original, &retained)); assert_eq!(transition.operation, Operation::SyncFailure); + assert_eq!( + machine.failure_stage(), + Some(crate::lifecycle::FailureStage::ProviderCall) + ); } } diff --git a/litellm-rust/crates/core/src/lifecycle/program.rs b/litellm-rust/crates/core/src/lifecycle/program.rs new file mode 100644 index 00000000000..97c2080d976 --- /dev/null +++ b/litellm-rust/crates/core/src/lifecycle/program.rs @@ -0,0 +1,271 @@ +use super::{ + ActionBinding, ActionKind, Delivery, ErrorDisposition, FailurePolicy, Outcome, Owner, + ResultPolicy, +}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Operation { + Setup, + DeploymentPre, + Prepare, + PreCall, + Send, + DeploymentSuccess, + DeploymentFailure, + SyncSuccess, + AsyncSuccess, + SyncSuccessIfNeeded, + SyncFailure, + AsyncFailure, + Restore, + Complete(Outcome), +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct Observations { + pub logger_available: bool, + pub has_fallbacks: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Commitment { + Replayable, + ProviderStarted, + ResponseReceived, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FailureStage { + BeforeProvider, + ProviderCall, + AfterProviderResponse, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Transition { + pub operation: Operation, + pub error: ErrorDisposition, + pub commitment: Commitment, + pub failure_stage: Option, +} + +#[derive(Clone, Copy, Debug)] +pub struct ProgramOptions { + pub asynchronous: bool, + pub internal_call: bool, + pub pre_call: bool, +} + +#[derive(Debug)] +pub struct CallProgram { + operation: Operation, + outcome: Outcome, + commitment: Commitment, + failure_stage: Option, + options: ProgramOptions, +} + +impl CallProgram { + pub fn new(options: ProgramOptions) -> Self { + Self { + operation: Operation::Setup, + outcome: Outcome::Success, + commitment: Commitment::Replayable, + failure_stage: None, + options, + } + } + + pub fn operation(&self) -> Operation { + self.operation + } + + pub fn commitment(&self) -> Commitment { + self.commitment + } + + pub fn failure_stage(&self) -> Option { + self.failure_stage + } + + pub fn advance(&mut self, outcome: Outcome, observations: Observations) -> Option { + use Operation::*; + + if matches!(self.operation, Complete(_)) { + return None; + } + let current = self.operation; + let failure = if observations.logger_available + && !(self.options.asynchronous && self.options.internal_call) + { + SyncFailure + } else { + Restore + }; + let error = if outcome != Outcome::Success && current != DeploymentFailure { + self.outcome = outcome; + ErrorDisposition::Replace + } else { + ErrorDisposition::Preserve + }; + let failure_stage = + (outcome != Outcome::Success).then(|| self.classify_failure_stage(current)); + if error == ErrorDisposition::Replace { + self.failure_stage = failure_stage; + } + if current == Send { + self.commitment = if outcome == Outcome::Success { + Commitment::ResponseReceived + } else { + Commitment::ProviderStarted + }; + } + self.operation = match (current, outcome) { + (Restore, _) => Complete(self.outcome), + (DeploymentFailure, _) => failure, + (_, Outcome::Abort) => Restore, + (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, + (Prepare | PreCall | Send, Outcome::Failure) if self.options.asynchronous => { + DeploymentFailure + } + (_, Outcome::Failure) => failure, + (Setup, Outcome::Success) if self.options.asynchronous => DeploymentPre, + (Setup | DeploymentPre, Outcome::Success) => Prepare, + (Prepare, Outcome::Success) if self.options.pre_call => PreCall, + (Prepare | PreCall, Outcome::Success) => Send, + (Send, Outcome::Success) if self.options.asynchronous => DeploymentSuccess, + (Send, Outcome::Success) => SyncSuccess, + (DeploymentSuccess, Outcome::Success) => { + if self.options.internal_call || observations.has_fallbacks { + SyncSuccessIfNeeded + } else { + AsyncSuccess + } + } + (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded, + (SyncFailure, Outcome::Success) if self.options.asynchronous => AsyncFailure, + (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => { + Restore + } + (Complete(_), _) => unreachable!(), + }; + Some(Transition { + operation: self.operation, + error, + commitment: self.commitment, + failure_stage: self.failure_stage, + }) + } + + fn classify_failure_stage(&self, operation: Operation) -> FailureStage { + match (self.commitment, operation) { + (Commitment::Replayable, Operation::Send) => FailureStage::ProviderCall, + (Commitment::Replayable, _) => FailureStage::BeforeProvider, + (Commitment::ProviderStarted, _) => FailureStage::ProviderCall, + (Commitment::ResponseReceived, _) => FailureStage::AfterProviderResponse, + } + } +} + +pub fn actions_for(operation: Operation) -> &'static [ActionBinding] { + match operation { + Operation::Prepare | Operation::Send => &PROVIDER_ACTION, + Operation::PreCall => &PRE_CALL_ACTION, + Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => { + &FAILURE_ACTION + } + Operation::Restore => &RESTORE_ACTION, + Operation::Complete(_) => &[], + _ => &CALLBACK_ACTION, + } +} + +const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::ProviderCall, + delivery: Delivery::InlineAwaited, + on_result: ResultPolicy::Replace, + on_error: FailurePolicy::Propagate, + owner: Owner::Core, +}]; + +const PRE_CALL_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::RequestPolicy, + delivery: Delivery::InlineDirect, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::Propagate, + owner: Owner::Route, +}]; + +const CALLBACK_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::TerminalSuccess, + delivery: Delivery::InlineAwaited, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::RecordAndContinue, + owner: Owner::Route, +}]; + +const FAILURE_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::TerminalFailure, + delivery: Delivery::InlineAwaited, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::PreserveOriginalFailure, + owner: Owner::Route, +}]; + +const RESTORE_ACTION: [ActionBinding; 1] = [ActionBinding { + kind: ActionKind::Restore, + delivery: Delivery::InlineDirect, + on_result: ResultPolicy::Continue, + on_error: FailurePolicy::Propagate, + owner: Owner::Core, +}]; + +#[cfg(test)] +mod tests { + use super::*; + + fn observations() -> Observations { + Observations { + logger_available: true, + has_fallbacks: false, + } + } + + #[test] + fn commitment_classifies_failures_without_host_inference() { + let mut before = CallProgram::new(ProgramOptions { + asynchronous: false, + internal_call: false, + pre_call: false, + }); + let failure = before.advance(Outcome::Failure, observations()).unwrap(); + assert_eq!(failure.commitment, Commitment::Replayable); + assert_eq!(failure.failure_stage, Some(FailureStage::BeforeProvider)); + + let mut provider = CallProgram::new(ProgramOptions { + asynchronous: false, + internal_call: false, + pre_call: false, + }); + provider.advance(Outcome::Success, observations()).unwrap(); + provider.advance(Outcome::Success, observations()).unwrap(); + let failure = provider.advance(Outcome::Failure, observations()).unwrap(); + assert_eq!(failure.commitment, Commitment::ProviderStarted); + assert_eq!(failure.failure_stage, Some(FailureStage::ProviderCall)); + + let mut after = CallProgram::new(ProgramOptions { + asynchronous: false, + internal_call: false, + pre_call: false, + }); + after.advance(Outcome::Success, observations()).unwrap(); + after.advance(Outcome::Success, observations()).unwrap(); + after.advance(Outcome::Success, observations()).unwrap(); + let failure = after.advance(Outcome::Failure, observations()).unwrap(); + assert_eq!(failure.commitment, Commitment::ResponseReceived); + assert_eq!( + failure.failure_stage, + Some(FailureStage::AfterProviderResponse) + ); + } +} diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 4cb34acbe32..13db329b663 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -13,6 +13,12 @@ pub(super) async fn execute_messages_provider_call( request: MessagesRequest, ) -> Result { let request = prepare_provider_request(request)?; + execute_prepared_messages_provider_call(request).await +} + +pub async fn execute_prepared_messages_provider_call( + request: super::types::ProviderMessagesRequest, +) -> Result { let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); diff --git a/litellm-rust/crates/core/src/messages/lifecycle.rs b/litellm-rust/crates/core/src/messages/lifecycle.rs index cd33e844c5d..bc651acb18f 100644 --- a/litellm-rust/crates/core/src/messages/lifecycle.rs +++ b/litellm-rust/crates/core/src/messages/lifecycle.rs @@ -4,32 +4,17 @@ use std::sync::Arc; use crate::Error; use crate::integrations::custom_logger::{LogError, LogFuture}; use crate::integrations::types::Usage; +use crate::lifecycle::program::{CallProgram, ProgramOptions, actions_for}; use crate::lifecycle::{ - ActionBinding, ActionKind, ActionResult, CallLifecycle, CallLifecycleContext, Clock, Delivery, - ErrorDisposition, ExecutedCall, FailurePolicy, Lifecycle, LifecycleRoute, Outcome, Owner, - RequestPolicy, ResultPolicy, StreamingCall, StreamingObserver, TerminalDispatcher, - TerminalRecord, + ActionBinding, ActionResult, CallLifecycle, CallLifecycleContext, Clock, ExecutedCall, + Lifecycle, LifecycleRoute, Outcome, RequestPolicy, StreamingCall, StreamingObserver, + TerminalDispatcher, TerminalRecord, }; use super::handler::{execute_messages_provider_call, execute_messages_provider_stream}; use super::types::{AnthropicMessagesResponse, MessagesRequest}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum Operation { - Setup, - DeploymentPre, - Prepare, - Send, - DeploymentSuccess, - DeploymentFailure, - SyncSuccess, - AsyncSuccess, - SyncSuccessIfNeeded, - SyncFailure, - AsyncFailure, - Restore, - Complete(Outcome), -} +pub use crate::lifecycle::program::{Observations, Operation, Transition}; #[derive(Clone, Debug, Default)] pub struct Options { @@ -39,24 +24,9 @@ pub struct Options { pub trace_id: Option, } -#[derive(Clone, Copy, Debug, Default)] -pub struct Observations { - pub logger_available: bool, - pub has_fallbacks: bool, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct Transition { - pub operation: Operation, - pub error: ErrorDisposition, -} - #[derive(Debug)] pub struct MessagesState { - operation: Operation, - outcome: Outcome, - asynchronous: bool, - internal_call: bool, + program: CallProgram, } #[derive(Debug)] @@ -76,15 +46,16 @@ impl LifecycleRoute for MessagesRoute { fn admit(_: &(), options: Options) -> Result, Error> { Ok(Ok(MessagesState { - operation: Operation::Setup, - outcome: Outcome::Success, - asynchronous: options.asynchronous, - internal_call: options.internal_call, + program: CallProgram::new(ProgramOptions { + asynchronous: options.asynchronous, + internal_call: options.internal_call, + pre_call: false, + }), })) } fn operation(state: &Self::State) -> Operation { - state.operation + state.program.operation() } fn advance( @@ -92,98 +63,26 @@ impl LifecycleRoute for MessagesRoute { outcome: Outcome, observations: Observations, ) -> Result { - use Operation::*; - - if matches!(state.operation, Complete(_)) { - return Err(Error::InvalidRequest( - "messages lifecycle is already complete".into(), - )); - } - let failure = - if observations.logger_available && !(state.asynchronous && state.internal_call) { - SyncFailure - } else { - Restore - }; - let error = if outcome != Outcome::Success && state.operation != DeploymentFailure { - state.outcome = outcome; - ErrorDisposition::Replace - } else { - ErrorDisposition::Preserve - }; - state.operation = match (state.operation, outcome) { - (Restore, _) => Complete(state.outcome), - (DeploymentFailure, _) => failure, - (_, Outcome::Abort) => Restore, - (SyncFailure | AsyncFailure, Outcome::Failure) => Restore, - (Prepare | Send, Outcome::Failure) if state.asynchronous => DeploymentFailure, - (_, Outcome::Failure) => failure, - (Setup, Outcome::Success) if state.asynchronous => DeploymentPre, - (Setup | DeploymentPre, Outcome::Success) => Prepare, - (Prepare, Outcome::Success) => Send, - (Send, Outcome::Success) if state.asynchronous => DeploymentSuccess, - (Send, Outcome::Success) => SyncSuccess, - (DeploymentSuccess, Outcome::Success) => { - if state.internal_call || observations.has_fallbacks { - SyncSuccessIfNeeded - } else { - AsyncSuccess - } - } - (AsyncSuccess, Outcome::Success) => SyncSuccessIfNeeded, - (SyncFailure, Outcome::Success) if state.asynchronous => AsyncFailure, - (SyncSuccess | SyncSuccessIfNeeded | SyncFailure | AsyncFailure, Outcome::Success) => { - Restore - } - (Complete(_), _) => unreachable!(), - }; - Ok(Transition { - operation: state.operation, - error, - }) + state + .program + .advance(outcome, observations) + .ok_or_else(|| Error::InvalidRequest("messages lifecycle is already complete".into())) } fn actions_for(operation: Operation, _: &Observations) -> &'static [ActionBinding] { - match operation { - Operation::Prepare | Operation::Send => &PROVIDER_ACTION, - Operation::SyncFailure | Operation::AsyncFailure | Operation::DeploymentFailure => { - &FAILURE_ACTION - } - Operation::Restore => &RESTORE_ACTION, - Operation::Complete(_) => &[], - _ => &CALLBACK_ACTION, - } + actions_for(operation) } } -const PROVIDER_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::ProviderCall, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Replace, - on_error: FailurePolicy::Propagate, - owner: Owner::Core, -}]; -const CALLBACK_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::TerminalSuccess, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::RecordAndContinue, - owner: Owner::Route, -}]; -const FAILURE_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::TerminalFailure, - delivery: Delivery::InlineAwaited, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::PreserveOriginalFailure, - owner: Owner::Route, -}]; -const RESTORE_ACTION: [ActionBinding; 1] = [ActionBinding { - kind: ActionKind::Restore, - delivery: Delivery::InlineDirect, - on_result: ResultPolicy::Continue, - on_error: FailurePolicy::Propagate, - owner: Owner::Core, -}]; +impl Lifecycle { + pub fn commitment(&self) -> crate::lifecycle::Commitment { + self.state.program.commitment() + } + + pub fn failure_stage(&self) -> Option { + self.state.program.failure_stage() + } +} pub trait MessagesServices: RequestPolicy + TerminalDispatcher + Clock @@ -241,17 +140,29 @@ pub async fn messages( context: CallLifecycleContext, ) -> ExecutedCall { CallLifecycle - .run( + .run_with_usage( context, request, services, services, services, |request| async move { execute_messages_provider_call(request).await }, + anthropic_response_usage, ) .await } +fn anthropic_response_usage(response: &AnthropicMessagesResponse) -> Option { + let usage = response.usage.as_ref()?; + let prompt_tokens = usage.get("input_tokens")?.as_u64()?; + let completion_tokens = usage.get("output_tokens")?.as_u64()?; + Some(Usage { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + }) +} + pub async fn messages_stream( services: Arc, request: MessagesRequest, diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 16cf64503c1..da9bc813f83 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -19,6 +19,8 @@ pub mod types; use std::sync::Arc; use crate::lifecycle::StreamingCall; +pub use handler::execute_prepared_messages_provider_call; +pub use prepare::prepare_provider_request; use types::{AnthropicMessagesResponse, MessagesRequest}; pub async fn messages(request: MessagesRequest) -> Result { diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 45569857882..e8b0f57a0be 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -6,7 +6,7 @@ use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrateg use super::types::{MessagesRequest, ProviderMessagesRequest}; use serde_json::{Map, Value}; -pub(super) fn prepare_provider_request( +pub fn prepare_provider_request( request: MessagesRequest, ) -> Result { let provider_info = diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index f55a9c9dfa7..a1da80b0afa 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -15,14 +15,14 @@ pub struct MessagesRequest { pub timeout: Option, } -pub(super) struct ProviderMessagesRequest { - pub(super) provider: String, - pub(super) model: String, +pub struct ProviderMessagesRequest { + pub provider: String, + pub model: String, pub(super) config: &'static dyn AnthropicMessagesProviderConfig, - pub(super) url: String, - pub(super) body: Value, - pub(super) upstream_headers: Vec<(String, String)>, - pub(super) timeout: Option, + pub url: String, + pub body: Value, + pub upstream_headers: Vec<(String, String)>, + pub timeout: Option, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index f39d90cc892..a7e89658b00 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -3,37 +3,49 @@ pub mod transformation; pub mod types; use serde_json::Value; +use std::future::Future; use crate::Error; use crate::error::json_type_name; use crate::http_utils::{buffered_post, has_header}; -pub use types::{OcrAdmissionRequest, OcrDraft, OcrEndpoint, OcrResponseData, SettledOcrRequest}; +pub use types::{ + OcrAdmissionRequest, OcrDraft, OcrEndpoint, OcrResponseData, OcrTransportRequest, + OcrTransportResponse, SettledOcrRequest, +}; use types::{OcrDocument, OcrDocumentProjection}; use crate::lifecycle::{ CallLifecycle, CallLifecycleContext, Clock, ExecutedCall, TerminalDispatcher, }; -pub trait OcrServices: TerminalDispatcher + Clock {} +pub trait OcrTransport { + type SendFuture<'a>: Future> + where + Self: 'a; -impl OcrServices for T where T: TerminalDispatcher + Clock {} + fn send(&self, request: OcrTransportRequest) -> Self::SendFuture<'_>; +} -pub struct NoopOcrServices; +pub trait OcrServices: TerminalDispatcher + Clock + OcrTransport {} -impl Default for NoopOcrServices { +impl OcrServices for T where T: TerminalDispatcher + Clock + OcrTransport {} + +pub struct DefaultOcrServices; + +impl Default for DefaultOcrServices { fn default() -> Self { Self } } -impl Clock for NoopOcrServices { +impl Clock for DefaultOcrServices { fn now(&self) -> f64 { crate::lifecycle::SystemClock.now() } } -impl TerminalDispatcher for NoopOcrServices { +impl TerminalDispatcher for DefaultOcrServices { fn dispatch<'a>( &'a self, _: &'a crate::lifecycle::TerminalRecord, @@ -42,6 +54,27 @@ impl TerminalDispatcher for NoopOcrServices { } } +impl OcrTransport for DefaultOcrServices { + type SendFuture<'a> = impl Future> + 'a; + + fn send(&self, request: OcrTransportRequest) -> Self::SendFuture<'_> { + async move { + let response = buffered_post::send(buffered_post::Request { + url: request.url, + headers: request.headers, + body: request.body, + timeout_seconds: request.timeout_seconds, + }) + .await?; + Ok(OcrTransportResponse { + status: response.status, + headers: response.headers, + content: response.content, + }) + } + } +} + pub async fn ocr( services: &S, request: SettledOcrRequest, @@ -55,7 +88,11 @@ pub async fn ocr( &SettledOcrPolicy, services, services, - |request| async move { send(request).await.map(OcrResponseData::into_json) }, + |request| async move { + send(services, request) + .await + .map(OcrResponseData::into_json) + }, ) .await } @@ -89,7 +126,10 @@ impl crate::lifecycle::RequestPolicy for S } } -pub(crate) async fn send(request: SettledOcrRequest) -> Result { +pub(crate) async fn send( + transport: &S, + request: SettledOcrRequest, +) -> Result { let SettledOcrRequest { endpoint, headers, @@ -132,13 +172,14 @@ pub(crate) async fn send(request: SettledOcrRequest) -> Result, Vec)>, + pub body: Vec, + pub timeout_seconds: f64, +} + +#[derive(Debug, PartialEq, Eq)] +pub struct OcrTransportResponse { + pub status: u16, + pub headers: Vec<(Vec, Vec)>, + pub content: Vec, +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct OcrRequestData { pub data: Value, diff --git a/litellm-rust/crates/core/src/realtime/mod.rs b/litellm-rust/crates/core/src/realtime/mod.rs index 5c36ed03829..6e34a25482b 100644 --- a/litellm-rust/crates/core/src/realtime/mod.rs +++ b/litellm-rust/crates/core/src/realtime/mod.rs @@ -3,6 +3,4 @@ mod streaming; pub mod transformation; pub mod types; -pub use streaming::{ - RealtimeConnectionSpec, RealtimeRequest, WarmConnection, realtime, warmup, -}; +pub use streaming::{RealtimeConnectionSpec, RealtimeRequest, WarmConnection, realtime, warmup}; diff --git a/litellm-rust/crates/core/src/realtime/streaming.rs b/litellm-rust/crates/core/src/realtime/streaming.rs index 05db2eeddfb..240e27f40da 100644 --- a/litellm-rust/crates/core/src/realtime/streaming.rs +++ b/litellm-rust/crates/core/src/realtime/streaming.rs @@ -20,8 +20,8 @@ use crate::Error; use crate::integrations::custom_logger::CallbackTiming; use crate::integrations::types::Usage; use crate::lifecycle::{ - CallLifecycleContext, Clock, CostInputs, ExecutedCall, RouteProjection, - TerminalClassification, TerminalDispatcher, TerminalRecord, + CallLifecycleContext, Clock, CostInputs, ExecutedCall, RouteProjection, TerminalClassification, + TerminalDispatcher, TerminalRecord, }; use crate::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; use crate::realtime::transformation::RealtimeProviderConfig; @@ -48,8 +48,9 @@ impl RealtimeConnectionSpec { api_key: Option<&str>, api_base: Option<&str>, ) -> Result { + let model = model.into(); Ok(Self { - model: model.into(), + model: openai_model(&model)?.to_string(), api_key: resolve_api_key(api_key)?, api_base: api_base.map(str::to_string), }) @@ -88,6 +89,7 @@ impl std::fmt::Debug for RealtimeConnectionSpec { } pub struct WarmConnection { + connection: RealtimeConnectionSpec, upstream: Upstream, session_created: RealtimeEvent, } @@ -95,12 +97,17 @@ pub struct WarmConnection { impl WarmConnection { pub fn is_live(&mut self) -> bool { let mut context = Context::from_waker(futures_util::task::noop_waker_ref()); - matches!(Pin::new(&mut self.upstream).poll_next(&mut context), Poll::Pending) + matches!( + Pin::new(&mut self.upstream).poll_next(&mut context), + Poll::Pending + ) } } pub struct RealtimeRequest { - pub connection: RealtimeConnectionSpec, + pub model: String, + pub api_key: Option, + pub api_base: Option, pub warm: Option, pub idle_timeout: Option, } @@ -115,6 +122,7 @@ pub async fn warmup(connection: &RealtimeConnectionSpec) -> Result Ok(warm), - None => dial_upstream(&request.connection) + let model = request.model; + let connection = RealtimeConnectionSpec::new( + model.clone(), + request.api_key.as_deref(), + request.api_base.as_deref(), + ); + let connection = match (connection, request.warm) { + (Ok(connection), Some(warm)) if warm.connection == connection => Ok(warm), + (Ok(_), Some(_)) => Err(Error::InvalidRequest( + "realtime warm connection does not match the requested provider connection".to_string(), + )), + (Ok(connection), None) => dial_upstream(&connection) .await .map(|upstream| WarmConnection { + connection, upstream, session_created: empty_event(), }), + (Err(error), _) => Err(error), }; let mut observation = RealtimeObservation::new(context.litellm_call_id.clone(), model.clone()); let result = match connection { @@ -157,13 +175,13 @@ where ) .await } - Err(error) => Err(error), + Err(error) => Err(error.into()), }; let classification = match &result { Ok(()) => TerminalClassification::Success, - Err(error) => TerminalClassification::Failure { - kind: error_kind(error).to_string(), - message: error.to_string(), + Err(failure) => TerminalClassification::Failure { + kind: failure.kind.to_string(), + message: failure.error.to_string(), }, }; let projection = match &classification { @@ -194,7 +212,10 @@ where response: (), terminal, }, - Err(error) => ExecutedCall::Failure { error, terminal }, + Err(failure) => ExecutedCall::Failure { + error: failure.error, + terminal, + }, } } @@ -205,13 +226,14 @@ async fn splice( observation: &mut RealtimeObservation, mut client_in: In, mut client_out: Out, -) -> Result<(), Error> +) -> Result<(), RealtimeFailure> where In: Stream + Unpin + Send, Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { let WarmConnection { + connection: _, upstream, session_created, } = connection; @@ -223,31 +245,106 @@ where loop { tokio::select! { event = client_in.next() => { - let Some(event) = event else { return Ok(()) }; + let Some(event) = event else { + return observation.settle( + "Cancelled", + "realtime client disconnected before provider completion", + ); + }; for outbound in OPENAI_REALTIME_CONFIG.transform_realtime_request(&event, model)?.events { let payload = serde_json::to_string(&outbound) .map_err(|error| Error::InvalidResponse(error.to_string()))?; upstream_tx.send(Message::Text(payload.into())).await.map_err(ws_transport_error)?; } + observation.observe_client(&event); } message = upstream_rx.next() => { - let Some(message) = message else { return Ok(()) }; + let Some(message) = message else { + return observation.settle( + "NetworkError", + "realtime provider closed before completion", + ); + }; match message.map_err(ws_transport_error)? { Message::Text(text) => { let event = serde_json::from_str::(&text) .map_err(|error| Error::InvalidResponse(error.to_string()))?; observation.observe(&event); send_client_event(&mut client_out, &event, model).await?; + if event.event_type == "error" || response_failed(&event) { + return Err(RealtimeFailure::new( + "ProviderError", + Error::InvalidResponse(provider_error_message(&event)), + )); + } + } + Message::Close(_) => { + return observation.settle( + "NetworkError", + "realtime provider closed before completion", + ); } - Message::Close(_) => return Ok(()), _ => {} } } - _ = tokio::time::sleep(idle_timeout) => return Ok(()), + _ = tokio::time::sleep(idle_timeout) => { + return Err(RealtimeFailure::new( + "Timeout", + Error::Network("realtime session idle timeout".to_string()), + )); + } } } } +struct RealtimeFailure { + kind: &'static str, + error: Error, +} + +impl RealtimeFailure { + fn new(kind: &'static str, error: Error) -> Self { + Self { kind, error } + } +} + +impl From for RealtimeFailure { + fn from(error: Error) -> Self { + Self { + kind: error_kind(&error), + error, + } + } +} + +fn response_failed(event: &RealtimeEvent) -> bool { + event.event_type == "response.done" + && event + .data + .get("response") + .and_then(|response| response.get("status")) + .and_then(Value::as_str) + .is_some_and(|status| status != "completed") +} + +fn provider_error_message(event: &RealtimeEvent) -> String { + event + .data + .get("error") + .and_then(|error| error.get("message")) + .or_else(|| { + event + .data + .get("response") + .and_then(|response| response.get("status_details")) + .and_then(|details| details.get("error")) + .and_then(|error| error.get("message")) + }) + .and_then(Value::as_str) + .unwrap_or("realtime provider reported an error") + .to_string() +} + async fn send_client_event( client_out: &mut Out, event: &RealtimeEvent, @@ -292,10 +389,8 @@ async fn read_event(upstream: &mut Upstream) -> Result { } async fn dial_upstream(connection: &RealtimeConnectionSpec) -> Result { - let url = OPENAI_REALTIME_CONFIG.complete_url( - connection.api_base.as_deref(), - connection.model.as_str(), - ); + let url = OPENAI_REALTIME_CONFIG + .complete_url(connection.api_base.as_deref(), connection.model.as_str()); let mut request = url.into_client_request().map_err(ws_transport_error)?; request.headers_mut().insert( AUTHORIZATION, @@ -328,13 +423,12 @@ fn tls_config() -> Result, Error> { native.errors ))); } - let config = ClientConfig::builder_with_provider(Arc::new( - rustls::crypto::ring::default_provider(), - )) - .with_safe_default_protocol_versions() - .map_err(|error| Error::Connect(error.to_string()))? - .with_root_certificates(roots) - .with_no_client_auth(); + let config = + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map_err(|error| Error::Connect(error.to_string()))? + .with_root_certificates(roots) + .with_no_client_auth(); let config = Arc::new(config); Ok(Arc::clone(TLS_CONFIG.get_or_init(|| config))) } @@ -352,6 +446,18 @@ fn resolve_api_key(api_key: Option<&str>) -> Result { .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) } +fn openai_model(model: &str) -> Result<&str, Error> { + if let Some((provider, provider_model)) = model.split_once('/') { + if provider != "openai" { + return Err(Error::InvalidProvider(format!( + "realtime route does not support provider '{provider}'" + ))); + } + return Ok(provider_model); + } + Ok(model) +} + fn ws_handshake_error(error: WsError) -> Error { match error { WsError::Http(response) => Error::Http { @@ -400,6 +506,8 @@ struct RealtimeObservation { call_id: String, model: String, usage: Usage, + completed_response: bool, + pending_responses: usize, } impl RealtimeObservation { @@ -408,6 +516,30 @@ impl RealtimeObservation { call_id, model, usage: Usage::default(), + completed_response: false, + pending_responses: 0, + } + } + + fn settle(&self, kind: &'static str, message: &str) -> Result<(), RealtimeFailure> { + if self.completed_response && self.pending_responses == 0 { + Ok(()) + } else if kind == "Cancelled" { + Err(RealtimeFailure::new( + kind, + Error::InvalidRequest(message.to_string()), + )) + } else { + Err(RealtimeFailure::new( + kind, + Error::Network(message.to_string()), + )) + } + } + + fn observe_client(&mut self, event: &RealtimeEvent) { + if event.event_type == "response.create" { + self.pending_responses += 1; } } @@ -433,6 +565,10 @@ impl RealtimeObservation { if event.event_type != "response.done" { return; } + if !response_failed(event) { + self.completed_response = true; + self.pending_responses = self.pending_responses.saturating_sub(1); + } let Some(usage) = event .data .get("response") @@ -442,7 +578,10 @@ impl RealtimeObservation { else { return; }; - let input = usage.get("input_tokens").and_then(Value::as_u64).unwrap_or(0); + let input = usage + .get("input_tokens") + .and_then(Value::as_u64) + .unwrap_or(0); let output = usage .get("output_tokens") .and_then(Value::as_u64) @@ -488,20 +627,30 @@ mod tests { } } - async fn provider() -> String { + async fn scripted_provider(events: Vec, close_after_events: bool) -> String { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); tokio::spawn(async move { while let Ok((stream, _)) = listener.accept().await { + let events = events.clone(); tokio::spawn(async move { let mut socket = accept_async(stream).await.unwrap(); socket.send(Message::Text(json!({"type":"session.created","session":{"id":"sess-core","model":"upstream-model"}}).to_string().into())).await.unwrap(); while let Some(Ok(Message::Text(text))) = socket.next().await { let event: RealtimeEvent = serde_json::from_str(&text).unwrap(); if event.event_type == "response.create" { - socket.send(Message::Text(json!({"type":"response.done","response":{"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}}).to_string().into())).await.unwrap(); - socket.send(Message::Text(json!({"type":"response.done","response":{"usage":{"input_tokens":7,"output_tokens":11}}}).to_string().into())).await.unwrap(); - socket.close(None).await.unwrap(); + for event in &events { + if socket + .send(Message::Text(event.to_string().into())) + .await + .is_err() + { + return; + } + } + if close_after_events { + let _ = socket.close(None).await; + } } } }); @@ -510,25 +659,88 @@ mod tests { format!("ws://{address}") } + async fn provider() -> String { + scripted_provider( + vec![ + json!({"type":"response.done","response":{"usage":{"input_tokens":2,"output_tokens":3,"total_tokens":5}}}), + json!({"type":"response.done","response":{"usage":{"input_tokens":7,"output_tokens":11}}}), + ], + true, + ) + .await + } + + async fn execute_scenario( + events: Vec, + close_after_events: bool, + send_response_create: bool, + disconnect_client: bool, + idle_timeout: Duration, + ) -> (ExecutedCall<(), Error>, Vec) { + let base = scripted_provider(events, close_after_events).await; + let services = Services::default(); + let (input_tx, input) = mpsc::unbounded(); + if send_response_create { + input_tx + .unbounded_send(serde_json::from_value(json!({"type":"response.create"})).unwrap()) + .unwrap(); + } + if disconnect_client { + drop(input_tx); + } + let (output, _output_rx) = mpsc::unbounded(); + let result = realtime( + &services, + RealtimeRequest { + model: "requested".to_string(), + api_key: Some("key".to_string()), + api_base: Some(base), + warm: None, + idle_timeout: Some(idle_timeout), + }, + CallLifecycleContext::new("realtime", "requested", "openai", "fallback"), + input, + output, + ) + .await; + let terminals = services.terminals.into_inner().unwrap(); + (result, terminals) + } + async fn execute(warm: bool) -> (ExecutedCall<(), Error>, Vec, usize) { let base = provider().await; let spec = RealtimeConnectionSpec::new("requested", Some("key"), Some(&base)).unwrap(); let services = Services::default(); - let warm = if warm { Some(warmup(&spec).await.unwrap()) } else { None }; + let warm = if warm { + Some(warmup(&spec).await.unwrap()) + } else { + None + }; assert!(services.terminals.lock().unwrap().is_empty()); let (input_tx, input) = mpsc::unbounded(); let (output, mut output_rx) = mpsc::unbounded(); input_tx.unbounded_send(serde_json::from_value(json!({"type":"response.done","response":{"usage":{"input_tokens":1000,"output_tokens":1000,"total_tokens":2000}}})).unwrap()).unwrap(); - input_tx.unbounded_send(serde_json::from_value(json!({"type":"response.create"})).unwrap()).unwrap(); + input_tx + .unbounded_send(serde_json::from_value(json!({"type":"response.create"})).unwrap()) + .unwrap(); let result = realtime( &services, - RealtimeRequest { connection: spec, warm, idle_timeout: Some(Duration::from_secs(1)) }, + RealtimeRequest { + model: spec.model.clone(), + api_key: Some("key".to_string()), + api_base: Some(base), + warm, + idle_timeout: Some(Duration::from_secs(1)), + }, CallLifecycleContext::new("realtime", "requested", "openai", "fallback"), input, output, - ).await; + ) + .await; let mut events = Vec::new(); - while let Ok(Some(event)) = tokio::time::timeout(Duration::from_millis(10), output_rx.next()).await { + while let Ok(Some(event)) = + tokio::time::timeout(Duration::from_millis(10), output_rx.next()).await + { events.push(event); } let count = services.terminals.lock().unwrap().len(); @@ -541,21 +753,120 @@ mod tests { let (result, events, count) = execute(warm).await; assert_eq!(count, 1); assert_eq!(events.first().unwrap().event_type, "session.created"); - let ExecutedCall::Success { terminal, .. } = result else { panic!("session failed") }; + let ExecutedCall::Success { terminal, .. } = result else { + panic!("session failed") + }; + assert_eq!(terminal.classification, TerminalClassification::Success); assert_eq!(terminal.call_id, "sess-core"); assert_eq!(terminal.model, "upstream-model"); - assert_eq!(terminal.usage, Usage { prompt_tokens: 9, completion_tokens: 14, total_tokens: 23 }); + assert_eq!( + terminal.usage, + Usage { + prompt_tokens: 9, + completion_tokens: 14, + total_tokens: 23 + } + ); } } + #[tokio::test] + async fn client_disconnect_before_provider_completion_is_cancelled_once() { + let (result, terminals) = + execute_scenario(Vec::new(), false, false, true, Duration::from_secs(1)).await; + + assert!(matches!(result, ExecutedCall::Failure { .. })); + assert_eq!(terminals.len(), 1); + assert!(matches!( + &terminals[0].classification, + TerminalClassification::Failure { kind, .. } if kind == "Cancelled" + )); + } + + #[tokio::test] + async fn idle_timeout_before_provider_completion_fails_once() { + let (result, terminals) = + execute_scenario(Vec::new(), false, false, false, Duration::from_millis(20)).await; + + assert!(matches!(result, ExecutedCall::Failure { .. })); + assert_eq!(terminals.len(), 1); + assert!(matches!( + &terminals[0].classification, + TerminalClassification::Failure { kind, .. } if kind == "Timeout" + )); + } + + #[tokio::test] + async fn provider_error_event_fails_once() { + let (result, terminals) = execute_scenario( + vec![json!({"type":"error","error":{"message":"provider rejected event"}})], + false, + true, + false, + Duration::from_secs(1), + ) + .await; + + assert!(matches!(result, ExecutedCall::Failure { .. })); + assert_eq!(terminals.len(), 1); + assert!(matches!( + &terminals[0].classification, + TerminalClassification::Failure { kind, message } + if kind == "ProviderError" && message.contains("provider rejected event") + )); + } + + #[tokio::test] + async fn provider_close_before_response_done_fails_once() { + let (result, terminals) = + execute_scenario(Vec::new(), true, true, false, Duration::from_secs(1)).await; + + assert!(matches!(result, ExecutedCall::Failure { .. })); + assert_eq!(terminals.len(), 1); + assert!(matches!( + &terminals[0].classification, + TerminalClassification::Failure { kind, .. } if kind == "NetworkError" + )); + } + #[tokio::test] async fn warmup_success_and_failure_dispatch_nothing() { let services = Services::default(); let base = provider().await; let good = RealtimeConnectionSpec::new("model", Some("key"), Some(&base)).unwrap(); assert!(warmup(&good).await.is_ok()); - let bad = RealtimeConnectionSpec::new("model", Some("key"), Some("ws://127.0.0.1:1")).unwrap(); + let bad = + RealtimeConnectionSpec::new("model", Some("key"), Some("ws://127.0.0.1:1")).unwrap(); assert!(warmup(&bad).await.is_err()); assert!(services.terminals.lock().unwrap().is_empty()); } + + #[tokio::test] + async fn fresh_dial_failure_returns_and_dispatches_one_terminal() { + let services = Services::default(); + let (_, input) = mpsc::unbounded(); + let (output, _) = mpsc::unbounded(); + let result = realtime( + &services, + RealtimeRequest { + model: "model".to_string(), + api_key: Some("key".to_string()), + api_base: Some("ws://127.0.0.1:1".to_string()), + warm: None, + idle_timeout: None, + }, + CallLifecycleContext::new("realtime", "model", "openai", "call-failure"), + input, + output, + ) + .await; + + assert!(matches!(result, ExecutedCall::Failure { .. })); + let terminals = services.terminals.lock().unwrap(); + assert_eq!(terminals.len(), 1); + assert!(matches!( + terminals[0].classification, + TerminalClassification::Failure { .. } + )); + } } diff --git a/litellm-rust/crates/core/src/responses/instrumentation.rs b/litellm-rust/crates/core/src/responses/instrumentation.rs index 59f12449dd0..8c49ede9893 100644 --- a/litellm-rust/crates/core/src/responses/instrumentation.rs +++ b/litellm-rust/crates/core/src/responses/instrumentation.rs @@ -1,4 +1,5 @@ use crate::integrations::types::Usage; +use crate::lifecycle::TerminalClassification; use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType}; use serde_json::Value; use std::sync::Mutex; @@ -69,6 +70,71 @@ impl ResponsesWsInstrumentation { .unwrap_or_else(|poisoned| poisoned.into_inner()) .clone() } + + pub(crate) fn terminal_classification( + &self, + event: &ResponsesWsEvent, + ) -> Option { + match event.event_type { + ResponsesWsEventType::ResponseCompleted => Some(TerminalClassification::Success), + ResponsesWsEventType::ResponseFailed => Some(provider_failure( + "ResponseFailed", + response_error_message(event).unwrap_or("provider response failed"), + )), + ResponsesWsEventType::ResponseIncomplete => Some(provider_failure( + "ResponseIncomplete", + incomplete_message(event).unwrap_or("provider response was incomplete"), + )), + ResponsesWsEventType::Error => Some(provider_failure( + "ProviderError", + top_level_error_message(event).unwrap_or("provider returned an error"), + )), + _ => None, + } + } +} + +fn response_error_message(event: &ResponsesWsEvent) -> Option<&str> { + event + .data + .get("response") + .and_then(Value::as_object) + .and_then(|response| response.get("error")) + .and_then(Value::as_object) + .and_then(|error| error.get("message")) + .and_then(Value::as_str) +} + +fn incomplete_message(event: &ResponsesWsEvent) -> Option<&str> { + event + .data + .get("response") + .and_then(Value::as_object) + .and_then(|response| response.get("incomplete_details")) + .and_then(Value::as_object) + .and_then(|details| details.get("reason")) + .and_then(Value::as_str) +} + +fn top_level_error_message(event: &ResponsesWsEvent) -> Option<&str> { + event + .data + .get("error") + .and_then(Value::as_object) + .and_then(|error| error.get("message")) + .and_then(Value::as_str) +} + +fn provider_failure(kind: &str, message: &str) -> TerminalClassification { + let message = message.trim().chars().take(512).collect::(); + TerminalClassification::Failure { + kind: kind.to_string(), + message: if message.is_empty() { + "provider returned an unspecified failure".to_string() + } else { + message + }, + } } #[cfg(test)] @@ -102,4 +168,36 @@ mod tests { assert_eq!(observation.usage.completion_tokens, 5); assert_eq!(observation.usage.total_tokens, 8); } + + #[test] + fn classifies_provider_terminal_frames_without_serializing_the_frame() { + let instrumentation = ResponsesWsInstrumentation::default(); + let cases = [ + ( + serde_json::json!({"type":"response.failed","response":{"error":{"message":"request rejected"}}}), + "ResponseFailed", + "request rejected", + ), + ( + serde_json::json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"max_output_tokens"}}}), + "ResponseIncomplete", + "max_output_tokens", + ), + ( + serde_json::json!({"type":"error","error":{"type":"server_error","message":"provider unavailable"}}), + "ProviderError", + "provider unavailable", + ), + ]; + + for (frame, expected_kind, expected_message) in cases { + assert_eq!( + instrumentation.terminal_classification(&event(frame)), + Some(TerminalClassification::Failure { + kind: expected_kind.to_string(), + message: expected_message.to_string(), + }) + ); + } + } } diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index c7816f17337..ce14ca9000b 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -67,7 +67,7 @@ pub struct ResponsesWebSocketRequest { } pub async fn responses_websocket( - services: &S, + services: Arc, request: ResponsesWebSocketRequest, context: crate::lifecycle::CallLifecycleContext, client_in: In, @@ -75,69 +75,164 @@ pub async fn responses_websocket( ) -> Result, Error> where S: TerminalDispatcher + crate::lifecycle::Clock, - In: Stream + Unpin + Send, + S: 'static, + In: Stream> + Unpin + Send, Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { let key = resolve_api_key(request.api_key.as_deref())?; let upstream = dial_upstream(&request.model, &key, request.api_base.as_deref()).await?; let start_time = services.now(); - let instrumentation = ResponsesWsInstrumentation::default(); + let instrumentation = Arc::new(ResponsesWsInstrumentation::default()); + let mut completion = + ResponsesWsCompletion::new(services, context, start_time, Arc::clone(&instrumentation)); let result = splice( upstream, &request.model, request.first_frame, request.idle_timeout.unwrap_or(IDLE_TIMEOUT), - &instrumentation, + instrumentation.as_ref(), client_in, client_out, ) .await; - let observation = instrumentation.snapshot(); - let model = if observation.model.is_empty() { - context.model.clone() - } else { - observation.model - }; let classification = match &result { - Ok(()) => TerminalClassification::Success, - Err(error) => TerminalClassification::Failure { - kind: error_kind(error).to_string(), - message: error.to_string(), - }, + Ok(classification) => classification.clone(), + Err(failure) => failure.classification.clone(), }; - let projection = match &classification { - TerminalClassification::Success => Value::Null, - TerminalClassification::Failure { kind, message } => { - json!({"kind": kind, "message": message}) - } - }; - let terminal = TerminalRecord { - call_id: context.litellm_call_id, - trace_id: context.trace_id, - attempt: context.attempt, - call_type: context.call_type, - model, - provider: context.custom_llm_provider, - timing: CallbackTiming::new(start_time, services.now()), - usage: observation.usage, - cost_inputs: CostInputs { - response_cost: context.response_cost, - metadata: context.metadata, - }, - classification, - projection: RouteProjection::ResponsesWs { value: projection }, - }; - let _ = services.dispatch(&terminal).await; + let terminal = completion.settle(classification).await; Ok(match result { - Ok(()) => ExecutedCall::Success { + Ok(TerminalClassification::Success) => ExecutedCall::Success { response: (), terminal, }, - Err(error) => ExecutedCall::Failure { error, terminal }, + Ok(TerminalClassification::Failure { message, .. }) => ExecutedCall::Failure { + error: Error::InvalidResponse(message), + terminal, + }, + Err(failure) => ExecutedCall::Failure { + error: failure.error, + terminal, + }, }) } +trait ResponsesCompletionServices: TerminalDispatcher + crate::lifecycle::Clock {} + +impl ResponsesCompletionServices for T where T: TerminalDispatcher + crate::lifecycle::Clock {} + +struct ResponsesWsCompletion { + services: Arc, + context: Option, + start_time: f64, + instrumentation: Arc, +} + +impl ResponsesWsCompletion { + fn new( + services: Arc, + context: crate::lifecycle::CallLifecycleContext, + start_time: f64, + instrumentation: Arc, + ) -> Self + where + S: ResponsesCompletionServices + 'static, + { + Self { + services, + context: Some(context), + start_time, + instrumentation, + } + } + + async fn settle(&mut self, classification: TerminalClassification) -> TerminalRecord { + let terminal = self.terminal(classification); + let dispatched = terminal.clone(); + let services = Arc::clone(&self.services); + let dispatch = tokio::spawn(async move { + let _ = services.dispatch(&dispatched).await; + }); + let _ = dispatch.await; + terminal + } + + fn terminal(&mut self, classification: TerminalClassification) -> TerminalRecord { + let context = self.context.take().expect("Responses session settled once"); + let observation = self.instrumentation.snapshot(); + let model = if observation.model.is_empty() { + context.model + } else { + observation.model + }; + let projection = match &classification { + TerminalClassification::Success => Value::Null, + TerminalClassification::Failure { kind, message } => { + json!({"kind": kind, "message": message}) + } + }; + TerminalRecord { + call_id: context.litellm_call_id, + trace_id: context.trace_id, + attempt: context.attempt, + call_type: context.call_type, + model, + provider: context.custom_llm_provider, + timing: CallbackTiming::new(self.start_time, self.services.now()), + usage: observation.usage, + cost_inputs: CostInputs { + response_cost: context.response_cost, + metadata: context.metadata, + }, + classification, + projection: RouteProjection::ResponsesWs { value: projection }, + } + } +} + +impl Drop for ResponsesWsCompletion { + fn drop(&mut self) { + if self.context.is_none() { + return; + } + let terminal = self.terminal(TerminalClassification::Failure { + kind: "Cancelled".to_string(), + message: "Responses WebSocket session was cancelled before completion".to_string(), + }); + let services = Arc::clone(&self.services); + tokio::spawn(async move { + let _ = services.dispatch(&terminal).await; + }); + } +} + +struct ResponsesWsFailure { + error: Error, + classification: TerminalClassification, +} + +impl ResponsesWsFailure { + fn new(error: Error) -> Self { + Self { + classification: TerminalClassification::Failure { + kind: error_kind(&error).to_string(), + message: error.to_string(), + }, + error, + } + } + + fn session(kind: &str, message: &str) -> Self { + Self { + error: Error::Network(message.to_string()), + classification: TerminalClassification::Failure { + kind: kind.to_string(), + message: message.to_string(), + }, + } + } +} + async fn splice( upstream: Upstream, model: &str, @@ -146,39 +241,66 @@ async fn splice( instrumentation: &ResponsesWsInstrumentation, mut client_in: In, mut client_out: Out, -) -> Result<(), Error> +) -> Result where - In: Stream + Unpin + Send, + In: Stream> + Unpin + Send, Out: Sink + Unpin + Send, Out::Error: std::fmt::Display, { let (mut upstream_tx, mut upstream_rx) = upstream.split(); if let Some(event) = first_frame { - send_provider_event(&mut upstream_tx, &event, model).await?; + send_provider_event(&mut upstream_tx, &event, model) + .await + .map_err(ResponsesWsFailure::new)?; } loop { tokio::select! { event = client_in.next() => { - let Some(event) = event else { return Ok(()) }; - send_provider_event(&mut upstream_tx, &event, model).await?; + let Some(event) = event else { + return Err(ResponsesWsFailure::session( + "ClientDisconnected", + "client disconnected before response.completed", + )); + }; + let event = event.map_err(ResponsesWsFailure::new)?; + send_provider_event(&mut upstream_tx, &event, model).await.map_err(ResponsesWsFailure::new)?; } message = upstream_rx.next() => { - let Some(message) = message else { return Ok(()) }; - match message.map_err(ws_transport_error)? { + let Some(message) = message else { + return Err(ResponsesWsFailure::session( + "ProviderDisconnected", + "provider disconnected before a terminal response frame", + )); + }; + match message.map_err(ws_transport_error).map_err(ResponsesWsFailure::new)? { Message::Text(text) => { let event = serde_json::from_str::(&text) - .map_err(|error| Error::InvalidResponse(error.to_string()))?; + .map_err(|error| ResponsesWsFailure::new(Error::InvalidResponse(error.to_string())))?; instrumentation.observe(&event); - for outbound in OPENAI_RESPONSES_WS_CONFIG.transform_ws_response(&event, model)?.events { + let terminal = instrumentation.terminal_classification(&event); + for outbound in OPENAI_RESPONSES_WS_CONFIG.transform_ws_response(&event, model) + .map_err(ResponsesWsFailure::new)?.events { client_out.send(outbound).await - .map_err(|error| Error::Network(error.to_string()))?; + .map_err(|error| ResponsesWsFailure::session( + "ClientDisconnected", + &format!("failed to deliver provider event to client: {error}"), + ))?; + } + if let Some(classification) = terminal { + return Ok(classification); } } - Message::Close(_) => return Ok(()), + Message::Close(_) => return Err(ResponsesWsFailure::session( + "ProviderDisconnected", + "provider closed before a terminal response frame", + )), _ => {} } } - _ = tokio::time::sleep(idle_timeout) => return Ok(()), + _ = tokio::time::sleep(idle_timeout) => return Err(ResponsesWsFailure::session( + "IdleTimeout", + "Responses WebSocket session timed out before a terminal response frame", + )), } } } @@ -361,8 +483,7 @@ pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool { matches!( event_type, - ResponsesWsEventType::ResponseCreated - | ResponsesWsEventType::ResponseCompleted + ResponsesWsEventType::ResponseCompleted | ResponsesWsEventType::ResponseFailed | ResponsesWsEventType::ResponseIncomplete | ResponsesWsEventType::Error @@ -417,33 +538,64 @@ mod tests { serde_json::from_value(value).expect("event") } - async fn mock_provider() -> (String, tokio::task::JoinHandle<()>) { + async fn provider_with_frames( + frames: Vec, + remain_open: bool, + ) -> (String, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let task = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let mut socket = accept_async(stream).await.unwrap(); - if let Some(Ok(Message::Text(text))) = socket.next().await { - let request: Value = serde_json::from_str(&text).unwrap(); - assert_eq!(request["model"], "authorized"); - socket.send(Message::Text(json!({"type":"response.completed","response":{"id":"resp-1","model":"authorized","usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}}).to_string().into())).await.unwrap(); + for frame in frames { + socket + .send(Message::Text(frame.to_string().into())) + .await + .unwrap(); + } + if remain_open { + futures_util::future::pending::<()>().await; + } else { socket.close(None).await.unwrap(); } }); (format!("http://{address}"), task) } + fn input() -> ( + futures_channel::mpsc::UnboundedSender>, + futures_channel::mpsc::UnboundedReceiver>, + ) { + futures_channel::mpsc::unbounded() + } + + fn assert_failure(services: &Services, kind: &str, message: &str) { + let terminals = services.terminals.lock().unwrap(); + assert_eq!(terminals.len(), 1); + assert_eq!( + terminals[0].classification, + TerminalClassification::Failure { + kind: kind.to_string(), + message: message.to_string(), + } + ); + } + #[tokio::test] - async fn core_owns_splice_transformation_and_one_terminal() { - let (api_base, server) = mock_provider().await; - let services = Services::default(); - let (client_tx, client_rx) = futures_channel::mpsc::unbounded(); + async fn completed_frame_is_delivered_and_settles_once_while_provider_remains_open() { + let (api_base, server) = provider_with_frames( + vec![json!({"type":"response.completed","response":{"id":"resp-1","model":"authorized","usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}})], + true, + ) + .await; + let services = Arc::new(Services::default()); + let (client_tx, client_rx) = input(); let (output_tx, mut output_rx) = futures_channel::mpsc::unbounded(); client_tx - .unbounded_send(event(json!({"type":"response.create","model":"wrong"}))) + .unbounded_send(Ok(event(json!({"type":"response.create","model":"wrong"})))) .unwrap(); let result = responses_websocket( - &services, + Arc::clone(&services), ResponsesWebSocketRequest { model: "authorized".into(), api_key: Some("key".into()), @@ -464,8 +616,59 @@ mod tests { ); let terminals = services.terminals.lock().unwrap(); assert_eq!(terminals.len(), 1); + assert_eq!(terminals[0].classification, TerminalClassification::Success); assert_eq!(terminals[0].usage.total_tokens, 3); - server.await.unwrap(); + server.abort(); + } + + #[tokio::test] + async fn provider_failure_terminals_are_failures_and_are_delivered_once() { + let cases = [ + ( + json!({"type":"response.failed","response":{"error":{"message":"request rejected"}}}), + "ResponseFailed", + "request rejected", + ), + ( + json!({"type":"response.incomplete","response":{"incomplete_details":{"reason":"max_output_tokens"}}}), + "ResponseIncomplete", + "max_output_tokens", + ), + ( + json!({"type":"error","error":{"message":"provider unavailable"}}), + "ProviderError", + "provider unavailable", + ), + ]; + for (frame, kind, message) in cases { + let expected_type = frame["type"].as_str().unwrap().to_string(); + let (api_base, server) = provider_with_frames(vec![frame], true).await; + let services = Arc::new(Services::default()); + let (_client_tx, client_rx) = input(); + let (output_tx, mut output_rx) = futures_channel::mpsc::unbounded(); + let result = responses_websocket( + Arc::clone(&services), + ResponsesWebSocketRequest { + model: "model".into(), + api_key: Some("key".into()), + api_base: Some(api_base), + first_frame: None, + idle_timeout: Some(Duration::from_secs(1)), + }, + CallLifecycleContext::new("responses_websocket", "model", "openai", "call-1"), + client_rx, + output_tx, + ) + .await + .unwrap(); + assert!(matches!(result, ExecutedCall::Failure { .. })); + assert_eq!( + output_rx.next().await.unwrap().event_type.as_str(), + expected_type + ); + assert_failure(services.as_ref(), kind, message); + server.abort(); + } } #[tokio::test] @@ -479,14 +682,14 @@ mod tests { .await .unwrap(); }); - let services = Services::default(); + let services = Arc::new(Services::default()); let (_, input): ( _, - futures_channel::mpsc::UnboundedReceiver, - ) = futures_channel::mpsc::unbounded(); + futures_channel::mpsc::UnboundedReceiver>, + ) = input(); let (output, _) = futures_channel::mpsc::unbounded(); let error = responses_websocket( - &services, + Arc::clone(&services), ResponsesWebSocketRequest { model: "model".into(), api_key: Some("key".into()), @@ -513,11 +716,11 @@ mod tests { let mut socket = accept_async(stream).await.unwrap(); socket.send(Message::Text("not-json".into())).await.unwrap(); }); - let services = Services::default(); - let (client_tx, input) = futures_channel::mpsc::unbounded(); + let services = Arc::new(Services::default()); + let (client_tx, input) = input(); let (output, _) = futures_channel::mpsc::unbounded(); let result = responses_websocket( - &services, + Arc::clone(&services), ResponsesWebSocketRequest { model: "model".into(), api_key: Some("key".into()), @@ -547,6 +750,135 @@ mod tests { )); } + #[tokio::test] + async fn client_drop_provider_close_and_idle_timeout_are_distinct_failures() { + let cases = [ + ( + true, + true, + "ClientDisconnected", + "client disconnected before response.completed", + ), + ( + false, + false, + "ProviderDisconnected", + "provider closed before a terminal response frame", + ), + ( + false, + true, + "IdleTimeout", + "Responses WebSocket session timed out before a terminal response frame", + ), + ]; + for (drop_client, remain_open, kind, message) in cases { + let (api_base, server) = provider_with_frames(Vec::new(), remain_open).await; + let services = Arc::new(Services::default()); + let (client_tx, client_rx) = input(); + if drop_client { + drop(client_tx); + } + let (output_tx, _) = futures_channel::mpsc::unbounded(); + let result = responses_websocket( + Arc::clone(&services), + ResponsesWebSocketRequest { + model: "model".into(), + api_key: Some("key".into()), + api_base: Some(api_base), + first_frame: None, + idle_timeout: Some(Duration::from_millis(20)), + }, + CallLifecycleContext::new("responses_websocket", "model", "openai", "call-1"), + client_rx, + output_tx, + ) + .await + .unwrap(); + assert!(matches!(result, ExecutedCall::Failure { .. })); + assert_failure(services.as_ref(), kind, message); + server.abort(); + } + } + + #[tokio::test] + async fn cancelling_a_committed_session_dispatches_one_failure() { + let (api_base, server) = provider_with_frames(Vec::new(), true).await; + let services = Arc::new(Services::default()); + let (_client_tx, client_rx) = input(); + let (output_tx, _) = futures_channel::mpsc::unbounded(); + let task = tokio::spawn(responses_websocket( + Arc::clone(&services), + ResponsesWebSocketRequest { + model: "model".into(), + api_key: Some("key".into()), + api_base: Some(api_base), + first_frame: None, + idle_timeout: Some(Duration::from_secs(60)), + }, + CallLifecycleContext::new("responses_websocket", "model", "openai", "call-1"), + client_rx, + output_tx, + )); + tokio::time::sleep(Duration::from_millis(20)).await; + task.abort(); + let _ = task.await; + tokio::time::timeout(Duration::from_secs(1), async { + while services.terminals.lock().unwrap().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_failure( + services.as_ref(), + "Cancelled", + "Responses WebSocket session was cancelled before completion", + ); + server.abort(); + } + + #[tokio::test] + async fn client_protocol_error_dispatches_one_failure() { + let (api_base, server) = provider_with_frames(Vec::new(), true).await; + let services = Arc::new(Services::default()); + let (client_tx, client_rx) = input(); + client_tx + .unbounded_send(Err(Error::InvalidRequest( + "invalid client frame".to_string(), + ))) + .unwrap(); + let (output_tx, _) = futures_channel::mpsc::unbounded(); + let result = responses_websocket( + Arc::clone(&services), + ResponsesWebSocketRequest { + model: "model".into(), + api_key: Some("key".into()), + api_base: Some(api_base), + first_frame: None, + idle_timeout: Some(Duration::from_secs(1)), + }, + CallLifecycleContext::new("responses_websocket", "model", "openai", "call-1"), + client_rx, + output_tx, + ) + .await + .unwrap(); + assert!(matches!( + result, + ExecutedCall::Failure { + error: Error::InvalidRequest(_), + .. + } + )); + assert_failure( + services.as_ref(), + "InvalidRequest", + "invalid request: invalid client frame", + ); + server.abort(); + } + #[test] fn url_and_model_behavior_match_the_public_protocol() { assert_eq!( diff --git a/litellm-rust/crates/core/tests/messages_lifecycle.rs b/litellm-rust/crates/core/tests/messages_lifecycle.rs index 1b40307e729..92fc9659416 100644 --- a/litellm-rust/crates/core/tests/messages_lifecycle.rs +++ b/litellm-rust/crates/core/tests/messages_lifecycle.rs @@ -93,7 +93,7 @@ async fn upstream(status: u16) -> (String, tokio::task::JoinHandle<()>) { let mut buffer = [0_u8; 4096]; let _ = socket.read(&mut buffer).await.unwrap(); let body = if status == 200 { - r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":"end_turn","stop_sequence":null}"# + r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":7}}"# } else { r#"{"error":"failed"}"# }; @@ -168,6 +168,9 @@ async fn success_dispatches_exactly_one_terminal() { let terminals = services.terminals.lock().unwrap(); assert_eq!(terminals.len(), 1); assert_eq!(terminals[0].classification, TerminalClassification::Success); + assert_eq!(terminals[0].usage.prompt_tokens, 11); + assert_eq!(terminals[0].usage.completion_tokens, 7); + assert_eq!(terminals[0].usage.total_tokens, 18); } #[tokio::test] diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index df7abfb05f4..e60cdef99e5 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -8,7 +8,7 @@ use litellm_core::Error; use litellm_core::lifecycle::CallLifecycleContext; use litellm_core::ocr::prepare::prepare; use litellm_core::ocr::types::{OcrDocument, OcrDocumentProjection}; -use litellm_core::ocr::{NoopOcrServices, OcrAdmissionRequest as OcrRequest, OcrDraft}; +use litellm_core::ocr::{DefaultOcrServices, OcrAdmissionRequest as OcrRequest, OcrDraft}; use serde_json::{Value, json}; fn request() -> OcrRequest { @@ -49,7 +49,7 @@ async fn ocr( let model = prepared.endpoint.model().to_string(); let provider = prepared.endpoint.custom_llm_provider().to_string(); let response = litellm_core::ocr::ocr( - &NoopOcrServices, + &DefaultOcrServices, prepared.endpoint.settle(headers, body), Default::default(), CallLifecycleContext::new("ocr", model, provider, "test-call"), diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index ec85252ffb4..a8ace3b12d6 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -63,6 +63,21 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { } } +pub(crate) fn messages_provider_error_to_pyerr(err: Error) -> PyErr { + match err { + Error::Http { status, .. } => { + RustUpstreamError::new_err((status, format!("Provider request failed (HTTP {status})"))) + } + Error::Network(_) | Error::Connect(_) => { + RustUpstreamError::new_err((0u16, "Provider transport failed")) + } + Error::InvalidResponse(_) => { + RustUpstreamError::new_err((0u16, "Invalid provider response")) + } + error => core_error_to_pyerr(error), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index a0bc9b79fe9..2a14beb41bd 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,3 +1,5 @@ +#![recursion_limit = "256"] + mod diagnostics; mod driver; mod errors; diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index af60515b0e2..5d9b36fdaf5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -10,6 +10,17 @@ use serde_json::Value; use crate::errors::core_error_to_pyerr; use crate::marshal::{RouteOptions, RouteOptionsInputs, object_or_empty}; +struct AudioTranscriptionInputs { + model: String, + audio: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option, + optional_params: Option, + timeout_seconds: Option, +} + fn prepare_transcription( inputs: AudioTranscriptionInputs, ) -> PyResult> + Send + 'static> { @@ -47,25 +58,141 @@ fn prepare_transcription( }) } -bridge_route! { - sync = transcription, - asynchronous = atranscription, - inputs = AudioTranscriptionInputs, - required = { +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn transcription( + py: Python<'_>, + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] audio: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, + timeout_seconds: Option, +) -> PyResult> { + let future = prepare_transcription(AudioTranscriptionInputs { + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + })?; + litellm_python_interop::run_sync(py, future, core_error_to_pyerr) +} + +#[pyfunction] +#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn atranscription( + py: Python<'_>, + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] audio: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, + timeout_seconds: Option, +) -> PyResult> { + let future = prepare_transcription(AudioTranscriptionInputs { + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + })?; + litellm_python_interop::run_async(py, future, core_error_to_pyerr) +} + +pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::definition::add_function(module, wrap_pyfunction!(transcription, module)?)?; + super::definition::add_function(module, wrap_pyfunction!(atranscription, module)?) +} + +#[cfg(feature = "trace-parity")] +mod trace { + use pyo3::prelude::*; + use serde_json::Value; + + use super::{AudioTranscriptionInputs, core_error_to_pyerr, prepare_transcription}; + + #[pyfunction] + #[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] + #[allow(clippy::too_many_arguments)] + fn transcription( + py: Python<'_>, model: String, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - audio: serde_json::Value, - }, - optional = { + #[pyo3(from_py_with = litellm_python_interop::from_py)] audio: Value, api_key: Option, api_base: Option, custom_llm_provider: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - extra_headers: Option, - #[pyo3(from_py_with = litellm_python_interop::from_py)] - optional_params: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, timeout_seconds: Option, - }, - prepare = prepare_transcription, - errors = core_error_to_pyerr, + ) -> PyResult> { + let future = prepare_transcription(AudioTranscriptionInputs { + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + })?; + litellm_python_interop::run_sync( + py, + crate::function_trace::capture(future), + core_error_to_pyerr, + ) + } + + #[pyfunction] + #[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] + #[allow(clippy::too_many_arguments)] + fn atranscription( + py: Python<'_>, + model: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] audio: Value, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, + timeout_seconds: Option, + ) -> PyResult> { + let future = prepare_transcription(AudioTranscriptionInputs { + model, + audio, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + })?; + litellm_python_interop::run_async( + py, + crate::function_trace::capture(future), + core_error_to_pyerr, + ) + } + + pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::super::definition::add_function(module, wrap_pyfunction!(transcription, module)?)?; + super::super::definition::add_function(module, wrap_pyfunction!(atranscription, module)?) + } +} + +#[cfg(feature = "trace-parity")] +pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + trace::register(module) } diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs index 1867d407439..6478d3f553a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions.rs @@ -158,6 +158,11 @@ fn invoke( Operation::Setup => ("setup", false), Operation::DeploymentPre => ("deployment_pre", true), Operation::Prepare => ("prepare", false), + Operation::PreCall => { + return Err(PyRuntimeError::new_err( + "chat completions lifecycle selected an unsupported pre-call operation", + )); + } Operation::Send if asynchronous => ("send", true), Operation::Send => ("send_sync", false), Operation::DeploymentSuccess => ("deployment_success", true), diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index e7fe7214e9e..4fd409b41e3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -2,129 +2,6 @@ use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyCFunction; -macro_rules! bridge_route { - ( - sync = $sync_name:ident, - asynchronous = $async_name:ident, - inputs = $inputs:ident, - required = { $($(#[$required_attr:meta])* $required_name:ident: $required_type:ty),+ $(,)? }, - optional = { $($(#[$optional_attr:meta])* $optional_name:ident: $optional_type:ty),* $(,)? }, - prepare = $prepare:path, - errors = $map_error:path - $(, extra = [$($extra:ident),* $(,)?])? - $(,)? - ) => { - struct $inputs { - $($required_name: $required_type,)* - $($optional_name: $optional_type),* - } - - #[pyfunction] - #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] - #[allow(clippy::too_many_arguments)] - fn $sync_name( - py: pyo3::Python<'_>, - $($(#[$required_attr])* $required_name: $required_type,)* - $($(#[$optional_attr])* $optional_name: $optional_type,)* - ) -> pyo3::PyResult> { - let future = $prepare($inputs { - $($required_name,)* - $($optional_name),* - })?; - litellm_python_interop::run_sync(py, future, $map_error) - } - - #[pyfunction] - #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] - #[allow(clippy::too_many_arguments)] - fn $async_name( - py: pyo3::Python<'_>, - $($(#[$required_attr])* $required_name: $required_type,)* - $($(#[$optional_attr])* $optional_name: $optional_type,)* - ) -> pyo3::PyResult> { - let future = $prepare($inputs { - $($required_name,)* - $($optional_name),* - })?; - litellm_python_interop::run_async(py, future, $map_error) - } - - pub(super) fn register( - module: &pyo3::Bound<'_, pyo3::types::PyModule>, - ) -> pyo3::PyResult<()> { - $($($crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($extra, module)?)?;)*)? - $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($sync_name, module)?)?; - $crate::routes::definition::add_function(module, pyo3::wrap_pyfunction!($async_name, module)?)?; - Ok(()) - } - - #[cfg(feature = "trace-parity")] - mod trace { - use pyo3::prelude::*; - use super::{$inputs, $map_error, $prepare}; - - #[pyfunction] - #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] - #[allow(clippy::too_many_arguments)] - fn $sync_name( - py: pyo3::Python<'_>, - $($(#[$required_attr])* $required_name: $required_type,)* - $($(#[$optional_attr])* $optional_name: $optional_type,)* - ) -> pyo3::PyResult> { - let future = $prepare($inputs { - $($required_name,)* - $($optional_name),* - })?; - litellm_python_interop::run_sync( - py, - $crate::function_trace::capture(future), - $map_error, - ) - } - - #[pyfunction] - #[pyo3(signature = ($($required_name),*, $($optional_name=None),*))] - #[allow(clippy::too_many_arguments)] - fn $async_name( - py: pyo3::Python<'_>, - $($(#[$required_attr])* $required_name: $required_type,)* - $($(#[$optional_attr])* $optional_name: $optional_type,)* - ) -> pyo3::PyResult> { - let future = $prepare($inputs { - $($required_name,)* - $($optional_name),* - })?; - litellm_python_interop::run_async( - py, - $crate::function_trace::capture(future), - $map_error, - ) - } - - pub(super) fn register( - module: &pyo3::Bound<'_, pyo3::types::PyModule>, - ) -> pyo3::PyResult<()> { - $crate::routes::definition::add_function( - module, - pyo3::wrap_pyfunction!($sync_name, module)?, - )?; - $crate::routes::definition::add_function( - module, - pyo3::wrap_pyfunction!($async_name, module)?, - )?; - Ok(()) - } - } - - #[cfg(feature = "trace-parity")] - pub(super) fn register_trace( - module: &pyo3::Bound<'_, pyo3::types::PyModule>, - ) -> pyo3::PyResult<()> { - trace::register(module) - } - }; -} - pub(super) fn add_function( module: &Bound<'_, PyModule>, function: Bound<'_, PyCFunction>, @@ -169,15 +46,63 @@ mod tests { FUTURE_DROPPED.load(Ordering::SeqCst) } - bridge_route! { - sync = echo, - asynchronous = aecho, - inputs = EchoInputs, - required = { value: String }, - optional = {}, - prepare = prepare_echo, - errors = map_error, - extra = [future_dropped], + struct EchoInputs { + value: String, + } + + #[pyfunction] + fn echo(py: Python<'_>, value: String) -> PyResult> { + let future = prepare_echo(EchoInputs { value })?; + litellm_python_interop::run_sync(py, future, map_error) + } + + #[pyfunction] + fn aecho(py: Python<'_>, value: String) -> PyResult> { + let future = prepare_echo(EchoInputs { value })?; + litellm_python_interop::run_async(py, future, map_error) + } + + pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::add_function(module, wrap_pyfunction!(future_dropped, module)?)?; + super::add_function(module, wrap_pyfunction!(echo, module)?)?; + super::add_function(module, wrap_pyfunction!(aecho, module)?) + } + + #[cfg(feature = "trace-parity")] + mod trace { + use pyo3::prelude::*; + + use super::{EchoInputs, map_error, prepare_echo}; + + #[pyfunction] + fn echo(py: Python<'_>, value: String) -> PyResult> { + let future = prepare_echo(EchoInputs { value })?; + litellm_python_interop::run_sync( + py, + crate::function_trace::capture(future), + map_error, + ) + } + + #[pyfunction] + fn aecho(py: Python<'_>, value: String) -> PyResult> { + let future = prepare_echo(EchoInputs { value })?; + litellm_python_interop::run_async( + py, + crate::function_trace::capture(future), + map_error, + ) + } + + pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::super::add_function(module, wrap_pyfunction!(echo, module)?)?; + super::super::add_function(module, wrap_pyfunction!(aecho, module)?) + } + } + + #[cfg(feature = "trace-parity")] + pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { + trace::register(module) } fn prepare_echo( @@ -229,11 +154,7 @@ mod tests { "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", ), ("messages", "amessages", "(arguments)"), - ( - "chat_completions", - "achat_completions", - "(model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)", - ), + ("chat_completions", "achat_completions", "(arguments)"), ]; for (sync_name, async_name, expected) in routes { @@ -262,13 +183,20 @@ mod tests { crate::routes::register(&module).expect("routes should register"); let invalid_messages = PyDict::new(py); + let invalid_chat_arguments = PyDict::new(py); + invalid_chat_arguments + .set_item("model", "model") + .expect("arguments should accept model"); + invalid_chat_arguments + .set_item("messages", &invalid_messages) + .expect("arguments should accept messages"); let sync_chat_error = module .getattr("chat_completions") - .and_then(|function| function.call1(("model", &invalid_messages))) + .and_then(|function| function.call1((&invalid_chat_arguments,))) .expect_err("sync chat should reject a non-list messages value"); let async_chat_error = module .getattr("achat_completions") - .and_then(|function| function.call1(("model", &invalid_messages))) + .and_then(|function| function.call1((&invalid_chat_arguments,))) .expect_err("async chat should reject a non-list messages value"); assert_eq!( @@ -338,32 +266,26 @@ mod tests { crate::routes::register(&module).expect("routes should register"); let invalid = PyList::empty(py); - let chat_kwargs = PyDict::new(py); - chat_kwargs + let chat_arguments = PyDict::new(py); + chat_arguments + .set_item("model", "model") + .expect("arguments should accept model"); + chat_arguments .set_item("optional_params", &invalid) - .expect("kwargs should accept optional_params"); - chat_kwargs + .expect("arguments should accept optional_params"); + chat_arguments .set_item("extra_headers", &invalid) - .expect("kwargs should accept extra_headers"); + .expect("arguments should accept extra_headers"); let invalid_messages = PyDict::new(py); + chat_arguments + .set_item("messages", &invalid_messages) + .expect("arguments should accept messages"); let error = module .getattr("chat_completions") - .and_then(|function| { - function.call(("model", &invalid_messages), Some(&chat_kwargs)) - }) + .and_then(|function| function.call1((&chat_arguments,))) .expect_err("messages should be validated first"); assert_eq!(error.to_string(), "TypeError: messages must be a list"); - let valid_messages = PyList::empty(py); - let error = module - .getattr("chat_completions") - .and_then(|function| function.call(("model", &valid_messages), Some(&chat_kwargs))) - .expect_err("optional_params should be validated before headers"); - assert_eq!( - error.to_string(), - "TypeError: optional_params must be a dict" - ); - let headers_kwargs = PyDict::new(py); headers_kwargs .set_item("extra_headers", &invalid) diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs index df663885e6e..b85057e9d75 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages.rs @@ -1,7 +1,9 @@ +use litellm_core::lifecycle::FailureStage; use litellm_core::lifecycle::{ErrorDisposition, Lifecycle, Outcome}; use litellm_core::messages::lifecycle::{MessagesRoute, Observations, Operation, Options, machine}; -use litellm_core::messages::types::MessagesRequest; -use litellm_python_interop::{Pythonized, from_py, run_async_value, run_sync_value}; +use litellm_core::messages::types::{MessagesRequest, ProviderMessagesRequest}; +use litellm_core::messages::{execute_prepared_messages_provider_call, prepare_provider_request}; +use litellm_python_interop::{Pythonized, from_py, run_async_value, run_sync_value, to_py}; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; use pyo3::pyclass::{PyTraverseError, PyVisit}; @@ -9,25 +11,34 @@ use pyo3::sync::PyOnceLock; use pyo3::types::PyDict; use serde_json::{Map, Value}; -use crate::errors::core_error_to_pyerr; +use crate::errors::{RustUpstreamError, core_error_to_pyerr, messages_provider_error_to_pyerr}; use crate::marshal::optional_timeout; #[pyclass] struct MessagesState { arguments: Option>, - request: Option, + body: Option>, + headers: Option>, + prepared: Option, } #[pymethods] impl MessagesState { fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.arguments) + visit.call(&self.arguments)?; + visit.call(&self.body)?; + visit.call(&self.headers) } fn __clear__(slf: &Bound<'_, Self>) { let roots = { let mut state = slf.borrow_mut(); - (state.arguments.take(), state.request.take()) + ( + state.arguments.take(), + state.body.take(), + state.headers.take(), + state.prepared.take(), + ) }; drop(roots); } @@ -41,33 +52,30 @@ fn scalar(arguments: &Bound<'_, PyDict>, name: &str) -> PyResult> .transpose() } -fn decode_state(py: Python<'_>, arguments: Py) -> PyResult { - let bag = arguments.bind(py); - let body = bag +fn decode_request(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + let body = arguments .get_item(pyo3::intern!(py, "body"))? .ok_or_else(|| PyValueError::new_err("messages requires body"))?; let timeout = optional_timeout( - bag.get_item(pyo3::intern!(py, "timeout_seconds"))? + arguments + .get_item(pyo3::intern!(py, "timeout_seconds"))? .filter(|value| !value.is_none()) .map(|value| value.extract::()) .transpose()?, )?; - Ok(MessagesState { - request: Some(MessagesRequest { - body: from_py(&body)?, - model: scalar(bag, "model")? - .ok_or_else(|| PyValueError::new_err("messages requires model"))?, - api_key: scalar(bag, "api_key")?, - api_base: scalar(bag, "api_base")?, - custom_llm_provider: scalar(bag, "custom_llm_provider")?, - extra_headers: bag - .get_item("extra_headers")? - .filter(|value| !value.is_none()) - .map(|value| from_py::>(&value)) - .transpose()?, - timeout, - }), - arguments: Some(arguments), + Ok(MessagesRequest { + body: from_py(&body)?, + model: scalar(arguments, "model")? + .ok_or_else(|| PyValueError::new_err("messages requires model"))?, + api_key: scalar(arguments, "api_key")?, + api_base: scalar(arguments, "api_base")?, + custom_llm_provider: scalar(arguments, "custom_llm_provider")?, + extra_headers: arguments + .get_item("extra_headers")? + .filter(|value| !value.is_none()) + .map(|value| from_py::>(&value)) + .transpose()?, + timeout, }) } @@ -121,6 +129,10 @@ impl MessagesLifecycle { _ => None, } } + + fn failed_after_provider_response(&self) -> bool { + self.machine.failure_stage() == Some(FailureStage::AfterProviderResponse) + } } #[pyfunction] @@ -137,6 +149,11 @@ fn invoke( Operation::Setup => ("setup", false), Operation::DeploymentPre => ("deployment_pre", true), Operation::Prepare => ("prepare", false), + Operation::PreCall => { + return Err(PyRuntimeError::new_err( + "messages lifecycle selected an unsupported pre-call operation", + )); + } Operation::Send if asynchronous => ("send", true), Operation::Send => ("send_sync", false), Operation::DeploymentSuccess => ("deployment_success", true), @@ -156,46 +173,78 @@ fn invoke( #[pyfunction] fn prepare(py: Python<'_>, arguments: Py) -> PyResult> { - let state = decode_state(py, arguments)?; - let bag = state.arguments.as_ref().unwrap().bind(py); + let bag = arguments.bind(py); + let request = decode_request(py, bag)?; + let prepared = py + .detach(|| prepare_provider_request(request)) + .map_err(core_error_to_pyerr)?; + let body = to_py(py, &prepared.body)? + .into_bound(py) + .cast_into::()?; + let headers = PyDict::new(py); + for (name, value) in &prepared.upstream_headers { + headers.set_item(name, value)?; + } let logging = bag .get_item("litellm_logging_obj")? .filter(|value| !value.is_none()) .map(|value| value.unbind()) .ok_or_else(|| PyRuntimeError::new_err("messages logging was not initialized"))?; let additional = PyDict::new(py); - additional.set_item( - pyo3::intern!(py, "complete_input_dict"), - bag.get_item(pyo3::intern!(py, "body"))?, - )?; - additional.set_item( - pyo3::intern!(py, "api_base"), - bag.get_item(pyo3::intern!(py, "api_base"))?, - )?; - additional.set_item( - pyo3::intern!(py, "headers"), - bag.get_item(pyo3::intern!(py, "extra_headers"))?, - )?; + additional.set_item(pyo3::intern!(py, "complete_input_dict"), &body)?; + additional.set_item(pyo3::intern!(py, "api_base"), &prepared.url)?; + additional.set_item(pyo3::intern!(py, "headers"), &headers)?; let kwargs = PyDict::new(py); - kwargs.set_item( - "input", - bag.get_item("messages")? - .unwrap_or_else(|| py.None().into_bound(py)), - )?; + let serialized = py.import("json")?.call_method1("dumps", (&body,))?; + let message = PyDict::new(py); + message.set_item("role", "user")?; + message.set_item("content", serialized)?; + let messages = vec![message]; + logging + .bind(py) + .call_method1("update_messages", (&messages,))?; + kwargs.set_item("input", messages)?; kwargs.set_item("api_key", "")?; kwargs.set_item("additional_args", additional)?; logging .bind(py) .call_method(pyo3::intern!(py, "pre_call"), (), Some(&kwargs))?; - Py::new(py, state) + Py::new( + py, + MessagesState { + arguments: Some(arguments), + body: Some(body.unbind()), + headers: Some(headers.unbind()), + prepared: Some(prepared), + }, + ) } -fn take_request(py: Python<'_>, state: &Py) -> PyResult { - state - .borrow_mut(py) - .request - .take() - .ok_or_else(|| PyRuntimeError::new_err("messages request was already sent or cleared")) +fn take_request(py: Python<'_>, state: &Py) -> PyResult { + let (mut prepared, body, headers) = { + let mut state = state.borrow_mut(py); + let prepared = state.prepared.take().ok_or_else(|| { + PyRuntimeError::new_err("messages request was already sent or cleared") + })?; + let body = state + .body + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("messages body was cleared"))? + .clone_ref(py); + let headers = state + .headers + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("messages headers were cleared"))? + .clone_ref(py); + (prepared, body, headers) + }; + prepared.body = from_py(body.bind(py).as_any())?; + prepared.upstream_headers = headers + .bind(py) + .iter() + .map(|(name, value)| Ok((name.extract()?, value.extract()?))) + .collect::>()?; + Ok(prepared) } fn validate_arguments(arguments: &Bound<'_, PyDict>) -> PyResult<()> { @@ -221,8 +270,8 @@ fn send(py: Python<'_>, state: Py) -> PyResult> litellm_python_interop::run_async_py(py, async move { let _state = state; let response = run_async_value( - litellm_core::messages::messages(request), - core_error_to_pyerr, + execute_prepared_messages_provider_call(request), + messages_provider_error_to_pyerr, ) .await?; Ok(Pythonized(response)) @@ -234,12 +283,20 @@ fn send_sync(py: Python<'_>, state: Py) -> PyResult> { let request = take_request(py, &state)?; let response = run_sync_value( py, - litellm_core::messages::messages(request), - core_error_to_pyerr, + execute_prepared_messages_provider_call(request), + messages_provider_error_to_pyerr, )?; Ok(Pythonized(response).into_pyobject(py)?.unbind().into_any()) } +#[pyfunction] +fn committed_failure() -> PyResult<()> { + Err(RustUpstreamError::new_err(( + 0u16, + "Messages lifecycle failed after the provider returned", + ))) +} + #[pyfunction] fn messages(py: Python<'_>, arguments: Py) -> PyResult> { validate_arguments(arguments.bind(py))?; @@ -263,6 +320,10 @@ fn driver(py: Python<'_>) -> PyResult<&Bound<'_, PyModule>> { module.add("_prepare", wrap_pyfunction!(prepare, &module)?)?; module.add("_send", wrap_pyfunction!(send, &module)?)?; module.add("_send_sync", wrap_pyfunction!(send_sync, &module)?)?; + module.add( + "_committed_failure", + wrap_pyfunction!(committed_failure, &module)?, + )?; Ok(DRIVER.get_or_init(py, || module.unbind()).bind(py)) } @@ -270,7 +331,7 @@ const HOST: &str = r#" from datetime import datetime from litellm import utils from litellm.types.utils import CallTypes -from litellm.rust_bridge.messages import initialize_logging, invoke_terminal +from litellm.rust_bridge.messages import initialize_logging, invoke_terminal, retain_stream_response class Host: def __init__(self, arguments, asynchronous): @@ -284,10 +345,12 @@ class Host: self.error = None self.start = datetime.now() self.end = None + self.streaming = False def setup(self): self.logger = initialize_logging(self.arguments, self.asynchronous) self.arguments['litellm_logging_obj'] = self.logger + self.streaming = self.logger.stream is True async def deployment_pre(self): modified = await utils.async_pre_call_deployment_hook(self.current, 'amessages') @@ -308,11 +371,20 @@ class Host: async def deployment_success(self): self.response = await utils.async_post_call_success_deployment_hook(self.current, self.response, CallTypes.aanthropic_messages) + if self.streaming: + self.response = retain_stream_response( + self.response, + (self.arguments, self.current, self.state), + self.logger, + self.start, + ) async def deployment_failure(self): await utils.async_post_call_failure_deployment_hook(self.current, self.error, 'amessages') def terminal(self, action, value): + if self.streaming: + return None return invoke_terminal(action, (self.arguments, self.current, self.state), self.logger, None, value, self.start, self.end) def sync_success(self): return self.terminal('sync_success', self.response) @@ -320,7 +392,9 @@ class Host: def sync_success_if_needed(self): return self.terminal('sync_success_if_needed', self.response) def sync_failure(self): return self.terminal('sync_failure', self.error) def async_failure(self): return self.terminal('async_failure', self.error) - def restore(self): utils._restore_correlation_context_if_supported(self.logger) + def restore(self): + if not self.streaming: + utils._restore_correlation_context_if_supported(self.logger) def advance(self, outcome, error=None): if error is not None and self.end is None: @@ -334,6 +408,8 @@ class Host: def result(self): if self.machine.complete(): return self.response + if self.machine.failed_after_provider_response(): + _committed_failure() raise self.error "#; @@ -361,7 +437,7 @@ mod tests { arguments.set_item("model", "model").unwrap(); arguments.set_item("body", PyDict::new(py)).unwrap(); arguments.set_item("timeout_seconds", timeout).unwrap(); - let error = match decode_state(py, arguments.unbind()) { + let error = match decode_request(py, &arguments) { Ok(_) => panic!("invalid timeout should fail normally"), Err(error) => error, }; diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 77eaa2bda91..9b6e0c8d159 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -1,6 +1,5 @@ use pyo3::prelude::*; -#[macro_use] mod definition; #[cfg(feature = "trace-parity")] diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index d08ac23a289..66e302e8424 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -7,7 +7,7 @@ use litellm_core::lifecycle::ocr::{NativeOutcome, Observations, OcrRoute, Operat use litellm_core::lifecycle::{ CallLifecycleContext, ErrorDisposition, ExecutedCall, Lifecycle, Outcome, TerminalRecord, }; -use litellm_core::ocr::NoopOcrServices; +use litellm_core::ocr::DefaultOcrServices; use litellm_core::ocr::types::{ OcrAdmissionRequest, OcrDocumentProjection, OcrDraft, OcrEndpoint, SettledOcrRequest, }; @@ -490,7 +490,7 @@ fn send(py: Python<'_>, state: Py) -> PyResult> { }); let executed = run_async_value( async move { - let services = NoopOcrServices; + let services = DefaultOcrServices; Ok::<_, std::convert::Infallible>( litellm_core::ocr::ocr( &services, @@ -552,7 +552,7 @@ fn send_sync(py: Python<'_>, state: Py) -> PyResult> { let executed = run_sync_value( py, async move { - let services = NoopOcrServices; + let services = DefaultOcrServices; Ok::<_, std::convert::Infallible>( litellm_core::ocr::ocr( &services, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py index d0fed3225af..0c3b686b7eb 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py @@ -9,7 +9,8 @@ the LLM doesn't make a tool call, and we need to return a stream to the user. """ import json -from collections.abc import Mapping +from collections.abc import Callable, Mapping +from types import TracebackType from typing import Any, Final, cast from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -34,10 +35,11 @@ class FakeAnthropicMessagesStreamIterator: - message_stop """ - def __init__(self, response: AnthropicMessagesResponse): + def __init__(self, response: AnthropicMessagesResponse, on_complete: Callable[[], None] | None = None): self.response = response self.chunks = self._create_streaming_chunks() self.current_index = 0 + self.on_complete = on_complete def _create_content_block_chunks(self, block_dict: Mapping[str, object], index: int) -> list[bytes]: """Build SSE chunks for a single content block.""" @@ -196,6 +198,7 @@ class FakeAnthropicMessagesStreamIterator: async def __anext__(self): if self.current_index >= len(self.chunks): + self.close() raise StopAsyncIteration chunk: Final = self.chunks[self.current_index] @@ -207,8 +210,46 @@ class FakeAnthropicMessagesStreamIterator: def __next__(self): if self.current_index >= len(self.chunks): + self.close() raise StopIteration chunk: Final = self.chunks[self.current_index] self.current_index += 1 return chunk + + def close(self) -> None: + if self.on_complete is None: + return + on_complete, self.on_complete = self.on_complete, None + on_complete() + + async def aclose(self) -> None: + self.close() + + def __enter__(self) -> "FakeAnthropicMessagesStreamIterator": + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.close() + + async def __aenter__(self) -> "FakeAnthropicMessagesStreamIterator": + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + await self.aclose() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 7d01aee5d98..859bea0b6b5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -2,6 +2,7 @@ import asyncio import json from collections.abc import AsyncIterator, Mapping, Sequence from datetime import datetime +from types import TracebackType from typing import Any, Final, Protocol, runtime_checkable import httpx @@ -373,6 +374,17 @@ class AnthropicMessagesStreamingResponse: async def aclose(self) -> None: await aclose_if_supported(self.completion_stream) + async def __aenter__(self) -> "AnthropicMessagesStreamingResponse": + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + await self.aclose() + class BaseAnthropicMessagesStreamingIterator: """ diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3844be6fe36..0dda5c7c6ba 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2455,6 +2455,15 @@ class BaseLLMHTTPHandler: timeout=timeout, ) except Exception as rust_error: # noqa: BLE001 # rollout-safety fallback: any Rust bridge failure must fall back to the Python path + from litellm.rust_bridge.bindings import native_exception_types + from litellm.rust_bridge.runtime import BridgeErrorContext, _raise_upstream + + exception_types: Final = native_exception_types() + if exception_types is not None and isinstance(rust_error, exception_types[1]): + _raise_upstream( + rust_error, + BridgeErrorContext(route="messages", provider=custom_llm_provider, model=model), + ) verbose_logger.debug( "Rust Anthropic messages bridge raised %s; falling back to Python path", type(rust_error).__name__, @@ -2463,7 +2472,7 @@ class BaseLLMHTTPHandler: if rust_response is None: return None - response_obj: Final = cast(AnthropicMessagesResponse, dict(rust_response)) + response_obj: Final = cast(AnthropicMessagesResponse, rust_response) response_obj["_hidden_params"] = {"additional_headers": {"x-litellm-rust": "true"}} return response_obj @@ -2479,7 +2488,13 @@ class BaseLLMHTTPHandler: AnthropicMessagesStreamingResponse, ) - completion_stream = cast(AsyncIterator[bytes], FakeAnthropicMessagesStreamIterator(response=rust_response)) + completion_stream = cast( + AsyncIterator[bytes], + FakeAnthropicMessagesStreamIterator( + response=rust_response, + on_complete=getattr(rust_response, "complete", None), + ), + ) hidden_params: Final = AnthropicMessagesStreamHiddenParams(additional_headers={"x-litellm-rust": "true"}) return AnthropicMessagesStreamingResponse( completion_stream=completion_stream, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index c94236f5c23..364f9e3ca69 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -240,7 +240,31 @@ async def aocr( ) ``` """ - return await _legacy_aocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) + from litellm.rust_bridge import ocr as rust_ocr_bridge + from litellm.rust_bridge.configuration import rust_enabled + + if not rust_enabled() or rust_ocr_bridge.load_rust_aocr() is None: + return await _legacy_aocr( + model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs + ) + arguments: Final[dict[str, object]] = { + **kwargs, + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + } + try: + return await rust_ocr_bridge.aocr(arguments) + except NotImplementedError: + if "litellm_logging_obj" in arguments: + raise + return await _legacy_aocr( + model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs + ) async def _legacy_aocr( @@ -510,7 +534,27 @@ def ocr( print(f"Page {page.index}: {page.markdown}") ``` """ - return _legacy_ocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) + from litellm.rust_bridge import ocr as rust_ocr_bridge + from litellm.rust_bridge.configuration import rust_enabled + + if not rust_enabled() or kwargs.get("aocr") is True or rust_ocr_bridge.load_rust_ocr() is None: + return _legacy_ocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) + arguments: Final[dict[str, object]] = { + **kwargs, + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + } + try: + return rust_ocr_bridge.ocr(arguments) + except NotImplementedError: + if "litellm_logging_obj" in arguments: + raise + return _legacy_ocr(model, document, api_key, api_base, timeout, custom_llm_provider, extra_headers, **kwargs) def _legacy_ocr( diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 80409688789..277adeb7222 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -2,8 +2,11 @@ from __future__ import annotations from collections.abc import Awaitable from dataclasses import dataclass +from datetime import datetime from typing import Final, Protocol, cast +import httpx + from litellm.rust_bridge._lifecycle import ( initialize_logging as initialize_lifecycle_logging, ) @@ -19,6 +22,12 @@ class RustAmessages(Protocol): def __call__(self, arguments: dict[str, object]) -> Awaitable[dict[str, object]]: ... +class _MessagesLogging(Protocol): + model_call_details: dict[str, object] + + def _handle_anthropic_messages_response_logging(self, result: object) -> object: ... + + class _Unset: pass @@ -68,6 +77,57 @@ def initialize_logging(arguments: dict[str, object], asynchronous: bool) -> obje return initialize_lifecycle_logging(arguments, asynchronous, "messages") +class _RetainedMessagesResponse(dict[str, object]): + def __init__(self, response: dict[str, object], roots: object, logger: object, start_time: datetime) -> None: + super().__init__(response) + self._roots = roots + self._logger = cast(_MessagesLogging, logger) + self._start_time = start_time + self._completed = False + + def complete(self) -> None: + if self._completed: + return + self._completed = True + roots, self._roots = self._roots, None + try: + complete_response = self._logger._handle_anthropic_messages_response_logging( # pyright: ignore[reportPrivateUsage] # existing Messages logging transform + self + ) + self._logger.model_call_details["complete_streaming_response"] = complete_response + end_time = datetime.now() + try: + invoke_terminal( + "async_success", + roots, + self._logger, + None, + complete_response, + self._start_time, + end_time, + ) + finally: + invoke_terminal( + "sync_success_if_needed", + roots, + self._logger, + None, + complete_response, + self._start_time, + end_time, + ) + finally: + from litellm import utils + + utils._restore_correlation_context_if_supported(self._logger) # pyright: ignore[reportPrivateUsage] # lifecycle cleanup has no public wrapper + + +def retain_stream_response( + response: dict[str, object], roots: object, logger: object, start_time: datetime +) -> dict[str, object]: + return _RetainedMessagesResponse(response, roots, logger, start_time) + + def _arguments( arguments: dict[str, object], model: str, @@ -76,7 +136,7 @@ def _arguments( api_base: str | None, custom_llm_provider: str | None, extra_headers: dict[str, object] | None, - timeout: object, + timeout: float | httpx.Timeout | None, ) -> dict[str, object]: return { **arguments, @@ -98,14 +158,16 @@ def messages( api_base: str | None, custom_llm_provider: str | None, extra_headers: dict[str, object] | None, - timeout: object, + timeout: float | httpx.Timeout | None, arguments: dict[str, object] | None = None, ) -> dict[str, object] | None: implementation: Final = load_rust_messages() if implementation is None: return None return implementation( - arguments=_arguments(arguments or {}, model, body, api_key, api_base, custom_llm_provider, extra_headers, timeout) + arguments=_arguments( + arguments or {}, model, body, api_key, api_base, custom_llm_provider, extra_headers, timeout + ) ) @@ -117,15 +179,26 @@ async def amessages( api_base: str | None, custom_llm_provider: str | None, extra_headers: dict[str, object] | None, - timeout: object, + timeout: float | httpx.Timeout | None, arguments: dict[str, object] | None = None, ) -> dict[str, object] | None: implementation: Final = load_rust_amessages() if implementation is None: return None return await implementation( - arguments=_arguments(arguments or {}, model, body, api_key, api_base, custom_llm_provider, extra_headers, timeout) + arguments=_arguments( + arguments or {}, model, body, api_key, api_base, custom_llm_provider, extra_headers, timeout + ) ) -__all__ = ["amessages", "initialize_logging", "invoke_terminal", "load_rust_amessages", "load_rust_messages", "messages", "set_rust_messages"] +__all__ = [ + "amessages", + "initialize_logging", + "invoke_terminal", + "load_rust_amessages", + "load_rust_messages", + "messages", + "retain_stream_response", + "set_rust_messages", +] diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 6f3839750ed..228827ad4fb 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -1,6 +1,9 @@ """Tests for the optional Rust-backed Anthropic Messages path.""" +import asyncio import importlib +import weakref +from datetime import datetime from typing import cast import httpx @@ -76,6 +79,19 @@ class RaisingAsyncMessages: raise RuntimeError("upstream request failed with status 400: bad request") +class _CommittedMessagesError(Exception): + pass + + +class _DeclinedMessagesError(Exception): + pass + + +class _NativeExceptions: + RustBridgeDeclined = _DeclinedMessagesError + RustUpstreamError = _CommittedMessagesError + + @pytest.fixture(autouse=True) def _reset_rust_flag(): rust_messages.set_rust_messages(messages=None, amessages=None) @@ -217,6 +233,26 @@ async def test_gate_falls_back_to_python_when_bridge_raises(): assert bridge.calls == 1 +@pytest.mark.asyncio +async def test_gate_does_not_fall_back_after_provider_commit(monkeypatch): + calls = 0 + + async def committed(**kwargs: object) -> dict[str, object]: + nonlocal calls + calls += 1 + raise _CommittedMessagesError(429, "rate limited") + + monkeypatch.setattr("litellm.rust_bridge.bindings.get_native_bridge", lambda: _NativeExceptions()) + rust_messages.set_rust_messages(amessages=committed) + litellm.rust(True) + + with pytest.raises(litellm.APIError) as raised: + await _gate() + + assert raised.value.status_code == 429 + assert calls == 1 + + @pytest.mark.asyncio async def test_gate_skips_rust_when_flag_absent(): bridge = ExplodingAsyncMessages() @@ -348,6 +384,111 @@ async def test_fake_stream_wraps_rust_response_as_anthropic_sse(): assert b"event: message_stop" in joined +def test_fake_stream_completes_retained_response_after_sync_exhaustion(): + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + + completed: list[bool] = [] + stream = FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)), + on_complete=lambda: completed.append(True), + ) + + assert list(stream) + assert completed == [True] + assert list(stream) == [] + assert completed == [True] + + +@pytest.mark.asyncio +async def test_fake_stream_async_exhaustion_and_close_settle_exactly_once(): + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + + completed: list[bool] = [] + stream = FakeAnthropicMessagesStreamIterator( + response=cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)), + on_complete=lambda: completed.append(True), + ) + + assert [chunk async for chunk in stream] + await stream.aclose() + stream.close() + assert completed == [True] + + +@pytest.mark.asyncio +async def test_fake_stream_context_closes_after_early_cancellation(): + from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( + FakeAnthropicMessagesStreamIterator, + ) + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamHiddenParams, + AnthropicMessagesStreamingResponse, + ) + + completed: list[bool] = [] + completion_stream = FakeAnthropicMessagesStreamIterator( + cast(AnthropicMessagesResponse, dict(FAKE_MESSAGES_RESPONSE)), + on_complete=lambda: completed.append(True), + ) + stream = AnthropicMessagesStreamingResponse( + completion_stream, + AnthropicMessagesStreamHiddenParams(additional_headers={}), + ) + + with pytest.raises(asyncio.CancelledError): + async with stream: + await anext(stream) + raise asyncio.CancelledError + + await stream.aclose() + assert completed == [True] + + +def test_retained_response_releases_roots_and_restores_when_terminal_fails(monkeypatch): + class Root: + pass + + class Logger: + def __init__(self) -> None: + self.model_call_details: dict[str, object] = {} + + def _handle_anthropic_messages_response_logging(self, result: object) -> object: + return result + + root = Root() + root_ref = weakref.ref(root) + restored: list[object] = [] + + terminal_calls: list[bool] = [] + + def terminal(*args: object) -> None: + terminal_calls.append(True) + if len(terminal_calls) == 1: + raise RuntimeError("terminal callback failed") + + monkeypatch.setattr(rust_messages, "invoke_terminal", terminal) + monkeypatch.setattr( + litellm.utils, + "_restore_correlation_context_if_supported", + lambda logger: restored.append(logger), + ) + logger = Logger() + response = rust_messages.retain_stream_response(dict(FAKE_MESSAGES_RESPONSE), root, logger, datetime.now()) + del root + + with pytest.raises(RuntimeError, match="terminal callback failed"): + response.complete() + response.complete() + + assert root_ref() is None + assert terminal_calls == [True, True] + assert restored == [logger] + + @pytest.mark.asyncio async def test_gate_falls_back_when_bridge_unavailable(monkeypatch): monkeypatch.setattr( diff --git a/tests/test_litellm_rust/test_sdk_dispatch.py b/tests/test_litellm_rust/test_sdk_dispatch.py index 8b0b79b1473..0fc626c3279 100644 --- a/tests/test_litellm_rust/test_sdk_dispatch.py +++ b/tests/test_litellm_rust/test_sdk_dispatch.py @@ -15,10 +15,6 @@ def ocr_server(recording_server: RecordingServer) -> RecordingServer: return recording_server -@pytest.mark.xfail( - strict=True, - reason="public litellm.ocr does not dispatch to the native OCR bridge yet", -) def test_public_ocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None: response: Final = litellm.ocr( model=OCR_MODEL, @@ -29,3 +25,42 @@ def test_public_ocr_entrypoint_uses_native_transport_when_enabled(ocr_server: Re assert response.pages[0].markdown == "native OCR response" assert ocr_server.requests[0].headers["accept-encoding"] == "identity" + + +@pytest.mark.asyncio +async def test_public_aocr_entrypoint_uses_native_transport_when_enabled(ocr_server: RecordingServer) -> None: + response: Final = await litellm.aocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].headers["accept-encoding"] == "identity" + + +def test_public_ocr_falls_back_when_native_transport_declines(ocr_server: RecordingServer) -> None: + response: Final = litellm.ocr( + model=OCR_MODEL, + document={"type": "file", "file": b"%PDF-1.4", "mime_type": "application/pdf"}, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].headers["accept-encoding"] != "identity" + + +def test_public_ocr_uses_python_transport_when_disabled(ocr_server: RecordingServer) -> None: + litellm.rust(False) + + response: Final = litellm.ocr( + model=OCR_MODEL, + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + ) + + assert response.pages[0].markdown == "native OCR response" + assert ocr_server.requests[0].headers["accept-encoding"] != "identity"