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)