mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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 <ishaanjaffer0324@gmail.com>
This commit is contained in:
ishaan-berri
GitHub
parent
fda08dd727
commit
1d5ab42e14
28 changed files with 1151 additions and 88 deletions
118
litellm-rust/Cargo.lock
generated
118
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
50
litellm-rust/crates/ai-gateway/AGENTS.md
Normal file
50
litellm-rust/crates/ai-gateway/AGENTS.md
Normal file
|
|
@ -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<Router> + 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<AppState>`; `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 <key>` 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.
|
||||
26
litellm-rust/crates/ai-gateway/Cargo.toml
Normal file
26
litellm-rust/crates/ai-gateway/Cargo.toml
Normal file
|
|
@ -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"]
|
||||
54
litellm-rust/crates/ai-gateway/src/auth/mod.rs
Normal file
54
litellm-rust/crates/ai-gateway/src/auth/mod.rs
Normal file
|
|
@ -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 <key>` 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<AppState> for RequireMasterKey {
|
||||
type Rejection = (StatusCode, String);
|
||||
|
||||
async fn from_request_parts(
|
||||
parts: &mut Parts,
|
||||
state: &AppState,
|
||||
) -> Result<Self, Self::Rejection> {
|
||||
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(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
58
litellm-rust/crates/ai-gateway/src/gil.rs
Normal file
58
litellm-rust/crates/ai-gateway/src/gil.rs
Normal file
|
|
@ -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<u64>,
|
||||
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,
|
||||
}
|
||||
}
|
||||
109
litellm-rust/crates/ai-gateway/src/main.rs
Normal file
109
litellm-rust/crates/ai-gateway/src/main.rs
Normal file
|
|
@ -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<Arc<str>> = 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])
|
||||
}
|
||||
27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
Normal file
27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
Normal file
|
|
@ -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.
|
||||
39
litellm-rust/crates/ai-gateway/src/python/config.rs
Normal file
39
litellm-rust/crates/ai-gateway/src/python/config.rs
Normal file
|
|
@ -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<Router> {
|
||||
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<Deployment> = serde_json::from_str(&model_list_json)
|
||||
.map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?;
|
||||
|
||||
Ok(Router::new(deployments))
|
||||
})
|
||||
}
|
||||
4
litellm-rust/crates/ai-gateway/src/python/mod.rs
Normal file
4
litellm-rust/crates/ai-gateway/src/python/mod.rs
Normal file
|
|
@ -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;
|
||||
38
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
Normal file
38
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
Normal file
|
|
@ -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<AppState>`.**
|
||||
> `routes/mod.rs::app` merges them all and applies state once. Adding a route is:
|
||||
> create the module, then add one `.merge(<name>::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<AppState> { 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.
|
||||
30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
Normal file
30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
Normal file
|
|
@ -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<AppState> {
|
||||
Router::new().route("/health/gil", get(status))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct GilStatusResponse {
|
||||
gil_acquired_last_30s: bool,
|
||||
total_acquisitions: u64,
|
||||
seconds_since_last: Option<u64>,
|
||||
}
|
||||
|
||||
async fn status() -> Json<GilStatusResponse> {
|
||||
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,
|
||||
})
|
||||
}
|
||||
24
litellm-rust/crates/ai-gateway/src/routes/health.rs
Normal file
24
litellm-rust/crates/ai-gateway/src/routes/health.rs
Normal file
|
|
@ -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<AppState> {
|
||||
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
|
||||
}
|
||||
23
litellm-rust/crates/ai-gateway/src/routes/mod.rs
Normal file
23
litellm-rust/crates/ai-gateway/src/routes/mod.rs
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
//! HTTP routes.
|
||||
//!
|
||||
//! **Template:** every route module exposes `pub fn router() -> Router<AppState>`
|
||||
//! 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)
|
||||
}
|
||||
81
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
Normal file
81
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
Normal file
|
|
@ -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<AppState> {
|
||||
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<AppState>,
|
||||
Query(query): Query<RealtimeQuery>,
|
||||
) -> Result<Response, (StatusCode, String)> {
|
||||
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<ModelRouter>, 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::<RealtimeEvent>(&text).ok(),
|
||||
_ => None,
|
||||
}
|
||||
});
|
||||
let client_out = ws_sink.with(|event: RealtimeEvent| async move {
|
||||
Ok::<Message, axum::Error>(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;
|
||||
}
|
||||
|
|
@ -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<In, Out>(
|
||||
router: &Router,
|
||||
model: &str,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::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
|
||||
}
|
||||
12
litellm-rust/crates/ai-gateway/src/state.rs
Normal file
12
litellm-rust/crates/ai-gateway/src/state.rs
Normal file
|
|
@ -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<Router>,
|
||||
/// 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<Arc<str>>,
|
||||
}
|
||||
|
|
@ -6,6 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
rand.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
pub mod error;
|
||||
pub mod ocr;
|
||||
pub mod realtime;
|
||||
pub mod router;
|
||||
|
||||
pub use error::{CoreError, CoreResult};
|
||||
|
|
|
|||
44
litellm-rust/crates/core/src/router/deployment.rs
Normal file
44
litellm-rust/crates/core/src/router/deployment.rs
Normal file
|
|
@ -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<String>,
|
||||
#[serde(default)]
|
||||
pub api_base: Option<String>,
|
||||
}
|
||||
|
||||
/// 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")
|
||||
);
|
||||
}
|
||||
}
|
||||
87
litellm-rust/crates/core/src/router/mod.rs
Normal file
87
litellm-rust/crates/core/src/router/mod.rs
Normal file
|
|
@ -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<Deployment>,
|
||||
routing_strategy: RoutingStrategy,
|
||||
}
|
||||
|
||||
impl Router {
|
||||
/// Build a router from a `model_list` using the default `simple-shuffle` strategy.
|
||||
pub fn new(model_list: Vec<Deployment>) -> 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());
|
||||
}
|
||||
}
|
||||
26
litellm-rust/crates/core/src/router/strategy/mod.rs
Normal file
26
litellm-rust/crates/core/src/router/strategy/mod.rs
Normal file
|
|
@ -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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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());
|
||||
}
|
||||
}
|
||||
|
|
@ -15,3 +15,4 @@ futures-util.workspace = true
|
|||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
futures-channel = "0.3"
|
||||
|
|
|
|||
|
|
@ -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<String> {
|
|||
.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<In, Out>(
|
||||
model: &str,
|
||||
input_events: Vec<RealtimeEvent>,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Vec<RealtimeEvent>> {
|
||||
idle_timeout: Option<Duration>,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::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<RealtimeEvent> = 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::<RealtimeEvent>();
|
||||
// provider -> client (we hold `backend_rx` to read backend events)
|
||||
let (client_out, mut backend_rx) = mpsc::unbounded::<RealtimeEvent>();
|
||||
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
|
|
|
|||
28
litellm/proxy/read_model_list.py
Normal file
28
litellm/proxy/read_model_list.py
Normal file
|
|
@ -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 []
|
||||
32
tests/test_litellm/proxy/test_read_model_list.py
Normal file
32
tests/test_litellm/proxy/test_read_model_list.py
Normal file
|
|
@ -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)) == []
|
||||
Loading…
Add table
Reference in a new issue