scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

wip

Browse source
This commit is contained in:
Yujong Lee 2026-09-07 23:03:27 -07:00
parent 94ce534b93
commit aa8d8bf84b
53 changed files with 2224 additions and 1278 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm-rust/Cargo.lock generated
View file

@ -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",

2
litellm-rust/Cargo.toml
View file

@ -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"

20
litellm-rust/crates/ai-gateway/ARCHITECTURE.md
View file

@ -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<br/>LLM inference]
G <--> O[OpenAI realtime]
C[client] <--> G[Rust ai-gateway<br/>Axum transport]
G <--> K[litellm-core<br/>route and provider I/O]
K <--> O[provider]
G -. spend tracking callback .-> P[litellm proxy]
F[litellm-config<br/>load-time only] --> G
F -. Python backend .-> P

5
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -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"] }

17
litellm-rust/crates/ai-gateway/README.md
View file

@ -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://<host>/v1/realtime?model=<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).

4
litellm-rust/crates/ai-gateway/src/constants.rs
View file

@ -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";

127
litellm-rust/crates/ai-gateway/src/integrations/README.md
View file

@ -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.

1
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -1,3 +1,2 @@
pub mod audio_transcription;
pub mod realtime_pool;
pub(crate) mod tls;

46
litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs
View file

@ -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]

80
litellm-rust/crates/ai-gateway/src/io/tls.rs
View file

@ -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<Arc<ClientConfig>> = OnceLock::new();
fn build_config() -> Result<ClientConfig, Box<Error>> {
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<Arc<ClientConfig>, Box<Error>> {
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<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
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());
}
}

10
litellm-rust/crates/ai-gateway/src/main.rs
View file

@ -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(&params.model);
if let Some(key) = upstream_key(
provider_model,
&params.model,
params.api_key.as_deref(),
params.api_base.as_deref(),
) {

15
litellm-rust/crates/ai-gateway/src/routes/realtime/README.md
View file

@ -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

56
litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs
View file

@ -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(&params.model);
let connection = upstream_key(
provider_model,
&params.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,
},

6
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs
View file

@ -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::<ResponsesWsEvent>(&text).ok(),
Ok(Message::Text(text)) => Some(
serde_json::from_str::<ResponsesWsEvent>(&text)
.map_err(|error| litellm_core::Error::InvalidRequest(error.to_string())),
),
Err(error) => Some(Err(litellm_core::Error::Network(error.to_string()))),
_ => None,
}
}));

4
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs
View file

@ -51,7 +51,7 @@ pub async fn run<In, Out>(
client_out: Out,
) -> Result<ExecutedCall<(), Error>, Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
In: Stream<Item = Result<ResponsesWsEvent, Error>> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + 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(),

28
litellm-rust/crates/core/AGENTS.md
View file

@ -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.

2
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -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<Value, Error> {
let body = serde_json::to_vec(&request.body)

2
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -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)]

2
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -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<ProviderAudioTranscriptionRequest, Error> {
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)

167
litellm-rust/crates/core/src/call_lifecycle/README.md
View file

