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..1f02643abe1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -0,0 +1,297 @@ +//! 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, + /// A new dynamic document (the same keys `RouterCoreHandle(static, dynamic)` takes). + #[serde(default)] + settings: Option, +} + +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: spec.settings.map(dynamic_config).transpose()?, + }) +} + +#[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..5f83f7be20b --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -0,0 +1,1025 @@ +//! 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 picked deployment as configured: `model_name`, `litellm_params`, `model_info` + /// (with its `id`). The host builds the request from it; it never looks the deployment up + /// in a Python router. + deployment: Value, + /// 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, + }, + /// `healthy`: the deployments a pick would choose among. + Healthy { deployment_ids: Vec }, + /// `report_failure`: whether the deployment was cooled down. + Reported { cooled: bool }, +} + +/// 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(()) + } +} + +/// A deployment as Python spells a `model_list` entry, with its id in `model_info`. +pub(crate) fn deployment_record(deployment: &Deployment) -> Value { + let mut model_info = deployment.model_info.raw.clone(); + model_info.insert("id".to_owned(), Value::String(deployment.id.clone())); + serde_json::json!({ + "model_name": deployment.model_name, + "litellm_params": deployment.raw_params, + "model_info": model_info, + }) +} + +/// 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(), + deployment: deployment_record(&attempt.deployment), + 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) + } +} + +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"); + } + } + + /// The host only runs the attempt; the core records its outcome (cooldowns, fail counters, + /// usage) through the host's cache. + #[rstest] + #[case::failure(rate_limited())] + #[case::success(Answer::Attempt(Ok(7)))] + fn attempt_outcomes_are_recorded_by_the_core(#[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(), "the core wrote no outcome"); + } + + #[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..4b32038a317 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -0,0 +1,671 @@ +//! `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, SystemTime}, +}; + +use litellm_error::Report; +use litellm_host_python::{from_py_argument, json_loads, to_py}; +use litellm_router_core::{ + AttemptError, Error, FailureKind, Op, Outcome, Router, RouterEvent, + health::{UsageMetric, current_minute, usage_key}, + registry::{DeploymentFilter, ResolveScope}, +}; +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(), + }) + })) + } + + /// List the deployments a pick for `model` would choose among, driven like `select`. + fn healthy(&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 deployments = router.healthy(&Op::Completion, &request).await?; + Ok(Finish::Healthy { + deployment_ids: deployments.iter().map(|d| d.id.clone()).collect(), + }) + })) + } + + /// Apply `{"upsert": [deployments], "delete": [ids], "settings": {dynamic document}}` 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() + } + + /// Record a failure the router did not see itself (a health check, an interceptor), driven + /// like `start`: its steps are cache ops, and it ends with `done` (`cooled`). + fn report_failure(&self, id: &str, failure: &Bound<'_, PyDict>) -> PyResult { + let failure = match attempt_outcome(failure)? { + Err(failure) => failure, + Ok(_) => { + return Err(PyValueError::new_err( + "a reported failure needs `ok: False`", + )); + } + }; + let router = self.router.clone(); + let id = id.to_owned(); + Ok(RouterCall::new(async move { + router + .report(&id, Outcome::Failure(failure)) + .await + .map(|cooled| Finish::Reported { cooled }) + })) + } + + /// This worker's (successes, failures) against a deployment this minute. + fn outcomes_this_minute(&self, id: &str) -> (u64, u64) { + self.router.health().outcomes_this_minute(id) + } + + /// The usage keys of this minute for every deployment of `model_group`: + /// `[(id, tpm_key, rpm_key)]`, or None when the group has no deployment. + fn usage_keys(&self, model_group: &str) -> Option> { + let minute = current_minute(SystemTime::now()); + let group = self.router.registry().group(model_group)?; + Some( + group + .iter() + .map(|deployment| { + let model = &deployment.params.model; + ( + deployment.id.clone(), + usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), + ) + }) + .collect(), + ) + } + + /// Every configured deployment, as `model_list` entries with their ids, in config order. + fn deployments<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments(&DeploymentFilter::default()) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployment with this id, or None. + fn deployment(&self, py: Python<'_>, id: &str) -> PyResult>> { + self.router + .registry() + .deployment(id) + .map(|deployment| to_py(py, &host::deployment_record(&deployment))) + .transpose() + } + + /// What `name` resolves to: `{"model_group", "kind", "ids"}`, or None for an unknown name. + #[pyo3(signature = (name, specific_deployment=false))] + fn resolve<'py>( + &self, + py: Python<'py>, + name: &str, + specific_deployment: bool, + ) -> PyResult>> { + let registry = self.router.registry(); + let Some(resolution) = registry.resolve( + name, + ResolveScope { + specific_deployment, + }, + ) else { + return Ok(None); + }; + let dict = PyDict::new(py); + dict.set_item("model_group", resolution.group)?; + dict.set_item("kind", format!("{:?}", resolution.kind))?; + dict.set_item( + "ids", + PyList::new(py, resolution.deployments.iter().map(|d| d.id.clone()))?, + )?; + Ok(Some(dict)) + } + + 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, + deployment, + model, + model_group, + target_order, + candidate_ids, + retry, + fallback_depth, + fallback_level, + timeout, + } => { + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("deployment", to_py(py, &deployment)?)?; + 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) + } + Finish::Reported { cooled } => dict.set_item("cooled", cooled), + Finish::Healthy { deployment_ids } => dict.set_item("deployment_ids", deployment_ids), + } +} + +#[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"}) +attempt = drive(call, None) +tried = attempt["deployment_id"] +assert attempt["deployment"]["model_info"]["id"] == tried, attempt +assert attempt["deployment"]["litellm_params"]["model"].startswith("openai/"), attempt +assert [d["model_info"]["id"] for d in cooled.deployments()] == ["a1", "a2"] +assert cooled.deployment("a2")["model_name"] == "primary" +assert cooled.resolve("primary")["ids"] == ["a1", "a2"] +assert cooled.resolve("nope") is None +assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed" +# The core records the failure itself: the 429 cools the deployment down in the host's store. +assert f"deployment:{tried}:cooldown" in store, store +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..12c3172f74e --- /dev/null +++ b/litellm-rust/crates/router-core/src/call.rs @@ -0,0 +1,583 @@ +//! 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, Unsupported}, + 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 through the host, 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(), + }); + if !self.host.supports(op, deployment.provider()) { + return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op { + op: op.clone(), + provider: deployment.provider().to_owned(), + })) + .into()); + } + 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, + }; + // Every attempt runs through the host; the core records every outcome itself. + match self.host.invoke(&attempt).await { + Ok(output) => { + if 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) => { + 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 { + // The deployment's own group decides, also when the call named the deployment by id. + single_deployment: snapshot + .registry + .group(&deployment.model_name) + .map_or(candidates.all.len() == 1, |group| group.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..0a1cda9e6ad --- /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: `Host::supports` must accept each one 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..32d2819f51c --- /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("the host cannot run operation `{op}` on provider `{provider}`")] + Op { 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/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..1c9e2e1ed1c --- /dev/null +++ b/litellm-rust/crates/router-core/src/host.rs @@ -0,0 +1,165 @@ +//! [`Host`]: the router core's only boundary. Everything the core needs from the embedding +//! process goes through it, and every attempt runs through [`Host::invoke`]: the core never runs +//! an operation itself. `PythonHost` serves it in a Python process (attempts call `litellm.`); +//! `NativeHost` serves it in the Rust server (attempts call the `inference-*` crates). + +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 one attempt against the picked deployment. The core records its outcome (cooldowns, + /// allowed-fails counters, usage) itself; the host only runs it. + 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 `op` on `provider`. Checked for every served op and + /// deployment when the router is built or changed, so an unsupported pair fails the config, + /// never a request. `PythonHost` runs every op; `NativeHost` only those with a Rust route. + fn supports(&self, _op: &Op, _provider: &str) -> bool { + true + } + + /// Whether [`Self::select`] can run strategies and plugins. + fn can_select(&self) -> bool { + true + } +} diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 05d95d02c84..e410269ef76 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -43,21 +43,320 @@ //! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the //! policy is exhausted. //! -//! # Execution -//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs -//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A -//! started attempt is never replayed on a different executor. +//! # Hosts: the only boundary +//! The core reaches outside itself only through [`Host`], and every attempt runs through +//! `Host::invoke`: the core picks the deployment, the host runs the call, the core records the +//! outcome. There is no second, "native" attempt path in the core. Which code runs an attempt is +//! the host's business: +//! - `PythonHost` (inside a Python process) builds the request for the picked deployment and calls +//! `litellm.`. It never calls back into a Python router. +//! - `NativeHost` (the Rust server) calls the `inference-*` crates. An op or provider it cannot run +//! fails `Host::supports` when the router is built or changed, never partway through a request. //! -//! # Hosts -//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select` -//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks, -//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all -//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost` -//! serves them from the existing Python code when the router runs inside Python; `NativeHost` -//! serves them in the Rust server, where an operation without a Rust implementation makes the -//! config fail at build time, never partway through a request. +//! Besides `invoke` and `supports`, a host provides `select` for strategies and plugins not yet +//! ported, `allow_fallback` for access and budget checks, `sleep`, `observe` for routing events, +//! and `cache`, a DualCache-style store that holds all routing state under the same keys, +//! encodings and TTLs the Python router uses. The core writes cooldowns, allowed-fails counters +//! and usage itself, through that cache. //! //! # Plugins //! 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`]: the `Host` boundary. + +pub mod cache; +pub mod call; +pub mod config; +pub mod error; +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}, + 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, + 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 { + if !self.host.supports(op, deployment.provider()) { + return Err(Unsupported::Op { + op: op.clone(), + provider: deployment.provider().to_owned(), + } + .into()); + } + } + } + Ok(()) + } +} + +/// Builds a [`Router`], with test seams. +pub struct RouterBuilder { + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + rng: Arc, + clock: Arc, +} + +impl RouterBuilder { + #[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, + 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 { + 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 + } + + /// The deployments a pick for `request` would choose among: resolved, team- and + /// cooldown-filtered, through the plugins and the lowest `order`. + pub async fn healthy( + &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?; + Ok(candidates.healthy) + } + + 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. Returns whether + /// the deployment was cooled down. + pub async fn report(&self, id: &str, outcome: Outcome) -> Result { + let snapshot = self.core.snapshot.load_full(); + let Some(deployment) = snapshot.registry.deployment(id) else { + return Ok(false); + }; + match outcome { + Outcome::Success { total_tokens } => self + .core + .health + .record_success(&deployment, total_tokens) + .await + .map(|()| false), + 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(|cooldown| cooldown.is_some()) + } + } + } + + /// 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..4f9f4eb480c --- /dev/null +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -0,0 +1,865 @@ +//! 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: upserts, then deletes (`apply_changes`). 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(); + for spec in &changes.upserts { + let deployment = Deployment::from_spec(spec.clone())?; + deployments.retain(|existing| existing.id != deployment.id); + deployments.push(Arc::new(deployment)); + } + deployments.retain(|deployment| !changes.deletes.contains(&deployment.id)); + 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)) + }) + } + + /// `resolve` for a caller: a team's own deployments under their public name win + /// (`_core_get_all_deployments` with `team_id`), then the name as [`Self::resolve`] reads it, + /// then deployments whose `litellm_params.model` is the name (`_get_deployment_by_litellm_model`). + pub fn resolve_for( + &self, + name: &str, + scope: ResolveScope, + team_id: Option<&str>, + ) -> Option { + if let Some(team) = team_id + && !scope.specific_deployment + { + let own: Vec> = self + .deployments + .iter() + .filter(|d| { + d.model_info.team_id.as_deref() == Some(team) + && d.model_info.team_public_model_name.as_deref() == Some(name) + }) + .cloned() + .collect(); + if !own.is_empty() { + return Some(Resolution { + kind: ResolveKind::Group, + group: name.to_owned(), + deployments: own, + }); + } + } + if let Some(resolution) = self.resolve(name, scope) { + return Some(resolution); + } + let by_model: Vec> = self + .deployments + .iter() + .filter(|d| d.params.model == name) + .cloned() + .collect(); + (!by_model.is_empty()).then(|| Resolution { + kind: ResolveKind::Group, + group: name.to_owned(), + deployments: by_model, + }) + } + + /// 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..d13859de1bf --- /dev/null +++ b/litellm-rust/crates/router-core/src/select.rs @@ -0,0 +1,326 @@ +//! 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_for(name, scope, request.team_id.as_deref()) + .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(); + + // A deployment named by its id is served even while it cools down + // (`get_deployment_by_model_id` returns before the cooldown filter). + let cooling = if resolution.kind == ResolveKind::Id { + Vec::new() + } else { + let ids: Vec = all.iter().map(|d| d.id.clone()).collect(); + 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_attempt.py b/litellm/router_attempt.py new file mode 100644 index 00000000000..b69dbb5d24c --- /dev/null +++ b/litellm/router_attempt.py @@ -0,0 +1,177 @@ +""" +One router attempt, as the host runs it: the request for the deployment the router core picked. + +The Rust router core picks a deployment and hands the attempt to its host; ``PythonHost`` builds +the keyword arguments for that one deployment here and calls ``litellm.`` with them. Nothing +here selects, retries, falls back or records cooldowns: the core does all of that. Nothing here +reads a Python router either, so this module outlives ``router_python_core.py``. +""" + +import re +import uuid +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter +from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs + +# Request keys the router reads to route the call; the deployment call never receives them. +ROUTING_KEYS: Final = frozenset( + { + "model", + "num_retries", + "disable_fallbacks", + "specific_deployment", + "_router_weights", + "_retry_skipped_deployment_ids", + "_target_order", + } +) + +# The positional parameters of each endpoint, in litellm's order: a call that passes them +# positionally is read back into keywords before it is routed. +POSITIONAL: Final = MappingProxyType( + { + "acompletion": ("model", "messages"), + "completion": ("model", "messages"), + "atext_completion": ("model", "prompt"), + "text_completion": ("model", "prompt"), + "aembedding": ("model", "input"), + "embedding": ("model", "input"), + "aimage_generation": ("prompt", "model"), + "image_generation": ("prompt", "model"), + "atranscription": ("model", "file"), + "aspeech": ("model", "input", "voice"), + "arerank": ("model", "query", "documents"), + "amoderation": ("input", "model"), + "aresponses": ("input", "model"), + "aanthropic_messages": (), + } +) + + +def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]: + """The call's arguments as one keyword request. Raises ``TypeError`` like the endpoint would.""" + names: Final = POSITIONAL.get(op, ()) + if len(args) > len(names): + raise TypeError(f"{op}() takes {len(names)} positional arguments but {len(args)} were given") + request: Final = dict(kwargs) + for name, value in zip(names, args, strict=False): + if name in request: + raise TypeError(f"{op}() got multiple values for argument {name!r}") + request[name] = value + return request + + +# Endpoints the Python router serves through ``_ageneric_api_call_with_fallbacks``, which writes +# the routing metadata under ``litellm_metadata``. +_GENERIC_API_OPS: Final = frozenset({"aanthropic_messages", "aresponses"}) + + +def metadata_field(op: str) -> str: + """``metadata`` or ``litellm_metadata``: where the router writes what it routed the call as.""" + return "litellm_metadata" if op in _GENERIC_API_OPS else get_router_metadata_variable_name(function_name=op) + + +def prepare_call(op: str, request: Mapping[str, object], model_group_alias: str | None) -> Mapping[str, object]: + """What the router adds to a request once, before its first attempt + (``_update_kwargs_before_fallbacks``): a trace id and the model group it was asked for.""" + field: Final = metadata_field(op) + metadata: Final = dict(_as_mapping(request.get(field))) + metadata.update({"model_group": request.get("model"), "model_group_alias": model_group_alias}) + return {**request, "litellm_trace_id": request.get("litellm_trace_id") or str(uuid.uuid4()), field: metadata} + + +def for_requested_model(deployment: Mapping[str, object], requested: str | None) -> Mapping[str, object]: + """A wildcard deployment as the router serves ``requested`` through it: its ``model`` with + the wildcards filled in from the requested name (``PatternMatchRouter.route``). Any other + deployment comes back unchanged.""" + model_name: Final = deployment.get("model_name") + params: Final = _as_mapping(deployment.get("litellm_params")) + model: Final = params.get("model") + if requested is None or not isinstance(model_name, str) or "*" not in model_name or not isinstance(model, str): + return deployment + regex: Final = PatternMatchRouter().pattern_to_regex(model_name) + match: Final = re.match(regex, requested) or re.match(regex, f"{inferred_provider(requested)}/{requested}") + if match is None: + return deployment + return { + **deployment, + "litellm_params": {**params, "model": PatternMatchRouter.set_deployment_model_name(match, model)}, + } + + +def attempt_kwargs( + op: str, + request: Mapping[str, object], + deployment: Mapping[str, object], + *, + model_group: str, + default_litellm_params: Mapping[str, object], + timeout: float | None, +) -> Mapping[str, object]: + """The keyword arguments for ``litellm.`` against one deployment + (``_update_kwargs_with_deployment`` and the call sites that merge ``litellm_params``).""" + params: Final = _as_mapping(deployment.get("litellm_params")) + model_info: Final = dict(_as_mapping(deployment.get("model_info"))) + field: Final = metadata_field(op) + metadata: Final = dict(_as_mapping(request.get(field))) + metadata.update( + { + "model_group": model_group, + "deployment": params.get("model"), + "model_info": model_info, + "api_base": params.get("api_base"), + "deployment_model_name": deployment.get("model_name"), + } + ) + tags: Final = _as_list(params.get("tags")) + if tags: + merged: Final = list(_as_list(metadata.get("tags"))) + merged.extend(tag for tag in tags if tag not in merged) + metadata["tags"] = merged + call: Final = {key: value for key, value in request.items() if key not in ROUTING_KEYS} + for key, value in default_litellm_params.items(): + call.setdefault(key, value) + call[field] = metadata + call["model_info"] = model_info + if "timeout" not in call: + deployment_timeout: Final = params.get("timeout") + chosen: Final = deployment_timeout if deployment_timeout is not None else timeout + if chosen is not None: + call["timeout"] = chosen + return without_router_only_kwargs({**params, **call, "model": params.get("model")}) + + +def stamp_attempt_error(error: BaseException, deployment: Mapping[str, object]) -> None: + """What the router stamps on a failed attempt's error: the deployment's own ``num_retries`` + (``_set_deployment_num_retries_on_exception``) and its id (``failed_deployment_id``).""" + params: Final = _as_mapping(deployment.get("litellm_params")) + if getattr(error, "num_retries", None) is None and params.get("num_retries") is not None: + try: + error.num_retries = int(str(params["num_retries"])) # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the router stamps the error in place + except ValueError: + pass + deployment_id: Final = _as_mapping(deployment.get("model_info")).get("id") + if deployment_id and not getattr(error, "failed_deployment_id", None): + try: + error.failed_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: see above + except AttributeError: + pass + + +def _as_mapping(value: object) -> Mapping[str, object]: + if isinstance(value, Mapping): + entries: Final[Mapping[object, object]] = value # pyright: ignore[reportUnknownVariableType] # a request's untyped mapping + return {str(key): item for key, item in entries.items()} + return MappingProxyType({}) + + +def _as_list(value: object) -> list[object]: # mutable-ok: a fresh list + if isinstance(value, list): + items: Final[list[object]] = value # pyright: ignore[reportUnknownVariableType] # a request's untyped list # mutable-ok: read, then copied + return list(items) + return [] diff --git a/litellm/router_core.py b/litellm/router_core.py index acbe82a0643..3ac3bf067f1 100644 --- a/litellm/router_core.py +++ b/litellm/router_core.py @@ -1,14 +1,16 @@ """ ``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one. -``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. The Rust core -(litellm-rust/crates/router-core) implements the same operations. The core is picked once, when -the router is built, from its ``RouterConfig``. +``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. ``RustCore`` +(``litellm/router_rust_core.py`` over litellm-rust/crates/router-core) is a separate implementation +that reaches outside itself only through its host; it does not use ``PythonCore``. The core is +picked once, when 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 +53,26 @@ 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 the config alone: the + Rust core never builds, wraps or falls back to a Python core. + """ 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 + from litellm.rust_bridge._native import RustBridgeDeclined + + try: + return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported + except RustBridgeDeclined as declined: + # A config the Rust core cannot take keeps the router on Python while both cores exist. + if chosen is Decision.RUST_REQUIRED: + raise + verbose_router_logger.info("Router served by the Python core: %s", declined) + return python_class(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature diff --git a/litellm/router_python_host.py b/litellm/router_python_host.py index 8156740144e..ff2c1aad21a 100644 --- a/litellm/router_python_host.py +++ b/litellm/router_python_host.py @@ -7,6 +7,7 @@ come from ``dynamic()`` (``litellm/router_config.py``). A Rust core gets the sam """ import asyncio +import time from collections.abc import Awaitable, Callable, Sequence from typing import ( TYPE_CHECKING, @@ -19,12 +20,13 @@ from typing import ( from typing_extensions import overload import litellm -from litellm.router_config import RouterDynamicConfig, live_dynamic_config +from litellm.router_config import RouterConfig, RouterDynamicConfig, live_dynamic_config if TYPE_CHECKING: from pydantic import BaseModel - from litellm.caching.caching import Cache + from litellm.caching.caching import Cache, RedisCache + from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper @@ -107,6 +109,10 @@ class PythonHost: def sleep(self, seconds: float) -> Awaitable[None]: return asyncio.sleep(seconds) + def sleep_sync(self, seconds: float) -> None: + """The backoff of a sync endpoint's retry.""" + time.sleep(seconds) + def dynamic(self) -> RouterDynamicConfig: """The process settings, read live on each access (``current_dynamic_config()`` takes a snapshot).""" return cast(RouterDynamicConfig, live_dynamic_config) # cast-ok: the live view reads the same fields @@ -116,7 +122,7 @@ class PythonHost: return litellm.logging_callback_manager def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself - value: Final = getattr(litellm, name) + value: Final[object] = getattr(litellm, name) # pyright: ignore[reportAny] # an untyped litellm global return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given @@ -131,5 +137,50 @@ class PythonHost: def suppress_debug_info(self) -> None: litellm.suppress_debug_info = True + def routing_cache(self, config: RouterConfig) -> "DualCache": + """The store a router keeps its routing state in (cooldowns, fail counts, usage): in memory, + plus Redis when the router's arguments name one. The host builds and owns it; a Rust core + reads and writes it through cache steps.""" + from litellm.caching.caching import DualCache, InMemoryCache + + return DualCache(redis_cache=redis_cache_for(config), in_memory_cache=InMemoryCache()) + + +def redis_cache_for(config: RouterConfig) -> "RedisCache | None": + """The Redis cache the router's ``redis_*`` and ``cache_kwargs`` arguments describe, or None.""" + import json + + from litellm.caching.caching import RedisCache, RedisClusterCache + from litellm.utils import get_secret + + url: Final = config.get("redis_url") + host: Final = config.get("redis_host") + port: Final = config.get("redis_port") + if url is None and (host is None or port is None): + return None + settings: Final[dict[str, object]] = {} # mutable-ok: built here, then passed as keywords + if url is not None: + settings["url"] = url + if host is not None: + settings["host"] = host + if port is not None: + settings["port"] = str(port) + if config.get("redis_password") is not None: + settings["password"] = config.get("redis_password") + if config.get("redis_db") is not None: + settings["db"] = str(config.get("redis_db")) + extra: Final = config.get("cache_kwargs") + if isinstance(extra, dict): + settings.update({str(key): value for key, value in extra.items()}) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the cache_kwargs argument is untyped + nodes: object = settings.get("startup_nodes") # rebind-ok: the env cluster nodes when the arguments name none + if not nodes: + env_nodes: Final[object] = get_secret("REDIS_CLUSTER_NODES") # pyright: ignore[reportUnknownVariableType] # get_secret is untyped + if isinstance(env_nodes, str): + parsed: Final[object] = json.loads(env_nodes) # pyright: ignore[reportAny] # plain JSON + nodes = parsed # rebind-ok: see above + if nodes: + return RedisClusterCache(**{**settings, "startup_nodes": nodes}) # pyright: ignore[reportArgumentType] # RedisCache's keyword arguments + return RedisCache(**settings) # pyright: ignore[reportArgumentType] # RedisCache's keyword arguments + python_host: Final = PythonHost() diff --git a/litellm/router_rust_core.py b/litellm/router_rust_core.py new file mode 100644 index 00000000000..2e37e17a9d1 --- /dev/null +++ b/litellm/router_rust_core.py @@ -0,0 +1,1498 @@ +""" +``RustCore``: the router core served by litellm-rust/crates/router-core through ``_native``. + +The Rust core owns the registry, selection, retries, fallbacks and cooldowns. It reaches outside +itself only through its host, and the host never calls back into a Python router: + +- every attempt runs through the host: ``router_attempt`` builds the request for the one + deployment Rust picked and ``PythonHost.invoke(op)`` calls ``litellm.`` with it +- the core records every outcome (cooldowns, allowed-fails counters, usage) itself, in the routing + cache the host builds and owns (``PythonHost.routing_cache``) +- sleeps, fallback access and budget checks, and cache reads and writes are host steps + +``RustCore`` does not build, wrap or read a ``PythonCore``. Settings are resolved from the +``RouterConfig`` and the host's process settings; the views (``registry()``, ``settings()``, +``health()``) are read from the Rust handle. A member the Rust core does not serve raises +``RustCoreUnsupported`` (an ``AttributeError``) instead of reaching a Python core, so deleting +``router_python_core.py`` cannot change what this core does. +""" + +import dataclasses +import inspect +import json +import math +from collections.abc import Awaitable, Callable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast + +from pydantic import TypeAdapter + +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, + apply_remaining_usage_headers, + apply_response_model_id, + ensure_response_additional_headers, + prepare_response_for_header_attachment, +) +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.router_api import ( + ApplyResult, + ChangeSet, + RouterHealth, + RouterRegistry, + RouterSettings, + RoutingContext, + RoutingObservers, +) +from litellm.router_attempt import ( + attempt_kwargs, + for_requested_model, + metadata_field, + prepare_call, + request_from_call, + stamp_attempt_error, +) +from litellm.router_config import RouterConfig, RouterDynamicConfig +from litellm.router_python_host import PythonHost +from litellm.types.router import Deployment, RouterErrors, RouterRateLimitError +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.router_utils.cooldown_cache import CooldownCache + from litellm.rust_bridge._native import ( # pyright: ignore[reportPrivateUsage] # the stub's step shapes + RouterCall, + RouterCoreHandle, + _AttemptFailed, # pyright: ignore[reportPrivateUsage] # the stub's step shape + ) + from litellm.types.router import CustomRoutingStrategyBase + +# Async endpoints the Rust core routes: select a deployment, then run the attempt on the host. +RUST_SERVED_OPS: Final[frozenset[str]] = frozenset( + { + "acompletion", + "atext_completion", + "aembedding", + "aresponses", + "aanthropic_messages", + "aimage_generation", + "atranscription", + "aspeech", + "arerank", + "amoderation", + } +) +# Sync endpoints, routed the same way with the host's sync endpoint and sync cache. +SYNC_OPS: Final = MappingProxyType( + { + "completion": "acompletion", + "text_completion": "atext_completion", + "embedding": "aembedding", + "image_generation": "aimage_generation", + } +) + +_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]) +# The litellm function an op's attempt calls, where it differs from the op's name. +_HOST_ENDPOINTS: Final = MappingProxyType({"aanthropic_messages": "anthropic_messages"}) +_PROVIDERS: Final = tuple(provider.value for provider in LlmProviders) +_DEFAULT_COOLDOWN_SECONDS: Final = 5.0 + +# Constructor arguments this core reads. Any other argument set to something other than its +# unset value keeps the router on the Python core while both exist (make_core decides that from +# the config alone, before any core is built). +_SERVED_ARGUMENTS: Final = frozenset( + { + "model_list", + "routing_strategy", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_after", + "timeout", + "default_litellm_params", + "fallback_access_check", + "fallback_budget_check", + "redis_url", + "redis_host", + "redis_port", + "redis_password", + "redis_db", + "cache_kwargs", + "set_verbose", + "debug_level", + } +) +# Arguments that may be passed at their default without asking for anything. +_UNSET_DEFAULTS: Final[Mapping[str, object]] = MappingProxyType( + { + "client_ttl": 3600, + "cache_responses": False, + "routing_strategy_args": {}, + "enable_pre_call_checks": False, + "enable_tag_filtering": False, + "retry_policy": None, + "model_group_retry_policy": {}, + } +) +# Per-request routing overrides this core does not read yet: such a request is refused, not +# rerouted to another core. +_UNSERVED_REQUEST_KEYS: Final = frozenset( + { + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "default_fallbacks", + "max_fallbacks", + "model_group_retry_policy", + "retry_policy", + } +) +# Settings a changeset may change. Static ones rebuild the Rust core (settings are fixed for a +# core's life); dynamic ones are handed to it as a new dynamic document. +_STATIC_SETTINGS: Final = frozenset( + {"routing_strategy", "cooldown_time", "disable_cooldowns", "allowed_fails_policy", "retry_after", "timeout"} +) +_DYNAMIC_SETTINGS: Final = frozenset( + { + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "model_group_alias", + "default_litellm_params", + "fallback_access_check", + "fallback_budget_check", + } +) + + +class RustCoreUnsupported(AttributeError): + """A member or request the Rust core does not serve. It is never forwarded to a Python core.""" + + +# --- failures ----------------------------------------------------------------------------------- + + +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[object] = getattr(error, "litellm_response_headers", None) or getattr(response, "headers", None) + get: Final[object] = getattr(headers, "get", None) + if headers is None or not callable(get): + return None + milliseconds: Final[object] = get("retry-after-ms") + seconds: Final[object] = get("retry-after") + try: + value: Final = float(str(milliseconds)) / 1000 if milliseconds is not None else float(str(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 + + +# --- settings ----------------------------------------------------------------------------------- + + +def _json_default(value: object) -> object: + dump: Final = getattr(value, "model_dump", None) + return dump(exclude_none=True) if callable(dump) else str(value) + + +def _plain(value: object) -> object: + """``value`` as plain JSON values, the shape the Rust core reads.""" + loaded: Final[object] = json.loads(json.dumps(value, default=_json_default)) # pyright: ignore[reportAny] # plain JSON + return loaded + + +def _document(value: object) -> Mapping[str, object]: + """A mapping as the plain JSON document the Rust core reads.""" + return as_str_mapping(_plain(value)) or {} + + +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 + + +@dataclasses.dataclass(frozen=True, slots=True) +class Resolved: + """The router's settings, resolved from its arguments and the process settings the way the + Python router's constructor resolves them. Replaced, never changed, by ``apply``.""" + + routing_strategy: str + num_retries: int + max_fallbacks: int + allowed_fails: int + explicit_allowed_fails: int | None + cooldown_time: float + disable_cooldowns: bool | None + allowed_fails_policy: object + retry_after: float + timeout: float | None + fallbacks: list[object] | None # mutable-ok: the setting as Python spells it + context_window_fallbacks: list[object] | None # mutable-ok: same + content_policy_fallbacks: list[object] | None # mutable-ok: same + model_group_alias: Mapping[str, object] + default_litellm_params: Mapping[str, object] + fallback_access_check: object + fallback_budget_check: object + assigned: frozenset[str] = frozenset() + + @classmethod + def of(cls, config: RouterConfig, dynamic: RouterDynamicConfig) -> "Resolved": + get: Final = config.get + configured: Final[Sequence[object] | None] = list(_as_list(get("fallbacks"))) or ( + list(dynamic.fallbacks) if dynamic.fallbacks else None + ) + default_fallbacks: Final = get("default_fallbacks") or dynamic.default_fallbacks + if isinstance(default_fallbacks, str): + raise TypeError( + f"default_fallbacks must be a list of model names, got '{default_fallbacks}'. Use ['{default_fallbacks}']" + ) + fallbacks: Final[Sequence[object] | None] = ( + configured if default_fallbacks is None else [*(configured or ()), {"*": list(_as_list(default_fallbacks))}] + ) + num_retries: Final = get("num_retries") + max_fallbacks: Final = get("max_fallbacks") + allowed_fails: Final = get("allowed_fails") + timeout: Final = get("timeout") + return cls( + routing_strategy=str(get("routing_strategy") or "simple-shuffle"), + num_retries=int(str(num_retries)) + if num_retries is not None + else dynamic.num_retries + if dynamic.num_retries is not None + else 2, + max_fallbacks=int(str(max_fallbacks)) + if max_fallbacks is not None + else dynamic.max_fallbacks + if dynamic.max_fallbacks is not None + else dynamic.router_max_fallbacks, + allowed_fails=int(str(allowed_fails)) if allowed_fails is not None else dynamic.allowed_fails, + explicit_allowed_fails=int(str(allowed_fails)) if allowed_fails is not None else None, + cooldown_time=float(str(get("cooldown_time") or _DEFAULT_COOLDOWN_SECONDS)), + disable_cooldowns=_optional_bool(get("disable_cooldowns")), + allowed_fails_policy=get("allowed_fails_policy"), + retry_after=float(str(get("retry_after") or 0)), + timeout=float(str(timeout)) if timeout else dynamic.request_timeout, + fallbacks=fallbacks, + context_window_fallbacks=list(_as_list(get("context_window_fallbacks"))) + or (list(dynamic.context_window_fallbacks) if dynamic.context_window_fallbacks else None), + content_policy_fallbacks=list(_as_list(get("content_policy_fallbacks"))) + or (list(dynamic.content_policy_fallbacks) if dynamic.content_policy_fallbacks else None), + model_group_alias=MappingProxyType(dict(as_str_mapping(get("model_group_alias")) or {})), + default_litellm_params=MappingProxyType(dict(as_str_mapping(get("default_litellm_params")) or {})), + fallback_access_check=get("fallback_access_check"), + fallback_budget_check=get("fallback_budget_check"), + ) + + def changed(self, values: Mapping[str, object]) -> "Resolved": + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + unknown: Final = sorted(set(values) - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES}) + if unknown: + raise ValueError(f"not assignable router settings: {unknown}") + unserved: Final = sorted(set(values) - _STATIC_SETTINGS - _DYNAMIC_SETTINGS) + if unserved: + raise RustCoreUnsupported(f"the Rust router core does not serve the setting(s) {', '.join(unserved)} yet") + updates: Final[dict[str, object]] = {} # mutable-ok: built here, then applied once + for name, value in values.items(): + if name == "model_group_alias" or name == "default_litellm_params": + updates[name] = MappingProxyType(dict(as_str_mapping(value) or {})) + elif name in ("num_retries", "max_fallbacks", "allowed_fails"): + updates[name] = int(str(value)) + if name == "allowed_fails": + updates["explicit_allowed_fails"] = int(str(value)) + elif name in ("cooldown_time", "retry_after"): + updates[name] = float(str(value or 0)) or ( + _DEFAULT_COOLDOWN_SECONDS if name == "cooldown_time" else 0.0 + ) + elif name == "timeout": + updates[name] = float(str(value)) if value is not None else None + elif name == "routing_strategy": + updates[name] = str(value) + elif name == "disable_cooldowns": + updates[name] = _optional_bool(value) + elif name in ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks"): + updates[name] = list(_as_list(value)) or None + else: + updates[name] = value + return dataclasses.replace(self, **updates, assigned=self.assigned | frozenset(values)) # pyright: ignore[reportArgumentType] # each value converted to its field's type above + + def static_document(self, model_list: Sequence[object]) -> Mapping[str, object]: + values: Final = { + "model_list": _plain(list(model_list)), + "routing_strategy": self.routing_strategy, + "cooldown_time": self.cooldown_time, + "disable_cooldowns": self.disable_cooldowns, + "allowed_fails_policy": _plain(self.allowed_fails_policy), + "retry_after": self.retry_after, + "timeout": self.timeout, + } + return {name: value for name, value in values.items() if value is not None} + + def dynamic_document(self, version: int, model_names: Sequence[str]) -> Mapping[str, object]: + values: Final = { + "version": version, + "model_group_alias": _plain(dict(self.model_group_alias)), + "fallbacks": _plain(self.fallbacks), + "context_window_fallbacks": _plain(self.context_window_fallbacks), + "content_policy_fallbacks": _plain(self.content_policy_fallbacks), + "num_retries": self.num_retries, + "max_fallbacks": self.max_fallbacks, + "allowed_fails": self._router_allowed_fails(), + "fallback_keys": self._fallback_keys(model_names), + } + return {name: value for name, value in values.items() if value is not None} + + def _router_allowed_fails(self) -> int | None: + """``allowed_fails`` when it counts as set on the router (``_is_allowed_fails_set_on_router``): + passed explicitly, or different from the process default.""" + import litellm + + if self.explicit_allowed_fails is not None: + return self.explicit_allowed_fails + if self.allowed_fails != litellm.allowed_fails: + return self.allowed_fails + return None + + def _fallback_keys(self, model_names: Sequence[str]) -> Mapping[str, object]: + """The spellings a fallback key matches besides the group name (`get_fallback_model_group`).""" + from litellm.router_utils.fallback_event_handlers import ( + _provider_prefixed_model_group, # pyright: ignore[reportPrivateUsage] # the fallback walk's own spelling rule + ) + + keys: Final = tuple( # comprehension-ok: the first group of each fallback entry, across the three lists + str(next(iter(mapping))) + for entries in (self.fallbacks, self.context_window_fallbacks, self.content_policy_fallbacks) + for entry in _as_list(entries) + if (mapping := as_str_mapping(entry)) + ) + names: Final = {*model_names, *self.model_group_alias} + 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 _optional_bool(value: object) -> bool | None: + return None if value is None else bool(value) + + +def unsupported_arguments(config: RouterConfig) -> frozenset[str]: + """Arguments the Rust core does not serve.""" + unserved: Final = frozenset( + name + for name, value in config.arguments.items() + if name not in _SERVED_ARGUMENTS and value not in (None, _UNSET_DEFAULTS.get(name)) + ) + strategy: Final = config.get("routing_strategy") or "simple-shuffle" + # `simple_shuffle` raises on a non-numeric weight when it picks; only 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(config.get("model_list")) + ) + return ( + unserved + | ({"routing_strategy"} if strategy != "simple-shuffle" else frozenset[str]()) + | ({"model_list"} if odd_weight else frozenset[str]()) + ) + + +def _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 _same_deployment(old: Deployment, new: Deployment) -> bool: + """An upsert that changes nothing (``_upsert_deployment``'s equality check).""" + return ( + old.model_name == new.model_name + and old.litellm_params == new.litellm_params + and old.model_info == new.model_info + ) + + +def _refuse_unserved_deployments(model_list: Sequence[object]) -> None: + """Deployment fields the Rust core does not read yet: refused, never silently ignored.""" + for entry in model_list: + model_info = as_str_mapping((as_str_mapping(_record(entry)) or {}).get("model_info")) or {} + if model_info.get("allowed_fails_policy") is not None: + raise RustCoreUnsupported( + "a deployment's model_info.allowed_fails_policy is not served by the Rust router core yet" + ) + + +def _start_fields(request: Mapping[str, object]) -> Mapping[str, object]: + """The request fields the Rust core routes on.""" + fields: Final = { + key: request[key] + for key in ("specific_deployment", "num_retries", "disable_fallbacks") + if request.get(key) is not None + } + weights: Final = request.get("_router_weights") + team_id: Final = _team_id(request) + return { + **fields, + **({"router_weights": weights} if weights is not None else {}), + **({"team_id": team_id} if team_id else {}), + } + + +def _deployment(record: object) -> Deployment: + return Deployment(**dict(as_str_mapping(record) or {})) # pyright: ignore[reportArgumentType] # a model_list entry + + +# --- views -------------------------------------------------------------------------------------- + + +class RustRegistry(RouterRegistry): + """``registry()`` read from the Rust handle. Queries it does not serve raise ``RustCoreUnsupported``.""" + + __slots__ = ("_core",) + + def __init__(self, core: "RustCore") -> None: # pyright: ignore[reportMissingSuperCall] # reads the Rust core, not a Router + self._core = core + + def __getattribute__(self, name: str) -> object: + if name in _REGISTRY_SERVED or name.startswith("__") or name == "_core": + found: Final[object] = object.__getattribute__(self, name) # pyright: ignore[reportAny] # the view's own member + return found + raise RustCoreUnsupported(f"registry().{name} is not served by the Rust router core", name=name, obj=self) + + @property + def version(self) -> int: + return self._core._version # pyright: ignore[reportPrivateUsage] # the registry reads its own core + + def _records(self) -> Sequence[Mapping[str, object]]: + return [as_str_mapping(record) or {} for record in self._core._handle.deployments()] # pyright: ignore[reportPrivateUsage] # see above + + def deployment(self, id: str) -> Deployment | None: + record: Final = self._core._handle.deployment(id) # pyright: ignore[reportPrivateUsage] # see above + return None if record is None else _deployment(record) + + def has_deployment(self, id: str) -> bool: + return self._core._handle.deployment(id) is not None # pyright: ignore[reportPrivateUsage] # see above + + def deployment_dict(self, id: str) -> dict[str, object] | None: # pyright: ignore[reportIncompatibleMethodOverride] # mutable-ok: a fresh dict + record: Final = self._core._handle.deployment(id) # pyright: ignore[reportPrivateUsage] # see above + return None if record is None else dict(as_str_mapping(record) or {}) + + def deployments( # pyright: ignore[reportIncompatibleMethodOverride] # the "group" source only + self, + model_name: str | None = None, + *, + source: str = "group", + team_id: str | None = None, + **_: object, # kwargs-ok: the RouterRegistry signature + ) -> list[Mapping[str, object]] | None: # mutable-ok: a fresh list, as get_model_list returns + if source != "group": + raise RustCoreUnsupported( + f"registry().deployments(source={source!r}) is not served by the Rust router core" + ) + records: Final = self._records() + if model_name is None: + return list(records) + target: Final = self._target(model_name) + found: Final = [ + record for record in records if record.get("model_name") == target and _usable_by(record, team_id) + ] + return found + + def all(self) -> list[Mapping[str, object]]: # pyright: ignore[reportIncompatibleMethodOverride] # mutable-ok: a fresh list + return list(self._records()) + + def names( # pyright: ignore[reportIncompatibleMethodOverride] # the served kinds only + self, + model_name: str | None = None, + *, + kind: str = "groups", + team_id: str | None = None, + exclude_team_models: bool = False, + ) -> list[str] | set[str]: # mutable-ok: fresh, as the Python core returns them + records: Final = self._records() + if kind == "ids": + return [ + str((as_str_mapping(record.get("model_info")) or {}).get("id")) + for record in records + if (model_name is None or record.get("model_name") == model_name) + and not (exclude_team_models and _team_of(record) is not None) + ] + if kind == "group_set": + return {str(record.get("model_name")) for record in records} + if kind == "groups": + names: Final[list[str]] = [] # mutable-ok: filled here, then returned + for record in records: + team: str | None = _team_of(record) + if team is None: + names.append(str(record.get("model_name"))) + elif team_id is not None and team == team_id: + public = (as_str_mapping(record.get("model_info")) or {}).get("team_public_model_name") + if isinstance(public, str): + names.append(public) + aliases: Final = self._core._resolved.model_group_alias # pyright: ignore[reportPrivateUsage] # see above + names.extend(alias for alias, target in aliases.items() if not _hidden(target)) + return names + raise RustCoreUnsupported(f"registry().names(kind={kind!r}) is not served by the Rust router core") + + def resolve( # pyright: ignore[reportIncompatibleMethodOverride] # the served kinds only + self, + name: str | None, + *, + kind: str = "group", + team_id: str | None = None, + **_: object, # kwargs-ok: the RouterRegistry signature + ) -> str | bool | None: + aliases: Final = self._core._resolved.model_group_alias # pyright: ignore[reportPrivateUsage] # see above + if kind == "from_id": + record: Final = self._core._handle.deployment(str(name)) # pyright: ignore[reportPrivateUsage] # see above + return None if record is None else str((as_str_mapping(record) or {}).get("model_name")) + if name is None: + raise TypeError(f"registry().resolve(kind={kind!r}) needs a name") + if kind == "alias": + return _alias_target(aliases.get(name)) if name in aliases else None + if kind == "group": + resolution: Final = self._core._handle.resolve(name) # pyright: ignore[reportPrivateUsage] # see above + if resolution is not None: + return str(resolution["model_group"]) + return _alias_target(aliases.get(name)) if name in aliases else name + if kind == "recognized": + return self._core._handle.resolve(name) is not None # pyright: ignore[reportPrivateUsage] # see above + raise RustCoreUnsupported(f"registry().resolve(kind={kind!r}) is not served by the Rust router core") + + def aliases(self) -> dict[str, object]: # pyright: ignore[reportIncompatibleMethodOverride] # mutable-ok: a fresh dict + return dict(self._core._resolved.model_group_alias) # pyright: ignore[reportPrivateUsage] # see above + + def _target(self, name: str) -> str: + aliases: Final = self._core._resolved.model_group_alias # pyright: ignore[reportPrivateUsage] # see above + return _alias_target(aliases.get(name)) or name if name in aliases else name + + +_REGISTRY_SERVED: Final = frozenset( + { + "version", + "deployment", + "has_deployment", + "deployment_dict", + "deployments", + "all", + "names", + "resolve", + "aliases", + "_records", + "_target", + } +) + + +def _team_of(record: Mapping[str, object]) -> str | None: + team: Final = (as_str_mapping(record.get("model_info")) or {}).get("team_id") + return team if isinstance(team, str) else None + + +def _usable_by(record: Mapping[str, object], team_id: str | None) -> bool: + team: Final = _team_of(record) + return team is None or team == team_id + + +def _hidden(target: object) -> bool: + mapping: Final = as_str_mapping(target) + return mapping is not None and mapping.get("hidden") is True + + +def _alias_target(target: object) -> str | None: + mapping: Final = as_str_mapping(target) + if mapping is not None: + model: Final = mapping.get("model") + return model if isinstance(model, str) else None + return target if isinstance(target, str) else None + + +class RustSettings(RouterSettings): + """``settings()``: the resolved settings as a frozen snapshot.""" + + def as_dict(self) -> Mapping[str, object]: + router: Final = self._router + if not isinstance(router, _SettingsSource): + return super().as_dict() + return router.settings_dict() + + +class _SettingsSource: + def settings_dict(self) -> Mapping[str, object]: + raise NotImplementedError + + +class RustHealth(RouterHealth): + """``health()``: cooldowns, fail counts and usage, read through the host's routing cache and + written through the Rust core.""" + + __slots__ = ("_core",) + + def __init__(self, core: "RustCore") -> None: # pyright: ignore[reportMissingSuperCall] # reads the Rust core, not a Router + self._core = core + + def __getattribute__(self, name: str) -> object: + if name in _HEALTH_SERVED or name.startswith("__") or name == "_core": + found: Final[object] = object.__getattribute__(self, name) # pyright: ignore[reportAny] # the view's own member + return found + raise RustCoreUnsupported(f"health().{name} is not served by the Rust router core", name=name, obj=self) + + @property + def store(self) -> "DualCache": + return self._core._cache # pyright: ignore[reportPrivateUsage] # the health view reads its own core + + @property + def cooldowns(self) -> "CooldownCache": + return self._core._cooldowns # pyright: ignore[reportPrivateUsage] # see above + + def failures_this_minute(self, deployment_id: str) -> int: + return int(self._core._handle.outcomes_this_minute(deployment_id)[1]) # pyright: ignore[reportPrivateUsage] # see above + + def record_failure(self, deployment_id: str, exception: Exception, exception_status: str | int) -> bool: + failure: Final = dict(attempt_failure(exception)) + if failure["status"] is None and str(exception_status).isdigit(): + failure["status"] = int(str(exception_status)) + handle: Final = self._core._handle # pyright: ignore[reportPrivateUsage] # see above + failed: Final = cast("_AttemptFailed", failure) # cast-ok: attempt_failure builds that shape + call: Final = handle.report_failure(deployment_id, failed) + done: Final = self._core._drive_sync_steps(call) # pyright: ignore[reportPrivateUsage] # see above + return bool(done.get("cooled")) + + async def usage( # pyright: ignore[reportIncompatibleMethodOverride] # the same two kinds + self, model_group: str, *, kind: Literal["current", "remaining"] = "current" + ) -> tuple[int | None, int | None] | Mapping[str, int]: + keys: Final = self._core._handle.usage_keys(model_group) # pyright: ignore[reportPrivateUsage] # see above + if keys is None: + if kind == "current": + return None, None + return {} + values: Final = await self.store.async_batch_get_cache( # pyright: ignore[reportUnknownMemberType] # DualCache's reads are untyped + keys=[tpm for _, tpm, _ in keys] + [rpm for _, _, rpm in keys] + ) + if kind == "current": + tpm_values: Final = (values or [])[: len(keys)] + rpm_values: Final = (values or [])[len(keys) :] + return _summed(tpm_values), _summed(rpm_values) + if kind != "remaining": + raise ValueError(f"unknown usage kind {kind!r}") + current: Final = await self.usage(model_group) + current_tpm, current_rpm = current if isinstance(current, tuple) else (None, None) + limits: Final = [ + as_str_mapping((as_str_mapping(self._core._handle.deployment(id)) or {}).get("litellm_params")) or {} # pyright: ignore[reportPrivateUsage] # see above + for id, _, _ in keys + ] + tpm_limit: Final = _limit(limits, "tpm") + rpm_limit: Final = _limit(limits, "rpm") + remaining: Final[dict[str, int]] = {} # mutable-ok: the header map, built here + if tpm_limit is not None: + remaining["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0) + remaining["x-ratelimit-limit-tokens"] = tpm_limit + if rpm_limit is not None: + remaining["x-ratelimit-remaining-requests"] = rpm_limit - (current_rpm or 0) + remaining["x-ratelimit-limit-requests"] = rpm_limit + return remaining + + +_HEALTH_SERVED: Final = frozenset({"store", "cooldowns", "failures_this_minute", "record_failure", "usage"}) + + +def _summed(values: Sequence[object]) -> int | None: + numbers: Final = [int(value) for value in values if isinstance(value, int | float)] + return sum(numbers) if numbers else None + + +def _limit(params: Sequence[Mapping[str, object]], key: str) -> int | None: + values: Final = [int(value) for p in params if isinstance((value := p.get(key)), int | float)] + return sum(values) if values else None + + +# --- the core ----------------------------------------------------------------------------------- + + +_EndpointImpl: TypeAlias = Callable[..., Awaitable[object]] + + +class RustCore(_SettingsSource): + """A ``RouterCore`` served by the Rust router core over a ``PythonHost``. It holds no Python core.""" + + def __init__(self, config: RouterConfig, host: PythonHost, resolved: Resolved) -> None: + from litellm.router_utils.cooldown_cache import CooldownCache + + self._config: Final = config + self._host: Final = host + self._resolved: Resolved = resolved + self._cache: Final = host.routing_cache(config) + self._cooldowns: CooldownCache = CooldownCache(cache=self._cache, default_cooldown_time=resolved.cooldown_time) + self._version: int = 1 + self._retired: bool = False + self._settings: RustSettings | None = None + self._handle: RouterCoreHandle = self._new_handle(list(_as_list(config.get("model_list")))) + + @classmethod + def build(cls, config: RouterConfig, host: PythonHost) -> "RustCore": + from litellm.rust_bridge._native import RustBridgeDeclined + + extra: Final = unsupported_arguments(config) + if extra: + raise RustBridgeDeclined(f"router arguments not served by the Rust core: {sorted(extra)}") + return cls(config, host, Resolved.of(config, host.dynamic())) + + def __getattr__(self, name: str) -> object: + # Only reached for members this class does not define: never forwarded anywhere. + raise RustCoreUnsupported(f"'{name}' is not served by the Rust router core", name=name, obj=self) + + # -- building -------------------------------------------------------------------------------- + + def _new_handle(self, model_list: Sequence[object]) -> "RouterCoreHandle": + from litellm.rust_bridge._native import RouterCoreHandle + + _refuse_unserved_deployments(model_list) + names: Final = [ + str(name) + for entry in model_list + if isinstance(name := (as_str_mapping(entry) or {}).get("model_name"), str) + ] + return RouterCoreHandle( + self._resolved.static_document(model_list), self._resolved.dynamic_document(self._version, names) + ) + + def _rebuild(self) -> None: + """A new Rust core for changed static settings, over the same deployments.""" + old: Final = self._handle + self._handle = self._new_handle(list(old.deployments())) + old.retire() + + # -- entry points ---------------------------------------------------------------------------- + + def registry(self) -> RustRegistry: + return RustRegistry(self) + + def settings(self) -> RustSettings: + """A cached snapshot, replaced when a changeset changes a setting.""" + if self._settings is None: + self._settings = self._settings_snapshot() + return self._settings + + def _settings_snapshot(self) -> RustSettings: + r: Final = self._resolved + return RustSettings( + routing_strategy=r.routing_strategy, + fallbacks=r.fallbacks, + context_window_fallbacks=r.context_window_fallbacks, + content_policy_fallbacks=r.content_policy_fallbacks, + default_litellm_params=dict(r.default_litellm_params), + cooldown_time=r.cooldown_time, + allowed_fails=r.allowed_fails, + allowed_fails_policy=r.allowed_fails_policy, # pyright: ignore[reportArgumentType] # the policy as passed + disable_cooldowns=r.disable_cooldowns, + enable_pre_call_checks=False, + enable_health_check_routing=False, + enable_tag_filtering=False, + fallback_access_check=r.fallback_access_check, # pyright: ignore[reportArgumentType] # the check as passed + fallback_budget_check=r.fallback_budget_check, # pyright: ignore[reportArgumentType] # the check as passed + cache_responses=False, + _router=self, # pyright: ignore[reportArgumentType] # the snapshot reads as_dict from this core + _unset=frozenset( + { + "default_deployment", + "default_max_parallel_requests", + "health_check_ignore_transient_errors", + "background_health_check_model_groups", + "tag_filtering_match_any", + "tag_routing_prefix", + "router_general_settings", + "provider_budget_config", + "provider_default_deployment_ids", + "kubernetes_pod_discovery", + "search_tools", + "guardrail_list", + } + ), + ) + + def settings_dict(self) -> Mapping[str, object]: + r: Final = self._resolved + return { + "routing_strategy_args": {}, + "routing_strategy": r.routing_strategy, + "allowed_fails": r.allowed_fails, + "cooldown_time": r.cooldown_time, + "num_retries": r.num_retries, + "timeout": r.timeout, + "retry_after": r.retry_after, + "fallbacks": r.fallbacks, + "context_window_fallbacks": r.context_window_fallbacks, + "model_group_alias": dict(r.model_group_alias), + "routing_groups": [], + } + + def health(self) -> RustHealth: + return RustHealth(self) + + def observe(self) -> RoutingObservers: + raise RustCoreUnsupported("observe() is not served by the Rust router core yet", name="observe", obj=self) + + def register_plugin(self, plugin: "CustomRoutingStrategyBase") -> None: + # Plugins will run through the host's select step; until then a plugin is refused, never + # silently served by another core. + raise RustCoreUnsupported( + f"register_plugin({type(plugin).__name__}) is not served by the Rust router core yet", + name="register_plugin", + obj=self, + ) + + def retire(self) -> None: + self._retired = True + self._handle.retire() + + def apply(self, changes: ChangeSet) -> ApplyResult: + if changes.redis_cache is not None or changes.refresh_model_info: + raise RustCoreUnsupported("apply(redis_cache / refresh_model_info) is not served by the Rust router core") + values: Final = {**changes.assign, **changes.settings} + if values: + resolved: Final = self._resolved.changed(values) + static_changed: Final = bool(_STATIC_SETTINGS & set(values)) + self._resolved = resolved + self._settings = None + if "cooldown_time" in values: + from litellm.router_utils.cooldown_cache import CooldownCache + + self._cooldowns = CooldownCache(cache=self._cache, default_cooldown_time=resolved.cooldown_time) + if static_changed: + self._rebuild() + else: + names: Final = [str(d.get("model_name")) for d in self._handle.deployments()] + self._handle.apply({"settings": resolved.dynamic_document(self._version, names)}) + upserted: Final[list[Deployment | None]] = [] # mutable-ok: filled per upsert + deleted: Final[list[Deployment | None]] = [] # mutable-ok: filled per delete + # The registry as each write leaves it, in order: upserts first, then deletes + # (`apply_changes`). Writes that change nothing are not sent and keep the version. + view: Final[dict[str, Deployment | None]] = {} # mutable-ok: filled per write + records: Final[list[object]] = [] # mutable-ok: the upserts that change a deployment + deletes: Final[list[str]] = [] # mutable-ok: the deletes of a deployment that exists + + def current(id: str) -> Deployment | None: + if id in view: + return view[id] + record: Final = self._handle.deployment(id) + return None if record is None else _deployment(record) + + for deployment in changes.upsert: + record = _plain(_record(deployment)) + new = _deployment(record) + id = str(new.model_info.id) + old = current(id) + if old is not None and _same_deployment(old, new): + upserted.append(None) + continue + view[id] = new + records.append(record) + upserted.append(new) + for id in changes.delete: + old = current(id) + deleted.append(old) + if old is not None: + view[id] = None + deletes.append(id) + if records or deletes: + self._handle.apply({"upsert": records, "delete": deletes}) + self._version += 1 + return ApplyResult(version=self._version, upserted=upserted, deleted=deleted) + + def upsert_deployment(self, deployment: Deployment) -> Deployment | None: + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice + ) + + _note("upsert_deployment", "apply(ChangeSet(upsert=[deployment]))") + return self.apply(ChangeSet(upsert=[deployment])).upserted[0] + + def delete_deployment(self, id: str) -> Deployment | None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see above + + _note("delete_deployment", "apply(ChangeSet(delete=[id]))") + return self.apply(ChangeSet(delete=[id])).deleted[0] + + def __setattr__(self, name: str, value: object) -> None: + # Settings are fixed after construction: only apply() changes them. A public attribute + # write would be silently ignored, so it is refused. + if not name.startswith("_"): + raise RustCoreUnsupported( + f"writing {name!r} is not served by the Rust router core; use apply(ChangeSet(assign=...))", + name=name, + obj=self, + ) + object.__setattr__(self, name, value) + + async def select( # pyright: ignore[reportIncompatibleMethodOverride] # the RouterCore selection + self, + model: str, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the RouterCore signature + *, + kind: str = "pick", + specific_deployment: bool | None = None, + **_: object, # kwargs-ok: the RouterCore select options this core ignores + ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + """Pick a deployment for ``model`` without calling it (``pick``), or list the deployments + a pick would choose among (``healthy``).""" + if kind not in ("pick", "healthy"): + raise RustCoreUnsupported(f"select(kind={kind!r}) is not served by the Rust router core") + fields: Final = dict(_start_fields(request_kwargs or {})) + if specific_deployment: + fields["specific_deployment"] = True + call: Final = (self._handle.select if kind == "pick" else self._handle.healthy)(model, _document(fields)) + step: Mapping[str, object] = {} # rebind-ok: one step per round + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome + match step["kind"]: + case "cache": + outcome = await self._answer_cache(step) + case "select": + outcome = step["candidates"] + case "done" if kind == "healthy": + records = (self._handle.deployment(str(i)) for i in _as_list(step.get("deployment_ids"))) + return [for_requested_model(as_str_mapping(record) or {}, model) for record in records] + case "done": + record = self._handle.deployment(str(step["deployment_id"])) + return for_requested_model(as_str_mapping(record) or {}, model) + case "failed": + raise self._failure_error(step, model, None) + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + # -- endpoints ------------------------------------------------------------------------------- + + async def call(self, op: str, request: Mapping[str, object]) -> object: + if op in SYNC_OPS: + return self._call_sync(op, request) + if op not in RUST_SERVED_OPS: + raise RustCoreUnsupported(f"endpoint {op!r} is not served by the Rust router core yet", name=op, obj=self) + return await self._drive(op, request) + + async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("acompletion", request_from_call("acompletion", args, kwargs)) + + async def atext_completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("atext_completion", request_from_call("atext_completion", args, kwargs)) + + async def aembedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("aembedding", request_from_call("aembedding", args, kwargs)) + + async def aresponses(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("aresponses", request_from_call("aresponses", args, kwargs)) + + async def aanthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("aanthropic_messages", kwargs) + + async def anthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload + """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async, as in PythonCore.""" + return await self.call("aanthropic_messages", kwargs) + + async def aimage_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("aimage_generation", request_from_call("aimage_generation", args, kwargs)) + + async def atranscription(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("atranscription", request_from_call("atranscription", args, kwargs)) + + async def aspeech(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("aspeech", request_from_call("aspeech", args, kwargs)) + + async def arerank(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("arerank", request_from_call("arerank", args, kwargs)) + + async def amoderation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self.call("amoderation", request_from_call("amoderation", args, kwargs)) + + def completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync("completion", request_from_call("completion", args, kwargs)) + + def text_completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync("text_completion", request_from_call("text_completion", args, kwargs)) + + def embedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync("embedding", request_from_call("embedding", args, kwargs)) + + def image_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync("image_generation", request_from_call("image_generation", args, kwargs)) + + # -- driving a call -------------------------------------------------------------------------- + + def _start(self, op: str, request: Mapping[str, object]) -> "tuple[RouterCall, Mapping[str, object]]": + if self._retired: + raise RuntimeError("this router core was retired") + unserved: Final = sorted( + key for key in request if key in _UNSERVED_REQUEST_KEYS or key.startswith("mock_testing_") + ) + if unserved: + raise RustCoreUnsupported(f"request keys not served by the Rust router core yet: {unserved}") + model: Final = request.get("model") + if not isinstance(model, str): + raise TypeError(f"{op}() needs a model name") + field: Final = metadata_field(op) + if request.get(field) is not None and not isinstance(request.get(field), Mapping): + raise TypeError(f"{field} must be a mapping") + aliases: Final = self._resolved.model_group_alias + prepared: Final = prepare_call(op, request, model if model in aliases else None) + return self._handle.start(op, _document({"model": model, **_start_fields(request)})), prepared + + async def _drive(self, op: str, request: Mapping[str, object]) -> object: + call, prepared = self._start(op, request) + driver: Final = _Driver(self, op, prepared) + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome against the step it answers + driver.take_events(step.get("events")) + match step["kind"]: + case "attempt": + outcome = await driver.attempt(step) + case "sleep": + await self._host.sleep(float(str(step["seconds"]))) + outcome = None + case "allow_fallback": + outcome = await driver.allow_fallback(str(step["target"])) + case "cache": + outcome = await self._answer_cache(step) + case "select": + outcome = step["candidates"] + case "done": + return await driver.finish_async(step) + case "failed": + raise driver.failed(step) + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + def _call_sync(self, op: str, request: Mapping[str, object]) -> object: + async_op: Final = SYNC_OPS[op] + call, prepared = self._start(async_op, request) + driver: Final = _Driver(self, op, prepared) + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # see _drive + driver.take_events(step.get("events")) + match step["kind"]: + case "attempt": + outcome = driver.attempt_sync(step) + case "sleep": + self._host.sleep_sync(float(str(step["seconds"]))) + outcome = None + case "allow_fallback": + # The fallback checks are awaited on the async path only, as in the Python router. + outcome = True + case "cache": + outcome = self._answer_cache_sync(step) + case "select": + outcome = step["candidates"] + case "done": + return driver.finish(step) + case "failed": + raise driver.failed(step) + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + def _drive_sync_steps(self, call: "RouterCall") -> Mapping[str, object]: + """Drive a call whose only steps are cache ops (``report_failure``) to its end.""" + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # see _drive + match step["kind"]: + case "cache": + outcome = self._answer_cache_sync(step) + case "done": + return step + case "failed": + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + case other: + raise RuntimeError(f"router-core asked for {other!r} while recording an outcome") + + # -- cache steps ----------------------------------------------------------------------------- + + def _store_for(self, key: str) -> "DualCache": + """Cooldowns live in the cooldown store; every other routing key in the routing cache.""" + return self._cooldowns.cooldown_store if key.endswith(":cooldown") else self._cache + + async def _answer_cache(self, step: Mapping[str, object]) -> object: + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + match step.get("op"): + case "get": + return _plain(await self._store_for(key).async_get_cache(key)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + return _plain( + [await self._store_for(str(k)).async_get_cache(str(k)) for k in _as_list(step.get("keys"))] # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + ) + case "set": + store: Final = self._store_for(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] # see above + return None + case "increment": + by: Final = _amount(step.get("value")) + counter: Final = self._store_for(key) + total: Final = ( + await counter.async_increment_cache(key, by) # pyright: ignore[reportUnknownMemberType] # see above + if ttl is None + else await counter.async_increment_cache(key, by, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + ) + return by if total is None else total + case "delete": + await self._store_for(key).async_delete_cache(key) + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + def _answer_cache_sync(self, step: Mapping[str, object]) -> object: + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + match step.get("op"): + case "get": + return _plain(self._store_for(key).get_cache(key)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + return _plain([self._store_for(str(k)).get_cache(str(k)) for k in _as_list(step.get("keys"))]) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + case "set": + store: Final = self._store_for(key) + if ttl is None: + store.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + else: + store.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + return None + case "increment": + by: Final = _amount(step.get("value")) + counter: Final = self._store_for(key) + # DualCache.increment_cache is annotated int but adds floats as given. + return ( + counter.increment_cache(key, by) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above + if ttl is None + else counter.increment_cache(key, by, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above + ) + case "delete": + self._store_for(key).delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + # -- errors the core raises when no attempt failed ------------------------------------------ + + def _failure_error(self, step: Mapping[str, object], model: str, last: BaseException | None) -> BaseException: + import litellm.exceptions as litellm_exceptions + + kind: Final = step.get("error_kind") + aliases: Final = self._resolved.model_group_alias + if kind == "unknown_model": + named: Final = _alias_target(aliases.get(model)) or model if model in aliases else model + return litellm_exceptions.BadRequestError( + message=f"You passed in model={named}. {RouterErrors.no_healthy_deployments.value}", + model=named, + llm_provider="", + ) + if kind == "no_healthy_deployments": + group: Final = self.registry().resolve(model) + ids: Final = self.registry().names(str(group), kind="ids") + active: Final = self._cooldowns.get_active_cooldowns(list(ids), None) # pyright: ignore[reportArgumentType] # a list of ids + cooling: Final = [id for id, _ in active] + min_cooldown: Final = min( + (_seconds(entry.get("cooldown_time"), self._resolved.cooldown_time) for _, entry in active), + default=self._resolved.cooldown_time, + ) + return RouterRateLimitError( + model=model, + cooldown_time=min_cooldown, + enable_pre_call_checks=False, + cooldown_list=cooling, + model_ids=list(ids), # pyright: ignore[reportArgumentType] # a list of ids + ) + if last is not None: + return last + return RuntimeError(f"router-core ({kind}): {step.get('error')}") + + +def _seconds(value: object, default: float) -> float | int: + """A cooldown time as stored: an int stays an int, so error texts read ``60`` like Python's.""" + if isinstance(value, int | float) and not isinstance(value, bool): + return value + return int(default) if float(default).is_integer() else default + + +def _ttl(step: Mapping[str, object]) -> float | None: + ttl: Final = step.get("ttl") + return ttl if isinstance(ttl, int | float) else None + + +def _amount(value: object) -> float | int: + return value if isinstance(value, int | float) and not isinstance(value, bool) else 1 + + +def _record(deployment: object) -> object: + dump: Final = getattr(deployment, "model_dump", None) + return dump(exclude_none=True) if callable(dump) else deployment + + +class _Driver: + """Answers the attempt and fallback steps of one call, and shapes its result. + + Each attempt is one host call against the deployment Rust picked. The driver keeps the error + each fallback level ended on: a level that fails raises its own error, whatever its fallbacks + raised, as the Python router does. + """ + + def __init__(self, core: RustCore, op: str, request: Mapping[str, object]) -> None: + self._core: Final = core + self._op: Final = op + self._request: Final = request + self._last_response: object = None + self._last_error: BaseException | None = None + self._level_errors: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call runs + self._served: tuple[str | None, Mapping[str, object]] = (None, MappingProxyType({})) + + def take_events(self, events: object) -> None: + for raw_event in _as_list(events): + event = as_str_mapping(raw_event) + if event is None: + continue + if event.get("kind") != "retries_exhausted": + continue + 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._core._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # the driver serves its own core + 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] # same + + def _kwargs(self, step: Mapping[str, object]) -> tuple[Mapping[str, object], Mapping[str, object]]: + name: Final = step.get("model") + requested: Final = str(name) if name is not None else None + deployment: Final = for_requested_model(as_str_mapping(step.get("deployment")) or {}, requested) + resolved: Final = self._core._resolved # pyright: ignore[reportPrivateUsage] # see above + group: Final = str(step["model_group"]) + kwargs: Final = attempt_kwargs( + self._op, + self._request, + deployment, + model_group=requested if requested is not None and "*" in group else group, + default_litellm_params=resolved.default_litellm_params, + timeout=resolved.timeout, + ) + self._served = (requested, kwargs) + return deployment, kwargs + + def _failed( + self, step: Mapping[str, object], deployment: Mapping[str, object], error: BaseException + ) -> Mapping[str, object]: + stamp_attempt_error(error, deployment) + self._last_error = error + level: Final = step.get("fallback_level") + self._level_errors[level if isinstance(level, int) else 0] = error + return attempt_failure(error) + + async def attempt(self, step: Mapping[str, object]) -> Mapping[str, object]: + deployment, kwargs = self._kwargs(step) + try: + host: Final = self._core._host # pyright: ignore[reportPrivateUsage] # the driver serves its own core + endpoint: Final = host.invoke(_HOST_ENDPOINTS.get(self._op, self._op)) + # Every op the Rust core serves names an async endpoint, so the call returns an awaitable. + response: Final[object] = await cast("Awaitable[object]", endpoint(**kwargs)) # cast-ok: see above + except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop + return self._failed(step, deployment, error) + self._last_response = response + return {"ok": True, "total_tokens": total_tokens(response)} + + def attempt_sync(self, step: Mapping[str, object]) -> Mapping[str, object]: + deployment, kwargs = self._kwargs(step) + try: + response: Final = self._core._host.invoke(self._op)(**kwargs) # pyright: ignore[reportPrivateUsage] # a sync endpoint + except Exception as error: # noqa: BLE001 # see attempt + return self._failed(step, deployment, error) + if inspect.isawaitable(response): + raise TypeError(f"the host's {self._op} endpoint is not synchronous") + 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`.""" + if target == self._request.get("model"): + return True + ctx: Final = RoutingContext(self._core) # pyright: ignore[reportArgumentType] # the checks read registry() and settings() + request: Final = dict(self._request) + if not await ctx.fallback_access_allowed(model=target, request_kwargs=request): + return False + return await ctx.fallback_budget_allowed(model=target, request_kwargs=request) + + def finish(self, step: Mapping[str, object]) -> object: + """The response with the router's headers: the deployment id, the model group it was + served as, and the retry and fallback counts (``set_response_headers``).""" + response: object = prepare_response_for_header_attachment( + self._last_response + ) # rebind-ok: gains its headers in turn + if response is None: + return self._last_response + name, kwargs = self._served + additional: Final = ensure_response_additional_headers(response) + apply_response_model_id(response, as_str_mapping(kwargs.get(metadata_field(self._op)))) + additional["x-litellm-model-group"] = name + retries: Final = step.get("attempted_retries") + if isinstance(retries, int): + max_retries: Final = step.get("max_retries") + response = add_retry_headers_to_response( # rebind-ok: see above + response=response, + attempted_retries=retries, + max_retries=max_retries if isinstance(max_retries, int) and retries > 0 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 + ) + + async def finish_async(self, step: Mapping[str, object]) -> object: + response: Final = self.finish(step) + name: Final = self._served[0] + if name is not None and response is not None: + remaining: Final = await self._core.health().usage(name, kind="remaining") + if isinstance(remaining, Mapping): + apply_remaining_usage_headers(ensure_response_additional_headers(response), dict(remaining)) + return response + + def failed(self, step: Mapping[str, object]) -> BaseException: + model: Final = str(self._request.get("model")) + primary: Final = self._level_errors.get(0) + if primary is None: + return self._core._failure_error(step, model, self._last_error) # pyright: ignore[reportPrivateUsage] # see above + if self._core._host.dynamic().expose_router_debug_in_errors and hasattr(primary, "message"): # pyright: ignore[reportPrivateUsage] # see above + primary.message += self._explanation(primary, model) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + return primary + + def _explanation(self, primary: BaseException, model: str) -> str: + """What ``_async_function_with_fallbacks_common_utils`` appends to the error it re-raises + under ``expose_router_debug_in_errors``, from the resolved settings and what the call did.""" + import litellm.exceptions as litellm_exceptions + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group_for_lookup_groups + + resolved: Final = self._core._resolved # pyright: ignore[reportPrivateUsage] # the driver serves its own core + lookup: Final = (model,) + context_window: Final = isinstance(primary, litellm_exceptions.ContextWindowExceededError) + content_policy: Final = isinstance(primary, litellm_exceptions.ContentPolicyViolationError) + # The error kind's own list when it is set, else the generic one with a hint to set it. + own: Final = ( + resolved.context_window_fallbacks + if context_window + else resolved.content_policy_fallbacks + if content_policy + else None + ) + hint: Final = ( + f"\nmodel={model}. context_window_fallbacks=None. fallbacks={resolved.fallbacks}.\n\n" + "Set 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + if context_window and own is None + else "\n" + format_content_policy_fallback_hint(model, None, resolved.fallbacks) + if content_policy and own is None + else "" + ) + source: Final = own if own is not None else resolved.fallbacks + entries: Final = [dict(m) for entry in _as_list(source) if (m := as_str_mapping(entry))] + found, generic = ( + get_fallback_model_group_for_lookup_groups(entries, lookup) # pyright: ignore[reportArgumentType] # the fallbacks config shape + if entries + else (None, None) + ) + chain: Final[Sequence[object] | None] = ( + found if found is not None else _as_list(entries[generic].get("*")) if generic is not None else None + ) + if chain is None and entries and own is None: + return hint + format_no_fallback_group_message(lookup, entries) # pyright: ignore[reportArgumentType] # the fallbacks config shape + if chain is None: + return hint + format_fallback_outcome_message(model, None, "") + return hint + format_fallback_outcome_message(model, chain, self._fallback_detail(chain)) + + def _fallback_detail(self, chain: Sequence[object]) -> str: + """The error the fallback walk ended on: the first fallback level's last error, or the + unknown-model error of a target with no deployment, which fails before any attempt.""" + last: Final = str(chain[-1]) if chain else None + attempted: Final = self._level_errors.get(1) + unknown: Final = last is not None and self._core._handle.resolve(last) is None # pyright: ignore[reportPrivateUsage] # the driver serves its own core + error: Final = ( + attempted + if attempted is not None or not unknown or last is None + else self._core._failure_error({"error_kind": "unknown_model"}, last, None) # pyright: ignore[reportPrivateUsage] # see above + ) + return truncate_fallback_error_detail(redact_string(str(error))) if error is not None else "" + + +__all__ = [ + "RUST_SERVED_OPS", + "Resolved", + "RustCore", + "RustCoreUnsupported", + "RustHealth", + "RustRegistry", + "RustSettings", + "attempt_failure", + "failure_kind", + "total_tokens", + "unsupported_arguments", +] diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..b5df74e22ce 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,152 @@ 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 + # The deployment as a `model_list` entry, with its id in `model_info`. + deployment: dict[str, JsonValue] + # The name this attempt serves: the requested model or the fallback target. + model: str + model_group: str + target_order: int | None + candidate_ids: list[str] + retry: int + fallback_depth: int + fallback_level: 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 + +# The end of a `RouterCoreHandle.healthy`. +class _HealthyStep(TypedDict): + kind: Literal["done"] + deployment_ids: list[str] + +# The end of a `RouterCoreHandle.report_failure`. +class _ReportedStep(TypedDict): + kind: Literal["done"] + cooled: bool + +# `RouterCoreHandle.resolve`: what a name resolves to. +class _Resolution(TypedDict): + model_group: str + kind: str + ids: list[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: ... + # Driven like `select`; its `done` step is a `_HealthyStep`. + def healthy(self, model: str, request: _Document) -> RouterCall: ... + # `{"upsert": [deployments], "delete": [ids], "settings": {dynamic document}}`. + def apply(self, changes: _Document) -> int: ... + def registry_version(self) -> int: ... + # A failure the router did not see itself, driven like `start`: cache steps, then a + # `_ReportedStep`. + def report_failure(self, id: str, failure: _AttemptFailed) -> RouterCall: ... + # (successes, failures) recorded for the deployment this minute. + def outcomes_this_minute(self, id: str) -> tuple[int, int]: ... + # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. + def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... + # Every deployment as a `model_list` entry, in config order. + def deployments(self) -> list[dict[str, JsonValue]]: ... + def deployment(self, id: str) -> dict[str, JsonValue] | None: ... + def resolve(self, name: str, specific_deployment: bool = False) -> _Resolution | None: ... + 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 | _HealthyStep | _ReportedStep: ... 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/conftest.py b/tests/unit/test_router/core/conftest.py index dffa938c441..6714e08c261 100644 --- a/tests/unit/test_router/core/conftest.py +++ b/tests/unit/test_router/core/conftest.py @@ -1,4 +1,5 @@ -from collections.abc import Callable +from collections.abc import Callable, Generator +from typing import Final import pytest @@ -15,12 +16,33 @@ def host() -> ScriptedHost: def make_core(request: pytest.FixtureRequest, host: ScriptedHost) -> Callable[..., RouterCore]: """``make_core(**router_arguments)``: the core under test, on the ``host`` fixture. Every test that uses it runs once per core in ``CORE_BUILDERS`` (``-k python`` or ``-k rust`` picks one).""" - builder = CORE_BUILDERS[request.param] + name: Final[str] = str(request.param) # pyright: ignore[reportAny] # the fixture parameter + builder: Final = CORE_BUILDERS[name] def build(**arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments try: return builder(host, **arguments) except CoreDeclined as declined: - pytest.skip(f"{request.param} core declines this configuration: {declined}") + pytest.skip(f"{name} core declines this configuration: {declined}") return build + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, None]: + """A test that reaches a member or request the Rust core does not serve yet is skipped with + that member named, so the run separates "not served yet" from a wrong answer.""" + outcome: Final = yield + excinfo: Final[object] = getattr(outcome, "excinfo", None) + if not isinstance(excinfo, tuple): + return + info: Final[tuple[object, ...]] = excinfo # pyright: ignore[reportUnknownVariableType] # pluggy's (type, value, traceback) + if len(info) < 2: + return + from litellm.router_rust_core import RustCoreUnsupported + + error: Final[object] = info[1] + cause: Final[object] = getattr(error, "__cause__", None) + force_exception: Final[object] = getattr(outcome, "force_exception", None) + if (isinstance(error, RustCoreUnsupported) or isinstance(cause, RustCoreUnsupported)) and callable(force_exception): + force_exception(pytest.skip.Exception(f"not served by the Rust core yet: {error}", _use_item_location=True)) diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index af28160b620..8814a738be8 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -324,6 +324,9 @@ class ScriptedHost(PythonHost): return done() + def sleep_sync(self, seconds: float) -> None: + self.sleeps.append(seconds) + def dynamic(self) -> RouterDynamicConfig: return self._dynamic @@ -452,25 +455,44 @@ 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} + + +def records_outcomes_through_callbacks(core: RouterCore) -> bool: + """Whether ``core`` learns attempt outcomes from callbacks it registers on its host. The Python + core does; the Rust core records every outcome from the attempt itself and registers none.""" + return isinstance(core, PythonCore) + + +"""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 +504,7 @@ __all__ = [ "embedding_response", "python_core", "rate_limit", + "rust_core", "server_error", "timeout", ] diff --git a/tests/unit/test_router/core/test_endpoints.py b/tests/unit/test_router/core/test_endpoints.py index c0c4c8cf6dd..ce56ae403ba 100644 --- a/tests/unit/test_router/core/test_endpoints.py +++ b/tests/unit/test_router/core/test_endpoints.py @@ -22,6 +22,7 @@ from tests.unit.test_router.core.scripted_host import ( completion_response, deployment, embedding_response, + records_outcomes_through_callbacks, ) MESSAGES: Final = [{"role": "user", "content": "hi"}] @@ -45,11 +46,16 @@ ENDPOINTS: Final[dict[str, tuple[str, EndpointCall]]] = { "anthropic_messages", lambda core: core.aanthropic_messages(model="g", messages=MESSAGES, max_tokens=5), ), + # The proxy's /v1/messages route calls the router by this name (route_request). + "anthropic_messages": ( + "anthropic_messages", + lambda core: core.anthropic_messages(model="g", messages=MESSAGES, max_tokens=5), + ), "call-acompletion": ("acompletion", lambda core: core.call("acompletion", {"model": "g", "messages": MESSAGES})), "call-completion": ("completion", lambda core: core.call("completion", {"model": "g", "messages": MESSAGES})), } SYNC_ENDPOINTS: Final = frozenset({"completion", "embedding", "image_generation"}) -LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages"}) +LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages", "anthropic_messages"}) EMBEDDING_ENDPOINTS: Final = frozenset({"aembedding", "embedding"}) @@ -276,14 +282,15 @@ async def test_request_level_api_key_replaces_the_deployment_key_on_the_host_cal async def test_construction_registers_one_success_and_one_failure_callback_per_mode( make_core: MakeCore, host: ScriptedHost ): - _one_group(make_core) + core: Final = _one_group(make_core) registered: Final = {name: len(callbacks) for name, callbacks in host.lists.items()} + one: Final = 1 if records_outcomes_through_callbacks(core) else 0 assert registered == { "callbacks": 0, - "success_callback": 1, - "_async_success_callback": 1, - "failure_callback": 1, - "_async_failure_callback": 1, + "success_callback": one, + "_async_success_callback": one, + "failure_callback": one, + "_async_failure_callback": one, "input_callback": 0, } assert all(callable(callbacks[0]) for callbacks in host.lists.values() if callbacks) @@ -416,7 +423,6 @@ def _scripted_stream(*texts: str) -> Callable[..., CustomStreamWrapper]: yield ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": text}}]) # A stream that ends without a finish reason fails since #47665. yield ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]) - yield ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]) logging_obj: Final = LiteLLMLogging( model="d1", diff --git a/tests/unit/test_router/core/test_registry_changesets.py b/tests/unit/test_router/core/test_registry_changesets.py index d4aa4d14a75..567b082607f 100644 --- a/tests/unit/test_router/core/test_registry_changesets.py +++ b/tests/unit/test_router/core/test_registry_changesets.py @@ -16,6 +16,7 @@ from tests.unit.test_router.core.scripted_host import ( ScriptedHost, context_window, deployment, + records_outcomes_through_callbacks, server_error, ) @@ -637,6 +638,8 @@ def test_retire_releases_every_callback_the_core_registered_on_its_host( make_core: MakeCore, host: ScriptedHost ) -> None: core = make_core(model_list=_model_list()) + if not records_outcomes_through_callbacks(core): + pytest.skip("this core registers no callbacks on its host") assert _registered(host) != {} core.retire() @@ -646,6 +649,8 @@ def test_retire_releases_every_callback_the_core_registered_on_its_host( def test_retire_leaves_another_cores_callbacks_registered(make_core: MakeCore, host: ScriptedHost) -> None: kept = make_core(model_list=_model_list()) + if not records_outcomes_through_callbacks(kept): + pytest.skip("this core registers no callbacks on its host") retired = make_core(model_list=[deployment("other", "o")]) registered_by_both = _registered(host) 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)