From a5533cbcfd4e27f1204297ebed2529976cba471b Mon Sep 17 00:00:00 2001 From: narmstrong-berri Date: Sat, 31 Oct 2026 17:47:09 +0000 Subject: [PATCH] feat(router): add the Rust router core and serve routing through it behind LITELLM_RUST litellm-router-core implements the router core in Rust: a versioned registry with Python's deployment ids, selection with native simple-shuffle, the retry, fallback and cooldown loop with Python's cache keys, and per-attempt executors with a native seam for Messages to Anthropic. The python bridge exposes it as RouterCoreHandle, whose call loop is step driven so Rust never awaits Python. RustCore drives that loop for completions, messages, embeddings and responses and serves everything else from a PythonCore built from the same config. A new router rule in the route catalog opts in with LITELLM_RUST, and any config the Rust core cannot serve gets PythonCore at construction The call loop follows Python's fallback walk (per-target depth, a level raising its own error, lookup by the current and requested group, "*" after specific chains), stamps retries the way an exhausted retry loop does, and leaves health writes to the Python attempt, which records them. RustCore builds its handle from the constructor arguments and the applied changesets and rebuilds it when either version moves. The core unit suite runs on RustCore too --- litellm-rust/Cargo.lock | 19 + litellm-rust/Cargo.toml | 3 + litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/lib.rs | 5 + .../crates/python-bridge/src/router/config.rs | 294 +++++ .../crates/python-bridge/src/router/host.rs | 1011 +++++++++++++++++ .../crates/python-bridge/src/router/mod.rs | 549 +++++++++ litellm-rust/crates/router-core/Cargo.toml | 17 + litellm-rust/crates/router-core/src/cache.rs | 130 +++ litellm-rust/crates/router-core/src/call.rs | 592 ++++++++++ litellm-rust/crates/router-core/src/config.rs | 433 +++++++ litellm-rust/crates/router-core/src/error.rs | 190 ++++ .../crates/router-core/src/executor.rs | 75 ++ litellm-rust/crates/router-core/src/health.rs | 475 ++++++++ litellm-rust/crates/router-core/src/host.rs | 168 +++ litellm-rust/crates/router-core/src/lib.rs | 282 +++++ .../crates/router-core/src/registry.rs | 820 +++++++++++++ litellm-rust/crates/router-core/src/select.rs | 320 ++++++ .../crates/router-core/src/settings.rs | 69 ++ .../crates/router-core/tests/call_loop.rs | 167 +++ litellm/router_core.py | 27 +- litellm/router_rust_core.py | 857 ++++++++++++++ litellm/rust_bridge/_native.pyi | 116 +- litellm/rust_bridge/catalog.py | 2 + .../router_litellm_globals_check.py | 1 + .../router_surface_check.py | 1 + tests/unit/test_router/core/scripted_host.py | 34 +- tests/unit/test_router/core/test_selection.py | 37 + 28 files changed, 6682 insertions(+), 14 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/router/config.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/host.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/mod.rs create mode 100644 litellm-rust/crates/router-core/src/cache.rs create mode 100644 litellm-rust/crates/router-core/src/call.rs create mode 100644 litellm-rust/crates/router-core/src/config.rs create mode 100644 litellm-rust/crates/router-core/src/error.rs create mode 100644 litellm-rust/crates/router-core/src/executor.rs create mode 100644 litellm-rust/crates/router-core/src/health.rs create mode 100644 litellm-rust/crates/router-core/src/host.rs create mode 100644 litellm-rust/crates/router-core/src/registry.rs create mode 100644 litellm-rust/crates/router-core/src/select.rs create mode 100644 litellm-rust/crates/router-core/src/settings.rs create mode 100644 litellm-rust/crates/router-core/tests/call_loop.rs create mode 100644 litellm/router_rust_core.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f91deaaa401..65ed2ecb0d6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4331,6 +4331,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "async-trait", "aws-sdk-secretsmanager", "bytes", "criterion", @@ -4357,6 +4358,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-router-core", "litellm-router-types", "litellm-secret-value", "litellm-secrets", @@ -4439,6 +4441,23 @@ dependencies = [ [[package]] name = "litellm-router-core" version = "0.1.0" +dependencies = [ + "arc-swap", + "async-trait", + "litellm-cache", + "litellm-router-types", + "parking_lot", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "serde_with", + "sha2 0.10.9", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", +] [[package]] name = "litellm-router-types" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 75cac2a5d63..dda216c4dcf 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" } litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } +litellm-router-core = { path = "crates/router-core" } litellm-tracing = { path = "crates/tracing" } litellm-spend-clickhouse = { path = "crates/spend-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } @@ -78,6 +79,8 @@ litellm-cost = { path = "crates/cost" } litellm-model-catalog = { path = "crates/model-catalog" } litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } +arc-swap = "1.9.2" +async-trait = "0.1.91" askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..5a63ecd0004 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -25,6 +25,7 @@ litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true litellm-storage-clickhouse.workspace = true litellm-host.workspace = true +async-trait.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true @@ -49,6 +50,7 @@ litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "has litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true +litellm-router-core.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index fdcb969a423..7ecbc6907ed 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -12,6 +12,7 @@ mod marshal; mod model_registry; mod preflight; mod python_settings; +mod router; mod routes; mod secrets; mod tokenizer; @@ -36,6 +37,8 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::router::{RouterCall, RouterCoreHandle}; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -120,6 +123,8 @@ mod tests { "bundled_model_cost_map", "TokenCounter", "Tokenizer", + "RouterCoreHandle", + "RouterCall", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs new file mode 100644 index 00000000000..f0ee786712b --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -0,0 +1,294 @@ +//! The documents `RouterCoreHandle` reads: the router's static and dynamic config, a call's +//! request and an `apply` change set. A config key router-core does not read declines the +//! handle, so Python keeps the router on its own core. + +use std::collections::BTreeMap; + +use litellm_router_core::{ + ChangeSet, ConfigError, DeploymentSpec, RouterConfig, RouterDynamicConfig, RouterRequest, +}; +use serde::{Deserialize, de::DeserializeOwned}; +use serde_json::{Map, Value}; + +/// The `RouterConfig` keys router-core reads. +const STATIC_KEYS: &[&str] = &[ + "model_list", + "routing_strategy", + "routing_groups", + "plugins", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_policy", + "retry_after", + "timeout", + "request_timeout", + "drop_params", + "enable_pre_call_checks", + "enable_health_check_routing", + "provider_default_deployment_ids", +]; + +/// The `RouterDynamicConfig` keys router-core reads. +const DYNAMIC_KEYS: &[&str] = &[ + "version", + "model_alias_map", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "fallback_keys", +]; + +/// Router arguments that do not change routing, accepted in either config and dropped. +const IGNORED_KEYS: &[&str] = &[ + "set_verbose", + "debug_level", + "alerting_config", + "client_ttl", +]; + +/// Fallback lists Python treats as unset when empty (`context_window_fallbacks or None`). +const EMPTY_MEANS_UNSET: &[&str] = &["context_window_fallbacks", "content_policy_fallbacks"]; + +/// Which document a failure is about. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub(crate) enum Document { + #[strum(serialize = "router config")] + Static, + #[strum(serialize = "dynamic router config")] + Dynamic, + #[strum(serialize = "router request")] + Request, + #[strum(serialize = "change set")] + Changes, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Error { + #[error("the {document} must be a JSON object")] + NotObject { document: Document }, + #[error("router-core does not read `{key}` in the {document}")] + UnknownKey { document: Document, key: String }, + #[error("router-core does not run `{key}` = {value}")] + UnsupportedValue { key: &'static str, value: Value }, + #[error("the {document} does not have the shape router-core reads")] + Shape { + document: Document, + #[source] + source: serde_json::Error, + }, + #[error("the request names no model")] + MissingModel, + #[error("router-core rejected the router config")] + Core(#[from] ConfigError), +} + +impl Error { + /// Whether router-core cannot serve this config, so Python keeps its own core. + pub(crate) fn declined(&self) -> bool { + match self { + Self::UnknownKey { .. } | Self::UnsupportedValue { .. } | Self::Core(_) => true, + Self::Shape { document, .. } => { + matches!(document, Document::Static | Document::Dynamic) + } + Self::NotObject { .. } | Self::MissingModel => false, + } + } +} + +fn object(document: Document, value: Value) -> Result, Error> { + match value { + Value::Object(object) => Ok(object), + _ => Err(Error::NotObject { document }), + } +} + +fn shaped(document: Document, object: Map) -> Result { + serde_json::from_value(Value::Object(object)) + .map_err(|source| Error::Shape { document, source }) +} + +/// Drop unset and ignored keys, then reject any key outside `known`. +fn routing_keys( + document: Document, + value: Value, + known: &[&str], +) -> Result, Error> { + let mut object = object(document, value)?; + object.retain(|key, value| { + let empty_unset = EMPTY_MEANS_UNSET.contains(&key.as_str()) + && value.as_array().is_some_and(Vec::is_empty); + !(value.is_null() || empty_unset || IGNORED_KEYS.contains(&key.as_str())) + }); + if let Some(key) = object.keys().find(|key| !known.contains(&key.as_str())) { + return Err(Error::UnknownKey { + document, + key: key.clone(), + }); + } + Ok(object) +} + +pub(crate) fn router_config(value: Value) -> Result { + let object = routing_keys(Document::Static, value, STATIC_KEYS)?; + // Pre-call checks reach router-core as plugins; the Python flag alone has no native filter. + if let Some(value) = object.get("enable_pre_call_checks") + && value != &Value::Bool(false) + { + return Err(Error::UnsupportedValue { + key: "enable_pre_call_checks", + value: value.clone(), + }); + } + shaped(Document::Static, object) +} + +pub(crate) fn dynamic_config(value: Value) -> Result { + let object = routing_keys(Document::Dynamic, value, DYNAMIC_KEYS)?; + shaped(Document::Dynamic, object) +} + +/// A call's request as Python spells it. +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestSpec { + model: Option, + team_id: Option, + specific_deployment: Option, + num_retries: Option, + disable_fallbacks: Option, + router_weights: Option>>, + body: Option, +} + +/// Read a request; `model` replaces the request's own when given. +pub(crate) fn request(value: Value, model: Option<&str>) -> Result { + let spec: RequestSpec = shaped(Document::Request, object(Document::Request, value)?)?; + let model = model + .map(str::to_owned) + .or(spec.model) + .ok_or(Error::MissingModel)?; + Ok(RouterRequest { + model, + body: spec.body.unwrap_or_default(), + team_id: spec.team_id, + specific_deployment: spec.specific_deployment.unwrap_or(false), + num_retries: spec.num_retries, + disable_fallbacks: spec.disable_fallbacks.unwrap_or(false), + router_weights: spec.router_weights, + payload: None, + }) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ChangesSpec { + #[serde(default)] + upsert: Vec, + #[serde(default)] + delete: Vec, +} + +pub(crate) fn changes(value: Value) -> Result { + let spec: ChangesSpec = shaped(Document::Changes, object(Document::Changes, value)?)?; + Ok(ChangeSet { + upserts: spec.upsert, + deletes: spec.delete, + settings: None, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::unknown_key(json!({"routing_strategy_args": {"ttl": 1}}), "routing_strategy_args")] + #[case::dynamic_key_in_static(json!({"num_retries": 2}), "num_retries")] + fn an_unread_static_key_declines_naming_the_key(#[case] json: Value, #[case] key: &str) { + let error = router_config(json).unwrap_err(); + assert!(error.declined()); + assert!(matches!(&error, Error::UnknownKey { key: named, .. } if named == key)); + } + + #[rstest] + #[case::null_unknown(json!({"routing_strategy_args": null}))] + #[case::ignored(json!({"set_verbose": true, "debug_level": "INFO"}))] + #[case::pre_call_checks_off(json!({"enable_pre_call_checks": false}))] + fn unset_and_ignored_keys_are_accepted(#[case] json: Value) { + router_config(json).unwrap(); + } + + #[rstest] + fn pre_call_checks_without_a_plugin_decline() { + let error = router_config(json!({"enable_pre_call_checks": true})).unwrap_err(); + assert!(error.declined()); + assert!(matches!( + error, + Error::UnsupportedValue { + key: "enable_pre_call_checks", + .. + } + )); + } + + #[rstest] + #[case::not_object(json!([]), false)] + #[case::bad_shape(json!({"retry_after": "soon"}), true)] + fn a_config_that_is_not_an_object_is_a_value_error_and_a_wrong_shape_declines( + #[case] json: Value, + #[case] declined: bool, + ) { + assert_eq!(router_config(json).unwrap_err().declined(), declined); + } + + #[rstest] + fn python_router_defaults_read_as_unset() { + let dynamic = dynamic_config(json!({ + "num_retries": null, "fallbacks": [], "context_window_fallbacks": [], + "content_policy_fallbacks": [], "model_group_alias": {"gpt": "gpt-4"}, + })) + .unwrap(); + assert_eq!(dynamic.num_retries, None); + assert!(dynamic.context_window_fallbacks.is_none()); + assert!(dynamic.content_policy_fallbacks.is_none()); + assert_eq!(dynamic.model_alias_map["gpt"], "gpt-4"); + } + + #[rstest] + #[case::own_model(json!({"model": "gpt", "num_retries": 1}), None, "gpt")] + #[case::override_model(json!({"model": "gpt"}), Some("other"), "other")] + fn a_request_reads_its_model( + #[case] json: Value, + #[case] model: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(request(json, model).unwrap().model, expected); + } + + #[rstest] + #[case::unknown_field(json!({"model": "gpt", "temperature": 1}))] + #[case::no_model(json!({"team_id": "t"}))] + fn a_bad_request_is_a_value_error(#[case] json: Value) { + assert!(!request(json, None).unwrap_err().declined()); + } + + #[rstest] + fn a_change_set_reads_upserts_and_deletes() { + let changes = changes(json!({ + "upsert": [{"model_name": "m", "litellm_params": {"model": "openai/x"}}], + "delete": ["old"], + })) + .unwrap(); + assert_eq!(changes.upserts.len(), 1); + assert_eq!(changes.deletes, ["old"]); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs new file mode 100644 index 00000000000..9dc116ab854 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -0,0 +1,1011 @@ +//! The host behind `RouterCoreHandle`. Rust never awaits Python: each host operation the call +//! loop needs becomes an [`Ask`] parked in the call's [`Exchange`], and the call future stays +//! pending until Python runs it and hands the [`Answer`] back through [`Driver::resume`]. +//! +//! The driver polls the call future with a no-op waker while the call's exchange is bound to the +//! current thread, so a host future finds its exchange without the router passing it through. +//! +//! The routing cache is a host operation too: [`StepCache`] parks every [`RouterCache`] op as a +//! `cache` step, and Python runs it against the router's own `DualCache` (memory plus Redis), so +//! Python and Rust read and write the same cooldowns, fail counters and usage. + +use std::{ + cell::RefCell, + future::{Future, poll_fn}, + pin::Pin, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + task::{Context, Poll, Waker}, + time::Duration, +}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Deployment, Error, Host, Output, Payload, RouterCache, RouterEvent, + RouterRequest, RouterResponse, RoutingContext, SelectStage, +}; +use serde_json::Value; + +/// The `kind` of a step dict. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum StepKind { + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "sleep")] + Sleep, + #[strum(serialize = "allow_fallback")] + AllowFallback, + #[strum(serialize = "select")] + Select, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "done")] + Done, + #[strum(serialize = "failed")] + Failed, +} + +/// Which selection stage Python is asked to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum Stage { + /// Pre-call checks and sub-routers: return the deployments that remain. + Plugins, + /// A strategy router-core does not implement: return the candidates with the pick first. + Strategy(String), +} + +/// The `op` of a `cache` step: one per [`RouterCache`] method. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum CacheOp { + /// `async_get_cache(key)`: answered with the stored JSON value, or `None`. + #[strum(serialize = "get")] + Get, + /// `async_set_cache(key, value, ttl)`: answered with `None`. + #[strum(serialize = "set")] + Set, + /// `async_increment(key, value, ttl)`: answered with the new number. The TTL applies when + /// the key is created. + #[strum(serialize = "increment")] + Increment, + /// `async_batch_get_cache(keys)`: answered with one JSON value or `None` per key, in order. + #[strum(serialize = "batch_get")] + BatchGet, + /// `async_delete_cache(key)`: answered with `None`. + #[strum(serialize = "delete")] + Delete, +} + +/// Python's outcome of a `cache` step. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum CacheOutcome { + /// `get`: the stored value (`None` when missing). + Value(Option), + /// `batch_get`: one entry per key. + Values(Vec>), + /// `increment`: the counter after the increment. + Number(f64), + /// `set` and `delete`. + Stored, +} + +impl CacheOutcome { + /// The op this outcome answers (`Stored` answers both `set` and `delete`). + fn answers(&self, op: CacheOp) -> bool { + matches!( + (self, op), + (Self::Value(_), CacheOp::Get) + | (Self::Values(_), CacheOp::BatchGet) + | (Self::Number(_), CacheOp::Increment) + | (Self::Stored, CacheOp::Set | CacheOp::Delete) + ) + } +} + +/// What a call is blocked on, for Python to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum Ask { + Attempt { + deployment_id: String, + /// The name the attempt resolved: the requested name or the fallback target. + model: String, + model_group: String, + /// The `order` level of a same-group order fallback. + target_order: Option, + /// Every deployment `model` resolved to for this caller, the picked one included. + candidate_ids: Vec, + retry: u32, + fallback_depth: u32, + fallback_level: u32, + timeout: Option, + }, + Sleep { + seconds: f64, + }, + AllowFallback { + target: String, + }, + Select { + stage: Stage, + model_group: String, + candidates: Vec, + }, + /// One [`RouterCache`] op against the host's cache. `key` is set for every op but + /// `batch_get`, which sets `keys`; `value` is set for `set` (the JSON to store) and + /// `increment` (the amount); `ttl` is in seconds. + Cache { + op: CacheOp, + key: Option, + keys: Option>, + value: Option, + ttl: Option, + }, +} + +impl Ask { + pub(crate) fn kind(&self) -> StepKind { + match self { + Self::Attempt { .. } => StepKind::Attempt, + Self::Sleep { .. } => StepKind::Sleep, + Self::AllowFallback { .. } => StepKind::AllowFallback, + Self::Select { .. } => StepKind::Select, + Self::Cache { .. } => StepKind::Cache, + } + } +} + +/// Python's answer to an [`Ask`]. +#[derive(Debug)] +pub(crate) enum Answer { + /// The attempt's total tokens, or its failure. + Attempt(Result), + Slept, + AllowFallback(bool), + /// The candidate ids that remain, pick first. + Select(Vec), + Cache(CacheOutcome), +} + +impl Answer { + fn kind(&self) -> StepKind { + match self { + Self::Attempt(_) => StepKind::Attempt, + Self::Slept => StepKind::Sleep, + Self::AllowFallback(_) => StepKind::AllowFallback, + Self::Select(_) => StepKind::Select, + Self::Cache(_) => StepKind::Cache, + } + } + + fn check(&self, ask: &Ask) -> Result<(), DriveError> { + if self.kind() != ask.kind() { + return Err(DriveError::WrongAnswer { + expected: ask.kind(), + got: self.kind(), + }); + } + if let (Self::Select(ids), Ask::Select { candidates, .. }) = (self, ask) + && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) + { + return Err(DriveError::UnknownCandidate { id: id.clone() }); + } + if let (Self::Cache(outcome), Ask::Cache { op, keys, .. }) = (self, ask) { + if !outcome.answers(*op) { + return Err(DriveError::WrongCacheOutcome { op: *op }); + } + if let (CacheOutcome::Values(values), Some(keys)) = (outcome, keys) + && values.len() != keys.len() + { + return Err(DriveError::BatchLength { + expected: keys.len(), + got: values.len(), + }); + } + } + Ok(()) + } +} + +/// How a driven future ended well. +#[derive(Debug)] +pub(crate) enum Finish { + /// `start`: the call's response. + Call(RouterResponse), + /// `select`: the deployment picked and the model group it serves. + Selected { + deployment_id: String, + model_group: String, + }, +} + +/// One step of a call: what it is blocked on, or how it ended. +#[derive(Debug)] +pub(crate) enum Step { + Ask(Ask), + Done(Finish), + Failed(Error), +} + +/// A misuse of the step protocol by the caller. +#[derive(Debug, thiserror::Error)] +pub(crate) enum DriveError { + #[error("the router call already finished")] + Finished, + #[error("the router call has not asked for anything yet, so the outcome must be None")] + Unasked, + #[error("the router call is waiting for an outcome of its `{expected}` step")] + MissingAnswer { expected: StepKind }, + #[error("the router call is waiting for an outcome of its `{expected}` step, not `{got}`")] + WrongAnswer { expected: StepKind, got: StepKind }, + #[error("deployment `{id}` was not one of the select step's candidates")] + UnknownCandidate { id: String }, + #[error("the outcome does not answer a `{op}` cache step")] + WrongCacheOutcome { op: CacheOp }, + #[error("a `batch_get` cache step for {expected} keys was answered with {got} values")] + BatchLength { expected: usize, got: usize }, + #[error("the router call is pending on something other than a step")] + Stalled, +} + +impl DriveError { + /// Whether the caller passed a bad outcome, rather than driving a call that cannot go on. + pub(crate) fn is_request(&self) -> bool { + match self { + Self::Unasked + | Self::MissingAnswer { .. } + | Self::WrongAnswer { .. } + | Self::UnknownCandidate { .. } + | Self::WrongCacheOutcome { .. } + | Self::BatchLength { .. } => true, + Self::Finished | Self::Stalled => false, + } + } +} + +/// The question a call is blocked on and the answer that releases it. +#[derive(Default)] +struct Exchange { + asked: Option, + answer: Option, + /// Events the call reported since the last step, handed to Python with the next one. + events: Vec, +} + +type Shared = Arc>; + +fn lock(exchange: &Mutex) -> MutexGuard<'_, Exchange> { + exchange.lock().unwrap_or_else(PoisonError::into_inner) +} + +thread_local! { + static CURRENT: RefCell> = const { RefCell::new(None) }; +} + +/// Binds an exchange to this thread for one poll and restores the previous one after. +struct Bound(Option); + +impl Bound { + fn new(exchange: &Shared) -> Self { + Self(CURRENT.replace(Some(exchange.clone()))) + } +} + +impl Drop for Bound { + fn drop(&mut self) { + CURRENT.set(self.0.take()); + } +} + +/// Park `question` in the bound exchange and wait for its answer. +async fn ask(question: Ask) -> Answer { + let mut question = Some(question); + poll_fn(move |_| { + CURRENT.with_borrow(|current| { + let Some(exchange) = current else { + return Poll::Pending; + }; + let mut exchange = lock(exchange); + if let Some(question) = question.take() { + exchange.asked = Some(question); + return Poll::Pending; + } + exchange.answer.take().map_or(Poll::Pending, Poll::Ready) + }) + }) + .await +} + +const CHECKED: &str = "the driver checks each answer against its ask"; + +fn ttl_seconds(ttl: Option) -> Option { + ttl.map(|ttl| ttl.as_secs_f64()) +} + +/// An increment amount as Python's counters store it: an int when it is whole. +fn amount_value(amount: f64) -> Value { + if amount.fract() == 0.0 && amount.abs() < i64::MAX as f64 { + Value::from(amount as i64) + } else { + Value::from(amount) + } +} + +/// The host's cache: every op is a `cache` step Python runs against the router's `DualCache`. +pub(crate) struct StepCache; + +async fn cache_step( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, +) -> CacheOutcome { + let question = Ask::Cache { + op, + key: key.map(str::to_owned), + keys: keys.map(<[String]>::to_vec), + value, + ttl: ttl_seconds(ttl), + }; + let Answer::Cache(outcome) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + outcome +} + +#[async_trait] +impl RouterCache for StepCache { + async fn async_get_cache(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + cache_step(CacheOp::Get, Some(key), None, None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Set, Some(key), None, Some(value), ttl).await; + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + cache_step(CacheOp::Increment, Some(key), None, value, ttl).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } + + async fn async_batch_get_cache( + &self, + keys: &[String], + ) -> Result>, litellm_cache::Error> { + let CacheOutcome::Values(values) = + cache_step(CacheOp::BatchGet, None, Some(keys), None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Delete, Some(key), None, None, None).await; + Ok(()) + } +} + +/// The [`Host`] every `RouterCoreHandle` router runs on. +pub(crate) struct StepHost; + +#[async_trait] +impl Host for StepHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + let question = Ask::Attempt { + deployment_id: attempt.deployment.id.clone(), + model: attempt.name.clone(), + model_group: attempt.model_group.clone(), + target_order: attempt.target_order, + candidate_ids: attempt.candidate_ids.clone(), + retry: attempt.retry, + fallback_depth: attempt.fallback_depth, + fallback_level: attempt.fallback_level, + timeout: attempt.timeout.map(|timeout| timeout.as_secs_f64()), + }; + let Answer::Attempt(result) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + result.map(|total_tokens| Output { + value: Payload(Arc::new(())), + total_tokens, + }) + } + + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + let question = Ask::Select { + stage: match stage { + SelectStage::Plugins => Stage::Plugins, + SelectStage::Strategy(strategy) => Stage::Strategy(strategy.to_string()), + }, + model_group: ctx.model_group.to_owned(), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let Answer::Select(ids) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + Ok(ids + .iter() + .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) + .collect()) + } + + async fn allow_fallback(&self, _request: &RouterRequest, target: &str) -> bool { + let question = Ask::AllowFallback { + target: target.to_owned(), + }; + let Answer::AllowFallback(allowed) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + allowed + } + + /// A zero sleep is asked too: Python's retry loop awaits `asyncio.sleep(0)` between sibling + /// retries, which yields to the event loop. + async fn sleep(&self, duration: Duration) { + let question = Ask::Sleep { + seconds: duration.as_secs_f64(), + }; + let Answer::Slept = ask(question).await else { + unreachable!("{CHECKED}"); + }; + } + + /// Kept for Python only when it acts on it: an exhausted retry loop is stamped on the + /// error Python raises, and a fallback hop tells it which target a failed walk ended on. + fn observe(&self, event: &RouterEvent) { + if !matches!( + event, + RouterEvent::RetriesExhausted { .. } | RouterEvent::Fallback { .. } + ) { + return; + } + CURRENT.with_borrow(|current| { + if let Some(exchange) = current { + lock(exchange).events.push(event.clone()); + } + }); + } + + fn cache(&self) -> Arc { + Arc::new(StepCache) + } + + /// Each attempt runs through the Python router, whose callbacks record cooldowns, + /// allowed-fails counters and usage. + fn records_outcomes(&self) -> bool { + true + } +} + +type CallFuture = Pin> + Send>>; + +/// Poll `future` once with `exchange` bound. +fn poll_bound(future: Pin<&mut F>, exchange: &Shared) -> Poll { + let _bound = Bound::new(exchange); + future.poll(&mut Context::from_waker(Waker::noop())) +} + +/// Drives one router call a step at a time. +pub(crate) struct Driver { + future: Option, + exchange: Shared, + waiting: Option, +} + +impl Driver { + pub(crate) fn new(call: impl Future> + Send + 'static) -> Self { + Self { + future: Some(Box::pin(call)), + exchange: Shared::default(), + waiting: None, + } + } + + /// The step the call is blocked on, if it asked for one. + pub(crate) fn waiting(&self) -> Option<&Ask> { + self.waiting.as_ref() + } + + /// The events reported since the last call. + pub(crate) fn take_events(&self) -> Vec { + std::mem::take(&mut lock(&self.exchange).events) + } + + /// Hand back the answer to the waiting step (`None` on the first call) and run the call to + /// its next step or its end. + pub(crate) fn resume(&mut self, answer: Option) -> Result { + let Some(future) = self.future.as_mut() else { + return Err(DriveError::Finished); + }; + match (&self.waiting, answer) { + (None, None) => {} + (None, Some(_)) => return Err(DriveError::Unasked), + (Some(ask), None) => { + return Err(DriveError::MissingAnswer { + expected: ask.kind(), + }); + } + (Some(ask), Some(answer)) => { + answer.check(ask)?; + lock(&self.exchange).answer = Some(answer); + } + } + self.waiting = None; + match poll_bound(future.as_mut(), &self.exchange) { + Poll::Ready(result) => { + self.future = None; + Ok(match result { + Ok(finish) => Step::Done(finish), + Err(error) => Step::Failed(error), + }) + } + Poll::Pending => { + let Some(ask) = lock(&self.exchange).asked.take() else { + self.future = None; + return Err(DriveError::Stalled); + }; + self.waiting = Some(ask.clone()); + Ok(Step::Ask(ask)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + time::{SystemTime, UNIX_EPOCH}, + }; + + use litellm_router_core::{ + FailureKind, Op, Router, RouterConfig, RouterDynamicConfig, SeededRandomness, + }; + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + fn router(config: Value, dynamic: Value) -> Router { + let config: RouterConfig = serde_json::from_value(config).unwrap(); + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, Arc::new(StepHost)) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() + } + + fn deployment(group: &str, id: &str) -> Value { + json!({"model_name": group, "litellm_params": {"model": format!("openai/{id}")}, "model_info": {"id": id}}) + } + + fn start(router: &Router, model: &str) -> Driver { + let router = router.clone(); + let request = RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + }; + Driver::new(async move { router.call(Op::Completion, request).await.map(Finish::Call) }) + } + + fn rate_limited() -> Answer { + Answer::Attempt(Err(AttemptError::new(FailureKind::RateLimit, Some(429)))) + } + + /// A dict-backed stand-in for Python's `DualCache` that answers `cache` steps and keeps + /// every one it answered. + #[derive(Default)] + struct Store { + entries: HashMap, + asked: Vec, + } + + impl Store { + fn answer(&mut self, ask: &Ask) -> Answer { + let Ask::Cache { + op, + key, + keys, + value, + .. + } = ask + else { + panic!("not a cache step: {ask:?}"); + }; + self.asked.push(ask.clone()); + let key = key.clone().unwrap_or_default(); + Answer::Cache(match op { + CacheOp::Get => CacheOutcome::Value(self.entries.get(&key).cloned()), + CacheOp::Set => { + self.entries.insert(key, value.clone().unwrap()); + CacheOutcome::Stored + } + CacheOp::Increment => { + let current = self.entries.get(&key).and_then(Value::as_f64); + let next = + current.unwrap_or(0.0) + value.as_ref().and_then(Value::as_f64).unwrap(); + self.entries.insert(key, amount_value(next)); + CacheOutcome::Number(next) + } + CacheOp::BatchGet => CacheOutcome::Values( + keys.iter() + .flatten() + .map(|key| self.entries.get(key).cloned()) + .collect(), + ), + CacheOp::Delete => { + self.entries.remove(&key); + CacheOutcome::Stored + } + }) + } + + /// Resume with `answer`, answering every cache step on the way, up to the next other + /// step. + fn resume(&mut self, driver: &mut Driver, answer: Option) -> Step { + let mut step = driver.resume(answer).unwrap(); + while let Step::Ask(ask @ Ask::Cache { .. }) = &step { + let answer = self.answer(ask); + step = driver.resume(Some(answer)).unwrap(); + } + step + } + } + + #[fixture] + fn pair() -> Router { + router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({}), + ) + } + + #[rstest] + fn a_first_attempt_success_is_done(pair: Router) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let Step::Ask(Ask::Attempt { + deployment_id, + model_group, + retry: 0, + fallback_depth: 0, + .. + }) = store.resume(&mut driver, None) + else { + panic!("the call should start with an attempt"); + }; + assert_eq!(model_group, "primary"); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(Answer::Attempt(Ok(12)))) + else { + panic!("a successful attempt should finish the call"); + }; + assert_eq!(response.deployment_id, deployment_id); + assert_eq!(response.output.total_tokens, 12); + assert_eq!( + (response.attempted_retries, response.attempted_fallbacks), + (0, 0) + ); + assert!(matches!(driver.resume(None), Err(DriveError::Finished))); + } + + #[rstest] + fn cooldown_reads_are_cache_steps(pair: Router) { + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::Ask(Ask::Cache { + op: CacheOp::BatchGet, + key: None, + keys: Some(keys), + value: None, + ttl: None, + }) = step + else { + panic!("the call should read cooldowns first, got {step:?}"); + }; + assert_eq!(keys, ["deployment:a1:cooldown", "deployment:a2:cooldown"]); + } + + #[rstest] + #[case::lone_deployment_backs_off(vec![deployment("primary", "a1")], true)] + #[case::sibling_retries_at_once( + vec![deployment("primary", "a1"), deployment("primary", "a2")], + false + )] + fn a_rate_limit_retries_after_the_backoff( + #[case] model_list: Vec, + #[case] backs_off: bool, + ) { + let router = router(json!({"model_list": model_list}), json!({"num_retries": 1})); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let step = store.resume(&mut driver, Some(rate_limited())); + let Step::Ask(Ask::Sleep { seconds }) = step else { + panic!("a retry should sleep first, got {step:?}"); + }; + assert_eq!(seconds > 0.0, backs_off, "slept {seconds}s"); + let step = store.resume(&mut driver, Some(Answer::Slept)); + assert!( + matches!(step, Step::Ask(Ask::Attempt { retry: 1, .. })), + "expected the retry, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(Answer::Attempt(Ok(1)))) + else { + panic!("the retry should finish the call"); + }; + assert_eq!(response.attempted_retries, 1); + } + + #[rstest] + fn a_cooldown_python_wrote_is_read_back_by_the_next_call() { + let pair_without_retries = router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + store.entries.insert( + "deployment:a1:cooldown".into(), + json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}), + ); + for _ in 0..8 { + let mut call = start(&pair_without_retries, "primary"); + let Step::Ask(Ask::Attempt { deployment_id, .. }) = store.resume(&mut call, None) + else { + panic!("the call should attempt the healthy deployment"); + }; + assert_eq!(deployment_id, "a2", "the cooled-down deployment was picked"); + } + } + + /// Each attempt runs `PythonCore.call`, which records its own cooldowns, fail counters and + /// usage, so the Rust call loop writes none of them. + #[rstest] + #[case::failure(rate_limited())] + #[case::success(Answer::Attempt(Ok(7)))] + fn attempt_outcomes_are_recorded_by_the_python_attempt(#[case] outcome: Answer) { + let tracked = |id: &str| { + json!({ + "model_name": "primary", + "litellm_params": {"model": format!("openai/{id}"), "tpm": 100, "cooldown_time": 7}, + "model_info": {"id": id, "allowed_fails": 3}, + }) + }; + let router = router( + json!({"model_list": [tracked("a1"), tracked("a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let step = store.resume(&mut driver, Some(outcome)); + assert!( + matches!(step, Step::Done(_) | Step::Failed(_)), + "got {step:?}" + ); + let writes: Vec<&Ask> = store + .asked + .iter() + .filter(|ask| { + matches!(ask, Ask::Cache { op, .. } if !matches!(op, CacheOp::Get | CacheOp::BatchGet)) + }) + .collect(); + assert!(writes.is_empty(), "{writes:?}"); + } + + #[rstest] + fn a_failed_group_falls_back_once_python_allows_it() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let bad_request = + Answer::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + let step = store.resume(&mut driver, Some(bad_request)); + assert!( + matches!(&step, Step::Ask(Ask::AllowFallback { target }) if target == "backup"), + "expected the fallback check, got {step:?}" + ); + let step = store.resume(&mut driver, Some(Answer::AllowFallback(true))); + assert!( + matches!(&step, Step::Ask(Ask::Attempt { deployment_id, fallback_depth: 1, .. }) if deployment_id == "b1"), + "expected the fallback attempt, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(Answer::Attempt(Ok(3)))) + else { + panic!("the fallback should finish the call"); + }; + assert_eq!(response.model_group, "backup"); + assert_eq!(response.attempted_fallbacks, 1); + } + + #[rstest] + fn a_refused_fallback_fails_with_the_attempt_error() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let bad_request = + Answer::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + store.resume(&mut driver, Some(bad_request)); + let Step::Failed(error) = store.resume(&mut driver, Some(Answer::AllowFallback(false))) + else { + panic!("a refused fallback should fail the call"); + }; + assert_eq!( + error.attempt().map(|a| a.kind), + Some(FailureKind::BadRequest) + ); + } + + #[rstest] + #[case::none_while_waiting(None, StepKind::Attempt)] + #[case::wrong_kind(Some(Answer::Slept), StepKind::Attempt)] + #[case::cache_outcome(Some(Answer::Cache(CacheOutcome::Stored)), StepKind::Attempt)] + fn a_mismatched_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] answer: Option, + #[case] expected: StepKind, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + store.resume(&mut driver, None); + let error = driver.resume(answer).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.waiting().map(Ask::kind), Some(expected)); + assert!(matches!( + store.resume(&mut driver, Some(Answer::Attempt(Ok(1)))), + Step::Done(_) + )); + } + + #[rstest] + #[case::wrong_op(Answer::Cache(CacheOutcome::Stored))] + #[case::number_for_batch(Answer::Cache(CacheOutcome::Number(1.0)))] + #[case::short_batch(Answer::Cache(CacheOutcome::Values(vec![None])))] + #[case::not_a_cache_outcome(Answer::Slept)] + fn a_mismatched_cache_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] answer: Answer, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::Ask( + ask @ Ask::Cache { + op: CacheOp::BatchGet, + .. + }, + ) = &step + else { + panic!("expected the cooldown read, got {step:?}"); + }; + let error = driver.resume(Some(answer)).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.waiting(), Some(ask)); + let answer = store.answer(ask); + assert_eq!( + store.resume(&mut driver, Some(answer)).unwrap_kind(), + StepKind::Attempt + ); + } + + #[rstest] + fn a_strategy_python_runs_asks_for_a_selection_and_keeps_the_pick() { + let router = router( + json!({ + "model_list": [deployment("primary", "a1"), deployment("primary", "a2")], + "routing_strategy": "least-busy", + }), + json!({}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + let step = store.resume(&mut driver, None); + let Step::Ask(Ask::Select { + stage: Stage::Strategy(strategy), + candidates, + .. + }) = step + else { + panic!("a host strategy should ask for a selection, got {step:?}"); + }; + assert_eq!(strategy, "least-busy"); + assert_eq!(candidates, ["a1", "a2"]); + let step = store.resume(&mut driver, Some(Answer::Select(vec!["a2".to_owned()]))); + assert!( + matches!(&step, Step::Ask(Ask::Attempt { deployment_id, .. }) if deployment_id == "a2"), + "expected the picked deployment, got {step:?}" + ); + } + + #[rstest] + fn a_selection_is_driven_through_cache_steps(pair: Router) { + let request = RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + }; + let mut store = Store::default(); + let mut driver = Driver::new(async move { + pair.select(&Op::Completion, &request) + .await + .map(|deployment| Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + }) + }); + let Step::Done(Finish::Selected { + deployment_id, + model_group, + }) = store.resume(&mut driver, None) + else { + panic!("a selection should finish with the pick"); + }; + assert!(["a1", "a2"].contains(&deployment_id.as_str())); + assert_eq!(model_group, "primary"); + assert!(matches!( + store.asked.as_slice(), + [Ask::Cache { + op: CacheOp::BatchGet, + .. + }] + )); + } + + trait StepKindOf { + fn unwrap_kind(self) -> StepKind; + } + + impl StepKindOf for Step { + fn unwrap_kind(self) -> StepKind { + match self { + Step::Ask(ask) => ask.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + } + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs new file mode 100644 index 00000000000..992dc58d546 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -0,0 +1,549 @@ +//! `RouterCoreHandle`: router-core behind a step-driven handle. Python drives each call through +//! [`RouterCall::next`], running every attempt, sleep, fallback check, selection and cache op the +//! call loop asks for and handing the outcome back; Rust never awaits Python. `select` is driven +//! the same way, since reading cooldowns is a cache step. + +mod config; +mod host; + +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_error::Report; +use litellm_host_python::{from_py_argument, json_loads, to_py}; +use litellm_router_core::{AttemptError, Error, FailureKind, Op, Router, RouterEvent}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::{PyBool, PyDict, PyList, PyString}, +}; +use serde_json::Value; + +use self::host::{ + Answer, Ask, CacheOp, CacheOutcome, DriveError, Driver, Finish, Stage, Step, StepHost, StepKind, +}; +use crate::errors::RustBridgeDeclined; + +/// The `error_kind` of a `failed` step. +#[derive(Clone, Copy, Debug, strum::IntoStaticStr)] +enum ErrorKind { + #[strum(serialize = "config")] + Config, + #[strum(serialize = "unknown_model")] + UnknownModel, + #[strum(serialize = "no_healthy_deployments")] + NoHealthyDeployments, + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "retired")] + Retired, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + match error { + Error::Config(_) => Self::Config, + Error::UnknownModel { .. } => Self::UnknownModel, + Error::NoHealthyDeployments { .. } => Self::NoHealthyDeployments, + Error::Attempt(_) => Self::Attempt, + Error::Cache(_) => Self::Cache, + Error::Retired => Self::Retired, + } + } +} + +fn config_error_to_pyerr(error: config::Error) -> PyErr { + let message = Report::new(&error).to_string(); + if error.declined() { + RustBridgeDeclined::new_err(message) + } else { + PyValueError::new_err(message) + } +} + +fn drive_error_to_pyerr(error: DriveError) -> PyErr { + let message = Report::new(&error).to_string(); + if error.is_request() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) + } +} + +/// A document Python passes as a mapping or as its JSON text. +fn document(value: &Bound<'_, PyAny>) -> PyResult { + match value.cast::() { + Ok(text) => { + let parsed = json_loads(value.py(), text.to_str()?.as_bytes())?; + from_py_argument(parsed.bind(value.py())) + } + Err(_) => from_py_argument(value), + } +} + +/// A router-core router whose host operations Python runs. +#[pyclass(frozen)] +pub(crate) struct RouterCoreHandle { + router: Router, +} + +#[pymethods] +impl RouterCoreHandle { + #[new] + fn new(config: &Bound<'_, PyAny>, dynamic: &Bound<'_, PyAny>) -> PyResult { + let (config, dynamic) = (document(config)?, document(dynamic)?); + let build = || { + let config = config::router_config(config)?; + let dynamic = config::dynamic_config(dynamic)?; + Ok(Router::new(config, dynamic, Arc::new(StepHost))?) + }; + build() + .map(|router| Self { router }) + .map_err(config_error_to_pyerr) + } + + /// Start a call; nothing runs until the first `next(None)`. + fn start(&self, op: &str, request: &Bound<'_, PyAny>) -> PyResult { + let op: Op = op.parse().unwrap_or_else(|_| Op::Other(op.to_owned())); + let request = config::request(document(request)?, None).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router.call(op, request).await.map(Finish::Call) + })) + } + + /// Pick the deployment `model` routes to now, without calling it. Driven like `start`: its + /// steps are the cache reads and selections the pick needs, and it ends with `done` + /// (`deployment_id`, `model_group`) or `failed`. + fn select(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = + config::request(document(request)?, Some(model)).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let deployment = router.select(&Op::Completion, &request).await?; + Ok(Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + }) + })) + } + + /// Apply `{"upsert": [deployments], "delete": [ids]}` and return the new registry version. + fn apply(&self, changes: &Bound<'_, PyAny>) -> PyResult { + let changes = config::changes(document(changes)?).map_err(config_error_to_pyerr)?; + self.router + .apply(changes) + .map_err(|error| config_error_to_pyerr(error.into())) + } + + fn registry_version(&self) -> u64 { + self.router.registry().version() + } + + fn retire(&self) { + self.router.retire(); + } +} + +/// One router call, driven a step at a time by Python. +#[pyclass(frozen)] +pub(crate) struct RouterCall { + driver: Mutex, +} + +impl RouterCall { + fn new( + future: impl std::future::Future> + Send + 'static, + ) -> Self { + Self { + driver: Mutex::new(Driver::new(future)), + } + } +} + +#[pymethods] +impl RouterCall { + /// Hand back the outcome of the last step (`None` on the first call) and return the next + /// step, or the call's end. + #[pyo3(signature = (outcome))] + fn next<'py>( + &self, + py: Python<'py>, + outcome: Option>, + ) -> PyResult> { + let mut driver = self.driver.lock().unwrap_or_else(PoisonError::into_inner); + let answer = match driver.waiting() { + None => match outcome { + None => None, + Some(_) => return Err(drive_error_to_pyerr(DriveError::Unasked)), + }, + Some(ask) => Some(answer(ask, outcome)?), + }; + let step = driver.resume(answer).map_err(drive_error_to_pyerr)?; + let dict = step_dict(py, step)?; + dict.set_item("events", events_list(py, driver.take_events())?)?; + Ok(dict) + } +} + +/// Read Python's outcome for `ask`. +fn answer(ask: &Ask, outcome: Option>) -> PyResult { + let missing = || { + drive_error_to_pyerr(DriveError::MissingAnswer { + expected: ask.kind(), + }) + }; + match ask { + Ask::Sleep { .. } => Ok(Answer::Slept), + Ask::AllowFallback { .. } => Ok(Answer::AllowFallback( + outcome.ok_or_else(missing)?.extract()?, + )), + Ask::Select { .. } => Ok(Answer::Select(outcome.ok_or_else(missing)?.extract()?)), + Ask::Attempt { .. } => { + let outcome = outcome.ok_or_else(missing)?; + attempt_outcome(outcome.cast::()?).map(Answer::Attempt) + } + Ask::Cache { op, .. } => cache_outcome(*op, outcome).map(Answer::Cache), + } +} + +/// `get`: the JSON value or None; `batch_get`: a list of JSON values or None, one per key; +/// `increment`: the new int or float; `set` and `delete`: None. +fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult { + let outcome = outcome.filter(|outcome| !outcome.is_none()); + let json = |value: Option| value.filter(|value| !value.is_null()); + match op { + CacheOp::Get => outcome + .map(|outcome| from_py_argument::(&outcome)) + .transpose() + .map(|value| CacheOutcome::Value(json(value))), + CacheOp::BatchGet => { + let outcome = outcome.ok_or_else(|| { + PyValueError::new_err("a `batch_get` cache outcome must be a list, not None") + })?; + if !outcome.is_instance_of::() { + return Err(PyValueError::new_err( + "a `batch_get` cache outcome must be a list with one entry per key", + )); + } + let values: Vec> = from_py_argument(&outcome)?; + Ok(CacheOutcome::Values(values.into_iter().map(json).collect())) + } + CacheOp::Increment => outcome + .filter(|outcome| !outcome.is_instance_of::()) + .and_then(|outcome| outcome.extract::().ok()) + .map(CacheOutcome::Number) + .ok_or_else(|| { + PyValueError::new_err("an `increment` cache outcome must be the new int or float") + }), + CacheOp::Set | CacheOp::Delete => match outcome { + None => Ok(CacheOutcome::Stored), + Some(_) => Err(PyValueError::new_err(format!( + "a `{op}` cache step has no outcome, so pass None" + ))), + }, + } +} + +/// `{"ok": True, "total_tokens": int}` or `{"ok": False, "kind": str, "status": int | None, +/// "retry_after": float | None, "num_retries": int | None, "detail": str | None}`. +fn attempt_outcome(outcome: &Bound<'_, PyDict>) -> PyResult> { + fn item<'py, T: FromPyObjectOwned<'py>>( + outcome: &Bound<'py, PyDict>, + key: &str, + ) -> PyResult> { + match outcome.get_item(key)? { + Some(value) if !value.is_none() => Ok(Some(value.extract().map_err(Into::into)?)), + Some(_) | None => Ok(None), + } + } + let ok: bool = item(outcome, "ok")? + .ok_or_else(|| PyValueError::new_err("an attempt outcome needs `ok`"))?; + if ok { + return Ok(Ok(item(outcome, "total_tokens")?.unwrap_or(0))); + } + let kind: String = item(outcome, "kind")? + .ok_or_else(|| PyValueError::new_err("a failed attempt outcome needs `kind`"))?; + let kind: FailureKind = kind + .parse() + .map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind")))?; + let retry_after = item::(outcome, "retry_after")? + .map(|seconds| { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("`retry_after` must be a non-negative number")) + }) + .transpose()?; + let mut failure = AttemptError::new(kind, item(outcome, "status")?); + failure.retry_after = retry_after; + failure.num_retries = item(outcome, "num_retries")?; + failure.detail = item::(outcome, "detail")?.map(Into::into); + Ok(Err(failure)) +} + +/// `[{"kind": "retries_exhausted", "model_group": str, "attempted": int, "max": int} +/// | {"kind": "fallback", "to": str, "order": int | None, "level": int}]`. +fn events_list(py: Python<'_>, events: Vec) -> PyResult> { + let list = PyList::empty(py); + for event in events { + let dict = PyDict::new(py); + match event { + RouterEvent::RetriesExhausted { + model_group, + attempted, + max, + } => { + dict.set_item("kind", "retries_exhausted")?; + dict.set_item("model_group", model_group)?; + dict.set_item("attempted", attempted)?; + dict.set_item("max", max)?; + } + RouterEvent::Fallback { + to, order, level, .. + } => { + dict.set_item("kind", "fallback")?; + dict.set_item("to", to)?; + dict.set_item("order", order)?; + dict.set_item("level", level)?; + } + _ => continue, + } + list.append(dict)?; + } + Ok(list) +} + +fn step_dict(py: Python<'_>, step: Step) -> PyResult> { + let dict = PyDict::new(py); + let kind = match &step { + Step::Ask(ask) => ask.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + }; + dict.set_item("kind", <&'static str>::from(kind))?; + match step { + Step::Ask(ask) => ask_items(py, &dict, ask)?, + Step::Done(finish) => done_items(&dict, finish)?, + Step::Failed(error) => { + let attempt = error.attempt(); + dict.set_item("error", Report::new(&error).to_string())?; + dict.set_item("error_kind", <&'static str>::from(ErrorKind::from(&error)))?; + dict.set_item("attempt_failed", attempt.is_some())?; + dict.set_item("status", attempt.and_then(|failure| failure.status))?; + dict.set_item( + "failure_kind", + attempt.map(|failure| <&'static str>::from(failure.kind)), + )?; + } + } + Ok(dict) +} + +fn ask_items(py: Python<'_>, dict: &Bound<'_, PyDict>, ask: Ask) -> PyResult<()> { + match ask { + Ask::Attempt { + deployment_id, + model, + model_group, + target_order, + candidate_ids, + retry, + fallback_depth, + fallback_level, + timeout, + } => { + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("model", model)?; + dict.set_item("model_group", model_group)?; + dict.set_item("target_order", target_order)?; + dict.set_item("candidate_ids", PyList::new(py, candidate_ids)?)?; + dict.set_item("retry", retry)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("fallback_level", fallback_level)?; + dict.set_item("timeout", timeout) + } + Ask::Sleep { seconds } => dict.set_item("seconds", seconds), + Ask::AllowFallback { target } => dict.set_item("target", target), + Ask::Select { + stage, + model_group, + candidates, + } => { + let (stage, strategy) = match stage { + Stage::Plugins => ("plugins", None), + Stage::Strategy(strategy) => ("strategy", Some(strategy)), + }; + dict.set_item("stage", stage)?; + dict.set_item("strategy", strategy)?; + dict.set_item("model_group", model_group)?; + dict.set_item("candidates", PyList::new(py, candidates)?) + } + Ask::Cache { + op, + key, + keys, + value, + ttl, + } => { + dict.set_item("op", <&'static str>::from(op))?; + dict.set_item("key", key)?; + dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?; + dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?; + dict.set_item("ttl", ttl) + } + } +} + +fn done_items(dict: &Bound<'_, PyDict>, finish: Finish) -> PyResult<()> { + match finish { + Finish::Call(response) => { + dict.set_item("deployment_id", response.deployment_id)?; + dict.set_item("model_group", response.model_group)?; + dict.set_item("attempted_retries", response.attempted_retries)?; + dict.set_item("max_retries", response.max_retries)?; + dict.set_item("attempted_fallbacks", response.attempted_fallbacks)?; + dict.set_item("total_tokens", response.output.total_tokens) + } + Finish::Selected { + deployment_id, + model_group, + } => { + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("model_group", model_group) + } + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ffi::c_str, prelude::*, types::PyDict}; + use rstest::rstest; + + const CONFIG: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}} + ]}"#; + + const PAIR: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}} + ]}"#; + + #[rstest] + fn python_drives_calls_and_selections_through_the_handle() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("native", crate::native_module(py)).unwrap(); + locals.set_item("config", CONFIG).unwrap(); + locals.set_item("pair", PAIR).unwrap(); + py.run( + c_str!( + r#" +import time +store = {} +asked = [] + +def answer_cache(step): + asked.append(step) + op, key = step["op"], step["key"] + if op == "get": + return store.get(key) + if op == "batch_get": + return [store.get(k) for k in step["keys"]] + if op == "set": + store[key] = step["value"] + return None + if op == "increment": + store[key] = store.get(key, 0) + step["value"] + return store[key] + if op == "delete": + store.pop(key, None) + return None + raise AssertionError(op) + +def drive(call, outcome): + step = call.next(outcome) + while step["kind"] == "cache": + step = call.next(answer_cache(step)) + return step + +handle = native.RouterCoreHandle(config, {"num_retries": 1}) +picked = drive(handle.select("primary", {}), None) +assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked +assert [step["op"] for step in asked] == ["batch_get"] +assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "events"], list(asked[0]) +call = handle.start("acompletion", '{"model": "primary"}') +version = handle.registry_version() +assert handle.apply({"delete": ["gone"]}) == version + 1 +steps = [drive(call, None)] +steps.append(drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})) +steps.append(drive(call, None)) +steps.append(drive(call, {"ok": True, "total_tokens": 5})) +kinds = [step["kind"] for step in steps] + +cooled = native.RouterCoreHandle(pair, {"num_retries": 0}) +call = cooled.start("acompletion", {"model": "primary"}) +tried = drive(call, None)["deployment_id"] +assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed" +assert f"deployment:{tried}:cooldown" not in store, store +store[f"deployment:{tried}:cooldown"] = {"exception_received": "x", "status_code": "429", "timestamp": time.time(), "cooldown_time": 60.0} +other = drive(cooled.select("primary", {}), None) +assert other["kind"] == "done" and other["deployment_id"] != tried, other + +call = handle.select("primary", {}) +step = call.next(None) +assert step["op"] == "batch_get" +for bad in ([], [None, None], "nope", None): + try: + call.next(bad) + except ValueError: + pass + else: + raise AssertionError(f"{bad!r} should not answer a batch_get of one key") +assert drive(call, [None])["kind"] == "done" +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let kinds: Vec = locals + .get_item("kinds") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!(kinds, ["attempt", "sleep", "attempt", "done"]); + }); + } + + #[rstest] + #[case::unknown_key(r#"{"routing_strategy_args": {}}"#, true)] + #[case::bad_json("{", false)] + #[case::not_an_object("[]", false)] + fn construction_failures_raise_the_class_python_reads( + #[case] config: &str, + #[case] declined: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let error = crate::native_module(py) + .getattr("RouterCoreHandle") + .and_then(|handle| handle.call1((config, "{}"))) + .unwrap_err(); + assert_eq!( + error.is_instance_of::(py), + declined, + "{error}" + ); + assert_eq!( + error.is_instance_of::(py), + !declined, + "{error}" + ); + }); + } +} diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml index af97bd52c57..b6bb4514dfb 100644 --- a/litellm-rust/crates/router-core/Cargo.toml +++ b/litellm-rust/crates/router-core/Cargo.toml @@ -7,3 +7,20 @@ repository.workspace = true description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop" [dependencies] +arc-swap.workspace = true +async-trait.workspace = true +litellm-cache.workspace = true +litellm-router-types.workspace = true +parking_lot.workspace = true +rand.workspace = true +serde.workspace = true +serde_with.workspace = true +serde_json = { workspace = true, features = ["preserve_order"] } +sha2.workspace = true +strum.workspace = true +thiserror.workspace = true +tracing.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/router-core/src/cache.rs b/litellm-rust/crates/router-core/src/cache.rs new file mode 100644 index 00000000000..da5e99c535e --- /dev/null +++ b/litellm-rust/crates/router-core/src/cache.rs @@ -0,0 +1,130 @@ +//! [`RouterCache`]: the DualCache-style store all routing state lives in, under the keys, +//! encodings and TTLs the Python router uses, so Python and Rust workers share it. Method names +//! follow `litellm-cache`'s `BaseCache`, `BatchCache`, `CounterCache` and `DeleteCache`. + +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +use async_trait::async_trait; +use litellm_cache::Error; +use parking_lot::Mutex; +use serde_json::Value; + +#[async_trait] +pub trait RouterCache: Send + Sync { + async fn async_get_cache(&self, key: &str) -> Result, Error>; + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error>; + + /// Add `amount` to the number at `key` (0 when missing) and return the new value. The TTL + /// applies when the key is created, as Redis `INCRBYFLOAT` + `EXPIRE NX` does. + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result; + + async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error>; + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error>; +} + +#[derive(Default)] +struct Entry { + value: Value, + expires: Option, +} + +/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`. +#[derive(Clone, Default)] +pub struct InMemoryRouterCache { + entries: Arc>>, +} + +impl InMemoryRouterCache { + fn live(entries: &mut HashMap, key: &str) -> Option { + let expired = entries + .get(key) + .is_some_and(|entry| entry.expires.is_some_and(|at| at <= Instant::now())); + if expired { + entries.remove(key); + } + entries.get(key).map(|entry| entry.value.clone()) + } + + pub fn get(&self, key: &str) -> Option { + Self::live(&mut self.entries.lock(), key) + } + + pub fn increment(&self, key: &str, amount: f64, ttl: Option) -> f64 { + let mut entries = self.entries.lock(); + let current = Self::live(&mut entries, key).and_then(|value| value.as_f64()); + let next = current.unwrap_or(0.0) + amount; + let value = serde_json::Number::from_f64(next) + .map(|number| { + if next.fract() == 0.0 && next.abs() < i64::MAX as f64 { + Value::from(next as i64) + } else { + Value::Number(number) + } + }) + .unwrap_or(Value::Null); + let expires = match (current, entries.get(key)) { + (Some(_), Some(entry)) => entry.expires, + _ => ttl.map(|ttl| Instant::now() + ttl), + }; + entries.insert(key.to_owned(), Entry { value, expires }); + next + } +} + +#[async_trait] +impl RouterCache for InMemoryRouterCache { + async fn async_get_cache(&self, key: &str) -> Result, Error> { + Ok(self.get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error> { + let expires = ttl.map(|ttl| Instant::now() + ttl); + self.entries + .lock() + .insert(key.to_owned(), Entry { value, expires }); + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + Ok(self.increment(key, amount, ttl)) + } + + async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error> { + let mut entries = self.entries.lock(); + Ok(keys + .iter() + .map(|key| Self::live(&mut entries, key)) + .collect()) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + self.entries.lock().remove(key); + Ok(()) + } +} diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs new file mode 100644 index 00000000000..e356ab1e694 --- /dev/null +++ b/litellm-rust/crates/router-core/src/call.rs @@ -0,0 +1,592 @@ +//! The call loop: select, run the attempt, then retry in the same group or fall back to another +//! until it succeeds or the policy is exhausted. Mirrors `Router.async_function_with_fallbacks`, +//! `async_function_with_fallbacks_common_utils`, `async_function_with_retries`, +//! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and +//! `run_async_fallback` in `router_utils/fallback_event_handlers.py`. + +use std::{collections::HashSet, future::Future, pin::Pin, sync::Arc, time::Duration}; + +use tracing::warn; + +use crate::{ + Core, Snapshot, + config::Op, + error::{AttemptError, ConfigError, Error, FailureKind}, + executor::Executor, + health::{GroupShape, should_retry_status}, + host::{Attempt, Output, RouterEvent, RouterRequest}, + registry::ResolveKind, + select::{Candidates, SelectInputs}, +}; + +/// `INITIAL_RETRY_DELAY`, `MAX_RETRY_DELAY` and `JITTER` in `litellm/constants.py`. +const INITIAL_RETRY_DELAY: f64 = 0.5; +const MAX_RETRY_DELAY: f64 = 8.0; +const JITTER: f64 = 0.75; +/// `calculate_retry_after` obeys a provider's retry-after up to a minute. +const MAX_HONOURED_RETRY_AFTER: Duration = Duration::from_secs(60); + +#[derive(Clone, Debug)] +pub struct RouterResponse { + pub output: Output, + pub deployment_id: String, + pub model_group: String, + pub attempted_retries: u32, + /// The retry budget of the loop a retry succeeded in (`x-litellm-max-retries`); `None` + /// when the first attempt succeeded. + pub max_retries: Option, + pub attempted_fallbacks: u32, +} + +/// A failed attempt with the deployment it ran on, when one was picked. +struct Failed { + error: Error, + deployment_id: Option, +} + +impl From for Failed { + fn from(error: Error) -> Self { + Self { + error, + deployment_id: None, + } + } +} + +/// Where a level of the fallback walk sits: `depth` is `run_async_fallback`'s counter, one more +/// for every target tried, and `level` how many fallbacks deep the walk is. +#[derive(Clone, Copy, Debug, Default)] +struct Hop { + depth: u32, + level: u32, +} + +/// A fallback target: a group, or the same group at a later `order`. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct Target { + group: String, + order: Option, +} + +type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// `should_retry_this_error`, as a predicate. +fn should_retry( + failure: &AttemptError, + healthy: usize, + all: usize, + has_context_window_fallbacks: bool, + has_content_policy_fallbacks: bool, + has_regular_fallbacks: bool, +) -> bool { + match failure.kind { + FailureKind::ContextWindowExceeded if has_context_window_fallbacks => return false, + FailureKind::ContentPolicyViolation if has_content_policy_fallbacks => return false, + _ => {} + } + if let Some(status) = failure.status + && !should_retry_status(status) + && !matches!(status, 401 | 403) + && failure.kind != FailureKind::Declined + { + return false; + } + match failure.kind { + FailureKind::NotFound => false, + FailureKind::RateLimit if healthy == 0 && has_regular_fallbacks => false, + FailureKind::Authentication | FailureKind::PermissionDenied if all <= 1 => false, + _ => healthy > 0, + } +} + +impl Core { + pub(crate) async fn call( + &self, + op: Op, + request: Arc, + ) -> Result { + let snapshot = self.snapshot.load_full(); + let mut attempted = HashSet::new(); + let target = Target { + group: request.model.clone(), + order: None, + }; + self.with_fallbacks( + &snapshot, + &op, + &request, + target, + Hop::default(), + &mut attempted, + ) + .await + } + + /// One level of `async_function_with_fallbacks`: the target's retry loop, then its + /// fallbacks. A level that fails raises its own error, whatever its fallbacks raised. + fn with_fallbacks<'a>( + &'a self, + snapshot: &'a Snapshot, + op: &'a Op, + request: &'a Arc, + target: Target, + at: Hop, + attempted: &'a mut HashSet, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let error = match self.with_retries(snapshot, op, request, &target, at).await { + Ok(response) => return Ok(response), + Err(error) => error, + }; + let Some(targets) = self.fallback_targets(snapshot, request, &target, &error) else { + return Err(error); + }; + // `run_async_fallback`: the depth is checked once on entry, then every target + // tried at this level counts one more. + if at.depth >= snapshot.settings.max_fallbacks { + return Err(error); + } + attempted.insert(Target { + group: target.group.clone(), + order: None, + }); + let mut hop = at.depth; + for next in targets { + if next.order.is_none() && next.group == target.group { + continue; + } + if next.group != target.group + && !self.host.allow_fallback(request, &next.group).await + { + continue; + } + if !attempted.insert(next.clone()) { + continue; + } + hop += 1; + self.host.observe(&RouterEvent::Fallback { + from: target.group.clone(), + to: next.group.clone(), + order: next.order, + depth: hop, + level: at.level + 1, + }); + if let Ok(mut response) = self + .with_fallbacks( + snapshot, + op, + request, + next, + Hop { + depth: hop, + level: at.level + 1, + }, + attempted, + ) + .await + { + response.attempted_fallbacks = hop; + return Ok(response); + } + } + Err(error) + }) + } + + /// The fallback targets for a failure, or `None` when it must be raised as is + /// (`async_function_with_fallbacks_common_utils`). + fn fallback_targets( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + target: &Target, + error: &Error, + ) -> Option> { + if request.disable_fallbacks { + return None; + } + let kind = match error { + Error::Attempt(failure) => failure.kind, + Error::NoHealthyDeployments { .. } => FailureKind::RateLimit, + Error::Config(_) | Error::UnknownModel { .. } | Error::Cache(_) | Error::Retired => { + return None; + } + }; + if kind == FailureKind::GuardrailIntervention { + return None; + } + let dynamic = &snapshot.settings.dynamic; + // `fallback_lookup_groups`: the group being served, then the group the request was for. + let mut lookup = vec![target.group.as_str()]; + if request.model != target.group { + lookup.push(request.model.as_str()); + } + let groups = |names: Vec| { + names + .into_iter() + .map(|group| Target { group, order: None }) + .collect::>() + }; + let skip_order = matches!( + kind, + FailureKind::ContextWindowExceeded | FailureKind::ContentPolicyViolation + ); + if !skip_order { + let mut orders: Vec = snapshot + .registry + .group(&target.group) + .unwrap_or_default() + .iter() + .filter(|d| d.usable_by_team(request.team_id.as_deref())) + .filter_map(|d| d.order()) + .collect(); + orders.sort_unstable(); + orders.dedup(); + if orders.len() > 1 { + let current = target.order.unwrap_or(orders[0]); + let mut targets: Vec = orders + .into_iter() + .filter(|order| *order > current) + .map(|order| Target { + group: target.group.clone(), + order: Some(order), + }) + .collect(); + targets.extend(groups( + dynamic + .fallbacks + .lookup_groups(&lookup, &dynamic.fallback_keys) + .unwrap_or_default(), + )); + if !targets.is_empty() { + return Some(targets); + } + } + } + let specific = match kind { + FailureKind::ContextWindowExceeded => dynamic.context_window_fallbacks.as_ref(), + FailureKind::ContentPolicyViolation => dynamic.content_policy_fallbacks.as_ref(), + _ => None, + }; + if let Some(fallbacks) = specific { + return fallbacks + .lookup_groups(&lookup, &dynamic.fallback_keys) + .map(groups); + } + dynamic.fallback_targets(&lookup).map(groups) + } + + async fn with_retries( + &self, + snapshot: &Snapshot, + op: &Op, + request: &Arc, + target: &Target, + at: Hop, + ) -> Result { + let mut inputs = SelectInputs { + target_order: target.order, + ..SelectInputs::default() + }; + let counting = SelectInputs { + count_only: true, + ..SelectInputs::default() + }; + let first = self + .attempt(snapshot, op, request, target, &inputs, 0, at) + .await; + let mut failed = match first { + Ok(response) => return Ok(response), + Err(failed) => failed, + }; + let Error::Attempt(failure) = &failed.error else { + return Err(failed.error); + }; + if failure.kind == FailureKind::GuardrailIntervention { + return Err(failed.error); + } + let settings = &snapshot.settings; + let mut num_retries = request + .num_retries + .or(failure.num_retries) + .unwrap_or(settings.num_retries); + let policy_retries = (request.num_retries != Some(0)) + .then(|| settings.retry_policy.as_ref()?.retries_for(failure.kind)) + .flatten(); + let policy_applies = policy_retries.is_some(); + num_retries = policy_retries.unwrap_or(num_retries); + // `_async_get_healthy_deployments` finds no group for a deployment id, so a call by id + // counts no deployments when deciding whether to retry. + let counted = |c: &Candidates, n: usize| if c.kind == ResolveKind::Id { 0 } else { n }; + let all = self + .candidates(snapshot, op, request, &target.group, &counting) + .await + .map(|c| counted(&c, c.all.len())) + .unwrap_or(0); + + for retry in 0..=num_retries { + let Error::Attempt(failure) = &failed.error else { + return Err(failed.error); + }; + let healthy = self + .candidates(snapshot, op, request, &target.group, &counting) + .await + .map(|c| counted(&c, c.healthy.len())) + .unwrap_or(0); + if !policy_applies && !self.should_retry(snapshot, failure, healthy, all) { + return Err(failed.error); + } + if num_retries == 0 { + return Err(failed.error); + } + if let (Some(status), Some(id)) = (failure.status, &failed.deployment_id) + && !should_retry_status(status) + && !inputs.skipped_ids.contains(id) + { + inputs.skipped_ids.push(id.clone()); + } + let delay = self.time_to_sleep( + snapshot, + failure, + num_retries - retry, + num_retries, + healthy, + all, + ); + self.host.sleep(delay).await; + if retry == num_retries { + self.host.observe(&RouterEvent::RetriesExhausted { + model_group: target.group.clone(), + attempted: num_retries, + max: num_retries, + }); + break; + } + self.host.observe(&RouterEvent::Retry { + model_group: target.group.clone(), + retry: retry + 1, + delay, + }); + match self + .attempt(snapshot, op, request, target, &inputs, retry + 1, at) + .await + { + Ok(mut response) => { + response.max_retries = Some(num_retries); + return Ok(response); + } + Err(next) => failed = next, + } + } + Err(failed.error) + } + + fn should_retry( + &self, + snapshot: &Snapshot, + failure: &AttemptError, + healthy: usize, + all: usize, + ) -> bool { + let dynamic = &snapshot.settings.dynamic; + should_retry( + failure, + healthy, + all, + dynamic.context_window_fallbacks.is_some(), + dynamic.content_policy_fallbacks.is_some(), + dynamic.has_regular_fallbacks(), + ) + } + + /// `_time_to_sleep_before_retry` and `calculate_retry_after`. + fn time_to_sleep( + &self, + snapshot: &Snapshot, + failure: &AttemptError, + remaining_retries: u32, + num_retries: u32, + healthy: usize, + all: usize, + ) -> Duration { + if all != 1 && healthy > 0 { + return Duration::ZERO; + } + let jitter = JITTER * self.rng.next_f64(); + if let Some(retry_after) = failure.retry_after + && !retry_after.is_zero() + && retry_after <= MAX_HONOURED_RETRY_AFTER + { + return retry_after + Duration::from_secs_f64(jitter); + } + let exponent = num_retries.saturating_sub(remaining_retries); + let backoff = (INITIAL_RETRY_DELAY * 2f64.powi(exponent as i32)) + .max(snapshot.settings.retry_after.as_secs_f64()) + .min(MAX_RETRY_DELAY); + Duration::from_secs_f64(backoff + jitter) + } + + /// One attempt: select a deployment, run it on its executor, record the outcome. + #[allow(clippy::too_many_arguments)] + async fn attempt( + &self, + snapshot: &Snapshot, + op: &Op, + request: &Arc, + target: &Target, + inputs: &SelectInputs, + retry: u32, + at: Hop, + ) -> Result { + let candidates = self + .candidates(snapshot, op, request, &target.group, inputs) + .await?; + if candidates.healthy.is_empty() { + return Err(Error::NoHealthyDeployments { + model: candidates.model_group, + } + .into()); + } + let deployment = self.pick(snapshot, op, request, &candidates).await?; + self.host.observe(&RouterEvent::Selected { + model_group: candidates.model_group.clone(), + deployment_id: deployment.id.clone(), + }); + let executor = self + .executors + .resolve(op, deployment.provider()) + .map_err(|unsupported| Error::Config(ConfigError::Unsupported(unsupported)))?; + let attempt = Attempt { + op: op.clone(), + request: request.clone(), + deployment: deployment.clone(), + name: target.group.clone(), + model_group: candidates.model_group.clone(), + target_order: target.order, + candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(), + retry, + fallback_depth: at.depth, + fallback_level: at.level, + timeout: snapshot.settings.timeout, + }; + let (result, recorded) = match executor { + Executor::Native(route) => (route.execute(&attempt).await, false), + Executor::Host => ( + self.host.invoke(&attempt).await, + self.host.records_outcomes(), + ), + }; + match result { + Ok(output) => { + if !recorded + && let Err(error) = self + .health + .record_success(&deployment, output.total_tokens) + .await + { + warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded"); + } + self.host.observe(&RouterEvent::Succeeded { + deployment_id: deployment.id.clone(), + }); + Ok(RouterResponse { + output, + deployment_id: deployment.id.clone(), + model_group: candidates.model_group, + attempted_retries: retry, + max_retries: None, + attempted_fallbacks: at.depth, + }) + } + Err(failure) => { + if recorded { + self.host.observe(&RouterEvent::AttemptFailed { + deployment_id: deployment.id.clone(), + kind: failure.kind, + status: failure.status, + }); + } else { + self.record_failure(snapshot, &candidates, &deployment, &failure) + .await; + } + Err(Failed { + error: Error::Attempt(failure), + deployment_id: Some(deployment.id.clone()), + }) + } + } + } + + async fn record_failure( + &self, + snapshot: &Snapshot, + candidates: &Candidates, + deployment: &crate::registry::Deployment, + failure: &AttemptError, + ) { + self.host.observe(&RouterEvent::AttemptFailed { + deployment_id: deployment.id.clone(), + kind: failure.kind, + status: failure.status, + }); + let shape = GroupShape { + single_deployment: candidates.all.len() == 1, + has_fallback_group: snapshot + .settings + .dynamic + .fallback_targets(&[candidates.model_group.as_str()]) + .is_some(), + }; + match self + .health + .record_failure(&snapshot.settings, deployment, shape, failure) + .await + { + Ok(Some(cooldown)) => self.host.observe(&RouterEvent::CooldownSet { + deployment_id: deployment.id.clone(), + cooldown, + }), + Ok(None) => {} + Err(error) => { + warn!(deployment_id = %deployment.id, error = %error, "cooldown was not recorded"); + } + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::rate_limit_with_healthy(FailureKind::RateLimit, Some(429), 1, 2, false, true)] + #[case::rate_limit_none_healthy_with_fallbacks( + FailureKind::RateLimit, + Some(429), + 0, + 2, + true, + false + )] + #[case::bad_request(FailureKind::BadRequest, Some(400), 2, 2, false, false)] + #[case::auth_lone_deployment(FailureKind::Authentication, Some(401), 1, 1, false, false)] + #[case::auth_with_siblings(FailureKind::Authentication, Some(401), 1, 2, false, true)] + #[case::not_found(FailureKind::NotFound, Some(404), 2, 2, false, false)] + #[case::server_error_no_healthy(FailureKind::InternalServer, Some(500), 0, 2, false, false)] + fn should_retry_follows_python( + #[case] kind: FailureKind, + #[case] status: Option, + #[case] healthy: usize, + #[case] all: usize, + #[case] regular_fallbacks: bool, + #[case] expected: bool, + ) { + let failure = AttemptError::new(kind, status); + assert_eq!( + should_retry(&failure, healthy, all, false, false, regular_fallbacks), + expected + ); + } +} diff --git a/litellm-rust/crates/router-core/src/config.rs b/litellm-rust/crates/router-core/src/config.rs new file mode 100644 index 00000000000..558b8d90844 --- /dev/null +++ b/litellm-rust/crates/router-core/src/config.rs @@ -0,0 +1,433 @@ +//! What the router is told: [`RouterConfig`], fixed for the router's life, and +//! [`RouterDynamicConfig`], a versioned snapshot read once per call. The router never reads +//! process globals. + +use std::{collections::BTreeMap, time::Duration}; + +use serde::Deserialize; + +use crate::{ + error::FailureKind, + registry::{DeploymentSpec, RoutingGroupSpec}, +}; + +/// `DEFAULT_COOLDOWN_TIME_SECONDS` in `litellm/constants.py`. +pub const DEFAULT_COOLDOWN_TIME: Duration = Duration::from_secs(5); +/// `ROUTER_MAX_FALLBACKS` in `litellm/constants.py`. +pub const DEFAULT_MAX_FALLBACKS: u32 = 5; + +/// The router operation a call runs, spelled as the Python router's endpoint method. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + Hash, + strum::EnumString, + strum::Display, + serde_with::DeserializeFromStr, +)] +pub enum Op { + #[strum(serialize = "acompletion")] + Completion, + #[strum(serialize = "atext_completion")] + TextCompletion, + #[strum(serialize = "aembedding")] + Embedding, + #[strum(serialize = "aresponses")] + Responses, + #[strum(serialize = "aanthropic_messages")] + Messages, + #[strum(serialize = "aimage_generation")] + ImageGeneration, + #[strum(serialize = "atranscription")] + Transcription, + #[strum(serialize = "aspeech")] + Speech, + #[strum(serialize = "aocr")] + Ocr, + #[strum(default, transparent)] + Other(String), +} + +/// `routing_strategy`. Only simple-shuffle is native; the rest go through `Host::select`. +#[derive( + Clone, + Debug, + Default, + PartialEq, + Eq, + strum::EnumString, + strum::Display, + serde_with::DeserializeFromStr, +)] +pub enum RoutingStrategy { + #[default] + #[strum(serialize = "simple-shuffle")] + SimpleShuffle, + #[strum(serialize = "least-busy")] + LeastBusy, + #[strum(serialize = "usage-based-routing")] + UsageBased, + #[strum(serialize = "usage-based-routing-v2")] + UsageBasedV2, + #[strum(serialize = "latency-based-routing")] + LatencyBased, + #[strum(serialize = "cost-based-routing")] + CostBased, + #[strum(default, transparent)] + Other(String), +} + +impl RoutingStrategy { + pub fn is_native(&self) -> bool { + matches!(self, Self::SimpleShuffle) + } +} + +/// `PluginKind` in `litellm/router_api.py`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::Display)] +pub enum PluginKind { + #[serde(rename = "auto")] + #[strum(serialize = "auto")] + Auto, + #[serde(rename = "complexity")] + #[strum(serialize = "complexity")] + Complexity, + #[serde(rename = "quality")] + #[strum(serialize = "quality")] + Quality, + #[serde(rename = "adaptive")] + #[strum(serialize = "adaptive")] + Adaptive, + #[serde(rename = "routing")] + #[strum(serialize = "routing")] + Routing, + #[serde(rename = "budget_limiter")] + #[strum(serialize = "budget_limiter")] + BudgetLimiter, + #[serde(rename = "prompt_caching_pins")] + #[strum(serialize = "prompt_caching_pins")] + PromptCachingPins, +} + +/// One registered plugin. Its behaviour lives in the host until it is ported. +#[derive(Clone, Debug, PartialEq, Deserialize)] +pub struct PluginSpec { + pub kind: PluginKind, + pub name: String, + #[serde(default)] + pub config: serde_json::Value, +} + +/// `RetryPolicy` in `litellm/types/router.py`: retries per exception type. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[allow(non_snake_case)] +pub struct RetryPolicy { + pub BadRequestErrorRetries: Option, + pub AuthenticationErrorRetries: Option, + pub TimeoutErrorRetries: Option, + pub RateLimitErrorRetries: Option, + pub ContentPolicyViolationErrorRetries: Option, + pub InternalServerErrorRetries: Option, +} + +impl RetryPolicy { + pub fn retries_for(&self, kind: FailureKind) -> Option { + match kind { + FailureKind::BadRequest | FailureKind::ContextWindowExceeded => { + self.BadRequestErrorRetries + } + FailureKind::Authentication => self.AuthenticationErrorRetries, + FailureKind::Timeout => self.TimeoutErrorRetries, + FailureKind::RateLimit => self.RateLimitErrorRetries, + FailureKind::ContentPolicyViolation => self.ContentPolicyViolationErrorRetries, + FailureKind::InternalServer => self.InternalServerErrorRetries, + FailureKind::PermissionDenied + | FailureKind::NotFound + | FailureKind::ApiConnection + | FailureKind::ServiceUnavailable + | FailureKind::BadGateway + | FailureKind::GuardrailIntervention + | FailureKind::Declined + | FailureKind::CallerCredential + | FailureKind::Other => None, + } + } +} + +/// `AllowedFailsPolicy` in `litellm/types/router.py`: allowed fails per minute per exception type. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[allow(non_snake_case)] +pub struct AllowedFailsPolicy { + pub BadRequestErrorAllowedFails: Option, + pub AuthenticationErrorAllowedFails: Option, + pub TimeoutErrorAllowedFails: Option, + pub RateLimitErrorAllowedFails: Option, + pub ContentPolicyViolationErrorAllowedFails: Option, + pub InternalServerErrorAllowedFails: Option, + pub ServiceUnavailableErrorAllowedFails: Option, + pub BadGatewayErrorAllowedFails: Option, +} + +impl AllowedFailsPolicy { + pub fn allowed_fails_for(&self, kind: FailureKind) -> Option { + match kind { + FailureKind::BadRequest | FailureKind::ContextWindowExceeded => { + self.BadRequestErrorAllowedFails + } + FailureKind::Authentication => self.AuthenticationErrorAllowedFails, + FailureKind::Timeout => self.TimeoutErrorAllowedFails, + FailureKind::RateLimit => self.RateLimitErrorAllowedFails, + FailureKind::ContentPolicyViolation => self.ContentPolicyViolationErrorAllowedFails, + FailureKind::InternalServer => self.InternalServerErrorAllowedFails, + FailureKind::ServiceUnavailable => self.ServiceUnavailableErrorAllowedFails, + FailureKind::BadGateway => self.BadGatewayErrorAllowedFails, + FailureKind::PermissionDenied + | FailureKind::NotFound + | FailureKind::ApiConnection + | FailureKind::GuardrailIntervention + | FailureKind::Declined + | FailureKind::CallerCredential + | FailureKind::Other => None, + } + } +} + +/// A fallback list as Python spells it, `[{"gpt-4": ["gpt-3.5"]}, {"*": ["default"]}]`, kept in +/// order. `"*"` is the default entry. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(from = "Vec>>")] +pub struct Fallbacks(Vec<(String, Vec)>); + +impl From>>> for Fallbacks { + fn from(entries: Vec>>) -> Self { + Self(entries.into_iter().flatten().collect()) + } +} + +impl Fallbacks { + pub fn new(entries: impl IntoIterator)>) -> Self { + Self(entries.into_iter().collect()) + } + + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// `_get_first_default_fallback`: the first target of the first `"*"` entry. + pub fn default_target(&self) -> Option<&str> { + self.0 + .iter() + .find(|(key, targets)| key == "*" && !targets.is_empty()) + .map(|(_, targets)| targets[0].as_str()) + } + + /// `get_fallback_model_group`: the group's own entry, else an entry keyed by a + /// provider-prefixed or provider-stripped spelling of it (the last such entry), else the + /// last `"*"` entry before the group's own. The flag says the targets came from `"*"`. + pub fn lookup(&self, group: &str, keys: &FallbackKeys) -> Option<(&[String], bool)> { + let prefixed = keys.provider_prefixed.get(group).map(String::as_str); + let mut spelled = None; + let mut generic = None; + for (key, targets) in &self.0 { + if key == group { + return Some((targets, false)); + } + if Some(key.as_str()) == prefixed || keys.is_stripped(group, key) { + spelled = Some(targets.as_slice()); + } else if key == "*" { + generic = Some(targets.as_slice()); + } + } + spelled + .map(|targets| (targets, false)) + .or_else(|| generic.map(|targets| (targets, true))) + } + + /// `get_fallback_model_group_for_lookup_groups`: the first lookup group with its own + /// chain wins; `"*"` applies only after every group missed. + pub fn lookup_groups(&self, groups: &[&str], keys: &FallbackKeys) -> Option> { + let results: Vec<_> = groups + .iter() + .map(|group| self.lookup(group, keys)) + .collect(); + results + .iter() + .flatten() + .find(|(_, generic)| !generic) + .or_else(|| results.iter().flatten().next()) + .map(|(targets, _)| targets.to_vec()) + } +} + +/// What fallback keys match besides the exact group name: Python's `provider_list`, for +/// `"openai/gpt-4o"` to use a `"gpt-4o"` entry, and the provider-prefixed spelling of each +/// bare group name (`_provider_prefixed_model_group`), for `"gpt-4o"` to use an +/// `"openai/gpt-4o"` entry. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(default)] +pub struct FallbackKeys { + pub providers: Vec, + pub provider_prefixed: BTreeMap, +} + +impl FallbackKeys { + /// `_check_stripped_model_group`. + fn is_stripped(&self, group: &str, key: &str) -> bool { + self.providers.iter().any(|provider| { + let prefix = format!("{provider}/"); + group.starts_with(&prefix) && group.replace(&prefix, "") == key + }) + } +} + +/// Static config, fixed for the router's life. +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(default)] +pub struct RouterConfig { + pub model_list: Vec, + pub routing_strategy: RoutingStrategy, + pub routing_groups: Vec, + pub plugins: Vec, + /// Router default `cooldown_time`; `None` is `DEFAULT_COOLDOWN_TIME_SECONDS`. + pub cooldown_time: Option, + pub disable_cooldowns: bool, + pub allowed_fails_policy: Option, + pub retry_policy: Option, + /// `retry_after`: the minimum backoff between retries, in seconds. + pub retry_after: f64, + pub timeout: Option, + pub request_timeout: Option, + pub drop_params: Option, + pub enable_pre_call_checks: bool, + pub enable_health_check_routing: bool, + /// Deployment ids that are a provider's default and never cool down + /// (`provider_default_deployment_ids`). + pub provider_default_deployment_ids: Vec, + /// The operations this router must serve; with a host that cannot invoke, each must have a + /// native executor for every configured provider or the router fails to build. + pub served_ops: Vec, +} + +/// Dynamic config: a versioned snapshot read once per call. +#[derive(Clone, Debug, Deserialize)] +#[serde(default)] +pub struct RouterDynamicConfig { + pub version: u64, + /// `model_group_alias` as Python spells it: each alias maps to a group name, or to + /// `{"model": , "hidden": }`. + #[serde(alias = "model_group_alias", deserialize_with = "alias_targets")] + pub model_alias_map: BTreeMap, + pub fallbacks: Fallbacks, + pub default_fallbacks: Vec, + pub context_window_fallbacks: Option, + pub content_policy_fallbacks: Option, + pub num_retries: Option, + pub max_fallbacks: u32, + /// `allowed_fails`. `Some` means it was set explicitly, which switches cooldowns from the + /// failure-ratio rule to the allowed-fails counter (`_is_allowed_fails_set_on_router`). + pub allowed_fails: Option, + /// Other spellings a fallback key matches. + pub fallback_keys: FallbackKeys, +} + +/// One `model_group_alias` value: a group name, or `{"model": , "hidden": }`. +#[derive(Deserialize)] +#[serde(untagged)] +enum AliasTarget { + Group(String), + Entry { model: String }, +} + +fn alias_targets<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let aliases = BTreeMap::::deserialize(deserializer)?; + Ok(aliases + .into_iter() + .map(|(alias, target)| match target { + AliasTarget::Group(group) | AliasTarget::Entry { model: group } => (alias, group), + }) + .collect()) +} + +impl Default for RouterDynamicConfig { + fn default() -> Self { + Self { + version: 0, + model_alias_map: BTreeMap::new(), + fallbacks: Fallbacks::default(), + default_fallbacks: Vec::new(), + context_window_fallbacks: None, + content_policy_fallbacks: None, + num_retries: None, + max_fallbacks: DEFAULT_MAX_FALLBACKS, + allowed_fails: None, + fallback_keys: FallbackKeys::default(), + } + } +} + +impl RouterDynamicConfig { + /// The regular fallback targets for the lookup groups (`get_fallback_model_group_for_lookup_groups` + /// over `fallbacks`), then `default_fallbacks`, which Python folds into `fallbacks` as + /// `{"*": [...]}`. + pub fn fallback_targets(&self, groups: &[&str]) -> Option> { + self.fallbacks + .lookup_groups(groups, &self.fallback_keys) + .or_else(|| { + (!self.default_fallbacks.is_empty()).then(|| self.default_fallbacks.clone()) + }) + } + + pub fn has_regular_fallbacks(&self) -> bool { + !self.fallbacks.is_empty() || !self.default_fallbacks.is_empty() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::own_entry("gpt-4", Some(vec!["gpt-3.5".to_owned()]))] + #[case::star_entry("other", Some(vec!["default".to_owned()]))] + fn fallback_targets_prefer_the_groups_own_entry( + #[case] group: &str, + #[case] expected: Option>, + ) { + let dynamic: RouterDynamicConfig = serde_json::from_value(serde_json::json!({ + "fallbacks": [{"gpt-4": ["gpt-3.5"]}, {"*": ["default"]}] + })) + .unwrap(); + assert_eq!(dynamic.fallback_targets(&[group]), expected); + } + + #[rstest] + #[case::group_name(json!({"model_group_alias": {"gpt": "gpt-4"}}))] + #[case::hidden_entry(json!({"model_group_alias": {"gpt": {"model": "gpt-4", "hidden": true}}}))] + #[case::native_name(json!({"model_alias_map": {"gpt": "gpt-4"}}))] + fn model_group_alias_reads_both_python_spellings(#[case] wire: serde_json::Value) { + let dynamic: RouterDynamicConfig = serde_json::from_value(wire).unwrap(); + assert_eq!( + dynamic.model_alias_map, + BTreeMap::from([("gpt".to_owned(), "gpt-4".to_owned())]) + ); + } + + #[rstest] + #[case::native("simple-shuffle", RoutingStrategy::SimpleShuffle)] + #[case::unknown("my-strategy", RoutingStrategy::Other("my-strategy".to_owned()))] + fn routing_strategy_parses_its_wire_name( + #[case] wire: &str, + #[case] expected: RoutingStrategy, + ) { + assert_eq!(wire.parse::().unwrap(), expected); + assert_eq!(expected.to_string(), wire); + } +} diff --git a/litellm-rust/crates/router-core/src/error.rs b/litellm-rust/crates/router-core/src/error.rs new file mode 100644 index 00000000000..429a82f94fe --- /dev/null +++ b/litellm-rust/crates/router-core/src/error.rs @@ -0,0 +1,190 @@ +//! The router's errors: [`ConfigError`] when a router is built or changed, [`Error`] when a call +//! or a selection fails, and [`AttemptError`], the typed facts of one failed attempt. + +use std::{any::Any, fmt, sync::Arc, time::Duration}; + +use crate::config::{Op, RoutingStrategy}; + +/// What a host cannot serve, found when the router is built rather than partway through a call. +#[derive(Clone, Debug, PartialEq, thiserror::Error)] +pub enum Unsupported { + #[error("routing strategy `{0}` needs a host that can select")] + Strategy(RoutingStrategy), + #[error("routing plugin `{0}` needs a host that can select")] + Plugin(String), + #[error("operation `{op}` on provider `{provider}` has no executor")] + Executor { op: Op, provider: String }, +} + +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("the router config asks for something this host cannot serve")] + Unsupported(#[from] Unsupported), + #[error("deployment `{model_name}` has invalid litellm_params")] + InvalidDeployment { + model_name: String, + #[source] + source: serde_json::Error, + }, + #[error("deployment id `{id}` is configured more than once")] + DuplicateId { id: String }, +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("the router config is invalid")] + Config(#[from] ConfigError), + #[error("no deployment serves model `{model}`")] + UnknownModel { model: String }, + #[error("no deployment of `{model}` is available")] + NoHealthyDeployments { model: String }, + #[error("the attempt failed")] + Attempt(#[from] AttemptError), + #[error("the routing cache failed")] + Cache(#[from] litellm_cache::Error), + #[error("the router was retired")] + Retired, +} + +impl Error { + /// The failed attempt behind this error, when a deployment was called. + pub fn attempt(&self) -> Option<&AttemptError> { + match self { + Self::Attempt(attempt) => Some(attempt), + Self::Config(_) + | Self::UnknownModel { .. } + | Self::NoHealthyDeployments { .. } + | Self::Cache(_) + | Self::Retired => None, + } + } +} + +/// The class of a failed attempt, as litellm's exception mapping names it. +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + strum::Display, + strum::EnumString, + strum::IntoStaticStr, + strum::VariantArray, +)] +pub enum FailureKind { + #[strum(serialize = "RateLimitError")] + RateLimit, + #[strum(serialize = "ContextWindowExceededError")] + ContextWindowExceeded, + #[strum(serialize = "ContentPolicyViolationError")] + ContentPolicyViolation, + #[strum(serialize = "AuthenticationError")] + Authentication, + #[strum(serialize = "PermissionDeniedError")] + PermissionDenied, + #[strum(serialize = "NotFoundError")] + NotFound, + #[strum(serialize = "BadRequestError")] + BadRequest, + #[strum(serialize = "Timeout")] + Timeout, + #[strum(serialize = "APIConnectionError")] + ApiConnection, + #[strum(serialize = "InternalServerError")] + InternalServer, + #[strum(serialize = "ServiceUnavailableError")] + ServiceUnavailable, + #[strum(serialize = "BadGatewayError")] + BadGateway, + /// A guardrail blocked the request: never retried, never a fallback. + #[strum(serialize = "GuardrailIntervention")] + GuardrailIntervention, + /// litellm rejected the request before calling the deployment (`MalformedRequestFieldError`, + /// `UnsupportedEndpointError`): never cools the deployment down. + #[strum(serialize = "Declined")] + Declined, + /// The caller's own credential failed, not the deployment's. + #[strum(serialize = "CallerCredentialError")] + CallerCredential, + #[strum(serialize = "APIError")] + Other, +} + +/// An opaque value the host hands through the router unchanged, such as the original Python +/// exception or response object. +#[derive(Clone)] +pub struct Payload(pub Arc); + +impl fmt::Debug for Payload { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Payload(..)") + } +} + +/// One failed attempt, as typed facts the retry, fallback and cooldown rules decide by. +#[derive(Clone, Debug, thiserror::Error)] +#[error("the deployment call failed with {kind}")] +pub struct AttemptError { + pub kind: FailureKind, + pub status: Option, + /// The provider's `retry-after` (or `retry-after-ms`) header. + pub retry_after: Option, + /// A deployment-level `num_retries` the failure carries (Python's `exception.num_retries`). + pub num_retries: Option, + /// The provider's error text, kept only for the cooldown record Python writes + /// (`exception_received`, masked past 50 characters). + pub detail: Option>, + pub payload: Option, +} + +impl AttemptError { + pub fn new(kind: FailureKind, status: Option) -> Self { + Self { + kind, + status, + retry_after: None, + num_retries: None, + detail: None, + payload: None, + } + } + + #[must_use] + pub fn with_retry_after(self, retry_after: Duration) -> Self { + Self { + retry_after: Some(retry_after), + ..self + } + } + + /// The status Python's cooldown rules read: the provider's, else 500. + pub fn cooldown_status(&self) -> u16 { + self.status.unwrap_or(500) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::rate_limit("RateLimitError", FailureKind::RateLimit)] + #[case::context_window("ContextWindowExceededError", FailureKind::ContextWindowExceeded)] + #[case::caller_credential("CallerCredentialError", FailureKind::CallerCredential)] + #[case::other("APIError", FailureKind::Other)] + fn failure_kind_round_trips_its_exception_name(#[case] wire: &str, #[case] kind: FailureKind) { + assert_eq!(wire.parse::().unwrap(), kind); + assert_eq!(<&'static str>::from(kind), wire); + } + + #[rstest] + #[case::empty("")] + #[case::lowercase("ratelimiterror")] + fn failure_kind_rejects_unknown_names(#[case] wire: &str) { + assert!(wire.parse::().is_err()); + } +} diff --git a/litellm-rust/crates/router-core/src/executor.rs b/litellm-rust/crates/router-core/src/executor.rs new file mode 100644 index 00000000000..1c895b0d90e --- /dev/null +++ b/litellm-rust/crates/router-core/src/executor.rs @@ -0,0 +1,75 @@ +//! Executors: each attempt resolves one after its deployment is picked. An (op, provider) pair in +//! the route-rule table with a registered [`NativeRoute`] runs in-process; anything else goes +//! through `Host::invoke`. An attempt runs on exactly one executor and is never replayed on +//! another: a native failure is that attempt's failure, and a retry selects afresh. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::{ + config::Op, + error::{AttemptError, Unsupported}, + host::{Attempt, Output}, +}; + +/// The (op, provider) pairs that have a Rust route machine. +pub const NATIVE_RULES: &[(Op, &str)] = &[(Op::Messages, "anthropic")]; + +/// A Rust-native route machine (an `inference-*` crate) adapted to an attempt. The embedder +/// registers one per native rule; for Messages on Anthropic that is an adapter over +/// `litellm_inference_messages::MessagesRoute::execute`. +#[async_trait] +pub trait NativeRoute: Send + Sync { + async fn execute(&self, attempt: &Attempt) -> Result; +} + +#[derive(Clone)] +pub enum Executor { + Native(Arc), + Host, +} + +#[derive(Clone, Default)] +pub struct Executors { + native: Vec<(Op, String, Arc)>, + host_invokes: bool, +} + +impl Executors { + pub fn new(host_invokes: bool) -> Self { + Self { + native: Vec::new(), + host_invokes, + } + } + + /// Register a native route for a rule in [`NATIVE_RULES`]. A pair outside the table stays + /// on the host, so a route cannot silently take over an operation not yet ported. + pub fn register(&mut self, op: Op, provider: &str, route: Arc) -> bool { + let ruled = NATIVE_RULES + .iter() + .any(|(rule_op, rule_provider)| *rule_op == op && *rule_provider == provider); + if ruled { + self.native.push((op, provider.to_owned(), route)); + } + ruled + } + + pub fn resolve(&self, op: &Op, provider: &str) -> Result { + if let Some((_, _, route)) = self + .native + .iter() + .find(|(rule_op, rule_provider, _)| rule_op == op && rule_provider == provider) + { + return Ok(Executor::Native(route.clone())); + } + if self.host_invokes { + return Ok(Executor::Host); + } + Err(Unsupported::Executor { + op: op.clone(), + provider: provider.to_owned(), + }) + } +} diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs new file mode 100644 index 00000000000..fd8d9367fb3 --- /dev/null +++ b/litellm-rust/crates/router-core/src/health.rs @@ -0,0 +1,475 @@ +//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the shared +//! [`RouterCache`] under the Python router's keys, encodings and TTLs: +//! +//! | state | key | value | TTL | +//! |---|---|---|---| +//! | cooldown | `deployment:{id}:cooldown` | `{exception_received, status_code, timestamp, cooldown_time}` | cooldown time | +//! | allowed fails | `deployment:{id}:allowed_fails[:{suffix}]` | int counter | cooldown time | +//! | usage | `global_router:{id}:{model}:tpm:{HH-MM}` and `:rpm:` | int counter | 60s | +//! | outcomes (local only) | `{id}:successes`, `{id}:fails` | int counter | 60s | +//! +//! From `router_utils/cooldown_cache.py`, `router_utils/cooldown_handlers.py`, +//! `router_utils/router_callbacks/track_deployment_metrics.py`, `RouterCacheEnum` in +//! `types/router.py` and `Router._increment_usage` / `deployment_callback_on_failure` in +//! `router.py`. + +use std::{ + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use serde::{Deserialize, Serialize}; +use tracing::debug; + +use crate::{ + cache::{InMemoryRouterCache, RouterCache}, + error::{AttemptError, Error, FailureKind}, + registry::Deployment, + settings::Settings, +}; + +/// `RoutingArgs.ttl`: per-minute usage keys expire after a minute. +pub const USAGE_TTL: Duration = Duration::from_secs(60); +/// `DEFAULT_FAILURE_THRESHOLD_PERCENT`. +pub const FAILURE_THRESHOLD_PERCENT: f64 = 0.5; +/// `DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS`. +pub const FAILURE_THRESHOLD_MINIMUM_REQUESTS: u64 = 5; +/// `SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD`. +pub const SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: u64 = 1000; +/// `exception_masker` in `CooldownCache`: the visible prefix of `exception_received`. +const EXCEPTION_VISIBLE_PREFIX: usize = 50; + +pub trait Clock: Send + Sync { + fn now(&self) -> SystemTime; +} + +pub struct SystemClock; + +impl Clock for SystemClock { + fn now(&self) -> SystemTime { + SystemTime::now() + } +} + +fn unix_seconds(time: SystemTime) -> f64 { + time.duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs_f64()) + .unwrap_or_default() +} + +/// `get_utc_datetime().strftime("%H-%M")`. +pub fn current_minute(time: SystemTime) -> String { + let seconds = time + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs()) + .unwrap_or_default() + % 86_400; + format!("{:02}-{:02}", seconds / 3600, (seconds % 3600) / 60) +} + +pub fn cooldown_key(id: &str) -> String { + format!("deployment:{id}:cooldown") +} + +pub fn allowed_fails_key(id: &str, suffix: Option<&str>) -> String { + match suffix { + Some(suffix) => format!("deployment:{id}:allowed_fails:{suffix}"), + None => format!("deployment:{id}:allowed_fails"), + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub enum UsageMetric { + #[strum(serialize = "tpm")] + Tpm, + #[strum(serialize = "rpm")] + Rpm, +} + +/// `RouterCacheEnum.TPM` / `RPM`: `model` is the deployment's `litellm_params.model`. +pub fn usage_key(id: &str, model: &str, metric: UsageMetric, minute: &str) -> String { + format!("global_router:{id}:{model}:{metric}:{minute}") +} + +/// `CooldownCacheValue`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CooldownValue { + pub exception_received: String, + pub status_code: String, + pub timestamp: f64, + /// Whole seconds are written as an int, as Python's default `cooldown_time` is, so the + /// "Try again in N seconds" text reads the same. + #[serde(serialize_with = "whole_as_int")] + pub cooldown_time: f64, +} + +fn whole_as_int(value: &f64, serializer: S) -> Result { + if value.fract() == 0.0 && value.abs() < i64::MAX as f64 { + serializer.serialize_i64(*value as i64) + } else { + serializer.serialize_f64(*value) + } +} + +/// `litellm.should_retry`. +pub fn should_retry_status(status: u16) -> bool { + matches!(status, 408 | 409 | 429) || status >= 500 +} + +/// `_is_cooldown_required`: a 4xx cools down only for 429, 401, 402, 404 and 408, every other +/// status does, and a connection error never does. +pub fn is_cooldown_required(failure: &AttemptError) -> bool { + if failure.kind == FailureKind::ApiConnection { + return false; + } + match failure.cooldown_status() { + status @ 400..500 => matches!(status, 429 | 401 | 402 | 408 | 404), + _ => true, + } +} + +/// `SensitiveDataMasker(visible_prefix=50, visible_suffix=0, mask_short_values=False)`. +fn mask_exception(text: &str) -> String { + let length = text.chars().count(); + if length <= EXCEPTION_VISIBLE_PREFIX { + return text.to_owned(); + } + let mut masked: String = text.chars().take(EXCEPTION_VISIBLE_PREFIX).collect(); + masked.extend(std::iter::repeat_n('*', length - EXCEPTION_VISIBLE_PREFIX)); + masked +} + +/// What the cooldown rules need to know about the deployment's group. +#[derive(Clone, Copy, Debug, Default)] +pub struct GroupShape { + /// The deployment is the only one in its model group. + pub single_deployment: bool, + /// The requested group has a fallback configured. + pub has_fallback_group: bool, +} + +/// An outcome the router did not see itself (`Health.report`). +#[derive(Clone, Debug)] +pub enum Outcome { + Success { total_tokens: u64 }, + Failure(AttemptError), +} + +#[derive(Clone)] +pub struct Health { + cache: Arc, + local: InMemoryRouterCache, + clock: Arc, +} + +impl Health { + pub fn new(cache: Arc, clock: Arc) -> Self { + Self { + cache, + local: InMemoryRouterCache::default(), + clock, + } + } + + pub fn cache(&self) -> &Arc { + &self.cache + } + + fn now(&self) -> f64 { + unix_seconds(self.clock.now()) + } + + /// The ids among `ids` with an unexpired cooldown (`async_get_active_cooldowns`). + pub async fn active_cooldowns( + &self, + ids: &[String], + ) -> Result, Error> { + let keys: Vec = ids.iter().map(|id| cooldown_key(id)).collect(); + let values = self.cache.async_batch_get_cache(&keys).await?; + let now = self.now(); + Ok(ids + .iter() + .zip(values) + .filter_map(|(id, value)| { + let value: CooldownValue = serde_json::from_value(value?).ok()?; + (value.timestamp + value.cooldown_time - now > 0.0).then(|| (id.clone(), value)) + }) + .collect()) + } + + /// `add_deployment_to_cooldown`. + pub async fn add_cooldown( + &self, + id: &str, + failure: &AttemptError, + cooldown: Duration, + ) -> Result<(), Error> { + let value = CooldownValue { + exception_received: mask_exception(failure.detail.as_deref().unwrap_or_default()), + status_code: failure.cooldown_status().to_string(), + timestamp: self.now(), + cooldown_time: cooldown.as_secs_f64(), + }; + let value = serde_json::to_value(value).unwrap_or_default(); + self.cache + .async_set_cache(&cooldown_key(id), value, Some(cooldown)) + .await?; + Ok(()) + } + + fn local_count(&self, key: &str) -> u64 { + self.local + .get(key) + .and_then(|value| value.as_f64()) + .map_or(0, |count| count as u64) + } + + /// The per-minute (successes, fails) this worker saw. + pub fn outcomes_this_minute(&self, id: &str) -> (u64, u64) { + ( + self.local_count(&format!("{id}:successes")), + self.local_count(&format!("{id}:fails")), + ) + } + + /// Cooldown duration precedence from `deployment_callback_on_failure`: the deployment's + /// `cooldown_time`, then the provider's retry-after, then the router default. + pub fn cooldown_duration( + settings: &Settings, + deployment: &Deployment, + failure: &AttemptError, + ) -> Duration { + Self::configured_cooldown(deployment) + .or(failure.retry_after) + .unwrap_or(settings.cooldown_time) + } + + /// The deployment's own `cooldown_time`: `model_info` first, then `litellm_params` + /// (`_first_present(model_info, litellm_params, key="cooldown_time")`). It is both the + /// cooldown and the TTL of the deployment's `allowed_fails:generic` counter. + fn configured_cooldown(deployment: &Deployment) -> Option { + deployment + .model_info + .cooldown_time + .or_else(|| { + deployment + .raw_params + .get("cooldown_time") + .and_then(|v| v.as_f64()) + }) + .filter(|seconds| *seconds >= 0.0) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + } + + /// Record a failed attempt and cool the deployment down when the rules say so. Returns the + /// cooldown applied, if any. + pub async fn record_failure( + &self, + settings: &Settings, + deployment: &Deployment, + shape: GroupShape, + failure: &AttemptError, + ) -> Result, Error> { + let cooldown = Self::cooldown_duration(settings, deployment, failure); + self.local + .increment(&format!("{}:fails", deployment.id), 1.0, Some(USAGE_TTL)); + if !Self::should_run_cooldown_logic(settings, deployment, failure, cooldown) { + return Ok(None); + } + if !self + .should_cooldown(settings, deployment, shape, failure) + .await? + { + return Ok(None); + } + debug!(deployment_id = %deployment.id, cooldown_seconds = cooldown.as_secs_f64(), "cooling deployment down"); + self.add_cooldown(&deployment.id, failure, cooldown).await?; + Ok(Some(cooldown)) + } + + /// `_should_run_cooldown_logic`. + fn should_run_cooldown_logic( + settings: &Settings, + deployment: &Deployment, + failure: &AttemptError, + cooldown: Duration, + ) -> bool { + !matches!( + failure.kind, + FailureKind::Declined | FailureKind::CallerCredential + ) && !cooldown.is_zero() + && !settings.disable_cooldowns + && is_cooldown_required(failure) + && !settings + .provider_default_deployment_ids + .contains(&deployment.id) + } + + /// `_should_cooldown_deployment`. + async fn should_cooldown( + &self, + settings: &Settings, + deployment: &Deployment, + shape: GroupShape, + failure: &AttemptError, + ) -> Result { + let status = failure.cooldown_status(); + let policy_fails = settings.allowed_fails_for(failure.kind); + let lone = shape.single_deployment && !shape.has_fallback_group; + if shape.single_deployment && status == 402 && policy_fails.is_none() { + return Ok(false); + } + if let Some(allowed) = deployment.model_info.allowed_fails { + // A deployment-wide `allowed_fails` keeps the lone-deployment safety net + // (`_should_cooldown_based_on_deployment_policy`). + if lone { + return Ok(false); + } + let ttl = Self::configured_cooldown(deployment).unwrap_or(settings.cooldown_time); + return self + .exceeds_allowed_fails(&deployment.id, Some("generic"), allowed, ttl) + .await; + } + if settings.allowed_fails_policy.is_none() && settings.allowed_fails.is_none() { + let (successes, fails) = self.outcomes_this_minute(&deployment.id); + let total = successes + fails; + let percent_fails = if total > 0 { + fails as f64 / total as f64 + } else { + 0.0 + }; + return Ok((status == 429 && !shape.single_deployment) + || (percent_fails == 1.0 && total >= SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD) + || (percent_fails > FAILURE_THRESHOLD_PERCENT + && total >= FAILURE_THRESHOLD_MINIMUM_REQUESTS + && !shape.single_deployment) + || !should_retry_status(status)); + } + if lone && policy_fails.is_none() { + return Ok(false); + } + let allowed = policy_fails.unwrap_or_else(|| settings.effective_allowed_fails()); + self.exceeds_allowed_fails(&deployment.id, None, allowed, settings.cooldown_time) + .await + } + + /// `should_cooldown_based_on_allowed_fails_policy`: count this failure, then compare. + async fn exceeds_allowed_fails( + &self, + id: &str, + suffix: Option<&str>, + allowed: u32, + ttl: Duration, + ) -> Result { + let fails = self + .cache + .async_increment(&allowed_fails_key(id, suffix), 1.0, Some(ttl)) + .await?; + Ok(fails > f64::from(allowed)) + } + + /// Count a success, and the tokens and request against the deployment's per-minute usage + /// when it has limits to track. + pub async fn record_success( + &self, + deployment: &Deployment, + total_tokens: u64, + ) -> Result<(), Error> { + self.local.increment( + &format!("{}:successes", deployment.id), + 1.0, + Some(USAGE_TTL), + ); + if !deployment.tracks_usage() { + return Ok(()); + } + let minute = current_minute(self.clock.now()); + let model = &deployment.params.model; + self.cache + .async_increment( + &usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + total_tokens as f64, + Some(USAGE_TTL), + ) + .await?; + self.cache + .async_increment( + &usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), + 1.0, + Some(USAGE_TTL), + ) + .await?; + Ok(()) + } + + /// This minute's (tpm, rpm) for a deployment. + pub async fn usage(&self, deployment: &Deployment) -> Result<(u64, u64), Error> { + let minute = current_minute(self.clock.now()); + let model = &deployment.params.model; + let keys = [UsageMetric::Tpm, UsageMetric::Rpm] + .map(|metric| usage_key(&deployment.id, model, metric, &minute)); + let values = self.cache.async_batch_get_cache(&keys).await?; + let count = |index: usize| { + values + .get(index) + .cloned() + .flatten() + .and_then(|value| value.as_f64()) + .map_or(0, |count| count as u64) + }; + Ok((count(0), count(1))) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::rate_limit(FailureKind::RateLimit, Some(429), true)] + #[case::auth(FailureKind::Authentication, Some(401), true)] + #[case::payment(FailureKind::Other, Some(402), true)] + #[case::not_found(FailureKind::NotFound, Some(404), true)] + #[case::timeout(FailureKind::Timeout, Some(408), true)] + #[case::bad_request(FailureKind::BadRequest, Some(400), false)] + #[case::forbidden(FailureKind::PermissionDenied, Some(403), false)] + #[case::server(FailureKind::InternalServer, Some(500), true)] + #[case::no_status(FailureKind::Other, None, true)] + #[case::connection(FailureKind::ApiConnection, Some(500), false)] + fn cooldown_required_follows_python( + #[case] kind: FailureKind, + #[case] status: Option, + #[case] expected: bool, + ) { + assert_eq!( + is_cooldown_required(&AttemptError::new(kind, status)), + expected + ); + } + + #[rstest] + #[case::midnight(0, "00-00")] + #[case::afternoon(13 * 3600 + 7 * 60 + 59, "13-07")] + fn current_minute_is_utc_hours_and_minutes(#[case] seconds: u64, #[case] expected: &str) { + let time = UNIX_EPOCH + Duration::from_secs(86_400 * 3 + seconds); + assert_eq!(current_minute(time), expected); + } + + #[rstest] + fn keys_match_python() { + assert_eq!( + ( + cooldown_key("abc"), + allowed_fails_key("abc", Some("generic")), + usage_key("abc", "openai/gpt-4o", UsageMetric::Tpm, "09-05"), + ), + ( + "deployment:abc:cooldown".to_owned(), + "deployment:abc:allowed_fails:generic".to_owned(), + "global_router:abc:openai/gpt-4o:tpm:09-05".to_owned(), + ) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs new file mode 100644 index 00000000000..c0403f15206 --- /dev/null +++ b/litellm-rust/crates/router-core/src/host.rs @@ -0,0 +1,168 @@ +//! [`Host`]: what the embedding process provides. `PythonHost` serves it from the existing Python +//! code; `NativeHost` serves it in the Rust server. + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use serde_json::Value; + +use crate::{ + cache::RouterCache, + config::{Op, RoutingStrategy}, + error::{AttemptError, Error, FailureKind, Payload}, + health::Health, + registry::{Deployment, Registry}, + settings::Settings, +}; + +/// One call into the router. +#[derive(Clone, Debug, Default)] +pub struct RouterRequest { + /// The requested name: a model group, alias, routing group, deployment id or wildcard match. + pub model: String, + /// The request params, for native executors and plugins that read them. + pub body: Value, + /// `user_api_key_team_id`. + pub team_id: Option, + pub specific_deployment: bool, + /// A per-request `num_retries`, which wins over the router's. + pub num_retries: Option, + pub disable_fallbacks: bool, + /// `_router_weights`: model group to deployment id to weight. + pub router_weights: Option>>, + /// What the host needs to run the call itself, such as the Python kwargs. + pub payload: Option, +} + +/// One attempt against one deployment. +#[derive(Clone, Debug)] +pub struct Attempt { + pub op: Op, + pub request: Arc, + pub deployment: Arc, + /// The name this attempt resolved: the requested name, or the fallback target. + pub name: String, + /// The model group being served, after fallbacks. + pub model_group: String, + /// The `order` level an order fallback targets in the same group. + pub target_order: Option, + /// Every deployment `name` resolved to for this caller, the picked one included. + pub candidate_ids: Vec, + /// 0 for the first attempt in this group, then the retry number. + pub retry: u32, + pub fallback_depth: u32, + /// How many fallbacks deep the walk is: 0 for the requested name. + pub fallback_level: u32, + pub timeout: Option, +} + +/// What a successful attempt returns. +#[derive(Clone, Debug)] +pub struct Output { + pub value: Payload, + /// Prompt plus completion tokens, counted against the deployment's per-minute usage. + pub total_tokens: u64, +} + +/// What a plugin or strategy sees: never the router itself. +pub struct RoutingContext<'a> { + pub registry: &'a Registry, + pub settings: &'a Settings, + pub health: &'a Health, + pub request: &'a RouterRequest, + pub op: &'a Op, + pub model_group: &'a str, +} + +/// Which selection step the host is asked to run. +#[derive(Clone, Debug, PartialEq)] +pub enum SelectStage { + /// Pre-call checks and sub-routers registered as plugins: return the deployments that remain. + Plugins, + /// A strategy not yet ported: return the candidates with the pick first. + Strategy(RoutingStrategy), +} + +/// Routing events, delivered through [`Host::observe`]. +#[derive(Clone, Debug, PartialEq)] +pub enum RouterEvent { + Selected { + model_group: String, + deployment_id: String, + }, + Succeeded { + deployment_id: String, + }, + AttemptFailed { + deployment_id: String, + kind: FailureKind, + status: Option, + }, + CooldownSet { + deployment_id: String, + cooldown: Duration, + }, + Retry { + model_group: String, + retry: u32, + delay: Duration, + }, + Fallback { + from: String, + to: String, + /// The order level of `to` a same-group hop walks to. + order: Option, + depth: u32, + /// The fallback level `to` is tried at: 1 for the request's own fallbacks. + level: u32, + }, + /// A group's retry loop ran every retry and the last one failed: Python stamps + /// `num_retries` and `max_retries` on that error. + RetriesExhausted { + model_group: String, + attempted: u32, + max: u32, + }, +} + +#[async_trait] +pub trait Host: Send + Sync { + /// Run a non-native attempt. + async fn invoke(&self, attempt: &Attempt) -> Result; + + /// Run a selection stage the core does not implement. + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error>; + + /// Access and budget checks before a cross-group fallback target. + async fn allow_fallback(&self, request: &RouterRequest, target: &str) -> bool; + + async fn sleep(&self, duration: Duration); + + fn observe(&self, event: &RouterEvent); + + /// The shared store routing state lives in. + fn cache(&self) -> Arc; + + /// Whether [`Self::invoke`] can run an attempt. A `NativeHost` cannot, so every served + /// operation needs a native executor. + fn can_invoke(&self) -> bool { + true + } + + /// Whether [`Self::select`] can run strategies and plugins. + fn can_select(&self) -> bool { + true + } + + /// Whether [`Self::invoke`] records each attempt's outcome itself (cooldowns, allowed-fails + /// counters, usage), as the Python router's own callbacks do. The core then leaves the + /// health store to the host for host-run attempts, so no failure is counted twice. + fn records_outcomes(&self) -> bool { + false + } +} diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 05d95d02c84..e41de752369 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -61,3 +61,285 @@ //! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive) //! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and //! the current request), never the router itself. +//! +//! # Modules +//! - [`registry`]: deployments and their indices, `deployment_id`, `ChangeSet`. +//! - [`config`] / [`settings`]: static and dynamic config, and the frozen `Settings` view. +//! - [`cache`] / [`health`]: the shared routing cache, cooldowns and usage. +//! - [`select`]: resolve, filter chain, simple-shuffle. +//! - [`call`]: retries and fallbacks. +//! - [`host`] / [`executor`]: the `Host` seam and per-attempt executors. + +pub mod cache; +pub mod call; +pub mod config; +pub mod error; +pub mod executor; +pub mod health; +pub mod host; +pub mod registry; +pub mod select; +pub mod settings; + +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; + +use arc_swap::ArcSwap; +use parking_lot::Mutex; + +pub use crate::{ + cache::{InMemoryRouterCache, RouterCache}, + call::RouterResponse, + config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy}, + error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported}, + executor::{Executor, Executors, NativeRoute}, + health::{Clock, Health, Outcome, SystemClock}, + host::{Attempt, Host, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage}, + registry::{ChangeSet, Deployment, DeploymentSpec, Registry, RegistryVersion}, + select::{Randomness, SeededRandomness, ThreadRandomness}, + settings::Settings, +}; +use crate::{health::GroupShape, registry::DeploymentFilter, select::SelectInputs}; + +/// The registry and settings one call reads, swapped atomically by `apply`. +pub(crate) struct Snapshot { + pub registry: Arc, + pub settings: Arc, +} + +pub(crate) struct Core { + host: Arc, + executors: Executors, + rng: Arc, + health: Health, + config: RouterConfig, + snapshot: ArcSwap, + writer: Mutex<()>, + retired: AtomicBool, +} + +impl Core { + /// Reject what this host cannot serve, before any request. + fn validate(&self, registry: &Registry, settings: &Settings) -> Result<(), ConfigError> { + let can_select = self.host.can_select(); + let strategies = std::iter::once(&settings.routing_strategy).chain( + registry + .routing_groups() + .iter() + .map(|group| &group.routing_strategy), + ); + for strategy in strategies { + if !strategy.is_native() && !can_select { + return Err(Unsupported::Strategy(strategy.clone()).into()); + } + } + if let Some(plugin) = settings.plugins.first() + && !can_select + { + return Err(Unsupported::Plugin(plugin.name.clone()).into()); + } + for deployment in registry.deployments(&DeploymentFilter::default()) { + for op in &self.config.served_ops { + self.executors.resolve(op, deployment.provider())?; + } + } + Ok(()) + } +} + +/// Builds a [`Router`], with native routes and test seams. +pub struct RouterBuilder { + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + executors: Executors, + rng: Arc, + clock: Arc, +} + +impl RouterBuilder { + /// Register a native route for an (op, provider) rule; a pair outside + /// [`executor::NATIVE_RULES`] is ignored and stays on the host. + #[must_use] + pub fn native_route(mut self, op: Op, provider: &str, route: Arc) -> Self { + self.executors.register(op, provider, route); + self + } + + #[must_use] + pub fn randomness(self, rng: Arc) -> Self { + Self { rng, ..self } + } + + #[must_use] + pub fn clock(self, clock: Arc) -> Self { + Self { clock, ..self } + } + + pub fn build(self) -> Result { + let registry = Registry::build( + self.config.model_list.clone(), + self.dynamic.model_alias_map.clone(), + self.config.routing_groups.clone(), + )?; + let settings = Settings::new(&self.config, Arc::new(self.dynamic)); + let core = Core { + health: Health::new(self.host.cache(), self.clock), + host: self.host, + executors: self.executors, + rng: self.rng, + config: self.config, + snapshot: ArcSwap::from_pointee(Snapshot { + registry: Arc::new(Registry::default()), + settings: Arc::new(settings.clone()), + }), + writer: Mutex::new(()), + retired: AtomicBool::new(false), + }; + core.validate(®istry, &settings)?; + core.snapshot.store(Arc::new(Snapshot { + registry: Arc::new(registry), + settings: Arc::new(settings), + })); + Ok(Router { + core: Arc::new(core), + }) + } +} + +/// The router: the entry points `litellm.Router` delegates to. +#[derive(Clone)] +pub struct Router { + core: Arc, +} + +impl Router { + pub fn new( + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + ) -> Result { + Self::builder(config, dynamic, host).build() + } + + pub fn builder( + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + ) -> RouterBuilder { + RouterBuilder { + executors: Executors::new(host.can_invoke()), + config, + dynamic, + host, + rng: Arc::new(ThreadRandomness), + clock: Arc::new(SystemClock), + } + } + + fn live(&self) -> Result<(), Error> { + if self.core.retired.load(Ordering::Acquire) { + return Err(Error::Retired); + } + Ok(()) + } + + /// Select, run the attempt, and retry or fall back until it succeeds or the policy is + /// exhausted. + pub async fn call(&self, op: Op, request: RouterRequest) -> Result { + self.live()?; + self.core.call(op, Arc::new(request)).await + } + + /// Resolve, filter and pick a deployment without calling it. + pub async fn select(&self, op: &Op, request: &RouterRequest) -> Result, Error> { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + let candidates = self + .core + .candidates( + &snapshot, + op, + request, + &request.model, + &SelectInputs::default(), + ) + .await?; + self.core.pick(&snapshot, op, request, &candidates).await + } + + pub fn registry(&self) -> Arc { + self.core.snapshot.load().registry.clone() + } + + pub fn settings(&self) -> Arc { + self.core.snapshot.load().settings.clone() + } + + pub fn health(&self) -> &Health { + &self.core.health + } + + /// Record an outcome the router did not see itself, such as a health check. + pub async fn report(&self, id: &str, outcome: Outcome) -> Result<(), Error> { + let snapshot = self.core.snapshot.load_full(); + let Some(deployment) = snapshot.registry.deployment(id) else { + return Ok(()); + }; + match outcome { + Outcome::Success { total_tokens } => { + self.core + .health + .record_success(&deployment, total_tokens) + .await + } + Outcome::Failure(failure) => { + let shape = GroupShape { + single_deployment: snapshot + .registry + .group(&deployment.model_name) + .is_some_and(|group| group.len() == 1), + has_fallback_group: snapshot + .settings + .dynamic + .fallback_targets(&[deployment.model_name.as_str()]) + .is_some(), + }; + self.core + .health + .record_failure(&snapshot.settings, &deployment, shape, &failure) + .await + .map(drop) + } + } + } + + /// Apply one change set atomically and return the new registry version. + pub fn apply(&self, changes: ChangeSet) -> Result { + let _writer = self.core.writer.lock(); + let current = self.core.snapshot.load_full(); + let aliases = changes + .settings + .as_ref() + .map(|dynamic| dynamic.model_alias_map.clone()); + let registry = current.registry.apply(&changes, aliases)?; + let settings = match changes.settings { + Some(dynamic) => Settings::new(&self.core.config, Arc::new(dynamic)), + None => Settings::clone(¤t.settings), + }; + self.core.validate(®istry, &settings)?; + let version = registry.version(); + self.core.snapshot.store(Arc::new(Snapshot { + registry: Arc::new(registry), + settings: Arc::new(settings), + })); + Ok(version) + } + + /// Stop serving: every later call and selection fails with [`Error::Retired`]. + pub fn retire(&self) { + self.core.retired.store(true, Ordering::Release); + } +} diff --git a/litellm-rust/crates/router-core/src/registry.rs b/litellm-rust/crates/router-core/src/registry.rs new file mode 100644 index 00000000000..eb2d3047c45 --- /dev/null +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -0,0 +1,820 @@ +//! The configured deployments, indexed by id, model group, alias, wildcard pattern, routing +//! group and team. A [`Registry`] is immutable; [`Registry::apply`] builds the next version. + +use std::{ + collections::{BTreeMap, BTreeSet, HashMap}, + fmt::Write as _, + sync::Arc, +}; + +use litellm_router_types::{LitellmParams, Spelled}; +use serde::Deserialize; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; + +use crate::error::ConfigError; + +/// A registry's version; every change produces a higher one. +pub type RegistryVersion = u64; + +/// A deployment as a config or an `apply` spells it. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +pub struct DeploymentSpec { + pub model_name: String, + pub litellm_params: Map, + #[serde(default)] + pub model_info: Map, +} + +/// `routing_groups`: a named subset of model groups sharing a routing strategy. +#[derive(Clone, Debug, PartialEq, Deserialize)] +pub struct RoutingGroupSpec { + pub group_name: String, + pub models: Vec, + pub routing_strategy: crate::config::RoutingStrategy, +} + +/// The `model_info` fields the router reads; every other key stays in [`Self::raw`]. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ModelInfo { + pub team_id: Option, + pub team_public_model_name: Option, + pub access_groups: Vec, + pub blocked: bool, + pub cooldown_time: Option, + pub allowed_fails: Option, + pub raw: Map, +} + +impl ModelInfo { + fn from_raw(raw: Map) -> Self { + let text = |key: &str| raw.get(key).and_then(Value::as_str).map(str::to_owned); + Self { + team_id: text("team_id"), + team_public_model_name: text("team_public_model_name"), + access_groups: raw + .get("access_groups") + .and_then(Value::as_array) + .map(|groups| { + groups + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default(), + blocked: raw.get("blocked").and_then(Value::as_bool) == Some(true), + cooldown_time: raw.get("cooldown_time").and_then(Value::as_f64), + allowed_fails: raw + .get("allowed_fails") + .and_then(Value::as_u64) + .and_then(|fails| u32::try_from(fails).ok()), + raw, + } + } +} + +/// One configured deployment. +#[derive(Clone, Debug, PartialEq)] +pub struct Deployment { + pub id: String, + pub model_name: String, + pub params: LitellmParams, + /// `litellm_params` as written, in order: what the id hashes and the host receives. + pub raw_params: Map, + pub model_info: ModelInfo, +} + +impl Deployment { + pub fn from_spec(spec: DeploymentSpec) -> Result { + let params: LitellmParams = serde_json::from_value(Value::Object( + spec.litellm_params.clone(), + )) + .map_err(|source| ConfigError::InvalidDeployment { + model_name: spec.model_name.clone(), + source, + })?; + let id = match spec.model_info.get("id") { + Some(Value::String(id)) => id.clone(), + Some(Value::Number(id)) => id.to_string(), + _ => deployment_id(&spec.model_name, &spec.litellm_params), + }; + Ok(Self { + id, + model_name: spec.model_name, + params, + raw_params: spec.litellm_params, + model_info: ModelInfo::from_raw(spec.model_info), + }) + } + + /// The provider: `custom_llm_provider`, else the `model` prefix before `/`. + pub fn provider(&self) -> &str { + self.params + .custom_llm_provider + .as_deref() + .or_else(|| { + self.params + .model + .split_once('/') + .map(|(provider, _)| provider) + }) + .unwrap_or("openai") + } + + /// A numeric `litellm_params` value, as `simple_shuffle._metric_weight` reads it. + pub fn metric(&self, key: &str) -> Option { + let spelled = |value: &Option>| match value { + Some(Spelled::Value(value)) => Some(*value), + Some(Spelled::Text(_)) | None => None, + }; + match key { + "rpm" => spelled(&self.params.rpm), + "tpm" => spelled(&self.params.tpm), + _ => self.raw_params.get(key).and_then(Value::as_f64), + } + } + + /// `litellm_params.order` (`get_deployment_order`). + pub fn order(&self) -> Option { + self.raw_params.get("order").and_then(Value::as_i64) + } + + /// Whether `team_id` may use this deployment (`deployment_usable_by_team`). + pub fn usable_by_team(&self, team_id: Option<&str>) -> bool { + self.model_info + .team_id + .as_deref() + .is_none_or(|owner| Some(owner) == team_id) + } + + /// Whether the router counts per-minute usage for this deployment + /// (`_deployment_tracks_usage`, without the model-catalog limits). + pub fn tracks_usage(&self) -> bool { + self.params.tpm.is_some() + || self.params.rpm.is_some() + || self.params.itpm.is_some() + || self.params.otpm.is_some() + } +} + +/// The connection params of one deployment (`credentials(id)`). +#[derive(Clone, Debug, PartialEq)] +pub struct Credentials { + pub provider: String, + pub api_base: Option, + pub api_version: Option, + pub credential_name: Option, + pub params: LitellmParams, +} + +/// What a name resolved through. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ResolveKind { + SpecificDeployment, + Id, + Alias, + RoutingGroup, + Group, + Wildcard, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Resolution { + pub kind: ResolveKind, + /// The model group the deployments are routed as. + pub group: String, + pub deployments: Vec>, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct ResolveScope { + /// The caller named a deployment's `litellm_params.model` (`specific_deployment=True`). + pub specific_deployment: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct DeploymentFilter<'a> { + pub model_name: Option<&'a str>, + pub team_id: Option<&'a str>, + pub provider: Option<&'a str>, +} + +#[derive(Clone, Debug, Default)] +pub struct NameFilter { + pub include_aliases: bool, +} + +/// One atomic change: upserts, deletes and, optionally, a new dynamic config. +#[derive(Clone, Debug, Default)] +pub struct ChangeSet { + pub upserts: Vec, + pub deletes: Vec, + pub settings: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct Registry { + version: RegistryVersion, + deployments: Vec>, + by_id: HashMap, + by_group: BTreeMap>, + by_team: HashMap>, + wildcards: Vec<(String, Vec)>, + aliases: BTreeMap, + routing_groups: Vec, +} + +impl Registry { + pub fn build( + specs: Vec, + aliases: BTreeMap, + routing_groups: Vec, + ) -> Result { + let deployments = specs + .into_iter() + .map(|spec| Deployment::from_spec(spec).map(Arc::new)) + .collect::, _>>()?; + let mut seen = BTreeSet::new(); + if let Some(duplicate) = deployments.iter().find(|d| !seen.insert(d.id.as_str())) { + return Err(ConfigError::DuplicateId { + id: duplicate.id.clone(), + }); + } + Ok(Self::index(0, deployments, aliases, routing_groups)) + } + + fn index( + version: RegistryVersion, + deployments: Vec>, + aliases: BTreeMap, + routing_groups: Vec, + ) -> Self { + let mut by_id = HashMap::new(); + let mut by_group: BTreeMap> = BTreeMap::new(); + let mut by_team: HashMap> = HashMap::new(); + for (index, deployment) in deployments.iter().enumerate() { + by_id.insert(deployment.id.clone(), index); + by_group + .entry(deployment.model_name.clone()) + .or_default() + .push(index); + if let Some(team) = &deployment.model_info.team_id { + by_team.entry(team.clone()).or_default().push(index); + } + } + // In the order each pattern first appears, then most specific first, as + // `PatternUtils.sorted_patterns` keeps them: a stable sort, longest regex first, + // then the most regex syntax. + let mut wildcards: Vec<(String, Vec)> = Vec::new(); + for deployment in &deployments { + let name = &deployment.model_name; + if name.contains('*') && !wildcards.iter().any(|(pattern, _)| pattern == name) { + wildcards.push(( + name.clone(), + by_group.get(name).cloned().unwrap_or_default(), + )); + } + } + wildcards.sort_by_key(|(pattern, _)| std::cmp::Reverse(pattern_specificity(pattern))); + Self { + version, + deployments, + by_id, + by_group, + by_team, + wildcards, + aliases, + routing_groups, + } + } + + pub fn version(&self) -> RegistryVersion { + self.version + } + + /// The next registry: deletes, then upserts. An upserted deployment whose id exists is + /// removed and appended at the end, as the Python router does on an edit. + pub fn apply( + &self, + changes: &ChangeSet, + aliases: Option>, + ) -> Result { + let mut deployments = self.deployments.clone(); + deployments.retain(|deployment| !changes.deletes.contains(&deployment.id)); + for spec in &changes.upserts { + let deployment = Deployment::from_spec(spec.clone())?; + deployments.retain(|existing| existing.id != deployment.id); + deployments.push(Arc::new(deployment)); + } + Ok(Self::index( + self.version + 1, + deployments, + aliases.unwrap_or_else(|| self.aliases.clone()), + self.routing_groups.clone(), + )) + } + + fn at(&self, indices: &[usize]) -> Vec> { + indices + .iter() + .map(|index| self.deployments[*index].clone()) + .collect() + } + + pub fn deployment(&self, id: &str) -> Option> { + self.by_id + .get(id) + .map(|index| self.deployments[*index].clone()) + } + + pub fn deployments(&self, filter: &DeploymentFilter<'_>) -> Vec> { + let candidates = match (filter.model_name, filter.team_id) { + (Some(name), _) => self + .by_group + .get(name) + .map(|i| self.at(i)) + .unwrap_or_default(), + (None, Some(team)) => self + .by_team + .get(team) + .map(|i| self.at(i)) + .unwrap_or_default(), + (None, None) => self.deployments.clone(), + }; + candidates + .into_iter() + .filter(|d| { + filter + .team_id + .is_none_or(|team| d.model_info.team_id.as_deref() == Some(team)) + }) + .filter(|d| { + filter + .provider + .is_none_or(|provider| d.provider() == provider) + }) + .collect() + } + + /// The model group names, in name order, with aliases when asked. + pub fn names(&self, filter: &NameFilter) -> Vec { + let mut names: Vec = self.by_group.keys().cloned().collect(); + if filter.include_aliases { + names.extend(self.aliases.keys().cloned()); + } + names + } + + pub fn group(&self, name: &str) -> Option>> { + self.by_group.get(name).map(|indices| self.at(indices)) + } + + /// The routing group a model group belongs to, if any. + pub fn routing_group_of(&self, model_name: &str) -> Option<&RoutingGroupSpec> { + self.routing_groups + .iter() + .find(|group| group.models.iter().any(|model| model == model_name)) + } + + /// The strategy a routing group named `group_name` declares. + pub fn resolve_routing_group_strategy( + &self, + group_name: &str, + ) -> Option { + self.routing_groups + .iter() + .find(|group| group.group_name == group_name) + .map(|group| group.routing_strategy.clone()) + } + + pub(crate) fn routing_groups(&self) -> &[RoutingGroupSpec] { + &self.routing_groups + } + + /// Resolve a requested name, in the order of `_common_checks_available_deployment`: + /// a specific deployment, an id, an alias, a routing group, a model group, a wildcard. + pub fn resolve(&self, name: &str, scope: ResolveScope) -> Option { + let resolution = |kind, group: &str, deployments: Vec>| { + (!deployments.is_empty()).then(|| Resolution { + kind, + group: group.to_owned(), + deployments, + }) + }; + if scope.specific_deployment { + let matching = self + .deployments + .iter() + .filter(|d| d.params.model == name) + .cloned() + .collect(); + return resolution(ResolveKind::SpecificDeployment, name, matching); + } + if !self.by_group.contains_key(name) + && let Some(deployment) = self.deployment(name) + { + let group = deployment.model_name.clone(); + return resolution(ResolveKind::Id, &group, vec![deployment]); + } + if let Some(target) = self.aliases.get(name) { + return self + .group(target) + .and_then(|deployments| resolution(ResolveKind::Alias, target, deployments)); + } + if let Some(routing_group) = self.routing_groups.iter().find(|g| g.group_name == name) { + let deployments = routing_group + .models + .iter() + .filter_map(|model| self.by_group.get(model)) + .flat_map(|indices| self.at(indices)) + .collect(); + return resolution(ResolveKind::RoutingGroup, name, deployments); + } + if let Some(indices) = self.by_group.get(name) { + return resolution(ResolveKind::Group, name, self.at(indices)); + } + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .and_then(|(pattern, indices)| { + resolution(ResolveKind::Wildcard, pattern, self.at(indices)) + }) + } + + /// Access group name to the model groups in it. + pub fn access_groups(&self) -> BTreeMap> { + let mut groups: BTreeMap> = BTreeMap::new(); + for deployment in &self.deployments { + for group in &deployment.model_info.access_groups { + groups + .entry(group.clone()) + .or_default() + .insert(deployment.model_name.clone()); + } + } + groups + } + + pub fn credentials(&self, id: &str) -> Option { + self.deployment(id).map(|deployment| Credentials { + provider: deployment.provider().to_owned(), + api_base: deployment.params.api_base.clone(), + api_version: deployment.params.api_version.clone(), + credential_name: deployment.params.litellm_credential_name.clone(), + params: deployment.params.clone(), + }) + } +} + +/// `PatternUtils.calculate_pattern_specificity` over the pattern's regex +/// (`re.escape(pattern).replace(r"\*", "(.*)")`): its length, then how much regex syntax it has. +fn pattern_specificity(pattern: &str) -> (usize, usize) { + let mut regex = String::new(); + for c in pattern.chars() { + match c { + '*' => regex.push_str("(.*)"), + '(' | ')' | '[' | ']' | '{' | '}' | '?' | '+' | '-' | '|' | '^' | '$' | '\\' | '.' + | '&' | '~' | '#' | ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c' => { + regex.push('\\'); + regex.push(c); + } + _ => regex.push(c), + } + } + let complexity = regex + .chars() + .filter(|c| matches!(c, '*' | '+' | '?' | '\\' | '^' | '$' | '|' | '(' | ')')) + .count(); + (regex.chars().count(), complexity) +} + +/// `*` matches any run of characters, as the pattern router's `(.*)` does under `re.match`: +/// anchored at the start only, so the name may run on past the pattern's last literal. +fn wildcard_matches(pattern: &str, name: &str) -> bool { + let mut parts = pattern.split('*'); + let first = parts.next().unwrap_or_default(); + let Some(mut rest) = name.strip_prefix(first) else { + return false; + }; + let parts: Vec<&str> = parts.collect(); + for part in parts { + match rest.find(part) { + Some(at) => rest = &rest[at + part.len()..], + None => return false, + } + } + true +} + +/// `deployment_id` in `litellm/router_api.py`: sha256 over the model group and every +/// `litellm_params` key and value in order. A string is taken as is, a dict as `json.dumps`, and +/// anything else as Python's `str()`. +pub fn deployment_id(model_group: &str, litellm_params: &Map) -> String { + let mut text = String::from(model_group); + for (key, value) in litellm_params { + text.push_str(key); + match value { + Value::String(value) => text.push_str(value), + Value::Object(_) => json_dumps(value, &mut text), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::Array(_) => { + python_str(value, &mut text); + } + } + } + let digest = Sha256::digest(text.as_bytes()); + let mut hex = String::with_capacity(64); + for byte in digest { + let _ = write!(hex, "{byte:02x}"); + } + hex +} + +/// Python's `repr(float)`: the shortest round-trip digits, positional for exponents in +/// `-4..16`, scientific otherwise. +fn python_float(value: f64, out: &mut String) { + if value.is_nan() { + out.push_str("nan"); + return; + } + if value.is_infinite() { + out.push_str(if value > 0.0 { "inf" } else { "-inf" }); + return; + } + let scientific = format!("{value:e}"); + let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); + let exponent: i32 = exponent.parse().unwrap_or(0); + let (sign, mantissa) = mantissa + .strip_prefix('-') + .map_or(("", mantissa), |m| ("-", m)); + let digits: String = mantissa.chars().filter(char::is_ascii_digit).collect(); + out.push_str(sign); + if (-4..16).contains(&exponent) { + if exponent < 0 { + out.push_str("0."); + out.extend(std::iter::repeat_n('0', (-exponent - 1) as usize)); + out.push_str(&digits); + } else { + let point = exponent as usize + 1; + if digits.len() <= point { + out.push_str(&digits); + out.extend(std::iter::repeat_n('0', point - digits.len())); + out.push_str(".0"); + } else { + out.push_str(&digits[..point]); + out.push('.'); + out.push_str(&digits[point..]); + } + } + } else { + out.push_str(&digits[..1]); + if digits.len() > 1 { + out.push('.'); + out.push_str(&digits[1..]); + } + let _ = write!( + out, + "e{}{:02}", + if exponent < 0 { '-' } else { '+' }, + exponent.abs() + ); + } +} + +fn number(value: &serde_json::Number, out: &mut String) { + match value.as_f64() { + Some(float) if !value.is_i64() && !value.is_u64() => python_float(float, out), + Some(_) | None => out.push_str(&value.to_string()), + } +} + +/// `json.dumps` with its defaults: `", "` and `": "` separators and `ensure_ascii`. +fn json_dumps(value: &Value, out: &mut String) { + match value { + Value::Null => out.push_str("null"), + Value::Bool(flag) => out.push_str(if *flag { "true" } else { "false" }), + Value::Number(value) => number(value, out), + Value::String(text) => json_string(text, out), + Value::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + json_dumps(item, out); + } + out.push(']'); + } + Value::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + json_string(key, out); + out.push_str(": "); + json_dumps(item, out); + } + out.push('}'); + } + } +} + +fn json_string(text: &str, out: &mut String) { + out.push('"'); + for ch in text.chars() { + match ch { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + '\u{08}' => out.push_str("\\b"), + '\u{0c}' => out.push_str("\\f"), + ch if ch.is_ascii() && !ch.is_ascii_control() => out.push(ch), + ch => { + let mut units = [0u16; 2]; + for unit in ch.encode_utf16(&mut units) { + let _ = write!(out, "\\u{unit:04x}"); + } + } + } + } + out.push('"'); +} + +/// Python's `str()` of a JSON value: `repr` for the items of a list or dict. +fn python_str(value: &Value, out: &mut String) { + match value { + Value::Null => out.push_str("None"), + Value::Bool(flag) => out.push_str(if *flag { "True" } else { "False" }), + Value::Number(value) => number(value, out), + Value::String(text) => python_repr_str(text, out), + Value::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + python_str(item, out); + } + out.push(']'); + } + Value::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + python_repr_str(key, out); + out.push_str(": "); + python_str(item, out); + } + out.push('}'); + } + } +} + +fn python_repr_str(text: &str, out: &mut String) { + let quote = if text.contains('\'') && !text.contains('"') { + '"' + } else { + '\'' + }; + out.push(quote); + for ch in text.chars() { + match ch { + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + ch if ch == quote => { + out.push('\\'); + out.push(ch); + } + ch if (ch as u32) < 0x20 || ch as u32 == 0x7f => { + let _ = write!(out, "\\x{:02x}", ch as u32); + } + ch => out.push(ch), + } + } + out.push(quote); +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn spec(model_name: &str, params: Value, info: Value) -> DeploymentSpec { + serde_json::from_value(json!({ + "model_name": model_name, "litellm_params": params, "model_info": info, + })) + .unwrap() + } + + /// Expected ids computed with `litellm.router_api.deployment_id`. + #[rstest] + #[case::strings( + "g", + json!({"model": "openai/gpt-4o", "api_key": "sk-1"}), + "30256ba697cc98ce2aff32c3b73c4b90514cb199876aef429ffcd91ae826d025" + )] + #[case::scalars( + "claude", + json!({"model": "anthropic/claude-3", "rpm": 10, "timeout": 1.5, "stream": true, "x": null}), + "645fd1504089c1084316573ccfc5dd1dd9622623b169026657484fb64a2c89a2" + )] + #[case::list_and_dict( + "m", + json!({"model": "azure/x", "tags": ["a", "b'c"], "extra_body": {"k": [1, 2.0, "é"], "z": {"q": false}}}), + "64f11e16dabed321268bd92132de2020290ef193c0cf554b2c74faee8badbb96" + )] + fn deployment_id_matches_python( + #[case] group: &str, + #[case] params: Value, + #[case] expected: &str, + ) { + let Value::Object(params) = params else { + unreachable!() + }; + assert_eq!(deployment_id(group, ¶ms), expected); + } + + #[rstest] + #[case::integral(2.0, "2.0")] + #[case::small(0.00001, "1e-05")] + #[case::large(1e16, "1e+16")] + #[case::fraction(0.1, "0.1")] + fn floats_print_as_python_repr(#[case] value: f64, #[case] expected: &str) { + let mut out = String::new(); + python_float(value, &mut out); + assert_eq!(out, expected); + } + + #[rstest] + #[case::group("gpt-4", ResolveKind::Group, "gpt-4")] + #[case::alias("smart", ResolveKind::Alias, "gpt-4")] + #[case::id("dep-b", ResolveKind::Id, "gpt-4")] + #[case::wildcard("anthropic/claude-x", ResolveKind::Wildcard, "anthropic/*")] + fn resolve_follows_python_order( + #[case] name: &str, + #[case] kind: ResolveKind, + #[case] group: &str, + ) { + let registry = Registry::build( + vec![ + spec( + "gpt-4", + json!({"model": "openai/gpt-4"}), + json!({"id": "dep-a"}), + ), + spec( + "gpt-4", + json!({"model": "azure/gpt-4"}), + json!({"id": "dep-b"}), + ), + spec("anthropic/*", json!({"model": "anthropic/*"}), json!({})), + ], + BTreeMap::from([("smart".to_owned(), "gpt-4".to_owned())]), + Vec::new(), + ) + .unwrap(); + let resolution = registry.resolve(name, ResolveScope::default()).unwrap(); + assert_eq!((resolution.kind, resolution.group.as_str()), (kind, group)); + } + + #[rstest] + fn apply_moves_an_edited_deployment_to_the_end() { + let registry = Registry::build( + vec![ + spec("a", json!({"model": "openai/a"}), json!({"id": "1"})), + spec("a", json!({"model": "openai/b"}), json!({"id": "2"})), + ], + BTreeMap::new(), + Vec::new(), + ) + .unwrap(); + let next = registry + .apply( + &ChangeSet { + upserts: vec![spec("a", json!({"model": "openai/c"}), json!({"id": "1"}))], + ..ChangeSet::default() + }, + None, + ) + .unwrap(); + let ids: Vec = next + .group("a") + .unwrap() + .iter() + .map(|d| d.id.clone()) + .collect(); + assert_eq!( + (next.version(), ids), + (1, vec!["2".to_owned(), "1".to_owned()]) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs new file mode 100644 index 00000000000..1d125451504 --- /dev/null +++ b/litellm-rust/crates/router-core/src/select.rs @@ -0,0 +1,320 @@ +//! Selection: resolve a name, filter its deployments in the order of +//! `Router._core_async_get_healthy_deployments` (team, cooldown, blocked, plugins, order, retry +//! skips), then pick one with the strategy. Simple-shuffle is native +//! (`router_strategy/simple_shuffle.py`); other strategies go through `Host::select`. + +use std::{collections::BTreeMap, sync::Arc}; + +use parking_lot::Mutex; +use rand::{Rng, SeedableRng, rngs::StdRng}; + +use crate::{ + Core, Snapshot, + config::{Op, RoutingStrategy}, + error::Error, + host::{RouterRequest, RoutingContext, SelectStage}, + registry::{Deployment, ResolveKind, ResolveScope}, +}; + +/// Randomness for the shuffle and retry jitter, injectable so tests are deterministic. +pub trait Randomness: Send + Sync { + /// A float in `[0, 1)`, as Python's `random.random()`. + fn next_f64(&self) -> f64; +} + +pub struct ThreadRandomness; + +impl Randomness for ThreadRandomness { + fn next_f64(&self) -> f64 { + rand::thread_rng().r#gen() + } +} + +pub struct SeededRandomness(Mutex); + +impl SeededRandomness { + pub fn new(seed: u64) -> Self { + Self(Mutex::new(StdRng::seed_from_u64(seed))) + } +} + +impl Randomness for SeededRandomness { + fn next_f64(&self) -> f64 { + self.0.lock().r#gen() + } +} + +/// A group's deployments before and after filtering. +#[derive(Clone, Debug)] +pub struct Candidates { + pub model_group: String, + pub kind: ResolveKind, + /// Every deployment the caller may use (after the team filter). + pub all: Vec>, + pub healthy: Vec>, +} + +/// Per-attempt selection inputs the call loop carries. +#[derive(Clone, Debug, Default)] +pub(crate) struct SelectInputs { + pub target_order: Option, + pub skipped_ids: Vec, + /// Stop before the plugins, order and skip filters: the group less cooldowns and blocked + /// deployments, which `_async_get_healthy_deployments` counts when deciding whether to retry. + pub count_only: bool, +} + +/// `simple_shuffle`: weighted by `_router_weights[group]`, else `weight`, else `rpm`, else +/// `tpm`, each normalized by its maximum, else uniform. +pub fn simple_shuffle( + deployments: &[Arc], + model_group: &str, + router_weights: Option<&BTreeMap>>, + rng: &dyn Randomness, +) -> Option> { + if deployments.is_empty() { + return None; + } + let scoped = router_weights + .and_then(|weights| weights.get(model_group)) + .filter(|weights| !weights.is_empty()) + .map(|weights| { + deployments + .iter() + .map(|d| weights.get(&d.id).copied().unwrap_or(0.0)) + .collect::>() + }); + let metric_sets = ["weight", "rpm", "tpm"].map(|metric| { + deployments + .iter() + .map(|d| d.metric(metric).unwrap_or(0.0)) + .collect::>() + }); + for weights in scoped.into_iter().chain(metric_sets) { + let largest = weights.iter().copied().fold(0.0, f64::max); + if largest <= 0.0 { + continue; + } + let normalized: Vec = weights.iter().map(|weight| weight / largest).collect(); + let total: f64 = normalized.iter().sum(); + if total <= 0.0 { + continue; + } + // `random.choices`: bisect_right over the cumulative weights. + let target = rng.next_f64() * total; + let mut cumulative = 0.0; + let index = normalized + .iter() + .position(|weight| { + cumulative += weight; + target < cumulative + }) + .unwrap_or(deployments.len() - 1); + return Some(deployments[index].clone()); + } + let index = ((rng.next_f64() * deployments.len() as f64) as usize).min(deployments.len() - 1); + Some(deployments[index].clone()) +} + +impl Core { + fn context<'a>( + &'a self, + snapshot: &'a Snapshot, + op: &'a Op, + request: &'a RouterRequest, + model_group: &'a str, + ) -> RoutingContext<'a> { + RoutingContext { + registry: &snapshot.registry, + settings: &snapshot.settings, + health: &self.health, + request, + op, + model_group, + } + } + + /// The eligible deployments for `name`, in filter-chain order. + pub(crate) async fn candidates( + &self, + snapshot: &Snapshot, + op: &Op, + request: &RouterRequest, + name: &str, + inputs: &SelectInputs, + ) -> Result { + let registry = &snapshot.registry; + let scope = ResolveScope { + specific_deployment: request.specific_deployment, + }; + let dynamic = &snapshot.settings.dynamic; + let resolution = registry + .resolve(name, scope) + .or_else(|| { + // No deployment for the name: route to the first default fallback. + let default = dynamic + .fallbacks + .default_target() + .or_else(|| dynamic.default_fallbacks.first().map(String::as_str))?; + registry.resolve(default, scope) + }) + .ok_or_else(|| Error::UnknownModel { + model: name.to_owned(), + })?; + let team_id = request.team_id.as_deref(); + let all: Vec> = resolution + .deployments + .into_iter() + .filter(|d| d.usable_by_team(team_id)) + .collect(); + + let ids: Vec = all.iter().map(|d| d.id.clone()).collect(); + let cooling = self.health.active_cooldowns(&ids).await?; + let mut healthy: Vec> = all + .iter() + .filter(|d| !cooling.iter().any(|(id, _)| *id == d.id)) + .cloned() + .collect(); + let settings = &snapshot.settings; + if healthy.is_empty() + && settings.enable_health_check_routing + && settings.allowed_fails_policy.is_some() + { + healthy = all.clone(); + } + healthy.retain(|d| !d.model_info.blocked); + if inputs.count_only { + return Ok(Candidates { + model_group: resolution.group, + kind: resolution.kind, + all, + healthy, + }); + } + + if !settings.plugins.is_empty() && !healthy.is_empty() { + let ctx = self.context(snapshot, op, request, &resolution.group); + healthy = self + .host + .select(&SelectStage::Plugins, &ctx, healthy) + .await?; + } + + let order = inputs + .target_order + .or_else(|| healthy.iter().filter_map(|d| d.order()).min()); + if let Some(order) = order { + healthy.retain(|d| d.order() == Some(order)); + } + + if !inputs.skipped_ids.is_empty() { + let kept: Vec> = healthy + .iter() + .filter(|d| !inputs.skipped_ids.contains(&d.id)) + .cloned() + .collect(); + if !kept.is_empty() { + healthy = kept; + } + } + + Ok(Candidates { + model_group: resolution.group, + kind: resolution.kind, + all, + healthy, + }) + } + + /// Pick one of `candidates.healthy` with the group's strategy. + pub(crate) async fn pick( + &self, + snapshot: &Snapshot, + op: &Op, + request: &RouterRequest, + candidates: &Candidates, + ) -> Result, Error> { + let no_deployment = || Error::NoHealthyDeployments { + model: candidates.model_group.clone(), + }; + let strategy = match candidates.kind { + ResolveKind::RoutingGroup => snapshot + .registry + .resolve_routing_group_strategy(&candidates.model_group), + ResolveKind::SpecificDeployment + | ResolveKind::Id + | ResolveKind::Alias + | ResolveKind::Group + | ResolveKind::Wildcard => None, + } + .unwrap_or_else(|| snapshot.settings.routing_strategy.clone()); + if strategy == RoutingStrategy::SimpleShuffle { + return simple_shuffle( + &candidates.healthy, + &candidates.model_group, + request.router_weights.as_ref(), + self.rng.as_ref(), + ) + .ok_or_else(no_deployment); + } + let ctx = self.context(snapshot, op, request, &candidates.model_group); + self.host + .select( + &SelectStage::Strategy(strategy), + &ctx, + candidates.healthy.clone(), + ) + .await? + .into_iter() + .next() + .ok_or_else(no_deployment) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::registry::DeploymentSpec; + + struct Fixed(f64); + + impl Randomness for Fixed { + fn next_f64(&self) -> f64 { + self.0 + } + } + + fn deployment(id: &str, params: serde_json::Value) -> Arc { + let spec: DeploymentSpec = serde_json::from_value(json!({ + "model_name": "g", "litellm_params": params, "model_info": {"id": id}, + })) + .unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + } + + #[rstest] + #[case::weight_low_draw(0.1, "a")] + #[case::weight_high_draw(0.9, "b")] + fn shuffle_uses_weight_before_rpm(#[case] draw: f64, #[case] expected: &str) { + let deployments = [ + deployment("a", json!({"model": "openai/a", "weight": 1, "rpm": 1000})), + deployment("b", json!({"model": "openai/b", "weight": 3})), + ]; + let picked = simple_shuffle(&deployments, "g", None, &Fixed(draw)).unwrap(); + assert_eq!(picked.id, expected); + } + + #[rstest] + fn shuffle_prefers_request_weights() { + let deployments = [ + deployment("a", json!({"model": "openai/a", "weight": 100})), + deployment("b", json!({"model": "openai/b"})), + ]; + let weights = BTreeMap::from([("g".to_owned(), BTreeMap::from([("b".to_owned(), 1.0)]))]); + let picked = simple_shuffle(&deployments, "g", Some(&weights), &Fixed(0.99)).unwrap(); + assert_eq!(picked.id, "b"); + } +} diff --git a/litellm-rust/crates/router-core/src/settings.rs b/litellm-rust/crates/router-core/src/settings.rs new file mode 100644 index 00000000000..bf8346b293f --- /dev/null +++ b/litellm-rust/crates/router-core/src/settings.rs @@ -0,0 +1,69 @@ +//! [`Settings`]: one typed, frozen view of the configuration a caller may read, built from the +//! static config and the current dynamic snapshot. + +use std::{sync::Arc, time::Duration}; + +use crate::{ + config::{ + AllowedFailsPolicy, DEFAULT_COOLDOWN_TIME, PluginSpec, RetryPolicy, RouterConfig, + RouterDynamicConfig, RoutingStrategy, + }, + error::FailureKind, +}; + +#[derive(Clone, Debug)] +pub struct Settings { + pub routing_strategy: RoutingStrategy, + pub plugins: Arc<[PluginSpec]>, + pub cooldown_time: Duration, + pub disable_cooldowns: bool, + pub allowed_fails: Option, + pub allowed_fails_policy: Option, + pub retry_policy: Option, + pub num_retries: u32, + pub max_fallbacks: u32, + pub retry_after: Duration, + pub timeout: Option, + pub drop_params: Option, + pub enable_health_check_routing: bool, + pub provider_default_deployment_ids: Arc<[String]>, + pub dynamic: Arc, +} + +fn seconds(value: f64) -> Duration { + Duration::try_from_secs_f64(value).unwrap_or_default() +} + +impl Settings { + pub fn new(config: &RouterConfig, dynamic: Arc) -> Self { + Self { + routing_strategy: config.routing_strategy.clone(), + plugins: config.plugins.clone().into(), + cooldown_time: config.cooldown_time.map_or(DEFAULT_COOLDOWN_TIME, seconds), + disable_cooldowns: config.disable_cooldowns, + allowed_fails: dynamic.allowed_fails, + allowed_fails_policy: config.allowed_fails_policy.clone(), + retry_policy: config.retry_policy.clone(), + num_retries: dynamic.num_retries.unwrap_or(0), + max_fallbacks: dynamic.max_fallbacks, + retry_after: seconds(config.retry_after), + timeout: config.request_timeout.or(config.timeout).map(seconds), + drop_params: config.drop_params, + enable_health_check_routing: config.enable_health_check_routing, + provider_default_deployment_ids: config.provider_default_deployment_ids.clone().into(), + dynamic, + } + } + + /// `allowed_fails_for`: the policy's count for this failure kind. + pub fn allowed_fails_for(&self, kind: FailureKind) -> Option { + self.allowed_fails_policy + .as_ref() + .and_then(|policy| policy.allowed_fails_for(kind)) + } + + /// Python's `allowed_fails` default when set without a value (`litellm.allowed_fails`). + pub fn effective_allowed_fails(&self) -> u32 { + self.allowed_fails.unwrap_or(3) + } +} diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs new file mode 100644 index 00000000000..1a4705b8dd8 --- /dev/null +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -0,0 +1,167 @@ +//! A call through the router end to end: selection, retries, fallbacks and cooldowns against a +//! scripted host. + +use std::{collections::VecDeque, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Deployment, Error, FailureKind, Host, InMemoryRouterCache, Op, Output, + Payload, Router, RouterCache, RouterConfig, RouterDynamicConfig, RouterEvent, RouterRequest, + RoutingContext, SeededRandomness, SelectStage, +}; +use parking_lot::Mutex; +use rstest::{fixture, rstest}; +use serde_json::json; + +#[derive(Default)] +struct ScriptedHost { + /// Failures handed out in order, keyed by model group; a group with none left succeeds. + script: Mutex)>>, + events: Mutex>, + cache: InMemoryRouterCache, +} + +impl ScriptedHost { + fn failing(group: &str, failures: Vec) -> Self { + Self { + script: Mutex::new(vec![(group.to_owned(), failures.into())]), + ..Self::default() + } + } +} + +#[async_trait] +impl Host for ScriptedHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + let next = self + .script + .lock() + .iter_mut() + .find(|(group, _)| *group == attempt.model_group) + .and_then(|(_, failures)| failures.pop_front()); + match next { + Some(failure) => Err(failure), + None => Ok(Output { + value: Payload(Arc::new(attempt.deployment.id.clone())), + total_tokens: 10, + }), + } + } + + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + + async fn allow_fallback(&self, _request: &RouterRequest, _target: &str) -> bool { + true + } + + async fn sleep(&self, _duration: Duration) {} + + fn observe(&self, event: &RouterEvent) { + self.events.lock().push(event.clone()); + } + + fn cache(&self) -> Arc { + Arc::new(self.cache.clone()) + } +} + +#[fixture] +fn config() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/c"}, "model_info": {"id": "b1"}}, + ] + })) + .unwrap() +} + +fn router(config: RouterConfig, dynamic: serde_json::Value, host: Arc) -> Router { + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, host) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() +} + +fn request() -> RouterRequest { + RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + } +} + +#[rstest] +#[tokio::test] +async fn a_rate_limit_retries_in_the_same_group(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router(config, json!({"num_retries": 2}), host.clone()); + let response = router.call(Op::Completion, request()).await.unwrap(); + assert_eq!( + ( + response.model_group.as_str(), + response.attempted_retries, + response.attempted_fallbacks + ), + ("primary", 1, 0) + ); +} + +#[rstest] +#[tokio::test] +async fn a_bad_request_falls_back_without_retrying(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::BadRequest, Some(400))], + )); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let response = router.call(Op::Completion, request()).await.unwrap(); + let retried = host + .events + .lock() + .iter() + .any(|event| matches!(event, RouterEvent::Retry { .. })); + assert_eq!( + ( + response.deployment_id.as_str(), + response.attempted_fallbacks, + retried + ), + ("b1", 1, false) + ); +} + +#[rstest] +#[tokio::test] +async fn a_rate_limit_cools_the_deployment_down(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router(config, json!({}), host.clone()); + let error = router.call(Op::Completion, request()).await.unwrap_err(); + let cooled = router + .health() + .active_cooldowns(&["a1".to_owned(), "a2".to_owned()]) + .await + .unwrap(); + assert_eq!( + (error.attempt().map(|failure| failure.kind), cooled.len()), + (Some(FailureKind::RateLimit), 1) + ); +} diff --git a/litellm/router_core.py b/litellm/router_core.py index acbe82a0643..e6b4014b473 100644 --- a/litellm/router_core.py +++ b/litellm/router_core.py @@ -7,8 +7,9 @@ the router is built, from its ``RouterConfig``. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable +from litellm._logging import verbose_router_logger from litellm.router_config import RouterConfig if TYPE_CHECKING: @@ -51,6 +52,28 @@ class RouterCore(Protocol): def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore: + """ + The Rust core when the route catalog's router rule allows it and the Rust core serves this + config, else the Python core. Decided once, for the router's life. + """ from litellm.router_python_core import PythonCore + from litellm.rust_bridge.catalog import Route, RouteContext, decision + from litellm.rust_bridge.configuration import Decision - return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature + python_class: Final = core_class or PythonCore + chosen: Final = decision(RouteContext(Route.ROUTER)) + if chosen is not Decision.PYTHON: + from litellm.router_rust_core import RustCore, RustCoreDeclined + from litellm.rust_bridge._native import RustBridgeDeclined + + try: + return RustCore.build(config, host, python_class) # pyright: ignore[reportReturnType] # RustCore serves the rest of RouterCore from its PythonCore, through __getattr__ + except (RustBridgeDeclined, RustCoreDeclined) as declined: + # A config the Rust core cannot take keeps the router on Python; an error building the + # Python core itself is the router's own error and is raised once, as it is. + if chosen is Decision.RUST_REQUIRED: + raise + verbose_router_logger.info("Router served by the Python core: %s", declined) + if isinstance(declined, RustCoreDeclined): + return declined.python + return python_class(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature diff --git a/litellm/router_rust_core.py b/litellm/router_rust_core.py new file mode 100644 index 00000000000..b5752e9d62c --- /dev/null +++ b/litellm/router_rust_core.py @@ -0,0 +1,857 @@ +""" +``RustCore``: the router core served by litellm-rust/crates/router-core through ``_native``. + +The Rust core owns the registry, selection, retries, fallbacks and cooldowns for the operations in +``RUST_SERVED_OPS``. It never awaits Python: ``RouterCall.next`` returns the step the call loop is +blocked on (run an attempt on a deployment, sleep, check a fallback target, run a plugin) and this +driver answers it. Each attempt runs through the Python core with that one deployment, so request +preparation, inference, logging and cost tracking are unchanged. + +Everything else (the views, runtime writes, every other endpoint) is served by the ``PythonCore`` +the Rust core wraps, built from the same config. A config the Rust core cannot serve fails at +construction and ``make_core`` builds a ``PythonCore`` instead. +""" + +import inspect +import json +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_guardrail import is_guardrail_intervention +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.routing_common_utils import ( + format_content_policy_fallback_hint, + format_fallback_outcome_message, + format_no_fallback_group_message, + truncate_fallback_error_detail, +) +from litellm.litellm_core_utils.secret_redaction import redact_string +from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure +from litellm.router_config import RouterConfig +from litellm.router_python_core import PythonCore +from litellm.router_python_host import PythonHost +from litellm.router_utils.fallback_event_handlers import ( + AttemptedFallbackTargets, + _provider_prefixed_model_group, # pyright: ignore[reportPrivateUsage] # the Python fallback walk's own spelling rule + fallback_lookup_groups, + get_fallback_model_group_for_lookup_groups, + is_bad_request_without_fallback, +) +from litellm.types.router import FallbackSkipReason +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.rust_bridge._native import RouterCall, RouterCoreHandle + +RustServedOp: TypeAlias = Literal["acompletion", "aanthropic_messages", "aembedding", "aresponses"] +RUST_SERVED_OPS: Final[frozenset[str]] = frozenset({"acompletion", "aanthropic_messages", "aembedding", "aresponses"}) + +_FAILURE_KINDS: Final = ( + "RateLimitError", + "ContextWindowExceededError", + "ContentPolicyViolationError", + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "BadRequestError", + "Timeout", + "APIConnectionError", + "InternalServerError", + "ServiceUnavailableError", + "BadGatewayError", + "GuardrailIntervention", + "CallerCredentialError", +) +_DECLINED: Final = frozenset({"MalformedRequestFieldError", "UnsupportedEndpointError"}) +_STEP: Final = TypeAdapter(Mapping[str, object]) +_REQUEST_FIELDS: Final = ("specific_deployment", "num_retries", "disable_fallbacks") +# Request keys router-core reads under another name. +_RENAMED_REQUEST_FIELDS: Final = MappingProxyType({"_router_weights": "router_weights"}) +_PROVIDERS: Final = tuple(provider.value for provider in LlmProviders) +_PYTHON_ONLY_REQUEST_KEYS: Final = frozenset( + { + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "default_fallbacks", + "max_fallbacks", + "model_group_retry_policy", + "retry_policy", + } +) +_STATIC_FIELDS: Final = ( + "routing_strategy", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_after", + "timeout", +) +_HANDLED_FIELDS: Final = frozenset( + { + *_STATIC_FIELDS, + "model_list", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "set_verbose", + "debug_level", + } +) +_PYTHON_DEFAULTS: Final = MappingProxyType( + { + name: parameter.default + for name, parameter in inspect.signature(PythonCore).parameters.items() + if parameter.default is not inspect.Parameter.empty + } +) + + +def failure_kind(error: BaseException) -> str: + names: Final = tuple(cls.__name__ for cls in type(error).__mro__) + if not _DECLINED.isdisjoint(names): + return "Declined" + if is_guardrail_intervention(error): # pyright: ignore[reportArgumentType] # every attempt error is an Exception + return "GuardrailIntervention" + return next((name for name in names if name in _FAILURE_KINDS), "APIError") + + +def _retry_after_seconds(error: BaseException) -> float | None: + response: Final = getattr(error, "response", None) + headers: Final = getattr(error, "litellm_response_headers", None) or getattr(response, "headers", None) + if headers is None or not hasattr(headers, "get"): + return None + milliseconds: Final = headers.get("retry-after-ms") + seconds: Final = headers.get("retry-after") + try: + value: Final = float(milliseconds) / 1000 if milliseconds is not None else float(seconds) + except (TypeError, ValueError): + return None + return value if math.isfinite(value) and value >= 0 else None + + +def attempt_failure(error: BaseException) -> Mapping[str, object]: + status: Final = getattr(error, "status_code", None) + num_retries: Final = getattr(error, "num_retries", None) + return { + "ok": False, + "kind": failure_kind(error), + "status": status if isinstance(status, int) and 100 <= status <= 599 else None, + "retry_after": _retry_after_seconds(error), + "num_retries": num_retries if isinstance(num_retries, int) and num_retries >= 0 else None, + "detail": str(error)[:500], + } + + +def total_tokens(response: object) -> int: + usage: Final = getattr(response, "usage", None) + tokens: Final = getattr(usage, "total_tokens", None) + return tokens if isinstance(tokens, int) else 0 + + +def _json_default(value: object) -> object: + return str(value) + + +def _plain(value: object) -> object: + """``value`` as plain JSON values, the shape the Rust core reads.""" + return json.loads(json.dumps(value, default=_json_default)) + + +def _plain_document(value: Mapping[str, object]) -> Mapping[str, object]: + """A mapping as plain JSON values: still a mapping, with string keys.""" + return as_str_mapping(_plain(value)) or MappingProxyType({}) + + +def _as_list(value: object) -> Sequence[object]: + return value if isinstance(value, list) else () # pyright: ignore[reportUnknownVariableType] # a list's items are read as objects + + +def _applied_static(python: PythonCore, name: str) -> object: + """Where an applied static setting lands on the Python core: the value the Rust core is rebuilt with.""" + match name: + case "routing_strategy": + return python.settings().routing_strategy + case "cooldown_time": + return python.settings().cooldown_time + case "disable_cooldowns": + return python.settings().disable_cooldowns + case "allowed_fails_policy": + return python.settings().allowed_fails_policy + case "retry_after": + return python.retry_after + case "timeout": + return python.timeout + case _: + raise KeyError(name) + + +def _static_value(value: object) -> object: + dump: Final = getattr(value, "model_dump", None) + return _plain(dump(exclude_none=True) if callable(dump) else value) + + +def static_settings(config: RouterConfig, python: PythonCore) -> Mapping[str, object]: + """The static settings: the constructor arguments, overlaid with every setting a changeset + (``apply``) has changed since, as the Python core resolved it.""" + applied: Final = python._core_applied # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + values: Final[dict[str, object]] = { # mutable-ok: the arguments, then the applied settings over them + name: config.arguments[name] for name in _STATIC_FIELDS if name in config.arguments + } + values.update({name: _applied_static(python, name) for name in _STATIC_FIELDS if name in applied}) + return {name: _static_value(value) for name, value in values.items() if value is not None} + + +def static_document(config: RouterConfig, python: PythonCore) -> Mapping[str, object]: + """The Python core's own normalized deployments (with their ids) and the static settings.""" + return { + "model_list": _plain(python._core_model_list), # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownArgumentType] # RustCore drives its own PythonCore, whose list is untyped + **static_settings(config, python), + } + + +def dynamic_document(python: PythonCore) -> Mapping[str, object]: + """The values the Python core resolved from its arguments and the process settings.""" + settings: Final = python.settings() + applied: Final = python._core_applied # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + resolved: Final[Mapping[str, object]] = { # pyright: ignore[reportUnknownVariableType] # the fallback settings are untyped lists + "version": python._registry_version, # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + "model_group_alias": python._core_model_group_alias, # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + "fallbacks": settings.fallbacks, # pyright: ignore[reportUnknownMemberType] # the setting's own untyped list + "context_window_fallbacks": settings.context_window_fallbacks, # pyright: ignore[reportUnknownMemberType] # the setting's own untyped list + "content_policy_fallbacks": settings.content_policy_fallbacks, # pyright: ignore[reportUnknownMemberType] # the setting's own untyped list + "num_retries": python.num_retries, + "max_fallbacks": python.max_fallbacks, + "allowed_fails": python._core_allowed_fails # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + if "allowed_fails" in applied + else python.explicit_allowed_fails, + "fallback_keys": fallback_keys(python, settings), + } + return {name: _plain(value) for name, value in resolved.items() if value is not None} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the fallback settings are untyped lists + + +def fallback_keys(python: PythonCore, settings: object) -> Mapping[str, object]: + """The spellings a fallback key matches besides the group name (`get_fallback_model_group`).""" + lists: Final = ( + getattr(settings, "fallbacks", None), + getattr(settings, "context_window_fallbacks", None), + getattr(settings, "content_policy_fallbacks", None), + ) + keys: Final = tuple( # comprehension-ok: the first group of each fallback entry, across the three lists + str(next(iter(mapping))) + for entries in lists + for entry in _as_list(entries) + if (mapping := as_str_mapping(entry)) and isinstance(mapping, dict) + ) + names: Final = { + *(python._core_model_names or ()), # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + *(python._core_model_group_alias or {}), # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + } + prefixed: Final = { + name: spelled for name in names if (spelled := _provider_prefixed_model_group(name, keys)) is not None + } + return {"providers": _PROVIDERS, "provider_prefixed": prefixed} + + +def unsupported_arguments(config: RouterConfig) -> frozenset[str]: + """Arguments the Rust core does not serve; any of them keeps the router on Python.""" + changed: Final = frozenset( + name + for name, value in config.arguments.items() + if name not in _HANDLED_FIELDS and value != _PYTHON_DEFAULTS.get(name, None) + ) + strategy: Final = config.arguments.get("routing_strategy", "simple-shuffle") + model_list: Final = config.arguments.get("model_list") + # `simple_shuffle` raises on a non-numeric weight when it picks; the Python core keeps that. + odd_weight: Final = any( + (params := as_str_mapping((as_str_mapping(deployment) or {}).get("litellm_params"))) is not None + and params.get("weight") is not None + and (isinstance(params.get("weight"), bool) or not isinstance(params.get("weight"), int | float)) + for deployment in _as_list(model_list) + ) + return ( + changed + | ({"routing_strategy"} if strategy != "simple-shuffle" else frozenset[str]()) + | ({"model_list"} if odd_weight else frozenset[str]()) + ) + + +def runs_on_python(request: Mapping[str, object]) -> bool: + """Per-request routing overrides the Rust core does not read yet.""" + return any(key in _PYTHON_ONLY_REQUEST_KEYS or key.startswith("mock_testing_") for key in request) + + +def request_team_id(request: Mapping[str, object]) -> str | None: + metadata: Final = as_str_mapping(request.get("metadata") or request.get("litellm_metadata")) + team_id: Final = metadata.get("user_api_key_team_id") if metadata is not None else None + return team_id if isinstance(team_id, str) else None + + +def _bound_request(bound: inspect.BoundArguments) -> Mapping[str, object]: + """The keyword payload of a call bound to the Python core's endpoint signature.""" + named: Final = {name: value for name, value in bound.arguments.items() if name not in ("self", "kwargs")} + extra: Final = bound.arguments.get("kwargs") + return {**named, **(extra if isinstance(extra, Mapping) else {})} + + +_ACOMPLETION: Final = inspect.signature(PythonCore.acompletion) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # only the parameter names are read +_AEMBEDDING: Final = inspect.signature(PythonCore.aembedding) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # only the parameter names are read +_OWN_ATTRIBUTES: Final = frozenset({"_handle", "_config", "_python", "_host", "_built_for"}) + + +class RustCoreDeclined(Exception): + """The Rust core cannot serve a router whose Python core is already built: the router keeps that one.""" + + def __init__(self, reason: str, python: PythonCore) -> None: + super().__init__(reason) + self.python = python + + +class RustCore: + """A ``RouterCore`` whose routing for ``RUST_SERVED_OPS`` runs in Rust, over a ``PythonCore``.""" + + def __init__(self, handle: "RouterCoreHandle", config: RouterConfig, python: PythonCore, host: PythonHost) -> None: + self._handle: RouterCoreHandle | None = handle + self._config: Final = config + self._python: Final = python + self._host: Final = host + self._built_for: tuple[int, int] = self._state() + + @classmethod + def build(cls, config: RouterConfig, host: PythonHost, python_class: type[PythonCore] = PythonCore) -> "RustCore": + from litellm.rust_bridge._native import RouterCoreHandle, RustBridgeDeclined + + extra: Final = unsupported_arguments(config) + if extra: + raise RustBridgeDeclined(f"router arguments not served by the Rust core: {sorted(extra)}") + python: Final = python_class(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature + try: + handle: Final = RouterCoreHandle(static_document(config, python), dynamic_document(python)) + except Exception as declined: + # The Python core is built already (its callbacks registered): keep it, don't build another. + raise RustCoreDeclined(str(declined), python) from declined + return cls(handle, config, python, host) + + def __getattr__(self, name: str) -> object: + return getattr(self._python, name) + + def __setattr__(self, name: str, value: object) -> None: + if name in _OWN_ATTRIBUTES: + object.__setattr__(self, name, value) + return + setattr(self._python, name, value) + + def __delattr__(self, name: str) -> None: + if name in _OWN_ATTRIBUTES: + object.__delattr__(self, name) + return + delattr(self._python, name) + + def _state(self) -> tuple[int, int]: + """What the handle was built from: the deployments, and the settings changesets applied.""" + return (self._python._registry_version, self._python._settings_version) # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + + def _current_handle(self) -> "RouterCoreHandle | None": + """The Rust handle, rebuilt after a deployment change or an applied changeset; ``None`` while + the Rust core cannot serve the router.""" + if "_core_async_get_available_deployment" in vars(self._python): + # `register_plugin` installed a custom strategy, which only the Python core runs. + return None + if self._state() == self._built_for: + return self._handle + from litellm.rust_bridge._native import RouterCoreHandle + + if self._handle is not None: + self._handle.retire() + strategy: Final = self._python.settings().routing_strategy + try: + if strategy not in (None, "simple-shuffle"): + raise ValueError(f"routing_strategy {strategy!r} is not served by the Rust core") + self._handle = RouterCoreHandle(static_document(self._config, self._python), dynamic_document(self._python)) + except Exception as declined: # noqa: BLE001 # a change the Rust core cannot take keeps the router on Python until the next one + verbose_router_logger.info("Router served by the Python core after a change: %s", declined) + self._handle = None + self._built_for = self._state() + return self._handle + + async def call(self, op: str, request: Mapping[str, object]) -> object: + if op not in RUST_SERVED_OPS: + return await self._python.call(op, request) + return await self._drive(op, request) + + async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + try: + bound: Final = _ACOMPLETION.bind(self._python, *args, **kwargs) + except TypeError: + # The Python core raises the caller's mistake in its own words. + return await self._python.acompletion(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType, reportUnknownMemberType, reportUnknownVariableType] # the call is expected to raise + return await self._drive("acompletion", _bound_request(bound)) + + async def aanthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: the endpoint's own keyword payload + return await self._drive("aanthropic_messages", kwargs) + + async def aembedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + try: + bound: Final = _AEMBEDDING.bind(self._python, *args, **kwargs) + except TypeError: + return await self._python.aembedding(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType, reportUnknownMemberType] # the call is expected to raise + return await self._drive("aembedding", _bound_request(bound)) + + async def aresponses(self, **kwargs: object) -> object: # kwargs-ok: the endpoint's own keyword payload + return await self._drive("aresponses", kwargs) + + def retire(self) -> None: + if self._handle is not None: + self._handle.retire() + self._python.retire() + + async def _drive(self, op: str, request: Mapping[str, object]) -> object: + handle: Final = self._current_handle() + if handle is None or runs_on_python(request) or not isinstance(request.get("model"), str): + return await self._python.call(op, request) + if not hasattr(request.get(_METADATA_FIELDS.get(op, "metadata"), {}), "update"): + # `_update_kwargs_before_fallbacks` raises on metadata it cannot update, before any attempt. + return await self._python.call(op, request) + fields: Final = { + **{key: request[key] for key in _REQUEST_FIELDS if request.get(key) is not None}, + **{name: request[key] for key, name in _RENAMED_REQUEST_FIELDS.items() if request.get(key) is not None}, + } + team_id: Final = request_team_id(request) + try: + call: Final = handle.start( + op, _plain_document({"model": request["model"], **fields, **({"team_id": team_id} if team_id else {})}) + ) + except ValueError: + return await self._python.call(op, request) + return await _Driver(call, self._python, op, request).run() + + +# A fallback chain list in the shape the Python core's fallback walk reads it. +_Chains: TypeAlias = Sequence[Mapping[str, list[str]]] | None + +# The request field `_update_kwargs_before_fallbacks` writes the model group into, per op. +_METADATA_FIELDS: Final = MappingProxyType( + {"aanthropic_messages": "litellm_metadata", "aresponses": "litellm_metadata"} +) + + +def _ttl(step: Mapping[str, object]) -> float | None: + ttl: Final = step.get("ttl") + return ttl if isinstance(ttl, int | float) else None + + +def _store_for(python: PythonCore, key: str) -> "DualCache": + """Cooldowns live in the router's cooldown store; every other routing key in its main DualCache.""" + health: Final = python.health() + return health.cooldowns.cooldown_store if key.endswith(":cooldown") else health.store + + +async def answer_cache_step(python: PythonCore, step: Mapping[str, object]) -> object: + """Run one cache op the Rust core asked for against the router's own DualCaches (memory plus Redis).""" + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + match step.get("op"): + case "get": + return _plain(await _store_for(python, key).async_get_cache(key)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + keys: Final = [str(k) for k in _as_list(step.get("keys"))] + values: Final[list[object]] = [ # mutable-ok: built here, then made plain + await _store_for(python, k).async_get_cache(k) # pyright: ignore[reportUnknownMemberType] # DualCache's reads are untyped + for k in keys + ] + return _plain(values) + case "set": + store: Final = _store_for(python, key) + if ttl is None: + await store.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + else: + await store.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + return None + case "increment": + amount: Final = step.get("value") + by: Final = amount if isinstance(amount, int | float) and not isinstance(amount, bool) else 1 + counter: Final = _store_for(python, key) + total: Final = ( + await counter.async_increment_cache(key, by) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + if ttl is None + else await counter.async_increment_cache(key, by, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + ) + return total if isinstance(total, int | float) else by + case "delete": + await _store_for(python, key).async_delete_cache(key) + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + +@dataclass(slots=True) +class _Hop: + """One target of a fallback walk and what its level ended on.""" + + group: str + pending: Mapping[str, object] | None + error: BaseException | None = None + failed: Mapping[str, object] | None = None + children: "list[_Hop]" = field(default_factory=list) # mutable-ok: the fallbacks tried from this target + + +class _Driver: + """Answers the steps one Rust call loop blocks on until it is done. + + Each attempt runs on the Python core with retries and fallbacks off, so the driver puts back + what Python's own retry and fallback layers add to the outcome: the retry stamp on an error + whose retry loop ran out, the retry and fallback headers on a response, and the explanation + the fallback layer appends to the error it finally raises. + """ + + def __init__(self, call: "RouterCall", python: PythonCore, op: str, request: Mapping[str, object]) -> None: + self._call: Final = call + self._python: Final = python + self._op: Final = op + self._request: Final = request + self._last_response: object = None + self._last_error: BaseException | None = None + # The last failure at each fallback level: a level raises its own last error. + self._level_errors: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call loop runs + self._attempted: bool = False + self._walk: Final = AttemptedFallbackTargets() + # Fallback targets in the order Rust walked them, and the groups an attempt was made in. + self._hop_targets: Final[list[str]] = [] # mutable-ok: filled as the call loop runs + self._attempted_groups: Final[set[str]] = set() # mutable-ok: filled as the call loop runs + # The walk in order: ("hop", (level, group, order)), ("attempt", level), ("error", (level, + # error, request)), ("sleep", None) before each retry, and ("exhausted", group) when a retry + # loop ran out. + self._timeline: Final[list[tuple[str, object]]] = [] # mutable-ok: filled as the call loop runs + + async def run(self) -> object: + outcome: object = None # rebind-ok: each step's outcome answers it + while True: + step = _STEP.validate_python(self._call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome against the step it answers + self._take_events(step.get("events")) + match step["kind"]: + case "attempt": + outcome = await self._attempt(step) + case "sleep": + self._timeline.append(("sleep", None)) + await self._python._host.sleep(float(str(step["seconds"]))) # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + outcome = None + case "allow_fallback": + outcome = await self._allow_fallback(str(step["target"])) + case "cache": + outcome = await answer_cache_step(self._python, step) + case "select": + outcome = step["candidates"] + case "done": + return self._with_headers(step) + case "failed": + primary = self._level_errors.get(0) + if primary is not None: + if self._walks_order_tiers(primary, self._request.get("model")): + # Python's order-tier walk raises what its last fallback raised, as is. + raise await self._walk_error(self._walk_tree(), primary) + await self._fill_unattempted_level() + self._explain(primary) + raise primary + if not self._attempted: + # No deployment was tried (unknown model, nothing healthy): Python raises its own error. + return await self._python.call(self._op, self._request) + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + def _take_events(self, events: object) -> None: + for raw_event in _as_list(events): + event = as_str_mapping(raw_event) + if event is not None and event.get("kind") == "fallback": + self._timeline.append(("hop", (event.get("level"), event.get("to"), event.get("order")))) + if event is None or event.get("kind") != "retries_exhausted": + continue + self._timeline.append(("exhausted", event.get("model_group"))) + error: BaseException | None = self._last_error + # `async_function_with_retries` stamps the error its exhausted retry loop raises. + if error is not None and type(error) in self._python._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + error.max_retries = event.get("max") # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes + error.num_retries = event.get("attempted") # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes + + def _walks_order_tiers(self, error: BaseException, group: object) -> bool: + """`async_function_with_fallbacks_common_utils` hands a group with more than one order level + to `run_async_fallback` outside its try, so the walk's own error is raised with no suffix.""" + import litellm.exceptions as litellm_exceptions + + python: Final = self._python + request: Final = self._request + if ( + request.get("disable_fallbacks") is True + or not isinstance(group, str) + or is_guardrail_intervention(error) # pyright: ignore[reportArgumentType] # every attempt error is an Exception + or is_bad_request_without_fallback(error, python.fallback_on_unsupported_params) # pyright: ignore[reportArgumentType] # every attempt error is an Exception + or isinstance( + error, + litellm_exceptions.ContextWindowExceededError | litellm_exceptions.ContentPolicyViolationError, + ) + ): + return False + metadata: Final = as_str_mapping(request.get("metadata")) + team: Final = metadata.get("user_api_key_team_id") if metadata is not None else None + deployments: Final = ( + python._core_get_model_list( # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + model_name=group, team_id=team if isinstance(team, str) else None + ) + or () + ) + orders: Final = { + order + for deployment in deployments + if (order := (deployment.get("litellm_params") or {}).get("order")) is not None # pyright: ignore[reportUnknownMemberType] # a deployment's untyped litellm_params + } + return len(orders) > 1 + + def _walk_tree(self) -> _Hop: + """The fallback walk as Rust ran it: each hop under the level it was tried from.""" + root: Final = _Hop(str(self._request.get("model")), None) + current: Final[dict[int, _Hop]] = {0: root} # mutable-ok: the newest hop at each level + recent: _Hop | None = None # rebind-ok: the hop whose attempt failed last + for kind, value in self._timeline: + if kind == "hop" and isinstance(value, tuple): + level, group, order = value # pyright: ignore[reportUnknownVariableType] # the timeline's tuples are untyped + hop = _Hop(str(group), {**self._request, "model": str(group)}) # pyright: ignore[reportUnknownArgumentType] # the timeline's tuples are untyped + if isinstance(order, int): + hop.pending = {**self._request, "model": str(group), "_target_order": order} # pyright: ignore[reportUnknownArgumentType] # the timeline's tuples are untyped + parent = current.get(level - 1 if isinstance(level, int) else -1) + if parent is not None and isinstance(level, int): + parent.children.append(hop) + current[level] = hop + recent = None + elif kind == "attempt" and isinstance(value, int) and value in current: + current[value].pending = None + elif kind == "error" and isinstance(value, tuple) and value[0] in current: + recent = current[value[0]] + recent.error, recent.failed = value[1], value[2] + elif kind == "sleep" and recent is not None: + recent.pending = recent.failed # pyright: ignore[reportUnknownMemberType] # the timeline's tuples are untyped + elif kind == "exhausted" and recent is not None and recent.pending is recent.failed: # pyright: ignore[reportUnknownMemberType] # the timeline's tuples are untyped + recent.pending = None + return root + + async def _walk_error(self, hop: _Hop, primary: BaseException) -> BaseException: + """The error a hop's level raises. Python raises its own error where Rust made no attempt: + for a target with no deployment, and for a retry that found no deployment left to try. A + group with more than one order level raises its walk's last error instead + (`run_async_fallback`), whatever that walk tried.""" + own: Final = await self._own_error(hop, primary) + if not hop.children or not self._walks_order_tiers(own, hop.group): + return own + last: BaseException = own # rebind-ok: each fallback tried replaces it + for child in hop.children: + last = await self._walk_error(child, primary) + return last + + async def _own_error(self, hop: _Hop, primary: BaseException) -> BaseException: + """The error the hop's own target raised: Python's, replayed, when Rust made no attempt.""" + if hop.pending is not None: + try: + await self._python.call(self._op, {**hop.pending, "num_retries": 0, "disable_fallbacks": True}) + except Exception as error: # noqa: BLE001 # the error Python's walk would end on + return error + return hop.error or primary + + async def _fill_unattempted_level(self) -> None: + """A first fallback that resolved no deployment made no attempt, so Rust holds no error for + it: ask Python, which raises its own no-deployment error without calling a provider.""" + if 1 in self._level_errors: + return + target: Final = next((t for t in self._hop_targets if t not in self._attempted_groups), None) + if target is None: + return + try: + await self._python.call( + self._op, {**self._request, "model": target, "num_retries": 0, "disable_fallbacks": True} + ) + except Exception as error: # noqa: BLE001 # the error is only quoted in the primary error's message + self._level_errors[1] = error + + def _with_headers(self, step: Mapping[str, object]) -> object: + response: object = self._last_response # rebind-ok: the response gains its headers in turn + retries: Final = step.get("attempted_retries") + if isinstance(retries, int) and retries > 0: + max_retries: Final = step.get("max_retries") + response = add_retry_headers_to_response( # rebind-ok: the response gains its headers in turn + response=response, + attempted_retries=retries, + max_retries=max_retries if isinstance(max_retries, int) else None, + ) + fallbacks: Final = step.get("attempted_fallbacks") + return add_fallback_headers_to_response( + response=response, attempted_fallbacks=fallbacks if isinstance(fallbacks, int) else 0 + ) + + def _explain(self, error: BaseException) -> None: + """Append what `async_function_with_fallbacks_common_utils` appends to the error it raises.""" + python: Final = self._python + request: Final = self._request + if ( + request.get("disable_fallbacks") is True + or is_guardrail_intervention(error) # pyright: ignore[reportArgumentType] # every attempt error is an Exception + or is_bad_request_without_fallback(error, python.fallback_on_unsupported_params) # pyright: ignore[reportArgumentType] # every attempt error is an Exception + or not hasattr(error, "message") + or not python._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + ): + return + model_group: Final = request.get("model") + model_name: Final = model_group if isinstance(model_group, str) else None + try: + self._explain_fallbacks(error, model_name) + except Exception: # noqa: BLE001 # Python looks the chains up inside its fallback try: a bad chain means none was attempted + error.message += format_fallback_outcome_message( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix + model_name, None, "", skipped=self._walk.skipped + ) + + def _explain_fallbacks(self, error: BaseException, model_name: str | None) -> None: + import litellm.exceptions as litellm_exceptions + + python: Final = self._python + request: Final = self._request + fallbacks: Final[_Chains] = python._core_fallbacks # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType, reportAssignmentType] # RustCore drives its own PythonCore, whose chains are untyped lists + context_window_fallbacks: Final[_Chains] = python._core_context_window_fallbacks # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType, reportAssignmentType] # RustCore drives its own PythonCore, whose chains are untyped lists + content_policy_fallbacks: Final[_Chains] = python._core_content_policy_fallbacks # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType, reportAssignmentType] # RustCore drives its own PythonCore, whose chains are untyped lists + lookup_groups: Final = fallback_lookup_groups(request, model_name) + context_window: Final = isinstance(error, litellm_exceptions.ContextWindowExceededError) + content_policy: Final = isinstance(error, litellm_exceptions.ContentPolicyViolationError) + attempted_group: Sequence[object] | None = None # rebind-ok: set below + orders: Final = sorted( + { + order + for deployment in python._core_get_model_list(model_name=model_name) or () # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + if (order := (deployment.get("litellm_params") or {}).get("order")) is not None # pyright: ignore[reportUnknownMemberType] # a deployment's untyped litellm_params + } + ) + if len(orders) > 1 and not (context_window or content_policy): + external, generic = ( + get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=lookup_groups) + if fallbacks is not None and lookup_groups + else (None, None) + ) + if external is None and generic is not None and fallbacks is not None: + external = fallbacks[generic]["*"] # rebind-ok: the generic chain when no specific one matched + combined: Final = [{"model": model_name, "_target_order": o} for o in orders if o > orders[0]] + list( + external or () + ) + if combined: + attempted_group = combined # rebind-ok: settled in the order Python's walk does + if attempted_group is None: + specific: Final = ( + context_window_fallbacks if context_window else content_policy_fallbacks if content_policy else None + ) + if (context_window or content_policy) and specific is not None: + attempted_group, _ = ( # rebind-ok: settled in the order Python's walk does + get_fallback_model_group_for_lookup_groups(fallbacks=specific, lookup_groups=lookup_groups) + ) + if attempted_group is None: + # Python raises before its generic fallbacks: "No fallback was attempted". + error.message += format_fallback_outcome_message( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix + model_name, None, "", skipped=self._walk.skipped + ) + return + elif context_window: + error.message += ( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix + f"\nmodel={model_name}. context_window_fallbacks={mask_sensitive_structure(context_window_fallbacks)}." + f" fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'context_window_fallback' -" + " https://docs.litellm.ai/docs/routing#fallbacks" + ) + elif content_policy: + error.message += ( # rebind-ok: the error Python raises carries the same suffix + f"\n{format_content_policy_fallback_hint(model_name, content_policy_fallbacks, fallbacks)}" # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + ) + if attempted_group is None and fallbacks is not None and lookup_groups: + attempted_group, generic_index = ( # rebind-ok: settled in the order Python's walk does + get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=lookup_groups) + ) + if attempted_group is None and generic_index is not None: + attempted_group = fallbacks[generic_index][ # rebind-ok: settled in the order Python's walk does + "*" + ] + if attempted_group is None: + error.message += format_no_fallback_group_message(lookup_groups, fallbacks) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix + return + detail: Final = ( + truncate_fallback_error_detail(redact_string(str(self._level_errors.get(1) or error))) + if attempted_group + else "" + ) + error.message += format_fallback_outcome_message( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix + model_name, attempted_group, detail, skipped=self._walk.skipped + ) + + def _pinned_request(self, step: Mapping[str, object]) -> Mapping[str, object]: + """The request as the Python router's own fallback walk would send it, narrowed to the one + deployment Rust picked: the name Rust resolved, with every other deployment it resolved to + skipped.""" + deployment_id: Final = str(step["deployment_id"]) + name: Final = str(step.get("model") or step["model_group"]) + candidates: Final = step.get("candidate_ids") + others: Final = tuple(str(i) for i in _as_list(candidates) if str(i) != deployment_id) + pinned: Final[dict[str, object]] = { # mutable-ok: built here, then handed to the Python core as its request + **self._request, + "model": name, + } + if others: + pinned["_retry_skipped_deployment_ids"] = others + target_order: Final = step.get("target_order") + if isinstance(target_order, int): + pinned["_target_order"] = target_order + return pinned + + async def _attempt(self, step: Mapping[str, object]) -> Mapping[str, object]: + self._attempted = True + self._attempted_groups.add(str(step["model_group"])) + self._timeline.append(("attempt", step.get("fallback_level"))) + request: Final = self._pinned_request(step) + try: + response: Final = await self._python.call( + self._op, {**request, "num_retries": 0, "disable_fallbacks": True} + ) + except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop + self._last_error = error + level: Final = step.get("fallback_level") + self._level_errors[level if isinstance(level, int) else 0] = error + self._timeline.append(("error", (level if isinstance(level, int) else 0, error, request))) + return attempt_failure(error) + self._last_response = response + return {"ok": True, "total_tokens": total_tokens(response)} + + async def _allow_fallback(self, target: str) -> bool: + """`_is_fallback_target_authorized` then `_is_fallback_target_within_budget`.""" + ctx: Final = self._python._core_ctx # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore + settings: Final = ctx.settings() + request: Final = dict(self._request) + original: Final = self._request.get("model") + if target == original: + return True + if settings.fallback_access_check is not None and not await ctx.fallback_access_allowed( + model=target, request_kwargs=request + ): + self._walk.record_skip(target, FallbackSkipReason.NOT_AUTHORIZED) + return False + if settings.fallback_budget_check is not None and not await ctx.fallback_budget_allowed( + model=target, request_kwargs=request + ): + self._walk.record_skip(target, FallbackSkipReason.OVER_BUDGET) + return False + self._hop_targets.append(target) + return True diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..87a02ecc288 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import Mapping, Sequence -from typing import Never, final +from typing import Literal, Never, NotRequired, TypedDict, final import httpx from pydantic import JsonValue @@ -218,6 +218,8 @@ __all__ = [ "NativeDiagnosticProcessor", "ProcessReservedForForking", "ResponsesWebSocketConnection", + "RouterCall", + "RouterCoreHandle", "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", @@ -325,3 +327,115 @@ class NativeCacheHandle: def ping(self) -> Future[bool]: ... def disconnect(self) -> Future[None]: ... def delete(self, keys: Sequence[str]) -> Future[None]: ... + +# The FailureKind names an attempt outcome's `kind` takes (router-core `FailureKind`). +_FailureKind = Literal[ + "RateLimitError", + "ContextWindowExceededError", + "ContentPolicyViolationError", + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "BadRequestError", + "Timeout", + "APIConnectionError", + "InternalServerError", + "ServiceUnavailableError", + "GuardrailIntervention", + "Declined", + "CallerCredentialError", + "APIError", +] + +class _AttemptStep(TypedDict): + kind: Literal["attempt"] + deployment_id: str + model_group: str + retry: int + fallback_depth: int + timeout: float | None + +class _SleepStep(TypedDict): + kind: Literal["sleep"] + seconds: float + +class _AllowFallbackStep(TypedDict): + kind: Literal["allow_fallback"] + target: str + +class _SelectStep(TypedDict): + kind: Literal["select"] + stage: Literal["plugins", "strategy"] + strategy: str | None + model_group: str + candidates: list[str] + +# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`, +# which sets `keys`; `value` is the JSON to store (`set`) or the amount (`increment`); `ttl` is +# in seconds. Outcomes: get -> the JSON value or None; batch_get -> one JSON value or None per +# key, in order; increment -> the new int or float; set and delete -> None. +class _CacheStep(TypedDict): + kind: Literal["cache"] + op: Literal["get", "set", "increment", "batch_get", "delete"] + key: str | None + keys: list[str] | None + value: JsonValue + ttl: float | None + +class _DoneStep(TypedDict): + kind: Literal["done"] + deployment_id: str + model_group: str + attempted_retries: int + attempted_fallbacks: int + total_tokens: int + +class _FailedStep(TypedDict): + kind: Literal["failed"] + error: str + error_kind: Literal["config", "unknown_model", "no_healthy_deployments", "attempt", "cache", "retired"] + attempt_failed: bool + status: int | None + failure_kind: _FailureKind | None + +# The end of a `RouterCoreHandle.select`. +class _SelectedStep(TypedDict): + kind: Literal["done"] + deployment_id: str + model_group: str + +_RouterStep = _AttemptStep | _SleepStep | _AllowFallbackStep | _SelectStep | _CacheStep | _DoneStep | _FailedStep + +class _AttemptSucceeded(TypedDict): + ok: Literal[True] + total_tokens: NotRequired[int] + +class _AttemptFailed(TypedDict): + ok: Literal[False] + kind: _FailureKind + status: NotRequired[int | None] + retry_after: NotRequired[float | None] + num_retries: NotRequired[int | None] + detail: NotRequired[str | None] + +# Each document is a mapping or its JSON text. +_Document = Mapping[str, object] | str + +@final +class RouterCoreHandle: + def __new__(cls, config: _Document, dynamic: _Document) -> RouterCoreHandle: ... + def start(self, op: str, request: _Document) -> RouterCall: ... + # Driven like `start`; its `done` step is a `_SelectedStep`. + def select(self, model: str, request: _Document) -> RouterCall: ... + def apply(self, changes: _Document) -> int: ... + def registry_version(self) -> int: ... + def retire(self) -> None: ... + +@final +class RouterCall: + # attempt: _AttemptSucceeded | _AttemptFailed; sleep: None; allow_fallback: bool; + # select: the remaining candidate ids, pick first; cache: see _CacheStep; first call: None. + # A call from `select` steps through select and cache only, and ends in _SelectedStep or failed. + def next( + self, outcome: _AttemptSucceeded | _AttemptFailed | bool | Sequence[str] | JsonValue | None + ) -> _RouterStep | _SelectedStep: ... diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 629d9098970..a080eb1c51d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -23,6 +23,7 @@ class Route(str, Enum): OCR = "ocr" TOKEN_COUNTER = "token_counter" TOKENIZER = "tokenizer" + ROUTER = "router" class Target(str, Enum): @@ -221,6 +222,7 @@ RULES: Final[Rules] = ( RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), + RouteRule(Route.ROUTER, Rollout.RUST_OPT_IN), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/tests/code_coverage_tests/router_litellm_globals_check.py b/tests/code_coverage_tests/router_litellm_globals_check.py index ff9fac0e3dc..479ccfb48d1 100644 --- a/tests/code_coverage_tests/router_litellm_globals_check.py +++ b/tests/code_coverage_tests/router_litellm_globals_check.py @@ -21,6 +21,7 @@ CHECKED = ( "litellm/router_core.py", "litellm/router_deprecated.py", "litellm/router_python_core.py", + "litellm/router_rust_core.py", ) diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py index a8e0d6fa0d7..a6b8150a921 100644 --- a/tests/code_coverage_tests/router_surface_check.py +++ b/tests/code_coverage_tests/router_surface_check.py @@ -42,6 +42,7 @@ ROUTER_FILES = { "litellm/router_core.py", "litellm/router_python_core.py", "litellm/router_python_host.py", + "litellm/router_rust_core.py", "litellm/router_deprecated.py", "litellm/router_host.py", } diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index af28160b620..1370ea50cd9 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -452,25 +452,36 @@ def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwarg return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds -CoreBuilder = Callable[..., RouterCore] -"""``builder(host, **router_arguments)``: a router core on ``host``.""" - -CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core} -"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a -configuration raises ``CoreDeclined``, and the test is skipped for that core.""" - - class CoreDeclined(Exception): """The core does not serve this configuration (the router would stay on the Python core).""" +CoreBuilder = Callable[..., RouterCore] +"""``builder(host, **router_arguments)``: a router core on ``host``.""" + + +def rust_core(host: ScriptedHost, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + """A ``RustCore`` on ``host``, built the way ``make_core`` builds one when the Rust router is on.""" + from litellm.router_rust_core import RustCore + + try: + return RustCore.build(build_router_config(arguments), host) + except Exception as declined: # noqa: BLE001 # make_core keeps such a router on the Python core + raise CoreDeclined(str(declined)) from declined + + +CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core, "rust": rust_core} +"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a +configuration raises ``CoreDeclined``, and the test is skipped for that core.""" + + __all__ = [ - "Call", - "HostCallbacks", - "INFERENCE_OPS", "CORE_BUILDERS", + "INFERENCE_OPS", + "Call", "CoreBuilder", "CoreDeclined", + "HostCallbacks", "ScriptedHost", "auth_error", "bad_request", @@ -482,6 +493,7 @@ __all__ = [ "embedding_response", "python_core", "rate_limit", + "rust_core", "server_error", "timeout", ] diff --git a/tests/unit/test_router/core/test_selection.py b/tests/unit/test_router/core/test_selection.py index 4b4c320c428..0d301128f0c 100644 --- a/tests/unit/test_router/core/test_selection.py +++ b/tests/unit/test_router/core/test_selection.py @@ -622,3 +622,40 @@ async def test_cooling_down_one_group_leaves_other_groups_pickable(make_core: Ma with pytest.raises(RouterRateLimitError): await core.select("g") assert await _picked_ids(core, "h") <= {"c", "d"} + + +# --- custom strategies --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_registered_strategy_picks_every_deployment_an_endpoint_call_reaches( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.types.router import CustomRoutingStrategyBase + + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b"), deployment("g", "c")]) + seen: Final[list[Mapping[str, object]]] = [] # mutable-ok: the request each pick was asked for + + class HighestId(CustomRoutingStrategyBase): + async def async_get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + seen.append(dict(request_kwargs or {})) + deployments: Final = core.get_model_list(model_name=model) # pyright: ignore[reportAttributeAccessIssue] # both cores serve the Python core's lookup + return max(deployments, key=lambda d: d["model_info"]["id"]) + + def get_available_deployment(self, model: str) -> dict: # mutable-ok: the strategy interface's own type + raise NotImplementedError + + core.register_plugin(HighestId()) + for _ in range(PICKS): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert set(host.deployments_called("acompletion")) == {"c"} + # The strategy picks from the whole group: no other core narrowed the pick before it. + assert len(seen) == PICKS + assert not any("_retry_skipped_deployment_ids" in request for request in seen)