From 1d5ab42e14fae6daea963e38633d67b7613eb302 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Tue, 23 Jun 2026 19:16:34 -0700 Subject: [PATCH] feat: add minimal rust router + axum ai-gateway calling router.realtime (2/2) (#31135) * add CoreError::Routing variant for deployment selection failures * add minimal Rust Router (simple-shuffle) mirroring router.py spec * add litellm-router crate manifest * add ai-gateway POST /v1/realtime handler calling router.realtime * add ai-gateway health routes * wire ai-gateway routes into the axum app * add ai-gateway AppState holding the shared router * add ai-gateway axum server entrypoint * add litellm-ai-gateway binary crate manifest * docs: add ai-gateway folder-architecture AGENTS.md * register router + ai-gateway crates and axum/rand deps in workspace * update Cargo.lock for router + ai-gateway crates * split router: extract model_list types into deployment module * split router: extract routing policy into strategy module * split router: move Router orchestration into router module * router lib: wire submodules and re-export public API * add read_model_list helper reusing ProxyConfig env/secret resolution * add GIL-activity tracker (records acquisitions, 30s window) * add GET /health/gil endpoint for polling GIL activity * add pyo3 load_router_from_config bridge (feature-gated, load-time only) * register /health/gil route in ai-gateway * wire build_router: load from python config when feature enabled * add optional pyo3 dep + python-config feature to ai-gateway * update Cargo.lock for optional pyo3 dependency * fix: satisfy strict ruff budget (FA100) in read_model_list * test: cover read_model_list env resolution + empty config * ai-gateway: bind localhost by default, warn on bad PORT/missing keys, wire gateway key * ai-gateway: add gateway_key to AppState for realtime auth * ai-gateway: require bearer auth + map unknown model to 404 on /v1/realtime * ai-gateway: move python interop into python/ with load-time-only AGENTS.md * ai-gateway: document auth, gil, and python folder in AGENTS.md * core: add router module (model_list types + simple-shuffle selection) * ai-gateway: dispatch realtime via core router + providers (drop router crate dep) * update Cargo.lock: fold router into core * workspace: drop crates/router member and litellm-router dep * read_model_list: reuse ProxyConfig.get_config (includes + os.environ + DB) instead of thin yaml read * ai-gateway: constant-time bearer compare + 500 (not 503) for unconfigured key * ai-gateway: trim stored gateway key to match trimmed bearer token * ai-gateway: add subtle dep for constant-time comparison * workspace: add subtle dependency * update Cargo.lock for subtle * core router: make strategy a folder (one module per strategy, simple_shuffle) * providers: make realtime() a streaming splice (client stream <-> OpenAI) instead of collect * providers: add futures-channel dev-dep for the streaming live test * ai-gateway: make /v1/realtime a WebSocket (auth before upgrade, splice typed events) * ai-gateway: dispatch realtime as a stream splice * ai-gateway: route /v1/realtime via GET (WebSocket), drop POST * ai-gateway: enable axum ws feature + futures-util * update Cargo.lock for ws feature + futures-channel * core router: add has_deployment() for pre-flight model checks * ai-gateway: extract auth into auth/ module (single master key, LITELLM_MASTER_KEY) * ai-gateway routes: adopt router()-per-module template + merge in app() * ai-gateway: document auth/ + routes template in AGENTS.md * ai-gateway: realtime route as thin handler + service + transport * ai-gateway: auth as a RequireMasterKey extractor (idiomatic axum FromRequestParts) * ai-gateway: docs for auth extractor + simplified route template * ai-gateway: collapse realtime route to mod.rs + service.rs; docs for extractor/template * providers realtime: enforce idle timeout around the splice (reap stalled sessions) * ai-gateway: rename realtime service timeout param to idle_timeout --------- Co-authored-by: Ishaan Jaffer --- litellm-rust/Cargo.lock | 118 +++++++++ litellm-rust/Cargo.toml | 4 + litellm-rust/crates/ai-gateway/AGENTS.md | 50 ++++ litellm-rust/crates/ai-gateway/Cargo.toml | 26 ++ .../crates/ai-gateway/src/auth/mod.rs | 54 +++++ litellm-rust/crates/ai-gateway/src/gil.rs | 58 +++++ litellm-rust/crates/ai-gateway/src/main.rs | 109 +++++++++ .../crates/ai-gateway/src/python/AGENTS.md | 27 +++ .../crates/ai-gateway/src/python/config.rs | 39 +++ .../crates/ai-gateway/src/python/mod.rs | 4 + .../crates/ai-gateway/src/routes/AGENTS.md | 38 +++ .../crates/ai-gateway/src/routes/gil.rs | 30 +++ .../crates/ai-gateway/src/routes/health.rs | 24 ++ .../crates/ai-gateway/src/routes/mod.rs | 23 ++ .../ai-gateway/src/routes/realtime/mod.rs | 81 +++++++ .../ai-gateway/src/routes/realtime/service.rs | 45 ++++ litellm-rust/crates/ai-gateway/src/state.rs | 12 + litellm-rust/crates/core/Cargo.toml | 1 + litellm-rust/crates/core/src/error.rs | 2 + litellm-rust/crates/core/src/lib.rs | 1 + .../crates/core/src/router/deployment.rs | 44 ++++ litellm-rust/crates/core/src/router/mod.rs | 87 +++++++ .../crates/core/src/router/strategy/mod.rs | 26 ++ .../src/router/strategy/simple_shuffle.rs | 47 ++++ litellm-rust/crates/providers/Cargo.toml | 1 + litellm-rust/crates/providers/src/realtime.rs | 228 +++++++++++------- litellm/proxy/read_model_list.py | 28 +++ .../proxy/test_read_model_list.py | 32 +++ 28 files changed, 1151 insertions(+), 88 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/AGENTS.md create mode 100644 litellm-rust/crates/ai-gateway/Cargo.toml create mode 100644 litellm-rust/crates/ai-gateway/src/auth/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/gil.rs create mode 100644 litellm-rust/crates/ai-gateway/src/main.rs create mode 100644 litellm-rust/crates/ai-gateway/src/python/AGENTS.md create mode 100644 litellm-rust/crates/ai-gateway/src/python/config.rs create mode 100644 litellm-rust/crates/ai-gateway/src/python/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/routes/AGENTS.md create mode 100644 litellm-rust/crates/ai-gateway/src/routes/gil.rs create mode 100644 litellm-rust/crates/ai-gateway/src/routes/health.rs create mode 100644 litellm-rust/crates/ai-gateway/src/routes/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs create mode 100644 litellm-rust/crates/ai-gateway/src/state.rs create mode 100644 litellm-rust/crates/core/src/router/deployment.rs create mode 100644 litellm-rust/crates/core/src/router/mod.rs create mode 100644 litellm-rust/crates/core/src/router/strategy/mod.rs create mode 100644 litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs create mode 100644 litellm/proxy/read_model_list.py create mode 100644 tests/test_litellm/proxy/test_read_model_list.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2ec86fc00e7..a269a224d97 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,17 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -14,6 +25,64 @@ version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "axum" +version = "0.7.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" +dependencies = [ + "async-trait", + "axum-core", + "base64", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "rustversion", + "serde", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sha1", + "sync_wrapper", + "tokio", + "tokio-tungstenite", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" +dependencies = [ + "async-trait", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "rustversion", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + [[package]] name = "base64" version = "0.22.1" @@ -283,6 +352,12 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + [[package]] name = "hyper" version = "1.10.1" @@ -296,6 +371,7 @@ dependencies = [ "http", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -483,10 +559,26 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "litellm-ai-gateway" +version = "0.1.0" +dependencies = [ + "axum", + "futures-util", + "litellm-core", + "litellm-providers", + "pyo3", + "serde", + "serde_json", + "subtle", + "tokio", +] + [[package]] name = "litellm-core" version = "0.1.0" dependencies = [ + "rand 0.8.6", "serde", "serde_json", "thiserror 2.0.18", @@ -496,6 +588,7 @@ dependencies = [ name = "litellm-providers" version = "0.1.0" dependencies = [ + "futures-channel", "futures-util", "litellm-core", "reqwest", @@ -532,6 +625,12 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "matchit" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" + [[package]] name = "memchr" version = "2.8.2" @@ -547,6 +646,12 @@ dependencies = [ "autocfg", ] +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + [[package]] name = "mio" version = "1.2.1" @@ -1001,6 +1106,17 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1231,6 +1347,7 @@ dependencies = [ "tokio", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -1269,6 +1386,7 @@ version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ + "log", "pin-project-lite", "tracing-core", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 0b1528672a7..06289e5a46f 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -3,6 +3,7 @@ members = [ "crates/core", "crates/providers", "crates/python-bridge", + "crates/ai-gateway", ] resolver = "2" @@ -14,10 +15,13 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } litellm-providers = { path = "crates/providers" } +axum = "0.7" pyo3 = "0.23.5" +rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +subtle = "2" thiserror = "2.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } diff --git a/litellm-rust/crates/ai-gateway/AGENTS.md b/litellm-rust/crates/ai-gateway/AGENTS.md new file mode 100644 index 00000000000..d9e6e1adde5 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/AGENTS.md @@ -0,0 +1,50 @@ +# ai-gateway — folder architecture + +The Axum server that fronts the Rust gateway. It owns transport + config + auth +only; deployment selection lives in `core::router`, transforms in `core`/`providers`. + +``` +src/ + main.rs # entrypoint: build AppState (router + master key), bind, serve + state.rs # AppState — shared Arc + master_key + gil.rs # GIL-activity tracker (records Python acquisitions) + auth/ # authentication as an axum extractor — added to handler args + mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY) + routes/ # one module per route, all matching the same template + AGENTS.md # ← the route template (read this before adding a route) + mod.rs # app(): merges every module's router() + health.rs # simple route (one file): router() + liveness/readiness + gil.rs # simple route (one file): router() + GET /health/gil + realtime/ # route with logic → axum surface + a no-axum service: + mod.rs # router() + handler + WS<->events adapter (the axum surface) + service.rs # business logic (select deployment, call provider) — no axum, testable + python/ # Python interop (feature: python-config) — load-time only + mod.rs, config.rs, AGENTS.md +``` + +## Rules + +- **Routes follow one template.** Each route module exposes + `pub fn router() -> Router`; `routes/mod.rs` only merges them. Simple + routes are one file; non-trivial routes are a folder (`handler`/`service`/ + `transport`). See `routes/AGENTS.md`. +- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's + args; it runs during extraction. Never re-implement the check per route. +- **Handlers are thin.** A handler validates and delegates to its `service`. No + business logic, no provider calls, no transforms in handlers. +- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in + `state.rs`; read env/config only in `main.rs` when building state. + +## Auth (interim) + +A single **master key** (`LITELLM_MASTER_KEY`), enforced by the +`auth::RequireMasterKey` extractor: any caller presenting it as +`Authorization: Bearer ` may invoke the gateway. Fails closed (500) when +unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to +override). Full per-key auth + budgets/rate-limits are delegated to the Python +proxy in a later phase. Health routes don't add the extractor (unauthenticated). + +## Python interop + +Anything that calls into Python lives in `python/` and is **load-time only** — see +`python/AGENTS.md`. The realtime data path never takes the GIL. diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml new file mode 100644 index 00000000000..79bdc4bdb26 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "litellm-ai-gateway" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[[bin]] +name = "litellm-ai-gateway" +path = "src/main.rs" + +[dependencies] +litellm-core.workspace = true +litellm-providers.workspace = true +axum = { workspace = true, features = ["ws"] } +futures-util.workspace = true +tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time"] } +serde.workspace = true +serde_json.workspace = true +subtle.workspace = true +pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } + +[features] +# Build the gateway's config from the proxy YAML via an embedded Python +# interpreter (links libpython; requires `litellm` importable at runtime). +python-config = ["dep:pyo3"] diff --git a/litellm-rust/crates/ai-gateway/src/auth/mod.rs b/litellm-rust/crates/ai-gateway/src/auth/mod.rs new file mode 100644 index 00000000000..e2dd51f656d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/auth/mod.rs @@ -0,0 +1,54 @@ +//! Gateway authentication, as an axum **extractor** (the idiomatic pattern — +//! keeps handlers clean and auth testable). +//! +//! For now this is a single **master key**: any caller presenting it as +//! `Authorization: Bearer ` may invoke the gateway. Per-key auth, budgets, +//! and rate limits are delegated to the Python proxy in a later phase. +//! +//! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then +//! runs during extraction, before the handler body. Routes never re-implement it. + +use axum::extract::FromRequestParts; +use axum::http::header::AUTHORIZATION; +use axum::http::request::Parts; +use axum::http::StatusCode; +use subtle::ConstantTimeEq; + +use crate::state::AppState; + +/// Extractor that requires the configured master key as a bearer token. +/// +/// Rejections: `500` when no master key is configured (permanent +/// misconfiguration, not a transient outage); `401` on a missing/incorrect +/// token. The comparison is constant-time. +pub struct RequireMasterKey; + +#[axum::async_trait] +impl FromRequestParts for RequireMasterKey { + type Rejection = (StatusCode, String); + + async fn from_request_parts( + parts: &mut Parts, + state: &AppState, + ) -> Result { + let Some(expected) = state.master_key.as_deref() else { + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(), + )); + }; + let provided = parts + .headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .map(str::trim); + match provided { + Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self), + _ => Err(( + StatusCode::UNAUTHORIZED, + "missing or invalid bearer token".to_string(), + )), + } + } +} diff --git a/litellm-rust/crates/ai-gateway/src/gil.rs b/litellm-rust/crates/ai-gateway/src/gil.rs new file mode 100644 index 00000000000..c749f722c73 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/gil.rs @@ -0,0 +1,58 @@ +//! GIL-activity tracking. +//! +//! Every acquisition of the Python GIL is recorded here so the `/health/gil` +//! endpoint can report whether Python was touched recently. The design goal is +//! that the GIL is acquired **only at load time** (config read) and never on the +//! realtime hot path — polling this endpoint during traffic should show the +//! count holding steady and `acquired_last_30s` falling to `false`. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +/// Window (seconds) for the "recently acquired" signal. +pub const RECENT_WINDOW_SECS: u64 = 30; + +static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0); +/// Unix seconds of the last acquisition; `0` means "never". +static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0); + +fn now_unix_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// Record that the GIL was just acquired. Call immediately before taking the GIL. +/// +/// Only invoked under the `python-config` feature; without it the gateway never +/// touches Python, so the recorder is unused (and the endpoint reports zero). +#[cfg_attr(not(feature = "python-config"), allow(dead_code))] +pub fn record_acquisition() { + GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed); + LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed); +} + +/// Point-in-time view of GIL activity. +pub struct GilSnapshot { + pub total_acquisitions: u64, + pub seconds_since_last: Option, + pub acquired_last_30s: bool, +} + +/// Read the current GIL-activity snapshot. +pub fn snapshot() -> GilSnapshot { + let total = GIL_ACQUISITIONS.load(Ordering::Relaxed); + let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed); + let seconds_since_last = if last == 0 { + None + } else { + Some(now_unix_secs().saturating_sub(last)) + }; + let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS); + GilSnapshot { + total_acquisitions: total, + seconds_since_last, + acquired_last_30s, + } +} diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs new file mode 100644 index 00000000000..944fa30881f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -0,0 +1,109 @@ +//! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router. +//! +//! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment +//! (simple-shuffle) → `providers::realtime::realtime()` invokes OpenAI. The +//! server owns transport + config; routing lives in the `router` crate. + +mod auth; +mod gil; +#[cfg(feature = "python-config")] +mod python; +mod routes; +mod state; + +use std::sync::Arc; + +use litellm_core::router::{Deployment, LiteLLMParams, Router}; + +use crate::state::AppState; + +/// Bind to localhost by default so the gateway is not a public, unauthenticated +/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`). +const DEFAULT_HOST: &str = "127.0.0.1"; +const DEFAULT_PORT: u16 = 4001; + +#[tokio::main] +async fn main() { + // Trim before storing so it matches the trimmed bearer token in `auth` + // (avoids a silent auth failure when the env var has surrounding whitespace). + let master_key: Option> = std::env::var("LITELLM_MASTER_KEY") + .ok() + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty()) + .map(Arc::from); + if master_key.is_none() { + eprintln!( + "warning: LITELLM_MASTER_KEY is not set; /v1/realtime will reject all requests (fail closed)" + ); + } + + let state = AppState { + router: Arc::new(build_router()), + master_key, + }; + + let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string()); + let port = resolve_port(); + + let listener = tokio::net::TcpListener::bind((host.as_str(), port)) + .await + .expect("failed to bind listener"); + eprintln!("litellm-ai-gateway listening on {host}:{port}"); + axum::serve(listener, routes::app(state)) + .await + .expect("server error"); +} + +/// Resolve `PORT`, warning (rather than silently defaulting) on an invalid value. +fn resolve_port() -> u16 { + match std::env::var("PORT") { + Ok(raw) => raw.parse().unwrap_or_else(|_| { + eprintln!("warning: PORT={raw:?} is not a valid port; using {DEFAULT_PORT}"); + DEFAULT_PORT + }), + Err(_) => DEFAULT_PORT, + } +} + +/// Build the router. With the `python-config` feature and `LITELLM_CONFIG_PATH` +/// set, load the resolved `model_list` from the proxy config via the embedded +/// Python reader (load time only). Otherwise fall back to the env stand-in. +fn build_router() -> Router { + #[cfg(feature = "python-config")] + if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") { + match python::config::load_router_from_config(&config_path) { + Ok(router) => { + eprintln!("loaded model_list from {config_path} via python config reader"); + return router; + } + Err(err) => { + eprintln!("config load failed ({err}); falling back to env deployment"); + } + } + } + build_router_from_env() +} + +/// Build a minimal single-deployment `model_list` from the environment. +/// +/// A real deployment loads `model_list` from config; this is the minimal stand-in +/// so the gateway has one OpenAI deployment to route to. +fn build_router_from_env() -> Router { + let model = + std::env::var("OPENAI_REALTIME_MODEL").unwrap_or_else(|_| "gpt-realtime".to_string()); + let api_key = std::env::var("OPENAI_API_KEY").ok(); + if api_key.is_none() { + eprintln!( + "warning: OPENAI_API_KEY is not set; realtime requests will fail with auth errors" + ); + } + let deployment = Deployment { + model_name: model.clone(), + litellm_params: LiteLLMParams { + model, + api_key, + api_base: None, + }, + }; + Router::new(vec![deployment]) +} diff --git a/litellm-rust/crates/ai-gateway/src/python/AGENTS.md b/litellm-rust/crates/ai-gateway/src/python/AGENTS.md new file mode 100644 index 00000000000..47aa117e0b9 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/python/AGENTS.md @@ -0,0 +1,27 @@ +# ai-gateway/src/python — Python interop (load-time only) + +Functions here embed the Python interpreter (pyo3) and take the GIL to call into +`litellm` (e.g. read the proxy `model_list`). Compiled only under the +`python-config` feature. + +## Hard rule: non-hot-path functions only + +Everything in this folder MUST run **at most once per process lifetime — at +startup / load time** (config read, warm-up). NEVER call into Python on the +request path: + +- No GIL acquisition per request, per connection, or per realtime event. +- No Python call inside a route handler, the router's hot path, or any loop that + scales with traffic. + +**Why:** the GIL serializes execution and would cap throughput; the realtime data +path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll +`GET /health/gil`, and `total_acquisitions` MUST stay flat under load. + +## How to add one + +Resolve whatever Python-derived data you need **once at boot** and hand the rest +of the gateway an owned, plain-Rust value (e.g. build a `Router` from the +resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()` +immediately before taking the GIL. If a function would need to run per request, +it does not belong here — move the work to Rust, or pre-resolve it at startup. diff --git a/litellm-rust/crates/ai-gateway/src/python/config.rs b/litellm-rust/crates/ai-gateway/src/python/config.rs new file mode 100644 index 00000000000..6ec9595469d --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/python/config.rs @@ -0,0 +1,39 @@ +//! Build the router by calling the Python proxy config reader (load time only). +//! +//! Embeds the interpreter via pyo3 and calls +//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's +//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot** +//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python. +//! +//! Compiled only under the `python-config` feature. + +use litellm_core::error::CoreError; +use litellm_core::router::{Deployment, Router}; +use litellm_core::CoreResult; +use pyo3::prelude::*; + +use crate::gil; + +/// Load the router's `model_list` from `config_path` via the Python reader. +pub fn load_router_from_config(config_path: &str) -> CoreResult { + gil::record_acquisition(); + Python::with_gil(|py| { + let model_list = py + .import("litellm.proxy.read_model_list") + .and_then(|module| module.getattr("read_model_list")) + .and_then(|reader| reader.call1((config_path,))) + .map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?; + + let model_list_json: String = py + .import("json") + .and_then(|json| json.getattr("dumps")) + .and_then(|dumps| dumps.call1((model_list,))) + .and_then(|encoded| encoded.extract()) + .map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?; + + let deployments: Vec = serde_json::from_str(&model_list_json) + .map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?; + + Ok(Router::new(deployments)) + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/python/mod.rs b/litellm-rust/crates/ai-gateway/src/python/mod.rs new file mode 100644 index 00000000000..a677bade676 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/python/mod.rs @@ -0,0 +1,4 @@ +//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path +//! only.** Compiled only under the `python-config` feature. + +pub mod config; diff --git a/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md new file mode 100644 index 00000000000..02c5f18c4f3 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/AGENTS.md @@ -0,0 +1,38 @@ +# routes/ — the route template + +Every route follows the **same shape** so the layout is predictable. The rule: + +> **Each route module exposes `pub fn router() -> Router`.** +> `routes/mod.rs::app` merges them all and applies state once. Adding a route is: +> create the module, then add one `.merge(::router())` line. + +## Default: one file +A route is a single file containing `router()` + its handler(s) (handlers stay +private). This is the norm — don't split until it hurts. +``` +pub fn router() -> Router { Router::new().route(PATH, get(handle)) } +async fn handle(...) -> impl IntoResponse { ... } +``` +`health.rs` and `gil.rs` are examples. + +## Split out `service` when there's real logic +When a route has business logic worth testing without axum, put it in a sibling +`service` (a file, or a folder if the route grows). The route file stays the +**axum surface** (router + handler + any socket/SSE adapter); `service` is plain +Rust with **no axum types**. `realtime/` is the example: +``` +realtime/ + mod.rs # axum surface: router() + handler + the WS<->events adapter + service.rs # pure logic: select deployment + call provider (no axum) — testable +``` +Split `service` further (or add `transport`, `repo`, …) only once a single file +genuinely gets hard to read. + +## Invariants +- **Auth is an extractor, not a manual call.** A handler requires auth by adding + `crate::auth::RequireMasterKey` to its arguments; it runs during extraction. + Never re-implement the check per route. +- **Handlers contain no business logic; `service` contains no axum types.** +- A route owns its paths in its own `router()`; `mod.rs` only merges. +- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`, + not duplicated in handlers. diff --git a/litellm-rust/crates/ai-gateway/src/routes/gil.rs b/litellm-rust/crates/ai-gateway/src/routes/gil.rs new file mode 100644 index 00000000000..0db0c6f0b14 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/gil.rs @@ -0,0 +1,30 @@ +//! `GET /health/gil` — poll to confirm Python is only touched at load time. +//! Simple-route template: a `router()` plus its handler, in one file. + +use axum::routing::get; +use axum::{Json, Router}; +use serde::Serialize; + +use crate::gil; +use crate::state::AppState; + +/// This route's contribution to the app router. +pub fn router() -> Router { + Router::new().route("/health/gil", get(status)) +} + +#[derive(Debug, Serialize)] +struct GilStatusResponse { + gil_acquired_last_30s: bool, + total_acquisitions: u64, + seconds_since_last: Option, +} + +async fn status() -> Json { + let snapshot = gil::snapshot(); + Json(GilStatusResponse { + gil_acquired_last_30s: snapshot.acquired_last_30s, + total_acquisitions: snapshot.total_acquisitions, + seconds_since_last: snapshot.seconds_since_last, + }) +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/health.rs b/litellm-rust/crates/ai-gateway/src/routes/health.rs new file mode 100644 index 00000000000..15c67fea325 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/health.rs @@ -0,0 +1,24 @@ +//! Health probes. Simple-route template: a `router()` plus its handlers, in one file. + +use axum::http::StatusCode; +use axum::routing::get; +use axum::Router; + +use crate::state::AppState; + +/// This route's contribution to the app router. +pub fn router() -> Router { + Router::new() + .route("/health/liveness", get(liveness)) + .route("/health/readiness", get(readiness)) +} + +/// The process is up. +async fn liveness() -> StatusCode { + StatusCode::OK +} + +/// The server is ready to accept traffic. +async fn readiness() -> StatusCode { + StatusCode::OK +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/mod.rs new file mode 100644 index 00000000000..c6b9573781a --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/mod.rs @@ -0,0 +1,23 @@ +//! HTTP routes. +//! +//! **Template:** every route module exposes `pub fn router() -> Router` +//! that mounts its own paths; [`app`] merges them. A trivial route is a single +//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with +//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md. + +pub mod gil; +pub mod health; +pub mod realtime; + +use axum::Router; + +use crate::state::AppState; + +/// Assemble the application router by merging every route module's `router()`. +pub fn app(state: AppState) -> Router { + Router::new() + .merge(health::router()) + .merge(gil::router()) + .merge(realtime::router()) + .with_state(state) +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs new file mode 100644 index 00000000000..1c3caa88e59 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs @@ -0,0 +1,81 @@ +//! `GET /v1/realtime` (WebSocket). +//! +//! This file is the **axum surface**: `router()`, the handler, and the small +//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is +//! the `RequireMasterKey` extractor, so the handler stays thin. + +mod service; + +use std::sync::Arc; + +use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::response::Response; +use axum::routing::get; +use axum::Router; +use futures_util::{SinkExt, StreamExt}; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::router::Router as ModelRouter; +use serde::Deserialize; + +use crate::auth::RequireMasterKey; +use crate::state::AppState; + +/// This route's contribution to the app router. +pub fn router() -> Router { + Router::new().route("/v1/realtime", get(handle)) +} + +#[derive(Debug, Deserialize)] +struct RealtimeQuery { + model: String, +} + +/// Auth runs via the `RequireMasterKey` extractor. We validate the model BEFORE +/// the upgrade so failures are clean HTTP (400/404), not a socket that opens then +/// closes, then hand the socket to `bridge`. +async fn handle( + _auth: RequireMasterKey, + ws: WebSocketUpgrade, + State(state): State, + Query(query): Query, +) -> Result { + if query.model.trim().is_empty() { + return Err(( + StatusCode::BAD_REQUEST, + "missing 'model' query param".to_string(), + )); + } + if !state.router.has_deployment(&query.model) { + return Err(( + StatusCode::NOT_FOUND, + format!("no deployment for model '{}'", query.model), + )); + } + + let router = state.router.clone(); + let model = query.model; + Ok(ws.on_upgrade(move |socket| bridge(socket, router, model))) +} + +/// Adapt the axum socket (text frames) to the typed-event `Stream`/`Sink` the +/// service wants, keeping axum types out of `service`. +async fn bridge(socket: WebSocket, router: Arc, model: String) { + let (ws_sink, ws_stream) = socket.split(); + + let client_in = ws_stream.filter_map(|message| async move { + match message { + Ok(Message::Text(text)) => serde_json::from_str::(&text).ok(), + _ => None, + } + }); + let client_out = ws_sink.with(|event: RealtimeEvent| async move { + Ok::(Message::Text( + serde_json::to_string(&event).unwrap_or_default(), + )) + }); + + futures_util::pin_mut!(client_in, client_out); + let _ = service::run(&router, &model, None, client_in, client_out).await; +} diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs new file mode 100644 index 00000000000..92eba2f2cdd --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -0,0 +1,45 @@ +//! Business logic: select a deployment with the (pure) core router, then call the +//! provider splice. The seam between `core::router` (selection only) and +//! `providers` (the actual WebSocket I/O). + +use std::time::Duration; + +use futures_util::{Sink, Stream}; +use litellm_core::error::CoreError; +use litellm_core::realtime::types::RealtimeEvent; +use litellm_core::router::Router; +use litellm_core::CoreResult; + +/// Select a deployment for `model` and splice the client stream to the provider. +pub async fn run( + router: &Router, + model: &str, + idle_timeout: Option, + client_in: In, + client_out: Out, +) -> CoreResult<()> +where + In: Stream + Unpin + Send, + Out: Sink + Unpin + Send, + >::Error: std::fmt::Display, +{ + let deployment = router.get_available_deployment(model).ok_or_else(|| { + CoreError::Routing(format!("no deployment available for model '{model}'")) + })?; + let params = &deployment.litellm_params; + // Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model. + let provider_model = params + .model + .strip_prefix("openai/") + .unwrap_or(¶ms.model); + + litellm_providers::realtime::realtime( + provider_model, + params.api_key.as_deref(), + params.api_base.as_deref(), + idle_timeout, + client_in, + client_out, + ) + .await +} diff --git a/litellm-rust/crates/ai-gateway/src/state.rs b/litellm-rust/crates/ai-gateway/src/state.rs new file mode 100644 index 00000000000..500d2da6750 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/state.rs @@ -0,0 +1,12 @@ +use std::sync::Arc; + +use litellm_core::router::Router; + +/// Shared application state handed to every route handler. +#[derive(Clone)] +pub struct AppState { + pub router: Arc, + /// The gateway master key. Any caller presenting it as a bearer token may + /// invoke the gateway. `None` → auth not configured (routes fail closed). + pub master_key: Option>, +} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index e54002fe5e8..1881bcfa602 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +rand.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 645e261f76d..9b29260cca4 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -19,6 +19,8 @@ pub enum CoreError { Http { status: u16, body: String }, #[error("OCR network error: {0}")] Network(String), + #[error("routing error: {0}")] + Routing(String), } pub fn json_type_name(value: &serde_json::Value) -> &'static str { diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 5d4d5bfd142..9d686626edc 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,5 +1,6 @@ pub mod error; pub mod ocr; pub mod realtime; +pub mod router; pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/router/deployment.rs b/litellm-rust/crates/core/src/router/deployment.rs new file mode 100644 index 00000000000..1ee88e682a3 --- /dev/null +++ b/litellm-rust/crates/core/src/router/deployment.rs @@ -0,0 +1,44 @@ +//! `model_list` data types, mirroring Python's deployment dict. Deserialize-ready +//! so a deployment can be loaded straight from the proxy config's `model_list`. + +use serde::Deserialize; + +/// Per-deployment call parameters, mirroring Python's `litellm_params`. +#[derive(Clone, Debug, Deserialize)] +pub struct LiteLLMParams { + /// Provider model, e.g. `gpt-realtime` or `openai/gpt-realtime`. + pub model: String, + #[serde(default)] + pub api_key: Option, + #[serde(default)] + pub api_base: Option, +} + +/// One entry of the `model_list`, mirroring Python's deployment dict. +#[derive(Clone, Debug, Deserialize)] +pub struct Deployment { + /// Public alias clients request, e.g. `gpt-realtime`. + pub model_name: String, + pub litellm_params: LiteLLMParams, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn deserializes_from_model_list_entry() { + let entry = r#"{ + "model_name": "gpt-realtime", + "litellm_params": {"model": "openai/gpt-realtime", "api_base": "https://x"} + }"#; + let deployment: Deployment = serde_json::from_str(entry).expect("valid entry"); + assert_eq!(deployment.model_name, "gpt-realtime"); + assert_eq!(deployment.litellm_params.model, "openai/gpt-realtime"); + assert_eq!(deployment.litellm_params.api_key, None); + assert_eq!( + deployment.litellm_params.api_base.as_deref(), + Some("https://x") + ); + } +} diff --git a/litellm-rust/crates/core/src/router/mod.rs b/litellm-rust/crates/core/src/router/mod.rs new file mode 100644 index 00000000000..101c3499fa3 --- /dev/null +++ b/litellm-rust/crates/core/src/router/mod.rs @@ -0,0 +1,87 @@ +//! Minimal Rust port of LiteLLM's `router.py` deployment selection. +//! +//! A [`Router`] is built from a `model_list` of [`Deployment`]s +//! (`{ model_name, litellm_params: { model, api_key, api_base } }`) and selects +//! one per request via a [`RoutingStrategy`]. For now the only strategy is +//! `simple-shuffle` — a uniform random pick within a `model_name` group. +//! +//! This stays pure (no I/O): it only *chooses* a deployment. The host (the +//! gateway) takes the chosen deployment and performs the actual provider call. +//! +//! - [`deployment`] — the `model_list` data types. +//! - [`strategy`] — how a deployment is chosen. + +mod deployment; +mod strategy; + +pub use deployment::{Deployment, LiteLLMParams}; +pub use strategy::RoutingStrategy; + +/// Load-balancing router over a `model_list`. +#[derive(Clone, Debug, Default)] +pub struct Router { + model_list: Vec, + routing_strategy: RoutingStrategy, +} + +impl Router { + /// Build a router from a `model_list` using the default `simple-shuffle` strategy. + pub fn new(model_list: Vec) -> Self { + Self { + model_list, + routing_strategy: RoutingStrategy::SimpleShuffle, + } + } + + /// Whether any deployment is registered under `model`. + pub fn has_deployment(&self, model: &str) -> bool { + self.model_list + .iter() + .any(|deployment| deployment.model_name == model) + } + + /// Pick a deployment for `model` per the routing strategy. Returns `None` + /// when no deployment is registered under that `model_name`. + pub fn get_available_deployment(&self, model: &str) -> Option<&Deployment> { + let candidates: Vec<&Deployment> = self + .model_list + .iter() + .filter(|deployment| deployment.model_name == model) + .collect(); + self.routing_strategy.select(&candidates) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn deployment(name: &str, model: &str) -> Deployment { + Deployment { + model_name: name.to_string(), + litellm_params: LiteLLMParams { + model: model.to_string(), + api_key: None, + api_base: None, + }, + } + } + + #[test] + fn selects_a_matching_deployment() { + let router = Router::new(vec![ + deployment("gpt-realtime", "gpt-realtime"), + deployment("other", "other-model"), + ]); + let chosen = router + .get_available_deployment("gpt-realtime") + .expect("a deployment should match"); + assert_eq!(chosen.model_name, "gpt-realtime"); + } + + #[test] + fn unknown_model_returns_none() { + let router = Router::new(vec![deployment("gpt-realtime", "gpt-realtime")]); + assert!(router.get_available_deployment("missing").is_none()); + } +} diff --git a/litellm-rust/crates/core/src/router/strategy/mod.rs b/litellm-rust/crates/core/src/router/strategy/mod.rs new file mode 100644 index 00000000000..7e8ac217db3 --- /dev/null +++ b/litellm-rust/crates/core/src/router/strategy/mod.rs @@ -0,0 +1,26 @@ +//! Routing policy: how the router picks one deployment from a model group. +//! +//! One module per strategy; [`RoutingStrategy::select`] dispatches to it. New +//! strategies (least-busy, latency-based, …) get their own file here. + +mod simple_shuffle; + +use super::Deployment; + +/// How the router chooses among the deployments sharing a `model_name`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RoutingStrategy { + /// Uniform random pick among the matching deployments. + #[default] + SimpleShuffle, +} + +impl RoutingStrategy { + /// Choose one deployment from `candidates` (all sharing the requested + /// `model_name`). Returns `None` when there are no candidates. + pub fn select<'a>(&self, candidates: &[&'a Deployment]) -> Option<&'a Deployment> { + match self { + RoutingStrategy::SimpleShuffle => simple_shuffle::select(candidates), + } + } +} diff --git a/litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs b/litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs new file mode 100644 index 00000000000..74ce0c21e80 --- /dev/null +++ b/litellm-rust/crates/core/src/router/strategy/simple_shuffle.rs @@ -0,0 +1,47 @@ +//! `simple-shuffle`: a uniform random pick among the candidate deployments. + +use rand::seq::SliceRandom; + +use crate::router::Deployment; + +/// Uniform random choice among `candidates` (all sharing the requested +/// `model_name`). Returns `None` when there are no candidates. +pub fn select<'a>(candidates: &[&'a Deployment]) -> Option<&'a Deployment> { + candidates.choose(&mut rand::thread_rng()).copied() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::router::{Deployment, LiteLLMParams}; + + fn deployment(model: &str) -> Deployment { + Deployment { + model_name: "gpt-realtime".to_string(), + litellm_params: LiteLLMParams { + model: model.to_string(), + api_key: None, + api_base: None, + }, + } + } + + #[test] + fn picks_from_candidates() { + let a = deployment("key-a"); + let b = deployment("key-b"); + let candidates = vec![&a, &b]; + for _ in 0..20 { + let chosen = select(&candidates).expect("non-empty"); + assert!(matches!( + chosen.litellm_params.model.as_str(), + "key-a" | "key-b" + )); + } + } + + #[test] + fn empty_candidates_select_none() { + assert!(select(&[]).is_none()); + } +} diff --git a/litellm-rust/crates/providers/Cargo.toml b/litellm-rust/crates/providers/Cargo.toml index cb8a91aa5c2..c5b41424d66 100644 --- a/litellm-rust/crates/providers/Cargo.toml +++ b/litellm-rust/crates/providers/Cargo.toml @@ -15,3 +15,4 @@ futures-util.workspace = true [dev-dependencies] serde_json.workspace = true +futures-channel = "0.3" diff --git a/litellm-rust/crates/providers/src/realtime.rs b/litellm-rust/crates/providers/src/realtime.rs index f2946db4137..95e8e26d45f 100644 --- a/litellm-rust/crates/providers/src/realtime.rs +++ b/litellm-rust/crates/providers/src/realtime.rs @@ -1,14 +1,14 @@ //! End-to-end OpenAI realtime invocation. //! //! The host-facing entry point, mirroring `providers::ocr::run_ocr`: open the -//! WebSocket to OpenAI, drive typed events through the pure -//! `OPENAI_REALTIME_CONFIG` transforms, and collect the response events. +//! WebSocket to OpenAI, then splice a client realtime stream to the upstream, +//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms. //! Network, auth header, key resolution, and wire (de)serialization live here so //! the `transformation` module stays pure and typed. use std::time::Duration; -use futures_util::{SinkExt, StreamExt}; +use futures_util::{Sink, SinkExt, Stream, StreamExt}; use litellm_core::error::CoreError; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; @@ -24,11 +24,14 @@ use crate::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; /// Environment variable holding the OpenAI API key (last-resort fallback). const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -/// Default overall ceiling for a single realtime invocation. -const DEFAULT_TIMEOUT_SECS: u64 = 60; - const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; +/// Default **idle** timeout: if neither side sends a frame for this long, the +/// session is reaped. It resets on any activity, so it does not cap a healthy +/// (continuously streaming) session — it only frees a stalled one (e.g. a +/// half-open upstream that keeps the socket open but stops sending). +const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 300; + /// Resolve the OpenAI API key from the explicit param or the environment. /// /// Blank/whitespace values are treated as absent (guard at resolution time). @@ -45,28 +48,25 @@ fn resolve_api_key(api_key: Option<&str>) -> CoreResult { .ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string())) } -/// True for events that end a realtime turn: a completed response or an error. -fn is_terminal_event(event: &RealtimeEvent) -> bool { - event.event_type == "response.done" || event.event_type == "error" -} - -/// Invoke the OpenAI realtime API end to end over a WebSocket. +/// Splice a client realtime stream to OpenAI: forward client events upstream +/// (via `transform_realtime_request`) and backend events downstream (via +/// `transform_realtime_response`). Returns when either side closes. /// -/// Sends each `input_events` entry after passing it through -/// `transform_realtime_request`, then collects backend events — each passed -/// through `transform_realtime_response` — until a terminal event -/// (`response.done` / `error`) arrives, the socket closes, or the `timeout` -/// elapses. Returns the transformed backend events in arrival order. -/// -/// Mirrors `run_ocr`: pure transforms come from `core`/`providers`; the network, -/// auth header, key resolution, and JSON (de)serialization are owned here. -pub async fn realtime( +/// Generic over the client transport (typed events) so this crate stays +/// framework-agnostic; the gateway adapts its axum socket to these. +pub async fn realtime( model: &str, - input_events: Vec, api_key: Option<&str>, api_base: Option<&str>, - timeout: Option, -) -> CoreResult> { + idle_timeout: Option, + mut client_in: In, + mut client_out: Out, +) -> CoreResult<()> +where + In: Stream + Unpin + Send, + Out: Sink + Unpin + Send, + >::Error: std::fmt::Display, +{ let config = &OPENAI_REALTIME_CONFIG; let api_key = resolve_api_key(api_key)?; let url = config.complete_url(api_base, model); @@ -75,58 +75,62 @@ pub async fn realtime( .as_str() .into_client_request() .map_err(|err| CoreError::Network(err.to_string()))?; - // GA realtime API: only Authorization is needed. The legacy - // `OpenAI-Beta: realtime=v1` header opts into the now-removed beta request - // shape and triggers `beta_api_shape_disabled`, so we do not send it. + // GA realtime: only Authorization. The legacy OpenAI-Beta header triggers + // beta_api_shape_disabled, so we do not send it. request.headers_mut().insert( AUTHORIZATION, HeaderValue::from_str(&format!("Bearer {api_key}")) .map_err(|err| CoreError::Auth(err.to_string()))?, ); - let (mut ws, _response) = connect_async(request) + let (upstream, _response) = connect_async(request) .await .map_err(|err| CoreError::Network(err.to_string()))?; + let (mut upstream_tx, mut upstream_rx) = upstream.split(); - for event in &input_events { - for outbound in config.transform_realtime_request(event, model)?.events { - let payload = serde_json::to_string(&outbound) - .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; - ws.send(Message::Text(payload)) - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - } - } + let idle = idle_timeout.unwrap_or(Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS)); - let deadline = timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_TIMEOUT_SECS)); - let mut received: Vec = Vec::new(); - - let collect = async { - while let Some(message) = ws.next().await { - match message.map_err(|err| CoreError::Network(err.to_string()))? { - Message::Text(text) => { - let event: RealtimeEvent = serde_json::from_str(text.as_str()) + // One loop forwarding both directions. The `sleep(idle)` arm is rebuilt every + // iteration, so any frame (either way) resets it — it fires only when the + // session has been fully idle for `idle`, reaping a stalled connection + // (task + upstream TCP socket) instead of leaking it. + loop { + tokio::select! { + // client -> upstream + client_event = client_in.next() => { + let Some(event) = client_event else { break }; // client disconnected + for outbound in config.transform_realtime_request(&event, model)?.events { + let payload = serde_json::to_string(&outbound) .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; - for outbound in config.transform_realtime_response(&event, model)?.events { - let terminal = is_terminal_event(&outbound); - received.push(outbound); - if terminal { - return Ok::<(), CoreError>(()); + upstream_tx + .send(Message::Text(payload)) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + } + } + // upstream -> client + upstream_message = upstream_rx.next() => { + let Some(message) = upstream_message else { break }; // upstream closed + match message.map_err(|err| CoreError::Network(err.to_string()))? { + Message::Text(text) => { + let event: RealtimeEvent = serde_json::from_str(&text) + .map_err(|err| CoreError::InvalidResponse(err.to_string()))?; + for outbound in config.transform_realtime_response(&event, model)?.events { + client_out + .send(outbound) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; } } + Message::Close(_) => break, + _ => {} } - Message::Close(_) => return Ok(()), - _ => {} } + // idle timeout: no activity from either side within `idle` + _ = tokio::time::sleep(idle) => break, } - Ok(()) - }; - - tokio::time::timeout(deadline, collect) - .await - .map_err(|_| CoreError::Network("realtime call timed out".to_string()))??; - - Ok(received) + } + Ok(()) } #[cfg(test)] @@ -146,48 +150,96 @@ mod tests { } } - #[test] - fn is_terminal_event_matches_done_and_error_only() { - assert!(is_terminal_event(&event(r#"{"type":"response.done"}"#))); - assert!(is_terminal_event(&event(r#"{"type":"error","error":{}}"#))); - assert!(!is_terminal_event(&event( - r#"{"type":"response.output_text.delta"}"# - ))); - } - /// Live end-to-end check against OpenAI. Ignored by default (CI never runs /// it); run explicitly with `OPENAI_API_KEY` set: /// `cargo test -p litellm-providers realtime_invokes_openai -- --ignored --nocapture` #[tokio::test] #[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"] async fn realtime_invokes_openai_and_responds() { + use futures_channel::mpsc; + let key = std::env::var(OPENAI_API_KEY_ENV).expect("set OPENAI_API_KEY to run this ignored test"); - let response_create = event( - r#"{"type":"response.create","response":{"output_modalities":["text"],"instructions":"Respond with exactly: hello world"}}"#, + // client -> provider (we hold `client_tx` to push events upstream) + let (mut client_tx, client_in) = mpsc::unbounded::(); + // provider -> client (we hold `backend_rx` to read backend events) + let (client_out, mut backend_rx) = mpsc::unbounded::(); + + // Clone the key so the spawned task owns its `String` (no borrow across await). + let key_owned = key.clone(); + let call = tokio::spawn(async move { + realtime( + "gpt-realtime", + Some(&key_owned), + None, + None, + client_in, + client_out, + ) + .await + }); + + // 1. First backend event should be session.created. + let first = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()) + .await + .expect("timed out waiting for session.created") + .expect("backend stream closed before session.created"); + assert_eq!( + first.event_type, "session.created", + "expected session.created, got: {}", + first.event_type ); - let events = realtime( - "gpt-realtime", - vec![response_create], - Some(&key), - None, - Some(Duration::from_secs(30)), - ) - .await - .expect("realtime call should succeed"); + // 2. Ask for a short audio response. + client_tx + .send(event( + r#"{"type":"conversation.item.create","item":{"type":"message","role":"user","content":[{"type":"input_text","text":"Say hi."}]}}"#, + )) + .await + .expect("send conversation.item.create"); + client_tx + .send(event(r#"{"type":"response.create"}"#)) + .await + .expect("send response.create"); - let types: Vec<&str> = events.iter().map(|e| e.event_type.as_str()).collect(); - eprintln!("received {} events: {:?}", events.len(), types); + // 3. Read backend events; require a non-empty audio delta, then response.done. + let mut saw_audio_delta = false; + let mut saw_done = false; + for _ in 0..500 { + let next = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()).await; + let event = match next { + Ok(Some(event)) => event, + Ok(None) => break, + Err(_) => panic!("timed out waiting for backend events"), + }; + match event.event_type.as_str() { + "response.output_audio.delta" => { + let delta = event + .data + .get("delta") + .and_then(|value| value.as_str()) + .unwrap_or(""); + if !delta.is_empty() { + saw_audio_delta = true; + } + } + "response.done" => { + saw_done = true; + break; + } + _ => {} + } + } assert!( - types.contains(&"response.done"), - "expected a response.done event, got: {types:?}" - ); - assert!( - types.contains(&"response.output_text.delta"), - "expected streamed text output, got: {types:?}" + saw_audio_delta, + "expected a response.output_audio.delta with non-empty delta" ); + assert!(saw_done, "expected a response.done event"); + + // Drop the client sender so the provider's to_upstream side finishes. + drop(client_tx); + let _ = call.await; } } diff --git a/litellm/proxy/read_model_list.py b/litellm/proxy/read_model_list.py new file mode 100644 index 00000000000..2dff8eaf698 --- /dev/null +++ b/litellm/proxy/read_model_list.py @@ -0,0 +1,28 @@ +"""Resolve a proxy config's ``model_list`` for the Rust AI gateway. + +The Rust gateway calls this once at load time (via an embedded interpreter) and +builds its own (Rust) router from the returned ``model_list``. We do NOT call +``ProxyConfig.load_config`` here: that returns a *Python* ``litellm.Router`` (not +usable from Rust) and boots the whole proxy (callbacks, cache, DB, auth) as side +effects. + +Instead we reuse ``ProxyConfig.get_config`` — the actual config reader — so the +gateway inherits the same heavy lifting the proxy does: ``include:`` merging, +``os.environ/`` + secret-manager resolution, and DB-stored models (when a DB is +configured). It has no proxy-setup side effects. Returns the resolved +``model_list``; the Rust side deserializes each entry into its ``Deployment``. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + + +def read_model_list(config_path: str) -> list[dict[str, Any]]: + """Load ``config_path`` via the proxy's own reader and return its + resolved ``model_list``.""" + from litellm.proxy.proxy_server import ProxyConfig + + config = asyncio.run(ProxyConfig().get_config(config_file_path=config_path)) + return config.get("model_list") or [] diff --git a/tests/test_litellm/proxy/test_read_model_list.py b/tests/test_litellm/proxy/test_read_model_list.py new file mode 100644 index 00000000000..7703d57e21d --- /dev/null +++ b/tests/test_litellm/proxy/test_read_model_list.py @@ -0,0 +1,32 @@ +"""Tests for litellm.proxy.read_model_list (Rust AI gateway config bridge).""" + +from litellm.proxy.read_model_list import read_model_list + + +def test_read_model_list_resolves_os_environ(monkeypatch, tmp_path): + """`os.environ/` markers in the model_list are resolved via ProxyConfig.""" + monkeypatch.setenv("OPENAI_API_KEY", "sk-resolved-123") + config = tmp_path / "config.yaml" + config.write_text( + "model_list:\n" + " - model_name: gpt-realtime\n" + " litellm_params:\n" + " model: openai/gpt-realtime\n" + " api_key: os.environ/OPENAI_API_KEY\n" + ) + + model_list = read_model_list(str(config)) + + assert len(model_list) == 1 + params = model_list[0]["litellm_params"] + assert model_list[0]["model_name"] == "gpt-realtime" + assert params["model"] == "openai/gpt-realtime" + assert params["api_key"] == "sk-resolved-123" + + +def test_read_model_list_missing_key_returns_empty(tmp_path): + """A config without a model_list yields an empty list, not an error.""" + config = tmp_path / "config.yaml" + config.write_text("general_settings: {}\n") + + assert read_model_list(str(config)) == []