@ -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/<call_type>/
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<MessagesResponse> {
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 `<call_type>/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 `<call_type>/prepare.rs`
Resolve model/provider once, generate or preserve `litellm_call_id`, construct
callback and guardrail runners, and return `Prepared<CallType>Call`.
4. Add `<call_type>/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 `<call_type>/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

146
litellm-rust/crates/core/src/chat_completions/lifecycle.rs
View file

@ -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<Transition, Error> {
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<ChatCompletionsRoute> {
pub fn commitment(&self) -> crate::lifecycle::Commitment {
self.state.program.commitment()
}
pub fn failure_stage(&self) -> Option<crate::lifecycle::FailureStage> {
self.state.program.failure_stage()
}
}
pub fn machine(
admission: &Admission,

44
litellm-rust/crates/core/src/lifecycle/execution.rs
View file

@ -188,6 +188,47 @@ impl CallLifecycle {
ClockImpl: Clock,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
{
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<Resp, Error>
where
Resp: Serialize,
Policy: RequestPolicy<InitialReq, ProviderReq>,
Dispatcher: TerminalDispatcher,
ClockImpl: Clock,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
ResponseUsage: FnOnce(&Resp) -> Option<Usage>,
{
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,

2
litellm-rust/crates/core/src/lifecycle/mod.rs
View file

@ -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,

172
litellm-rust/crates/core/src/lifecycle/ocr.rs
View file

@ -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<T> {
@ -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<super::FailureStage> {
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<Self::Transition, Self::Error> {
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::<u128>() & !(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)
);
}
}

271
litellm-rust/crates/core/src/lifecycle/program.rs Normal file
View file

@ -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<FailureStage>,
}
#[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<FailureStage>,
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<FailureStage> {
self.failure_stage
}
pub fn advance(&mut self, outcome: Outcome, observations: Observations) -> Option<Transition> {
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)
);
}
}

6
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -13,6 +13,12 @@ pub(super) async fn execute_messages_provider_call(
request: MessagesRequest,
) -> Result<AnthropicMessagesResponse, Error> {
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<AnthropicMessagesResponse, Error> {
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);

167
litellm-rust/crates/core/src/messages/lifecycle.rs
View file

@ -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<String>,
}
#[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<Result<Self::State, Self::Decline>, 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<Transition, Error> {
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<MessagesRoute> {
pub fn commitment(&self) -> crate::lifecycle::Commitment {
self.state.program.commitment()
}
pub fn failure_stage(&self) -> Option<crate::lifecycle::FailureStage> {
self.state.program.failure_stage()
}
}
pub trait MessagesServices:
RequestPolicy<MessagesRequest, MessagesRequest> + TerminalDispatcher + Clock
@ -241,17 +140,29 @@ pub async fn messages<S: MessagesServices>(
context: CallLifecycleContext,
) -> ExecutedCall<AnthropicMessagesResponse, Error> {
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<Usage> {
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<S: MessagesServices + 'static>(
services: Arc<S>,
request: MessagesRequest,

2
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -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<AnthropicMessagesResponse, Error> {

2
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -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<ProviderMessagesRequest, Error> {
let provider_info =

14
litellm-rust/crates/core/src/messages/types.rs
View file

@ -15,14 +15,14 @@ pub struct MessagesRequest {
pub timeout: Option<Duration>,
}
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<Duration>,
pub url: String,
pub body: Value,
pub upstream_headers: Vec<(String, String)>,
pub timeout: Option<Duration>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

65
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -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<Output = Result<OcrTransportResponse, Error>>
where
Self: 'a;
impl<T> 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<T> 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<Output = Result<OcrTransportResponse, Error>> + '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<S: OcrServices>(
services: &S,
request: SettledOcrRequest,
@ -55,7 +88,11 @@ pub async fn ocr<S: OcrServices>(
&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<SettledOcrRequest, SettledOcrRequest> for S
}
}
pub(crate) async fn send(request: SettledOcrRequest) -> Result<OcrResponseData, Error> {
pub(crate) async fn send<S: OcrTransport>(
transport: &S,
request: SettledOcrRequest,
) -> Result<OcrResponseData, Error> {
let SettledOcrRequest {
endpoint,
headers,
@ -132,13 +172,14 @@ pub(crate) async fn send(request: SettledOcrRequest) -> Result<OcrResponseData,
.collect();
let body = serde_json::to_vec(&body)
.map_err(|_| Error::InvalidRequest("could not encode OCR request".into()))?;
let response = buffered_post::send(buffered_post::Request {
let response = transport
.send(OcrTransportRequest {
url: endpoint.url,
headers,
body,
timeout_seconds: endpoint.timeout_seconds,
})
.await?;
})
.await?;
if !(200..300).contains(&response.status) {
return Err(Error::Http {
status: response.status,

15
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -106,6 +106,21 @@ pub struct SettledOcrRequest {
pub(super) body: Value,
}
#[derive(Debug, PartialEq, Eq)]
pub struct OcrTransportRequest {
pub url: String,
pub headers: Vec<(Vec<u8>, Vec<u8>)>,
pub body: Vec<u8>,
pub timeout_seconds: f64,
}
#[derive(Debug, PartialEq, Eq)]
pub struct OcrTransportResponse {
pub status: u16,
pub headers: Vec<(Vec<u8>, Vec<u8>)>,
pub content: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrRequestData {
pub data: Value,

4
litellm-rust/crates/core/src/realtime/mod.rs
View file

@ -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};

397
litellm-rust/crates/core/src/realtime/streaming.rs
View file

@ -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<Self, Error> {
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<String>,
pub api_base: Option<String>,
pub warm: Option<WarmConnection>,
pub idle_timeout: Option<Duration>,
}
@ -115,6 +122,7 @@ pub async fn warmup(connection: &RealtimeConnectionSpec) -> Result<WarmConnectio
)));
}
Ok(WarmConnection {
connection: connection.clone(),
upstream,
session_created,
})
@ -134,15 +142,25 @@ where
Out::Error: std::fmt::Display,
{
let start_time = services.now();
let model = request.connection.model.clone();
let connection = match request.warm {
Some(warm) => 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<In, Out>(
observation: &mut RealtimeObservation,
mut client_in: In,
mut client_out: Out,
) -> Result<(), Error>
) -> Result<(), RealtimeFailure>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + 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::<RealtimeEvent>(&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<Error> 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<Out>(
client_out: &mut Out,
event: &RealtimeEvent,
@ -292,10 +389,8 @@ async fn read_event(upstream: &mut Upstream) -> Result<RealtimeEvent, Error> {
}
async fn dial_upstream(connection: &RealtimeConnectionSpec) -> Result<Upstream, Error> {
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<Arc<ClientConfig>, 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<String, Error> {
.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<Value>, 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<Value>,
close_after_events: bool,
send_response_create: bool,
disconnect_client: bool,
idle_timeout: Duration,
) -> (ExecutedCall<(), Error>, Vec<TerminalRecord>) {
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<RealtimeEvent>, 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 { .. }
));
}
}

98
litellm-rust/crates/core/src/responses/instrumentation.rs
View file

@ -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<TerminalClassification> {
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::<String>();
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(),
})
);
}
}
}

478
litellm-rust/crates/core/src/responses/websocket.rs
View file

@ -67,7 +67,7 @@ pub struct ResponsesWebSocketRequest {
}
pub async fn responses_websocket<S, In, Out>(
services: &S,
services: Arc<S>,
request: ResponsesWebSocketRequest,
context: crate::lifecycle::CallLifecycleContext,
client_in: In,
@ -75,69 +75,164 @@ pub async fn responses_websocket<S, In, Out>(
) -> Result<ExecutedCall<(), Error>, Error>
where
S: TerminalDispatcher + crate::lifecycle::Clock,
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
S: 'static,
In: Stream<Item = Result<ResponsesWsEvent, Error>> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + 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<T> ResponsesCompletionServices for T where T: TerminalDispatcher + crate::lifecycle::Clock {}
struct ResponsesWsCompletion {
services: Arc<dyn ResponsesCompletionServices>,
context: Option<crate::lifecycle::CallLifecycleContext>,
start_time: f64,
instrumentation: Arc<ResponsesWsInstrumentation>,
}
impl ResponsesWsCompletion {
fn new<S>(
services: Arc<S>,
context: crate::lifecycle::CallLifecycleContext,
start_time: f64,
instrumentation: Arc<ResponsesWsInstrumentation>,
) -> 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<In, Out>(
upstream: Upstream,
model: &str,
@ -146,39 +241,66 @@ async fn splice<In, Out>(
instrumentation: &ResponsesWsInstrumentation,
mut client_in: In,
mut client_out: Out,
) -> Result<(), Error>
) -> Result<TerminalClassification, ResponsesWsFailure>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
In: Stream<Item = Result<ResponsesWsEvent, Error>> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + 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::<ResponsesWsEvent>(&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<Value>,
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<Result<ResponsesWsEvent, Error>>,
futures_channel::mpsc::UnboundedReceiver<Result<ResponsesWsEvent, Error>>,
) {
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<ResponsesWsEvent>,
) = futures_channel::mpsc::unbounded();
futures_channel::mpsc::UnboundedReceiver<Result<ResponsesWsEvent, Error>>,
) = 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!(

5
litellm-rust/crates/core/tests/messages_lifecycle.rs
View file

@ -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]

4
litellm-rust/crates/core/tests/ocr.rs
View file

@ -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"),

15
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -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::*;

2
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,3 +1,5 @@
#![recursion_limit = "256"]
mod diagnostics;
mod driver;
mod errors;

159
litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs
View file

@ -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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Value>,
optional_params: Option<Value>,
timeout_seconds: Option<f64>,
}
fn prepare_transcription(
inputs: AudioTranscriptionInputs,
) -> PyResult<impl Future<Output = Result<Value, Error>> + 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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option<Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option<Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option<Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_transcription,
errors = core_error_to_pyerr,
) -> PyResult<Py<PyAny>> {
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<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] extra_headers: Option<Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option<Value>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
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)
}

5
litellm-rust/crates/python-bridge/src/routes/chat_completions.rs
View file

@ -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),

236
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -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<pyo3::Py<pyo3::PyAny>> {
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<pyo3::Bound<'_, pyo3::PyAny>> {
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<pyo3::Py<pyo3::PyAny>> {
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<pyo3::Bound<'_, pyo3::PyAny>> {
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<Py<PyAny>> {
let future = prepare_echo(EchoInputs { value })?;
litellm_python_interop::run_sync(py, future, map_error)
}
#[pyfunction]
fn aecho(py: Python<'_>, value: String) -> PyResult<Bound<'_, PyAny>> {
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<Py<PyAny>> {
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<Bound<'_, PyAny>> {
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)

194
litellm-rust/crates/python-bridge/src/routes/messages.rs
View file

@ -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<Py<PyDict>>,
request: Option<MessagesRequest>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
prepared: Option<ProviderMessagesRequest>,
}
#[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<Option<String>>
.transpose()
}
fn decode_state(py: Python<'_>, arguments: Py<PyDict>) -> PyResult<MessagesState> {
let bag = arguments.bind(py);
let body = bag
fn decode_request(py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult<MessagesRequest> {
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::<f64>())
.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::<Map<String, Value>>(&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::<Map<String, Value>>(&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<PyDict>) -> PyResult<Py<MessagesState>> {
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::<PyDict>()?;
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<MessagesState>) -> PyResult<MessagesRequest> {
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<MessagesState>) -> PyResult<ProviderMessagesRequest> {
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::<PyResult<_>>()?;
Ok(prepared)
}
fn validate_arguments(arguments: &Bound<'_, PyDict>) -> PyResult<()> {
@ -221,8 +270,8 @@ fn send(py: Python<'_>, state: Py<MessagesState>) -> PyResult<Bound<'_, PyAny>>
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<MessagesState>) -> PyResult<Py<PyAny>> {
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<PyDict>) -> PyResult<Bound<'_, PyAny>> {
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,
};

1
litellm-rust/crates/python-bridge/src/routes/mod.rs
View file

@ -1,6 +1,5 @@
use pyo3::prelude::*;
#[macro_use]
mod definition;
#[cfg(feature = "trace-parity")]

6
litellm-rust/crates/python-bridge/src/routes/ocr.rs
View file

@ -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<OcrState>) -> PyResult<Bound<'_, PyAny>> {
});
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<OcrState>) -> PyResult<Py<PyAny>> {
let executed = run_sync_value(
py,
async move {
let services = NoopOcrServices;
let services = DefaultOcrServices;
Ok::<_, std::convert::Infallible>(
litellm_core::ocr::ocr(
&services,

45
litellm/llms/anthropic/experimental_pass_through/messages/fake_stream_iterator.py
View file

@ -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

12
litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py
View file

@ -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:
"""

19
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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,

48
litellm/ocr/main.py
View file

@ -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(

85
litellm/rust_bridge/messages.py
View file

@ -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",
]

141
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -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(

43
tests/test_litellm_rust/test_sdk_dispatch.py
View file

@ -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"