feat(router): add the Rust router core and serve routing through it behind LITELLM_RUST
Some checks are pending
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run

litellm-router-core implements the router core in Rust: a versioned registry with Python's
deployment ids, selection with native simple-shuffle, the retry, fallback and cooldown loop with
Python's cache keys, and per-attempt executors with a native seam for Messages to Anthropic. The
python bridge exposes it as RouterCoreHandle, whose call loop is step driven so Rust never awaits
Python. RustCore drives that loop for completions, messages, embeddings and responses and serves
everything else from a PythonCore built from the same config. A new router rule in the route
catalog opts in with LITELLM_RUST, and any config the Rust core cannot serve gets PythonCore at
construction

The call loop follows Python's fallback walk (per-target depth, a level raising its own error,
lookup by the current and requested group, "*" after specific chains), stamps retries the way an
exhausted retry loop does, and leaves health writes to the Python attempt, which records them.
RustCore builds its handle from the constructor arguments and the applied changesets and rebuilds
it when either version moves. The core unit suite runs on RustCore too
This commit is contained in:
narmstrong-berri 2026-10-31 17:47:09 +00:00 • committed by Nate Armstrong
parent c43a1addf2
commit a5533cbcfd
28 changed files with 6682 additions and 14 deletions

View file

@ -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"

View file

@ -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"] }

View file

@ -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 }

View file

@ -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",

View file

@ -0,0 +1,294 @@
//! The documents `RouterCoreHandle` reads: the router's static and dynamic config, a call's
//! request and an `apply` change set. A config key router-core does not read declines the
//! handle, so Python keeps the router on its own core.
use std::collections::BTreeMap;
use litellm_router_core::{
ChangeSet, ConfigError, DeploymentSpec, RouterConfig, RouterDynamicConfig, RouterRequest,
};
use serde::{Deserialize, de::DeserializeOwned};
use serde_json::{Map, Value};
/// The `RouterConfig` keys router-core reads.
const STATIC_KEYS: &[&str] = &[
"model_list",
"routing_strategy",
"routing_groups",
"plugins",
"cooldown_time",
"disable_cooldowns",
"allowed_fails_policy",
"retry_policy",
"retry_after",
"timeout",
"request_timeout",
"drop_params",
"enable_pre_call_checks",
"enable_health_check_routing",
"provider_default_deployment_ids",
];
/// The `RouterDynamicConfig` keys router-core reads.
const DYNAMIC_KEYS: &[&str] = &[
"version",
"model_alias_map",
"model_group_alias",
"fallbacks",
"default_fallbacks",
"context_window_fallbacks",
"content_policy_fallbacks",
"num_retries",
"max_fallbacks",
"allowed_fails",
"fallback_keys",
];
/// Router arguments that do not change routing, accepted in either config and dropped.
const IGNORED_KEYS: &[&str] = &[
"set_verbose",
"debug_level",
"alerting_config",
"client_ttl",
];
/// Fallback lists Python treats as unset when empty (`context_window_fallbacks or None`).
const EMPTY_MEANS_UNSET: &[&str] = &["context_window_fallbacks", "content_policy_fallbacks"];
/// Which document a failure is about.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)]
pub(crate) enum Document {
#[strum(serialize = "router config")]
Static,
#[strum(serialize = "dynamic router config")]
Dynamic,
#[strum(serialize = "router request")]
Request,
#[strum(serialize = "change set")]
Changes,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Error {
#[error("the {document} must be a JSON object")]
NotObject { document: Document },
#[error("router-core does not read `{key}` in the {document}")]
UnknownKey { document: Document, key: String },
#[error("router-core does not run `{key}` = {value}")]
UnsupportedValue { key: &'static str, value: Value },
#[error("the {document} does not have the shape router-core reads")]
Shape {
document: Document,
#[source]
source: serde_json::Error,
},
#[error("the request names no model")]
MissingModel,
#[error("router-core rejected the router config")]
Core(#[from] ConfigError),
}
impl Error {
/// Whether router-core cannot serve this config, so Python keeps its own core.
pub(crate) fn declined(&self) -> bool {
match self {
Self::UnknownKey { .. } | Self::UnsupportedValue { .. } | Self::Core(_) => true,
Self::Shape { document, .. } => {
matches!(document, Document::Static | Document::Dynamic)
}
Self::NotObject { .. } | Self::MissingModel => false,
}
}
}
fn object(document: Document, value: Value) -> Result<Map<String, Value>, Error> {
match value {
Value::Object(object) => Ok(object),
_ => Err(Error::NotObject { document }),
}
}
fn shaped<T: DeserializeOwned>(document: Document, object: Map<String, Value>) -> Result<T, Error> {
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<Map<String, Value>, 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<RouterConfig, Error> {
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<RouterDynamicConfig, Error> {
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<String>,
team_id: Option<String>,
specific_deployment: Option<bool>,
num_retries: Option<u32>,
disable_fallbacks: Option<bool>,
router_weights: Option<BTreeMap<String, BTreeMap<String, f64>>>,
body: Option<Value>,
}
/// Read a request; `model` replaces the request's own when given.
pub(crate) fn request(value: Value, model: Option<&str>) -> Result<RouterRequest, Error> {
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<DeploymentSpec>,
#[serde(default)]
delete: Vec<String>,
}
pub(crate) fn changes(value: Value) -> Result<ChangeSet, Error> {
let spec: ChangesSpec = shaped(Document::Changes, object(Document::Changes, value)?)?;
Ok(ChangeSet {
upserts: spec.upsert,
deletes: spec.delete,
settings: None,
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::unknown_key(json!({"routing_strategy_args": {"ttl": 1}}), "routing_strategy_args")]
#[case::dynamic_key_in_static(json!({"num_retries": 2}), "num_retries")]
fn an_unread_static_key_declines_naming_the_key(#[case] json: Value, #[case] key: &str) {
let error = router_config(json).unwrap_err();
assert!(error.declined());
assert!(matches!(&error, Error::UnknownKey { key: named, .. } if named == key));
}
#[rstest]
#[case::null_unknown(json!({"routing_strategy_args": null}))]
#[case::ignored(json!({"set_verbose": true, "debug_level": "INFO"}))]
#[case::pre_call_checks_off(json!({"enable_pre_call_checks": false}))]
fn unset_and_ignored_keys_are_accepted(#[case] json: Value) {
router_config(json).unwrap();
}
#[rstest]
fn pre_call_checks_without_a_plugin_decline() {
let error = router_config(json!({"enable_pre_call_checks": true})).unwrap_err();
assert!(error.declined());
assert!(matches!(
error,
Error::UnsupportedValue {
key: "enable_pre_call_checks",
..
}
));
}
#[rstest]
#[case::not_object(json!([]), false)]
#[case::bad_shape(json!({"retry_after": "soon"}), true)]
fn a_config_that_is_not_an_object_is_a_value_error_and_a_wrong_shape_declines(
#[case] json: Value,
#[case] declined: bool,
) {
assert_eq!(router_config(json).unwrap_err().declined(), declined);
}
#[rstest]
fn python_router_defaults_read_as_unset() {
let dynamic = dynamic_config(json!({
"num_retries": null, "fallbacks": [], "context_window_fallbacks": [],
"content_policy_fallbacks": [], "model_group_alias": {"gpt": "gpt-4"},
}))
.unwrap();
assert_eq!(dynamic.num_retries, None);
assert!(dynamic.context_window_fallbacks.is_none());
assert!(dynamic.content_policy_fallbacks.is_none());
assert_eq!(dynamic.model_alias_map["gpt"], "gpt-4");
}
#[rstest]
#[case::own_model(json!({"model": "gpt", "num_retries": 1}), None, "gpt")]
#[case::override_model(json!({"model": "gpt"}), Some("other"), "other")]
fn a_request_reads_its_model(
#[case] json: Value,
#[case] model: Option<&str>,
#[case] expected: &str,
) {
assert_eq!(request(json, model).unwrap().model, expected);
}
#[rstest]
#[case::unknown_field(json!({"model": "gpt", "temperature": 1}))]
#[case::no_model(json!({"team_id": "t"}))]
fn a_bad_request_is_a_value_error(#[case] json: Value) {
assert!(!request(json, None).unwrap_err().declined());
}
#[rstest]
fn a_change_set_reads_upserts_and_deletes() {
let changes = changes(json!({
"upsert": [{"model_name": "m", "litellm_params": {"model": "openai/x"}}],
"delete": ["old"],
}))
.unwrap();
assert_eq!(changes.upserts.len(), 1);
assert_eq!(changes.deletes, ["old"]);
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,549 @@
//! `RouterCoreHandle`: router-core behind a step-driven handle. Python drives each call through
//! [`RouterCall::next`], running every attempt, sleep, fallback check, selection and cache op the
//! call loop asks for and handing the outcome back; Rust never awaits Python. `select` is driven
//! the same way, since reading cooldowns is a cache step.
mod config;
mod host;
use std::{
sync::{Arc, Mutex, PoisonError},
time::Duration,
};
use litellm_error::Report;
use litellm_host_python::{from_py_argument, json_loads, to_py};
use litellm_router_core::{AttemptError, Error, FailureKind, Op, Router, RouterEvent};
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
types::{PyBool, PyDict, PyList, PyString},
};
use serde_json::Value;
use self::host::{
Answer, Ask, CacheOp, CacheOutcome, DriveError, Driver, Finish, Stage, Step, StepHost, StepKind,
};
use crate::errors::RustBridgeDeclined;
/// The `error_kind` of a `failed` step.
#[derive(Clone, Copy, Debug, strum::IntoStaticStr)]
enum ErrorKind {
#[strum(serialize = "config")]
Config,
#[strum(serialize = "unknown_model")]
UnknownModel,
#[strum(serialize = "no_healthy_deployments")]
NoHealthyDeployments,
#[strum(serialize = "attempt")]
Attempt,
#[strum(serialize = "cache")]
Cache,
#[strum(serialize = "retired")]
Retired,
}
impl From<&Error> for ErrorKind {
fn from(error: &Error) -> Self {
match error {
Error::Config(_) => Self::Config,
Error::UnknownModel { .. } => Self::UnknownModel,
Error::NoHealthyDeployments { .. } => Self::NoHealthyDeployments,
Error::Attempt(_) => Self::Attempt,
Error::Cache(_) => Self::Cache,
Error::Retired => Self::Retired,
}
}
}
fn config_error_to_pyerr(error: config::Error) -> PyErr {
let message = Report::new(&error).to_string();
if error.declined() {
RustBridgeDeclined::new_err(message)
} else {
PyValueError::new_err(message)
}
}
fn drive_error_to_pyerr(error: DriveError) -> PyErr {
let message = Report::new(&error).to_string();
if error.is_request() {
PyValueError::new_err(message)
} else {
PyRuntimeError::new_err(message)
}
}
/// A document Python passes as a mapping or as its JSON text.
fn document(value: &Bound<'_, PyAny>) -> PyResult<Value> {
match value.cast::<PyString>() {
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<Self> {
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<RouterCall> {
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<RouterCall> {
let request =
config::request(document(request)?, Some(model)).map_err(config_error_to_pyerr)?;
let router = self.router.clone();
Ok(RouterCall::new(async move {
let deployment = router.select(&Op::Completion, &request).await?;
Ok(Finish::Selected {
deployment_id: deployment.id.clone(),
model_group: deployment.model_name.clone(),
})
}))
}
/// Apply `{"upsert": [deployments], "delete": [ids]}` and return the new registry version.
fn apply(&self, changes: &Bound<'_, PyAny>) -> PyResult<u64> {
let changes = config::changes(document(changes)?).map_err(config_error_to_pyerr)?;
self.router
.apply(changes)
.map_err(|error| config_error_to_pyerr(error.into()))
}
fn registry_version(&self) -> u64 {
self.router.registry().version()
}
fn retire(&self) {
self.router.retire();
}
}
/// One router call, driven a step at a time by Python.
#[pyclass(frozen)]
pub(crate) struct RouterCall {
driver: Mutex<Driver>,
}
impl RouterCall {
fn new(
future: impl std::future::Future<Output = Result<Finish, Error>> + 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<Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyDict>> {
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<Bound<'_, PyAny>>) -> PyResult<Answer> {
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::<PyDict>()?).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<Bound<'_, PyAny>>) -> PyResult<CacheOutcome> {
let outcome = outcome.filter(|outcome| !outcome.is_none());
let json = |value: Option<Value>| value.filter(|value| !value.is_null());
match op {
CacheOp::Get => outcome
.map(|outcome| from_py_argument::<Value>(&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::<PyList>() {
return Err(PyValueError::new_err(
"a `batch_get` cache outcome must be a list with one entry per key",
));
}
let values: Vec<Option<Value>> = from_py_argument(&outcome)?;
Ok(CacheOutcome::Values(values.into_iter().map(json).collect()))
}
CacheOp::Increment => outcome
.filter(|outcome| !outcome.is_instance_of::<PyBool>())
.and_then(|outcome| outcome.extract::<f64>().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<Result<u64, AttemptError>> {
fn item<'py, T: FromPyObjectOwned<'py>>(
outcome: &Bound<'py, PyDict>,
key: &str,
) -> PyResult<Option<T>> {
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::<f64>(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::<String>(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<RouterEvent>) -> PyResult<Bound<'_, PyList>> {
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<Bound<'_, PyDict>> {
let dict = PyDict::new(py);
let kind = match &step {
Step::Ask(ask) => ask.kind(),
Step::Done(_) => StepKind::Done,
Step::Failed(_) => StepKind::Failed,
};
dict.set_item("kind", <&'static str>::from(kind))?;
match step {
Step::Ask(ask) => ask_items(py, &dict, ask)?,
Step::Done(finish) => done_items(&dict, finish)?,
Step::Failed(error) => {
let attempt = error.attempt();
dict.set_item("error", Report::new(&error).to_string())?;
dict.set_item("error_kind", <&'static str>::from(ErrorKind::from(&error)))?;
dict.set_item("attempt_failed", attempt.is_some())?;
dict.set_item("status", attempt.and_then(|failure| failure.status))?;
dict.set_item(
"failure_kind",
attempt.map(|failure| <&'static str>::from(failure.kind)),
)?;
}
}
Ok(dict)
}
fn ask_items(py: Python<'_>, dict: &Bound<'_, PyDict>, ask: Ask) -> PyResult<()> {
match ask {
Ask::Attempt {
deployment_id,
model,
model_group,
target_order,
candidate_ids,
retry,
fallback_depth,
fallback_level,
timeout,
} => {
dict.set_item("deployment_id", deployment_id)?;
dict.set_item("model", model)?;
dict.set_item("model_group", model_group)?;
dict.set_item("target_order", target_order)?;
dict.set_item("candidate_ids", PyList::new(py, candidate_ids)?)?;
dict.set_item("retry", retry)?;
dict.set_item("fallback_depth", fallback_depth)?;
dict.set_item("fallback_level", fallback_level)?;
dict.set_item("timeout", timeout)
}
Ask::Sleep { seconds } => dict.set_item("seconds", seconds),
Ask::AllowFallback { target } => dict.set_item("target", target),
Ask::Select {
stage,
model_group,
candidates,
} => {
let (stage, strategy) = match stage {
Stage::Plugins => ("plugins", None),
Stage::Strategy(strategy) => ("strategy", Some(strategy)),
};
dict.set_item("stage", stage)?;
dict.set_item("strategy", strategy)?;
dict.set_item("model_group", model_group)?;
dict.set_item("candidates", PyList::new(py, candidates)?)
}
Ask::Cache {
op,
key,
keys,
value,
ttl,
} => {
dict.set_item("op", <&'static str>::from(op))?;
dict.set_item("key", key)?;
dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?;
dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?;
dict.set_item("ttl", ttl)
}
}
}
fn done_items(dict: &Bound<'_, PyDict>, finish: Finish) -> PyResult<()> {
match finish {
Finish::Call(response) => {
dict.set_item("deployment_id", response.deployment_id)?;
dict.set_item("model_group", response.model_group)?;
dict.set_item("attempted_retries", response.attempted_retries)?;
dict.set_item("max_retries", response.max_retries)?;
dict.set_item("attempted_fallbacks", response.attempted_fallbacks)?;
dict.set_item("total_tokens", response.output.total_tokens)
}
Finish::Selected {
deployment_id,
model_group,
} => {
dict.set_item("deployment_id", deployment_id)?;
dict.set_item("model_group", model_group)
}
}
}
#[cfg(test)]
mod tests {
use pyo3::{ffi::c_str, prelude::*, types::PyDict};
use rstest::rstest;
const CONFIG: &str = r#"{"model_list": [
{"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}
]}"#;
const PAIR: &str = r#"{"model_list": [
{"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}},
{"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}
]}"#;
#[rstest]
fn python_drives_calls_and_selections_through_the_handle() {
Python::initialize();
Python::attach(|py| {
let locals = PyDict::new(py);
locals.set_item("native", crate::native_module(py)).unwrap();
locals.set_item("config", CONFIG).unwrap();
locals.set_item("pair", PAIR).unwrap();
py.run(
c_str!(
r#"
import time
store = {}
asked = []
def answer_cache(step):
asked.append(step)
op, key = step["op"], step["key"]
if op == "get":
return store.get(key)
if op == "batch_get":
return [store.get(k) for k in step["keys"]]
if op == "set":
store[key] = step["value"]
return None
if op == "increment":
store[key] = store.get(key, 0) + step["value"]
return store[key]
if op == "delete":
store.pop(key, None)
return None
raise AssertionError(op)
def drive(call, outcome):
step = call.next(outcome)
while step["kind"] == "cache":
step = call.next(answer_cache(step))
return step
handle = native.RouterCoreHandle(config, {"num_retries": 1})
picked = drive(handle.select("primary", {}), None)
assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked
assert [step["op"] for step in asked] == ["batch_get"]
assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "events"], list(asked[0])
call = handle.start("acompletion", '{"model": "primary"}')
version = handle.registry_version()
assert handle.apply({"delete": ["gone"]}) == version + 1
steps = [drive(call, None)]
steps.append(drive(call, {"ok": False, "kind": "RateLimitError", "status": 429}))
steps.append(drive(call, None))
steps.append(drive(call, {"ok": True, "total_tokens": 5}))
kinds = [step["kind"] for step in steps]
cooled = native.RouterCoreHandle(pair, {"num_retries": 0})
call = cooled.start("acompletion", {"model": "primary"})
tried = drive(call, None)["deployment_id"]
assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed"
assert f"deployment:{tried}:cooldown" not in store, store
store[f"deployment:{tried}:cooldown"] = {"exception_received": "x", "status_code": "429", "timestamp": time.time(), "cooldown_time": 60.0}
other = drive(cooled.select("primary", {}), None)
assert other["kind"] == "done" and other["deployment_id"] != tried, other
call = handle.select("primary", {})
step = call.next(None)
assert step["op"] == "batch_get"
for bad in ([], [None, None], "nope", None):
try:
call.next(bad)
except ValueError:
pass
else:
raise AssertionError(f"{bad!r} should not answer a batch_get of one key")
assert drive(call, [None])["kind"] == "done"
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
let kinds: Vec<String> = 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::<crate::errors::RustBridgeDeclined>(py),
declined,
"{error}"
);
assert_eq!(
error.is_instance_of::<pyo3::exceptions::PyValueError>(py),
!declined,
"{error}"
);
});
}
}

View file

@ -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

View file

@ -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<Option<Value>, Error>;
async fn async_set_cache(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> 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<Duration>,
) -> Result<f64, Error>;
async fn async_batch_get_cache(&self, keys: &[String]) -> Result<Vec<Option<Value>>, Error>;
async fn async_delete_cache(&self, key: &str) -> Result<(), Error>;
}
#[derive(Default)]
struct Entry {
value: Value,
expires: Option<Instant>,
}
/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`.
#[derive(Clone, Default)]
pub struct InMemoryRouterCache {
entries: Arc<Mutex<HashMap<String, Entry>>>,
}
impl InMemoryRouterCache {
fn live(entries: &mut HashMap<String, Entry>, key: &str) -> Option<Value> {
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<Value> {
Self::live(&mut self.entries.lock(), key)
}
pub fn increment(&self, key: &str, amount: f64, ttl: Option<Duration>) -> 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<Option<Value>, Error> {
Ok(self.get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> 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<Duration>,
) -> Result<f64, Error> {
Ok(self.increment(key, amount, ttl))
}
async fn async_batch_get_cache(&self, keys: &[String]) -> Result<Vec<Option<Value>>, 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(())
}
}

View file

@ -0,0 +1,592 @@
//! The call loop: select, run the attempt, then retry in the same group or fall back to another
//! until it succeeds or the policy is exhausted. Mirrors `Router.async_function_with_fallbacks`,
//! `async_function_with_fallbacks_common_utils`, `async_function_with_retries`,
//! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and
//! `run_async_fallback` in `router_utils/fallback_event_handlers.py`.
use std::{collections::HashSet, future::Future, pin::Pin, sync::Arc, time::Duration};
use tracing::warn;
use crate::{
Core, Snapshot,
config::Op,
error::{AttemptError, ConfigError, Error, FailureKind},
executor::Executor,
health::{GroupShape, should_retry_status},
host::{Attempt, Output, RouterEvent, RouterRequest},
registry::ResolveKind,
select::{Candidates, SelectInputs},
};
/// `INITIAL_RETRY_DELAY`, `MAX_RETRY_DELAY` and `JITTER` in `litellm/constants.py`.
const INITIAL_RETRY_DELAY: f64 = 0.5;
const MAX_RETRY_DELAY: f64 = 8.0;
const JITTER: f64 = 0.75;
/// `calculate_retry_after` obeys a provider's retry-after up to a minute.
const MAX_HONOURED_RETRY_AFTER: Duration = Duration::from_secs(60);
#[derive(Clone, Debug)]
pub struct RouterResponse {
pub output: Output,
pub deployment_id: String,
pub model_group: String,
pub attempted_retries: u32,
/// The retry budget of the loop a retry succeeded in (`x-litellm-max-retries`); `None`
/// when the first attempt succeeded.
pub max_retries: Option<u32>,
pub attempted_fallbacks: u32,
}
/// A failed attempt with the deployment it ran on, when one was picked.
struct Failed {
error: Error,
deployment_id: Option<String>,
}
impl From<Error> 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<i64>,
}
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + 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<RouterRequest>,
) -> Result<RouterResponse, Error> {
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<RouterRequest>,
target: Target,
at: Hop,
attempted: &'a mut HashSet<Target>,
) -> BoxFuture<'a, Result<RouterResponse, Error>> {
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<Vec<Target>> {
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<String>| {
names
.into_iter()
.map(|group| Target { group, order: None })
.collect::<Vec<_>>()
};
let skip_order = matches!(
kind,
FailureKind::ContextWindowExceeded | FailureKind::ContentPolicyViolation
);
if !skip_order {
let mut orders: Vec<i64> = 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<Target> = 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<RouterRequest>,
target: &Target,
at: Hop,
) -> Result<RouterResponse, Error> {
let mut inputs = SelectInputs {
target_order: target.order,
..SelectInputs::default()
};
let counting = SelectInputs {
count_only: true,
..SelectInputs::default()
};
let first = self
.attempt(snapshot, op, request, target, &inputs, 0, at)
.await;
let mut failed = match first {
Ok(response) => return Ok(response),
Err(failed) => failed,
};
let Error::Attempt(failure) = &failed.error else {
return Err(failed.error);
};
if failure.kind == FailureKind::GuardrailIntervention {
return Err(failed.error);
}
let settings = &snapshot.settings;
let mut num_retries = request
.num_retries
.or(failure.num_retries)
.unwrap_or(settings.num_retries);
let policy_retries = (request.num_retries != Some(0))
.then(|| settings.retry_policy.as_ref()?.retries_for(failure.kind))
.flatten();
let policy_applies = policy_retries.is_some();
num_retries = policy_retries.unwrap_or(num_retries);
// `_async_get_healthy_deployments` finds no group for a deployment id, so a call by id
// counts no deployments when deciding whether to retry.
let counted = |c: &Candidates, n: usize| if c.kind == ResolveKind::Id { 0 } else { n };
let all = self
.candidates(snapshot, op, request, &target.group, &counting)
.await
.map(|c| counted(&c, c.all.len()))
.unwrap_or(0);
for retry in 0..=num_retries {
let Error::Attempt(failure) = &failed.error else {
return Err(failed.error);
};
let healthy = self
.candidates(snapshot, op, request, &target.group, &counting)
.await
.map(|c| counted(&c, c.healthy.len()))
.unwrap_or(0);
if !policy_applies && !self.should_retry(snapshot, failure, healthy, all) {
return Err(failed.error);
}
if num_retries == 0 {
return Err(failed.error);
}
if let (Some(status), Some(id)) = (failure.status, &failed.deployment_id)
&& !should_retry_status(status)
&& !inputs.skipped_ids.contains(id)
{
inputs.skipped_ids.push(id.clone());
}
let delay = self.time_to_sleep(
snapshot,
failure,
num_retries - retry,
num_retries,
healthy,
all,
);
self.host.sleep(delay).await;
if retry == num_retries {
self.host.observe(&RouterEvent::RetriesExhausted {
model_group: target.group.clone(),
attempted: num_retries,
max: num_retries,
});
break;
}
self.host.observe(&RouterEvent::Retry {
model_group: target.group.clone(),
retry: retry + 1,
delay,
});
match self
.attempt(snapshot, op, request, target, &inputs, retry + 1, at)
.await
{
Ok(mut response) => {
response.max_retries = Some(num_retries);
return Ok(response);
}
Err(next) => failed = next,
}
}
Err(failed.error)
}
fn should_retry(
&self,
snapshot: &Snapshot,
failure: &AttemptError,
healthy: usize,
all: usize,
) -> bool {
let dynamic = &snapshot.settings.dynamic;
should_retry(
failure,
healthy,
all,
dynamic.context_window_fallbacks.is_some(),
dynamic.content_policy_fallbacks.is_some(),
dynamic.has_regular_fallbacks(),
)
}
/// `_time_to_sleep_before_retry` and `calculate_retry_after`.
fn time_to_sleep(
&self,
snapshot: &Snapshot,
failure: &AttemptError,
remaining_retries: u32,
num_retries: u32,
healthy: usize,
all: usize,
) -> Duration {
if all != 1 && healthy > 0 {
return Duration::ZERO;
}
let jitter = JITTER * self.rng.next_f64();
if let Some(retry_after) = failure.retry_after
&& !retry_after.is_zero()
&& retry_after <= MAX_HONOURED_RETRY_AFTER
{
return retry_after + Duration::from_secs_f64(jitter);
}
let exponent = num_retries.saturating_sub(remaining_retries);
let backoff = (INITIAL_RETRY_DELAY * 2f64.powi(exponent as i32))
.max(snapshot.settings.retry_after.as_secs_f64())
.min(MAX_RETRY_DELAY);
Duration::from_secs_f64(backoff + jitter)
}
/// One attempt: select a deployment, run it on its executor, record the outcome.
#[allow(clippy::too_many_arguments)]
async fn attempt(
&self,
snapshot: &Snapshot,
op: &Op,
request: &Arc<RouterRequest>,
target: &Target,
inputs: &SelectInputs,
retry: u32,
at: Hop,
) -> Result<RouterResponse, Failed> {
let candidates = self
.candidates(snapshot, op, request, &target.group, inputs)
.await?;
if candidates.healthy.is_empty() {
return Err(Error::NoHealthyDeployments {
model: candidates.model_group,
}
.into());
}
let deployment = self.pick(snapshot, op, request, &candidates).await?;
self.host.observe(&RouterEvent::Selected {
model_group: candidates.model_group.clone(),
deployment_id: deployment.id.clone(),
});
let executor = self
.executors
.resolve(op, deployment.provider())
.map_err(|unsupported| Error::Config(ConfigError::Unsupported(unsupported)))?;
let attempt = Attempt {
op: op.clone(),
request: request.clone(),
deployment: deployment.clone(),
name: target.group.clone(),
model_group: candidates.model_group.clone(),
target_order: target.order,
candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(),
retry,
fallback_depth: at.depth,
fallback_level: at.level,
timeout: snapshot.settings.timeout,
};
let (result, recorded) = match executor {
Executor::Native(route) => (route.execute(&attempt).await, false),
Executor::Host => (
self.host.invoke(&attempt).await,
self.host.records_outcomes(),
),
};
match result {
Ok(output) => {
if !recorded
&& let Err(error) = self
.health
.record_success(&deployment, output.total_tokens)
.await
{
warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded");
}
self.host.observe(&RouterEvent::Succeeded {
deployment_id: deployment.id.clone(),
});
Ok(RouterResponse {
output,
deployment_id: deployment.id.clone(),
model_group: candidates.model_group,
attempted_retries: retry,
max_retries: None,
attempted_fallbacks: at.depth,
})
}
Err(failure) => {
if recorded {
self.host.observe(&RouterEvent::AttemptFailed {
deployment_id: deployment.id.clone(),
kind: failure.kind,
status: failure.status,
});
} else {
self.record_failure(snapshot, &candidates, &deployment, &failure)
.await;
}
Err(Failed {
error: Error::Attempt(failure),
deployment_id: Some(deployment.id.clone()),
})
}
}
}
async fn record_failure(
&self,
snapshot: &Snapshot,
candidates: &Candidates,
deployment: &crate::registry::Deployment,
failure: &AttemptError,
) {
self.host.observe(&RouterEvent::AttemptFailed {
deployment_id: deployment.id.clone(),
kind: failure.kind,
status: failure.status,
});
let shape = GroupShape {
single_deployment: candidates.all.len() == 1,
has_fallback_group: snapshot
.settings
.dynamic
.fallback_targets(&[candidates.model_group.as_str()])
.is_some(),
};
match self
.health
.record_failure(&snapshot.settings, deployment, shape, failure)
.await
{
Ok(Some(cooldown)) => self.host.observe(&RouterEvent::CooldownSet {
deployment_id: deployment.id.clone(),
cooldown,
}),
Ok(None) => {}
Err(error) => {
warn!(deployment_id = %deployment.id, error = %error, "cooldown was not recorded");
}
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::rate_limit_with_healthy(FailureKind::RateLimit, Some(429), 1, 2, false, true)]
#[case::rate_limit_none_healthy_with_fallbacks(
FailureKind::RateLimit,
Some(429),
0,
2,
true,
false
)]
#[case::bad_request(FailureKind::BadRequest, Some(400), 2, 2, false, false)]
#[case::auth_lone_deployment(FailureKind::Authentication, Some(401), 1, 1, false, false)]
#[case::auth_with_siblings(FailureKind::Authentication, Some(401), 1, 2, false, true)]
#[case::not_found(FailureKind::NotFound, Some(404), 2, 2, false, false)]
#[case::server_error_no_healthy(FailureKind::InternalServer, Some(500), 0, 2, false, false)]
fn should_retry_follows_python(
#[case] kind: FailureKind,
#[case] status: Option<u16>,
#[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
);
}
}

View file

@ -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<u32>,
pub AuthenticationErrorRetries: Option<u32>,
pub TimeoutErrorRetries: Option<u32>,
pub RateLimitErrorRetries: Option<u32>,
pub ContentPolicyViolationErrorRetries: Option<u32>,
pub InternalServerErrorRetries: Option<u32>,
}
impl RetryPolicy {
pub fn retries_for(&self, kind: FailureKind) -> Option<u32> {
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<u32>,
pub AuthenticationErrorAllowedFails: Option<u32>,
pub TimeoutErrorAllowedFails: Option<u32>,
pub RateLimitErrorAllowedFails: Option<u32>,
pub ContentPolicyViolationErrorAllowedFails: Option<u32>,
pub InternalServerErrorAllowedFails: Option<u32>,
pub ServiceUnavailableErrorAllowedFails: Option<u32>,
pub BadGatewayErrorAllowedFails: Option<u32>,
}
impl AllowedFailsPolicy {
pub fn allowed_fails_for(&self, kind: FailureKind) -> Option<u32> {
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<BTreeMap<String, Vec<String>>>")]
pub struct Fallbacks(Vec<(String, Vec<String>)>);
impl From<Vec<BTreeMap<String, Vec<String>>>> for Fallbacks {
fn from(entries: Vec<BTreeMap<String, Vec<String>>>) -> Self {
Self(entries.into_iter().flatten().collect())
}
}
impl Fallbacks {
pub fn new(entries: impl IntoIterator<Item = (String, Vec<String>)>) -> 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<Vec<String>> {
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<String>,
pub provider_prefixed: BTreeMap<String, String>,
}
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<DeploymentSpec>,
pub routing_strategy: RoutingStrategy,
pub routing_groups: Vec<RoutingGroupSpec>,
pub plugins: Vec<PluginSpec>,
/// Router default `cooldown_time`; `None` is `DEFAULT_COOLDOWN_TIME_SECONDS`.
pub cooldown_time: Option<f64>,
pub disable_cooldowns: bool,
pub allowed_fails_policy: Option<AllowedFailsPolicy>,
pub retry_policy: Option<RetryPolicy>,
/// `retry_after`: the minimum backoff between retries, in seconds.
pub retry_after: f64,
pub timeout: Option<f64>,
pub request_timeout: Option<f64>,
pub drop_params: Option<bool>,
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<String>,
/// The operations this router must serve; with a host that cannot invoke, each must have a
/// native executor for every configured provider or the router fails to build.
pub served_ops: Vec<Op>,
}
/// 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": <group>, "hidden": <bool>}`.
#[serde(alias = "model_group_alias", deserialize_with = "alias_targets")]
pub model_alias_map: BTreeMap<String, String>,
pub fallbacks: Fallbacks,
pub default_fallbacks: Vec<String>,
pub context_window_fallbacks: Option<Fallbacks>,
pub content_policy_fallbacks: Option<Fallbacks>,
pub num_retries: Option<u32>,
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<u32>,
/// Other spellings a fallback key matches.
pub fallback_keys: FallbackKeys,
}
/// One `model_group_alias` value: a group name, or `{"model": <group>, "hidden": <bool>}`.
#[derive(Deserialize)]
#[serde(untagged)]
enum AliasTarget {
Group(String),
Entry { model: String },
}
fn alias_targets<'de, D>(deserializer: D) -> Result<BTreeMap<String, String>, D::Error>
where
D: serde::Deserializer<'de>,
{
let aliases = BTreeMap::<String, AliasTarget>::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<Vec<String>> {
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<Vec<String>>,
) {
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::<RoutingStrategy>().unwrap(), expected);
assert_eq!(expected.to_string(), wire);
}
}

View file

@ -0,0 +1,190 @@
//! The router's errors: [`ConfigError`] when a router is built or changed, [`Error`] when a call
//! or a selection fails, and [`AttemptError`], the typed facts of one failed attempt.
use std::{any::Any, fmt, sync::Arc, time::Duration};
use crate::config::{Op, RoutingStrategy};
/// What a host cannot serve, found when the router is built rather than partway through a call.
#[derive(Clone, Debug, PartialEq, thiserror::Error)]
pub enum Unsupported {
#[error("routing strategy `{0}` needs a host that can select")]
Strategy(RoutingStrategy),
#[error("routing plugin `{0}` needs a host that can select")]
Plugin(String),
#[error("operation `{op}` on provider `{provider}` has no executor")]
Executor { op: Op, provider: String },
}
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
#[error("the router config asks for something this host cannot serve")]
Unsupported(#[from] Unsupported),
#[error("deployment `{model_name}` has invalid litellm_params")]
InvalidDeployment {
model_name: String,
#[source]
source: serde_json::Error,
},
#[error("deployment id `{id}` is configured more than once")]
DuplicateId { id: String },
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("the router config is invalid")]
Config(#[from] ConfigError),
#[error("no deployment serves model `{model}`")]
UnknownModel { model: String },
#[error("no deployment of `{model}` is available")]
NoHealthyDeployments { model: String },
#[error("the attempt failed")]
Attempt(#[from] AttemptError),
#[error("the routing cache failed")]
Cache(#[from] litellm_cache::Error),
#[error("the router was retired")]
Retired,
}
impl Error {
/// The failed attempt behind this error, when a deployment was called.
pub fn attempt(&self) -> Option<&AttemptError> {
match self {
Self::Attempt(attempt) => Some(attempt),
Self::Config(_)
| Self::UnknownModel { .. }
| Self::NoHealthyDeployments { .. }
| Self::Cache(_)
| Self::Retired => None,
}
}
}
/// The class of a failed attempt, as litellm's exception mapping names it.
#[derive(
Clone,
Copy,
Debug,
PartialEq,
Eq,
Hash,
strum::Display,
strum::EnumString,
strum::IntoStaticStr,
strum::VariantArray,
)]
pub enum FailureKind {
#[strum(serialize = "RateLimitError")]
RateLimit,
#[strum(serialize = "ContextWindowExceededError")]
ContextWindowExceeded,
#[strum(serialize = "ContentPolicyViolationError")]
ContentPolicyViolation,
#[strum(serialize = "AuthenticationError")]
Authentication,
#[strum(serialize = "PermissionDeniedError")]
PermissionDenied,
#[strum(serialize = "NotFoundError")]
NotFound,
#[strum(serialize = "BadRequestError")]
BadRequest,
#[strum(serialize = "Timeout")]
Timeout,
#[strum(serialize = "APIConnectionError")]
ApiConnection,
#[strum(serialize = "InternalServerError")]
InternalServer,
#[strum(serialize = "ServiceUnavailableError")]
ServiceUnavailable,
#[strum(serialize = "BadGatewayError")]
BadGateway,
/// A guardrail blocked the request: never retried, never a fallback.
#[strum(serialize = "GuardrailIntervention")]
GuardrailIntervention,
/// litellm rejected the request before calling the deployment (`MalformedRequestFieldError`,
/// `UnsupportedEndpointError`): never cools the deployment down.
#[strum(serialize = "Declined")]
Declined,
/// The caller's own credential failed, not the deployment's.
#[strum(serialize = "CallerCredentialError")]
CallerCredential,
#[strum(serialize = "APIError")]
Other,
}
/// An opaque value the host hands through the router unchanged, such as the original Python
/// exception or response object.
#[derive(Clone)]
pub struct Payload(pub Arc<dyn Any + Send + Sync>);
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<u16>,
/// The provider's `retry-after` (or `retry-after-ms`) header.
pub retry_after: Option<Duration>,
/// A deployment-level `num_retries` the failure carries (Python's `exception.num_retries`).
pub num_retries: Option<u32>,
/// The provider's error text, kept only for the cooldown record Python writes
/// (`exception_received`, masked past 50 characters).
pub detail: Option<Arc<str>>,
pub payload: Option<Payload>,
}
impl AttemptError {
pub fn new(kind: FailureKind, status: Option<u16>) -> 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::<FailureKind>().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::<FailureKind>().is_err());
}
}

View file

@ -0,0 +1,75 @@
//! Executors: each attempt resolves one after its deployment is picked. An (op, provider) pair in
//! the route-rule table with a registered [`NativeRoute`] runs in-process; anything else goes
//! through `Host::invoke`. An attempt runs on exactly one executor and is never replayed on
//! another: a native failure is that attempt's failure, and a retry selects afresh.
use std::sync::Arc;
use async_trait::async_trait;
use crate::{
config::Op,
error::{AttemptError, Unsupported},
host::{Attempt, Output},
};
/// The (op, provider) pairs that have a Rust route machine.
pub const NATIVE_RULES: &[(Op, &str)] = &[(Op::Messages, "anthropic")];
/// A Rust-native route machine (an `inference-*` crate) adapted to an attempt. The embedder
/// registers one per native rule; for Messages on Anthropic that is an adapter over
/// `litellm_inference_messages::MessagesRoute::execute`.
#[async_trait]
pub trait NativeRoute: Send + Sync {
async fn execute(&self, attempt: &Attempt) -> Result<Output, AttemptError>;
}
#[derive(Clone)]
pub enum Executor {
Native(Arc<dyn NativeRoute>),
Host,
}
#[derive(Clone, Default)]
pub struct Executors {
native: Vec<(Op, String, Arc<dyn NativeRoute>)>,
host_invokes: bool,
}
impl Executors {
pub fn new(host_invokes: bool) -> Self {
Self {
native: Vec::new(),
host_invokes,
}
}
/// Register a native route for a rule in [`NATIVE_RULES`]. A pair outside the table stays
/// on the host, so a route cannot silently take over an operation not yet ported.
pub fn register(&mut self, op: Op, provider: &str, route: Arc<dyn NativeRoute>) -> bool {
let ruled = NATIVE_RULES
.iter()
.any(|(rule_op, rule_provider)| *rule_op == op && *rule_provider == provider);
if ruled {
self.native.push((op, provider.to_owned(), route));
}
ruled
}
pub fn resolve(&self, op: &Op, provider: &str) -> Result<Executor, Unsupported> {
if let Some((_, _, route)) = self
.native
.iter()
.find(|(rule_op, rule_provider, _)| rule_op == op && rule_provider == provider)
{
return Ok(Executor::Native(route.clone()));
}
if self.host_invokes {
return Ok(Executor::Host);
}
Err(Unsupported::Executor {
op: op.clone(),
provider: provider.to_owned(),
})
}
}

View file

@ -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<S: serde::Serializer>(value: &f64, serializer: S) -> Result<S::Ok, S::Error> {
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<dyn RouterCache>,
local: InMemoryRouterCache,
clock: Arc<dyn Clock>,
}
impl Health {
pub fn new(cache: Arc<dyn RouterCache>, clock: Arc<dyn Clock>) -> Self {
Self {
cache,
local: InMemoryRouterCache::default(),
clock,
}
}
pub fn cache(&self) -> &Arc<dyn RouterCache> {
&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<Vec<(String, CooldownValue)>, Error> {
let keys: Vec<String> = 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<Duration> {
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<Option<Duration>, 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<bool, Error> {
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<bool, Error> {
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<u16>,
#[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(),
)
);
}
}

View file

@ -0,0 +1,168 @@
//! [`Host`]: what the embedding process provides. `PythonHost` serves it from the existing Python
//! code; `NativeHost` serves it in the Rust server.
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use async_trait::async_trait;
use serde_json::Value;
use crate::{
cache::RouterCache,
config::{Op, RoutingStrategy},
error::{AttemptError, Error, FailureKind, Payload},
health::Health,
registry::{Deployment, Registry},
settings::Settings,
};
/// One call into the router.
#[derive(Clone, Debug, Default)]
pub struct RouterRequest {
/// The requested name: a model group, alias, routing group, deployment id or wildcard match.
pub model: String,
/// The request params, for native executors and plugins that read them.
pub body: Value,
/// `user_api_key_team_id`.
pub team_id: Option<String>,
pub specific_deployment: bool,
/// A per-request `num_retries`, which wins over the router's.
pub num_retries: Option<u32>,
pub disable_fallbacks: bool,
/// `_router_weights`: model group to deployment id to weight.
pub router_weights: Option<BTreeMap<String, BTreeMap<String, f64>>>,
/// What the host needs to run the call itself, such as the Python kwargs.
pub payload: Option<Payload>,
}
/// One attempt against one deployment.
#[derive(Clone, Debug)]
pub struct Attempt {
pub op: Op,
pub request: Arc<RouterRequest>,
pub deployment: Arc<Deployment>,
/// 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<i64>,
/// Every deployment `name` resolved to for this caller, the picked one included.
pub candidate_ids: Vec<String>,
/// 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<Duration>,
}
/// 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<u16>,
},
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<i64>,
depth: u32,
/// The fallback level `to` is tried at: 1 for the request's own fallbacks.
level: u32,
},
/// A group's retry loop ran every retry and the last one failed: Python stamps
/// `num_retries` and `max_retries` on that error.
RetriesExhausted {
model_group: String,
attempted: u32,
max: u32,
},
}
#[async_trait]
pub trait Host: Send + Sync {
/// Run a non-native attempt.
async fn invoke(&self, attempt: &Attempt) -> Result<Output, AttemptError>;
/// Run a selection stage the core does not implement.
async fn select(
&self,
stage: &SelectStage,
ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, 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<dyn RouterCache>;
/// Whether [`Self::invoke`] can run an attempt. A `NativeHost` cannot, so every served
/// operation needs a native executor.
fn can_invoke(&self) -> bool {
true
}
/// Whether [`Self::select`] can run strategies and plugins.
fn can_select(&self) -> bool {
true
}
/// Whether [`Self::invoke`] records each attempt's outcome itself (cooldowns, allowed-fails
/// counters, usage), as the Python router's own callbacks do. The core then leaves the
/// health store to the host for host-run attempts, so no failure is counted twice.
fn records_outcomes(&self) -> bool {
false
}
}

View file

@ -61,3 +61,285 @@
//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive)
//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and
//! the current request), never the router itself.
//!
//! # Modules
//! - [`registry`]: deployments and their indices, `deployment_id`, `ChangeSet`.
//! - [`config`] / [`settings`]: static and dynamic config, and the frozen `Settings` view.
//! - [`cache`] / [`health`]: the shared routing cache, cooldowns and usage.
//! - [`select`]: resolve, filter chain, simple-shuffle.
//! - [`call`]: retries and fallbacks.
//! - [`host`] / [`executor`]: the `Host` seam and per-attempt executors.
pub mod cache;
pub mod call;
pub mod config;
pub mod error;
pub mod executor;
pub mod health;
pub mod host;
pub mod registry;
pub mod select;
pub mod settings;
use std::sync::{
Arc,
atomic::{AtomicBool, Ordering},
};
use arc_swap::ArcSwap;
use parking_lot::Mutex;
pub use crate::{
cache::{InMemoryRouterCache, RouterCache},
call::RouterResponse,
config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy},
error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported},
executor::{Executor, Executors, NativeRoute},
health::{Clock, Health, Outcome, SystemClock},
host::{Attempt, Host, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage},
registry::{ChangeSet, Deployment, DeploymentSpec, Registry, RegistryVersion},
select::{Randomness, SeededRandomness, ThreadRandomness},
settings::Settings,
};
use crate::{health::GroupShape, registry::DeploymentFilter, select::SelectInputs};
/// The registry and settings one call reads, swapped atomically by `apply`.
pub(crate) struct Snapshot {
pub registry: Arc<Registry>,
pub settings: Arc<Settings>,
}
pub(crate) struct Core {
host: Arc<dyn Host>,
executors: Executors,
rng: Arc<dyn Randomness>,
health: Health,
config: RouterConfig,
snapshot: ArcSwap<Snapshot>,
writer: Mutex<()>,
retired: AtomicBool,
}
impl Core {
/// Reject what this host cannot serve, before any request.
fn validate(&self, registry: &Registry, settings: &Settings) -> Result<(), ConfigError> {
let can_select = self.host.can_select();
let strategies = std::iter::once(&settings.routing_strategy).chain(
registry
.routing_groups()
.iter()
.map(|group| &group.routing_strategy),
);
for strategy in strategies {
if !strategy.is_native() && !can_select {
return Err(Unsupported::Strategy(strategy.clone()).into());
}
}
if let Some(plugin) = settings.plugins.first()
&& !can_select
{
return Err(Unsupported::Plugin(plugin.name.clone()).into());
}
for deployment in registry.deployments(&DeploymentFilter::default()) {
for op in &self.config.served_ops {
self.executors.resolve(op, deployment.provider())?;
}
}
Ok(())
}
}
/// Builds a [`Router`], with native routes and test seams.
pub struct RouterBuilder {
config: RouterConfig,
dynamic: RouterDynamicConfig,
host: Arc<dyn Host>,
executors: Executors,
rng: Arc<dyn Randomness>,
clock: Arc<dyn Clock>,
}
impl RouterBuilder {
/// Register a native route for an (op, provider) rule; a pair outside
/// [`executor::NATIVE_RULES`] is ignored and stays on the host.
#[must_use]
pub fn native_route(mut self, op: Op, provider: &str, route: Arc<dyn NativeRoute>) -> Self {
self.executors.register(op, provider, route);
self
}
#[must_use]
pub fn randomness(self, rng: Arc<dyn Randomness>) -> Self {
Self { rng, ..self }
}
#[must_use]
pub fn clock(self, clock: Arc<dyn Clock>) -> Self {
Self { clock, ..self }
}
pub fn build(self) -> Result<Router, ConfigError> {
let registry = Registry::build(
self.config.model_list.clone(),
self.dynamic.model_alias_map.clone(),
self.config.routing_groups.clone(),
)?;
let settings = Settings::new(&self.config, Arc::new(self.dynamic));
let core = Core {
health: Health::new(self.host.cache(), self.clock),
host: self.host,
executors: self.executors,
rng: self.rng,
config: self.config,
snapshot: ArcSwap::from_pointee(Snapshot {
registry: Arc::new(Registry::default()),
settings: Arc::new(settings.clone()),
}),
writer: Mutex::new(()),
retired: AtomicBool::new(false),
};
core.validate(&registry, &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<Core>,
}
impl Router {
pub fn new(
config: RouterConfig,
dynamic: RouterDynamicConfig,
host: Arc<dyn Host>,
) -> Result<Self, ConfigError> {
Self::builder(config, dynamic, host).build()
}
pub fn builder(
config: RouterConfig,
dynamic: RouterDynamicConfig,
host: Arc<dyn Host>,
) -> RouterBuilder {
RouterBuilder {
executors: Executors::new(host.can_invoke()),
config,
dynamic,
host,
rng: Arc::new(ThreadRandomness),
clock: Arc::new(SystemClock),
}
}
fn live(&self) -> Result<(), Error> {
if self.core.retired.load(Ordering::Acquire) {
return Err(Error::Retired);
}
Ok(())
}
/// Select, run the attempt, and retry or fall back until it succeeds or the policy is
/// exhausted.
pub async fn call(&self, op: Op, request: RouterRequest) -> Result<RouterResponse, Error> {
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<Arc<Deployment>, Error> {
self.live()?;
let snapshot = self.core.snapshot.load_full();
let candidates = self
.core
.candidates(
&snapshot,
op,
request,
&request.model,
&SelectInputs::default(),
)
.await?;
self.core.pick(&snapshot, op, request, &candidates).await
}
pub fn registry(&self) -> Arc<Registry> {
self.core.snapshot.load().registry.clone()
}
pub fn settings(&self) -> Arc<Settings> {
self.core.snapshot.load().settings.clone()
}
pub fn health(&self) -> &Health {
&self.core.health
}
/// Record an outcome the router did not see itself, such as a health check.
pub async fn report(&self, id: &str, outcome: Outcome) -> Result<(), Error> {
let snapshot = self.core.snapshot.load_full();
let Some(deployment) = snapshot.registry.deployment(id) else {
return Ok(());
};
match outcome {
Outcome::Success { total_tokens } => {
self.core
.health
.record_success(&deployment, total_tokens)
.await
}
Outcome::Failure(failure) => {
let shape = GroupShape {
single_deployment: snapshot
.registry
.group(&deployment.model_name)
.is_some_and(|group| group.len() == 1),
has_fallback_group: snapshot
.settings
.dynamic
.fallback_targets(&[deployment.model_name.as_str()])
.is_some(),
};
self.core
.health
.record_failure(&snapshot.settings, &deployment, shape, &failure)
.await
.map(drop)
}
}
}
/// Apply one change set atomically and return the new registry version.
pub fn apply(&self, changes: ChangeSet) -> Result<RegistryVersion, ConfigError> {
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(&current.settings),
};
self.core.validate(&registry, &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);
}
}

View file

@ -0,0 +1,820 @@
//! The configured deployments, indexed by id, model group, alias, wildcard pattern, routing
//! group and team. A [`Registry`] is immutable; [`Registry::apply`] builds the next version.
use std::{
collections::{BTreeMap, BTreeSet, HashMap},
fmt::Write as _,
sync::Arc,
};
use litellm_router_types::{LitellmParams, Spelled};
use serde::Deserialize;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::error::ConfigError;
/// A registry's version; every change produces a higher one.
pub type RegistryVersion = u64;
/// A deployment as a config or an `apply` spells it.
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
pub struct DeploymentSpec {
pub model_name: String,
pub litellm_params: Map<String, Value>,
#[serde(default)]
pub model_info: Map<String, Value>,
}
/// `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<String>,
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<String>,
pub team_public_model_name: Option<String>,
pub access_groups: Vec<String>,
pub blocked: bool,
pub cooldown_time: Option<f64>,
pub allowed_fails: Option<u32>,
pub raw: Map<String, Value>,
}
impl ModelInfo {
fn from_raw(raw: Map<String, Value>) -> 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<String, Value>,
pub model_info: ModelInfo,
}
impl Deployment {
pub fn from_spec(spec: DeploymentSpec) -> Result<Self, ConfigError> {
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<f64> {
let spelled = |value: &Option<Spelled<f64>>| 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<i64> {
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<String>,
pub api_version: Option<String>,
pub credential_name: Option<String>,
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<Arc<Deployment>>,
}
#[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<DeploymentSpec>,
pub deletes: Vec<String>,
pub settings: Option<crate::config::RouterDynamicConfig>,
}
#[derive(Clone, Debug, Default)]
pub struct Registry {
version: RegistryVersion,
deployments: Vec<Arc<Deployment>>,
by_id: HashMap<String, usize>,
by_group: BTreeMap<String, Vec<usize>>,
by_team: HashMap<String, Vec<usize>>,
wildcards: Vec<(String, Vec<usize>)>,
aliases: BTreeMap<String, String>,
routing_groups: Vec<RoutingGroupSpec>,
}
impl Registry {
pub fn build(
specs: Vec<DeploymentSpec>,
aliases: BTreeMap<String, String>,
routing_groups: Vec<RoutingGroupSpec>,
) -> Result<Self, ConfigError> {
let deployments = specs
.into_iter()
.map(|spec| Deployment::from_spec(spec).map(Arc::new))
.collect::<Result<Vec<_>, _>>()?;
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<Arc<Deployment>>,
aliases: BTreeMap<String, String>,
routing_groups: Vec<RoutingGroupSpec>,
) -> Self {
let mut by_id = HashMap::new();
let mut by_group: BTreeMap<String, Vec<usize>> = BTreeMap::new();
let mut by_team: HashMap<String, Vec<usize>> = 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<usize>)> = Vec::new();
for deployment in &deployments {
let name = &deployment.model_name;
if name.contains('*') && !wildcards.iter().any(|(pattern, _)| pattern == name) {
wildcards.push((
name.clone(),
by_group.get(name).cloned().unwrap_or_default(),
));
}
}
wildcards.sort_by_key(|(pattern, _)| std::cmp::Reverse(pattern_specificity(pattern)));
Self {
version,
deployments,
by_id,
by_group,
by_team,
wildcards,
aliases,
routing_groups,
}
}
pub fn version(&self) -> RegistryVersion {
self.version
}
/// The next registry: deletes, then upserts. An upserted deployment whose id exists is
/// removed and appended at the end, as the Python router does on an edit.
pub fn apply(
&self,
changes: &ChangeSet,
aliases: Option<BTreeMap<String, String>>,
) -> Result<Self, ConfigError> {
let mut deployments = self.deployments.clone();
deployments.retain(|deployment| !changes.deletes.contains(&deployment.id));
for spec in &changes.upserts {
let deployment = Deployment::from_spec(spec.clone())?;
deployments.retain(|existing| existing.id != deployment.id);
deployments.push(Arc::new(deployment));
}
Ok(Self::index(
self.version + 1,
deployments,
aliases.unwrap_or_else(|| self.aliases.clone()),
self.routing_groups.clone(),
))
}
fn at(&self, indices: &[usize]) -> Vec<Arc<Deployment>> {
indices
.iter()
.map(|index| self.deployments[*index].clone())
.collect()
}
pub fn deployment(&self, id: &str) -> Option<Arc<Deployment>> {
self.by_id
.get(id)
.map(|index| self.deployments[*index].clone())
}
pub fn deployments(&self, filter: &DeploymentFilter<'_>) -> Vec<Arc<Deployment>> {
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<String> {
let mut names: Vec<String> = self.by_group.keys().cloned().collect();
if filter.include_aliases {
names.extend(self.aliases.keys().cloned());
}
names
}
pub fn group(&self, name: &str) -> Option<Vec<Arc<Deployment>>> {
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<crate::config::RoutingStrategy> {
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<Resolution> {
let resolution = |kind, group: &str, deployments: Vec<Arc<Deployment>>| {
(!deployments.is_empty()).then(|| Resolution {
kind,
group: group.to_owned(),
deployments,
})
};
if scope.specific_deployment {
let matching = self
.deployments
.iter()
.filter(|d| d.params.model == name)
.cloned()
.collect();
return resolution(ResolveKind::SpecificDeployment, name, matching);
}
if !self.by_group.contains_key(name)
&& let Some(deployment) = self.deployment(name)
{
let group = deployment.model_name.clone();
return resolution(ResolveKind::Id, &group, vec![deployment]);
}
if let Some(target) = self.aliases.get(name) {
return self
.group(target)
.and_then(|deployments| resolution(ResolveKind::Alias, target, deployments));
}
if let Some(routing_group) = self.routing_groups.iter().find(|g| g.group_name == name) {
let deployments = routing_group
.models
.iter()
.filter_map(|model| self.by_group.get(model))
.flat_map(|indices| self.at(indices))
.collect();
return resolution(ResolveKind::RoutingGroup, name, deployments);
}
if let Some(indices) = self.by_group.get(name) {
return resolution(ResolveKind::Group, name, self.at(indices));
}
self.wildcards
.iter()
.find(|(pattern, _)| wildcard_matches(pattern, name))
.and_then(|(pattern, indices)| {
resolution(ResolveKind::Wildcard, pattern, self.at(indices))
})
}
/// Access group name to the model groups in it.
pub fn access_groups(&self) -> BTreeMap<String, BTreeSet<String>> {
let mut groups: BTreeMap<String, BTreeSet<String>> = 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<Credentials> {
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, Value>) -> 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, &params), 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<String> = next
.group("a")
.unwrap()
.iter()
.map(|d| d.id.clone())
.collect();
assert_eq!(
(next.version(), ids),
(1, vec!["2".to_owned(), "1".to_owned()])
);
}
}

View file

@ -0,0 +1,320 @@
//! Selection: resolve a name, filter its deployments in the order of
//! `Router._core_async_get_healthy_deployments` (team, cooldown, blocked, plugins, order, retry
//! skips), then pick one with the strategy. Simple-shuffle is native
//! (`router_strategy/simple_shuffle.py`); other strategies go through `Host::select`.
use std::{collections::BTreeMap, sync::Arc};
use parking_lot::Mutex;
use rand::{Rng, SeedableRng, rngs::StdRng};
use crate::{
Core, Snapshot,
config::{Op, RoutingStrategy},
error::Error,
host::{RouterRequest, RoutingContext, SelectStage},
registry::{Deployment, ResolveKind, ResolveScope},
};
/// Randomness for the shuffle and retry jitter, injectable so tests are deterministic.
pub trait Randomness: Send + Sync {
/// A float in `[0, 1)`, as Python's `random.random()`.
fn next_f64(&self) -> f64;
}
pub struct ThreadRandomness;
impl Randomness for ThreadRandomness {
fn next_f64(&self) -> f64 {
rand::thread_rng().r#gen()
}
}
pub struct SeededRandomness(Mutex<StdRng>);
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<Arc<Deployment>>,
pub healthy: Vec<Arc<Deployment>>,
}
/// Per-attempt selection inputs the call loop carries.
#[derive(Clone, Debug, Default)]
pub(crate) struct SelectInputs {
pub target_order: Option<i64>,
pub skipped_ids: Vec<String>,
/// 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<Deployment>],
model_group: &str,
router_weights: Option<&BTreeMap<String, BTreeMap<String, f64>>>,
rng: &dyn Randomness,
) -> Option<Arc<Deployment>> {
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::<Vec<_>>()
});
let metric_sets = ["weight", "rpm", "tpm"].map(|metric| {
deployments
.iter()
.map(|d| d.metric(metric).unwrap_or(0.0))
.collect::<Vec<_>>()
});
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<f64> = 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<Candidates, Error> {
let registry = &snapshot.registry;
let scope = ResolveScope {
specific_deployment: request.specific_deployment,
};
let dynamic = &snapshot.settings.dynamic;
let resolution = registry
.resolve(name, scope)
.or_else(|| {
// No deployment for the name: route to the first default fallback.
let default = dynamic
.fallbacks
.default_target()
.or_else(|| dynamic.default_fallbacks.first().map(String::as_str))?;
registry.resolve(default, scope)
})
.ok_or_else(|| Error::UnknownModel {
model: name.to_owned(),
})?;
let team_id = request.team_id.as_deref();
let all: Vec<Arc<Deployment>> = resolution
.deployments
.into_iter()
.filter(|d| d.usable_by_team(team_id))
.collect();
let ids: Vec<String> = all.iter().map(|d| d.id.clone()).collect();
let cooling = self.health.active_cooldowns(&ids).await?;
let mut healthy: Vec<Arc<Deployment>> = 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<Arc<Deployment>> = 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<Arc<Deployment>, 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<Deployment> {
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");
}
}

View file

@ -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<u32>,
pub allowed_fails_policy: Option<AllowedFailsPolicy>,
pub retry_policy: Option<RetryPolicy>,
pub num_retries: u32,
pub max_fallbacks: u32,
pub retry_after: Duration,
pub timeout: Option<Duration>,
pub drop_params: Option<bool>,
pub enable_health_check_routing: bool,
pub provider_default_deployment_ids: Arc<[String]>,
pub dynamic: Arc<RouterDynamicConfig>,
}
fn seconds(value: f64) -> Duration {
Duration::try_from_secs_f64(value).unwrap_or_default()
}
impl Settings {
pub fn new(config: &RouterConfig, dynamic: Arc<RouterDynamicConfig>) -> 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<u32> {
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)
}
}

View file

@ -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<Vec<(String, VecDeque<AttemptError>)>>,
events: Mutex<Vec<RouterEvent>>,
cache: InMemoryRouterCache,
}
impl ScriptedHost {
fn failing(group: &str, failures: Vec<AttemptError>) -> 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<Output, AttemptError> {
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<Arc<Deployment>>,
) -> Result<Vec<Arc<Deployment>>, 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<dyn RouterCache> {
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<ScriptedHost>) -> 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)
);
}

View file

@ -7,8 +7,9 @@ the router is built, from its ``RouterConfig``.
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
from litellm._logging import verbose_router_logger
from litellm.router_config import RouterConfig
if TYPE_CHECKING:
@ -51,6 +52,28 @@ class RouterCore(Protocol):
def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore:
"""
The Rust core when the route catalog's router rule allows it and the Rust core serves this
config, else the Python core. Decided once, for the router's life.
"""
from litellm.router_python_core import PythonCore
from litellm.rust_bridge.catalog import Route, RouteContext, decision
from litellm.rust_bridge.configuration import Decision
return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature
python_class: Final = core_class or PythonCore
chosen: Final = decision(RouteContext(Route.ROUTER))
if chosen is not Decision.PYTHON:
from litellm.router_rust_core import RustCore, RustCoreDeclined
from litellm.rust_bridge._native import RustBridgeDeclined
try:
return RustCore.build(config, host, python_class) # pyright: ignore[reportReturnType] # RustCore serves the rest of RouterCore from its PythonCore, through __getattr__
except (RustBridgeDeclined, RustCoreDeclined) as declined:
# A config the Rust core cannot take keeps the router on Python; an error building the
# Python core itself is the router's own error and is raised once, as it is.
if chosen is Decision.RUST_REQUIRED:
raise
verbose_router_logger.info("Router served by the Python core: %s", declined)
if isinstance(declined, RustCoreDeclined):
return declined.python
return python_class(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature

857
litellm/router_rust_core.py Normal file
View file

@ -0,0 +1,857 @@
"""
``RustCore``: the router core served by litellm-rust/crates/router-core through ``_native``.
The Rust core owns the registry, selection, retries, fallbacks and cooldowns for the operations in
``RUST_SERVED_OPS``. It never awaits Python: ``RouterCall.next`` returns the step the call loop is
blocked on (run an attempt on a deployment, sleep, check a fallback target, run a plugin) and this
driver answers it. Each attempt runs through the Python core with that one deployment, so request
preparation, inference, logging and cost tracking are unchanged.
Everything else (the views, runtime writes, every other endpoint) is served by the ``PythonCore``
the Rust core wraps, built from the same config. A config the Rust core cannot serve fails at
construction and ``make_core`` builds a ``PythonCore`` instead.
"""
import inspect
import json
import math
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import TypeAdapter
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_guardrail import is_guardrail_intervention
from litellm.litellm_core_utils.add_retry_fallback_headers import (
add_fallback_headers_to_response,
add_retry_headers_to_response,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.routing_common_utils import (
format_content_policy_fallback_hint,
format_fallback_outcome_message,
format_no_fallback_group_message,
truncate_fallback_error_detail,
)
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
from litellm.router_config import RouterConfig
from litellm.router_python_core import PythonCore
from litellm.router_python_host import PythonHost
from litellm.router_utils.fallback_event_handlers import (
AttemptedFallbackTargets,
_provider_prefixed_model_group, # pyright: ignore[reportPrivateUsage] # the Python fallback walk's own spelling rule
fallback_lookup_groups,
get_fallback_model_group_for_lookup_groups,
is_bad_request_without_fallback,
)
from litellm.types.router import FallbackSkipReason
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.rust_bridge._native import RouterCall, RouterCoreHandle
RustServedOp: TypeAlias = Literal["acompletion", "aanthropic_messages", "aembedding", "aresponses"]
RUST_SERVED_OPS: Final[frozenset[str]] = frozenset({"acompletion", "aanthropic_messages", "aembedding", "aresponses"})
_FAILURE_KINDS: Final = (
"RateLimitError",
"ContextWindowExceededError",
"ContentPolicyViolationError",
"AuthenticationError",
"PermissionDeniedError",
"NotFoundError",
"BadRequestError",
"Timeout",
"APIConnectionError",
"InternalServerError",
"ServiceUnavailableError",
"BadGatewayError",
"GuardrailIntervention",
"CallerCredentialError",
)
_DECLINED: Final = frozenset({"MalformedRequestFieldError", "UnsupportedEndpointError"})
_STEP: Final = TypeAdapter(Mapping[str, object])
_REQUEST_FIELDS: Final = ("specific_deployment", "num_retries", "disable_fallbacks")
# Request keys router-core reads under another name.
_RENAMED_REQUEST_FIELDS: Final = MappingProxyType({"_router_weights": "router_weights"})
_PROVIDERS: Final = tuple(provider.value for provider in LlmProviders)
_PYTHON_ONLY_REQUEST_KEYS: Final = frozenset(
{
"fallbacks",
"context_window_fallbacks",
"content_policy_fallbacks",
"default_fallbacks",
"max_fallbacks",
"model_group_retry_policy",
"retry_policy",
}
)
_STATIC_FIELDS: Final = (
"routing_strategy",
"cooldown_time",
"disable_cooldowns",
"allowed_fails_policy",
"retry_after",
"timeout",
)
_HANDLED_FIELDS: Final = frozenset(
{
*_STATIC_FIELDS,
"model_list",
"model_group_alias",
"fallbacks",
"default_fallbacks",
"context_window_fallbacks",
"content_policy_fallbacks",
"num_retries",
"max_fallbacks",
"allowed_fails",
"set_verbose",
"debug_level",
}
)
_PYTHON_DEFAULTS: Final = MappingProxyType(
{
name: parameter.default
for name, parameter in inspect.signature(PythonCore).parameters.items()
if parameter.default is not inspect.Parameter.empty
}
)
def failure_kind(error: BaseException) -> str:
names: Final = tuple(cls.__name__ for cls in type(error).__mro__)
if not _DECLINED.isdisjoint(names):
return "Declined"
if is_guardrail_intervention(error): # pyright: ignore[reportArgumentType] # every attempt error is an Exception
return "GuardrailIntervention"
return next((name for name in names if name in _FAILURE_KINDS), "APIError")
def _retry_after_seconds(error: BaseException) -> float | None:
response: Final = getattr(error, "response", None)
headers: Final = getattr(error, "litellm_response_headers", None) or getattr(response, "headers", None)
if headers is None or not hasattr(headers, "get"):
return None
milliseconds: Final = headers.get("retry-after-ms")
seconds: Final = headers.get("retry-after")
try:
value: Final = float(milliseconds) / 1000 if milliseconds is not None else float(seconds)
except (TypeError, ValueError):
return None
return value if math.isfinite(value) and value >= 0 else None
def attempt_failure(error: BaseException) -> Mapping[str, object]:
status: Final = getattr(error, "status_code", None)
num_retries: Final = getattr(error, "num_retries", None)
return {
"ok": False,
"kind": failure_kind(error),
"status": status if isinstance(status, int) and 100 <= status <= 599 else None,
"retry_after": _retry_after_seconds(error),
"num_retries": num_retries if isinstance(num_retries, int) and num_retries >= 0 else None,
"detail": str(error)[:500],
}
def total_tokens(response: object) -> int:
usage: Final = getattr(response, "usage", None)
tokens: Final = getattr(usage, "total_tokens", None)
return tokens if isinstance(tokens, int) else 0
def _json_default(value: object) -> object:
return str(value)
def _plain(value: object) -> object:
"""``value`` as plain JSON values, the shape the Rust core reads."""
return json.loads(json.dumps(value, default=_json_default))
def _plain_document(value: Mapping[str, object]) -> Mapping[str, object]:
"""A mapping as plain JSON values: still a mapping, with string keys."""
return as_str_mapping(_plain(value)) or MappingProxyType({})
def _as_list(value: object) -> Sequence[object]:
return value if isinstance(value, list) else () # pyright: ignore[reportUnknownVariableType] # a list's items are read as objects
def _applied_static(python: PythonCore, name: str) -> object:
"""Where an applied static setting lands on the Python core: the value the Rust core is rebuilt with."""
match name:
case "routing_strategy":
return python.settings().routing_strategy
case "cooldown_time":
return python.settings().cooldown_time
case "disable_cooldowns":
return python.settings().disable_cooldowns
case "allowed_fails_policy":
return python.settings().allowed_fails_policy
case "retry_after":
return python.retry_after
case "timeout":
return python.timeout
case _:
raise KeyError(name)
def _static_value(value: object) -> object:
dump: Final = getattr(value, "model_dump", None)
return _plain(dump(exclude_none=True) if callable(dump) else value)
def static_settings(config: RouterConfig, python: PythonCore) -> Mapping[str, object]:
"""The static settings: the constructor arguments, overlaid with every setting a changeset
(``apply``) has changed since, as the Python core resolved it."""
applied: Final = python._core_applied # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
values: Final[dict[str, object]] = { # mutable-ok: the arguments, then the applied settings over them
name: config.arguments[name] for name in _STATIC_FIELDS if name in config.arguments
}
values.update({name: _applied_static(python, name) for name in _STATIC_FIELDS if name in applied})
return {name: _static_value(value) for name, value in values.items() if value is not None}
def static_document(config: RouterConfig, python: PythonCore) -> Mapping[str, object]:
"""The Python core's own normalized deployments (with their ids) and the static settings."""
return {
"model_list": _plain(python._core_model_list), # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownArgumentType] # RustCore drives its own PythonCore, whose list is untyped
**static_settings(config, python),
}
def dynamic_document(python: PythonCore) -> Mapping[str, object]:
"""The values the Python core resolved from its arguments and the process settings."""
settings: Final = python.settings()
applied: Final = python._core_applied # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
resolved: Final[Mapping[str, object]] = { # pyright: ignore[reportUnknownVariableType] # the fallback settings are untyped lists
"version": python._registry_version, # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
"model_group_alias": python._core_model_group_alias, # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
"fallbacks": settings.fallbacks, # pyright: ignore[reportUnknownMemberType] # the setting's own untyped list
"context_window_fallbacks": settings.context_window_fallbacks, # pyright: ignore[reportUnknownMemberType] # the setting's own untyped list
"content_policy_fallbacks": settings.content_policy_fallbacks, # pyright: ignore[reportUnknownMemberType] # the setting's own untyped list
"num_retries": python.num_retries,
"max_fallbacks": python.max_fallbacks,
"allowed_fails": python._core_allowed_fails # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
if "allowed_fails" in applied
else python.explicit_allowed_fails,
"fallback_keys": fallback_keys(python, settings),
}
return {name: _plain(value) for name, value in resolved.items() if value is not None} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the fallback settings are untyped lists
def fallback_keys(python: PythonCore, settings: object) -> Mapping[str, object]:
"""The spellings a fallback key matches besides the group name (`get_fallback_model_group`)."""
lists: Final = (
getattr(settings, "fallbacks", None),
getattr(settings, "context_window_fallbacks", None),
getattr(settings, "content_policy_fallbacks", None),
)
keys: Final = tuple( # comprehension-ok: the first group of each fallback entry, across the three lists
str(next(iter(mapping)))
for entries in lists
for entry in _as_list(entries)
if (mapping := as_str_mapping(entry)) and isinstance(mapping, dict)
)
names: Final = {
*(python._core_model_names or ()), # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
*(python._core_model_group_alias or {}), # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
}
prefixed: Final = {
name: spelled for name in names if (spelled := _provider_prefixed_model_group(name, keys)) is not None
}
return {"providers": _PROVIDERS, "provider_prefixed": prefixed}
def unsupported_arguments(config: RouterConfig) -> frozenset[str]:
"""Arguments the Rust core does not serve; any of them keeps the router on Python."""
changed: Final = frozenset(
name
for name, value in config.arguments.items()
if name not in _HANDLED_FIELDS and value != _PYTHON_DEFAULTS.get(name, None)
)
strategy: Final = config.arguments.get("routing_strategy", "simple-shuffle")
model_list: Final = config.arguments.get("model_list")
# `simple_shuffle` raises on a non-numeric weight when it picks; the Python core keeps that.
odd_weight: Final = any(
(params := as_str_mapping((as_str_mapping(deployment) or {}).get("litellm_params"))) is not None
and params.get("weight") is not None
and (isinstance(params.get("weight"), bool) or not isinstance(params.get("weight"), int | float))
for deployment in _as_list(model_list)
)
return (
changed
| ({"routing_strategy"} if strategy != "simple-shuffle" else frozenset[str]())
| ({"model_list"} if odd_weight else frozenset[str]())
)
def runs_on_python(request: Mapping[str, object]) -> bool:
"""Per-request routing overrides the Rust core does not read yet."""
return any(key in _PYTHON_ONLY_REQUEST_KEYS or key.startswith("mock_testing_") for key in request)
def request_team_id(request: Mapping[str, object]) -> str | None:
metadata: Final = as_str_mapping(request.get("metadata") or request.get("litellm_metadata"))
team_id: Final = metadata.get("user_api_key_team_id") if metadata is not None else None
return team_id if isinstance(team_id, str) else None
def _bound_request(bound: inspect.BoundArguments) -> Mapping[str, object]:
"""The keyword payload of a call bound to the Python core's endpoint signature."""
named: Final = {name: value for name, value in bound.arguments.items() if name not in ("self", "kwargs")}
extra: Final = bound.arguments.get("kwargs")
return {**named, **(extra if isinstance(extra, Mapping) else {})}
_ACOMPLETION: Final = inspect.signature(PythonCore.acompletion) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # only the parameter names are read
_AEMBEDDING: Final = inspect.signature(PythonCore.aembedding) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # only the parameter names are read
_OWN_ATTRIBUTES: Final = frozenset({"_handle", "_config", "_python", "_host", "_built_for"})
class RustCoreDeclined(Exception):
"""The Rust core cannot serve a router whose Python core is already built: the router keeps that one."""
def __init__(self, reason: str, python: PythonCore) -> None:
super().__init__(reason)
self.python = python
class RustCore:
"""A ``RouterCore`` whose routing for ``RUST_SERVED_OPS`` runs in Rust, over a ``PythonCore``."""
def __init__(self, handle: "RouterCoreHandle", config: RouterConfig, python: PythonCore, host: PythonHost) -> None:
self._handle: RouterCoreHandle | None = handle
self._config: Final = config
self._python: Final = python
self._host: Final = host
self._built_for: tuple[int, int] = self._state()
@classmethod
def build(cls, config: RouterConfig, host: PythonHost, python_class: type[PythonCore] = PythonCore) -> "RustCore":
from litellm.rust_bridge._native import RouterCoreHandle, RustBridgeDeclined
extra: Final = unsupported_arguments(config)
if extra:
raise RustBridgeDeclined(f"router arguments not served by the Rust core: {sorted(extra)}")
python: Final = python_class(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature
try:
handle: Final = RouterCoreHandle(static_document(config, python), dynamic_document(python))
except Exception as declined:
# The Python core is built already (its callbacks registered): keep it, don't build another.
raise RustCoreDeclined(str(declined), python) from declined
return cls(handle, config, python, host)
def __getattr__(self, name: str) -> object:
return getattr(self._python, name)
def __setattr__(self, name: str, value: object) -> None:
if name in _OWN_ATTRIBUTES:
object.__setattr__(self, name, value)
return
setattr(self._python, name, value)
def __delattr__(self, name: str) -> None:
if name in _OWN_ATTRIBUTES:
object.__delattr__(self, name)
return
delattr(self._python, name)
def _state(self) -> tuple[int, int]:
"""What the handle was built from: the deployments, and the settings changesets applied."""
return (self._python._registry_version, self._python._settings_version) # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
def _current_handle(self) -> "RouterCoreHandle | None":
"""The Rust handle, rebuilt after a deployment change or an applied changeset; ``None`` while
the Rust core cannot serve the router."""
if "_core_async_get_available_deployment" in vars(self._python):
# `register_plugin` installed a custom strategy, which only the Python core runs.
return None
if self._state() == self._built_for:
return self._handle
from litellm.rust_bridge._native import RouterCoreHandle
if self._handle is not None:
self._handle.retire()
strategy: Final = self._python.settings().routing_strategy
try:
if strategy not in (None, "simple-shuffle"):
raise ValueError(f"routing_strategy {strategy!r} is not served by the Rust core")
self._handle = RouterCoreHandle(static_document(self._config, self._python), dynamic_document(self._python))
except Exception as declined: # noqa: BLE001 # a change the Rust core cannot take keeps the router on Python until the next one
verbose_router_logger.info("Router served by the Python core after a change: %s", declined)
self._handle = None
self._built_for = self._state()
return self._handle
async def call(self, op: str, request: Mapping[str, object]) -> object:
if op not in RUST_SERVED_OPS:
return await self._python.call(op, request)
return await self._drive(op, request)
async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload
try:
bound: Final = _ACOMPLETION.bind(self._python, *args, **kwargs)
except TypeError:
# The Python core raises the caller's mistake in its own words.
return await self._python.acompletion(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType, reportUnknownMemberType, reportUnknownVariableType] # the call is expected to raise
return await self._drive("acompletion", _bound_request(bound))
async def aanthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: the endpoint's own keyword payload
return await self._drive("aanthropic_messages", kwargs)
async def aembedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload
try:
bound: Final = _AEMBEDDING.bind(self._python, *args, **kwargs)
except TypeError:
return await self._python.aembedding(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType, reportUnknownMemberType] # the call is expected to raise
return await self._drive("aembedding", _bound_request(bound))
async def aresponses(self, **kwargs: object) -> object: # kwargs-ok: the endpoint's own keyword payload
return await self._drive("aresponses", kwargs)
def retire(self) -> None:
if self._handle is not None:
self._handle.retire()
self._python.retire()
async def _drive(self, op: str, request: Mapping[str, object]) -> object:
handle: Final = self._current_handle()
if handle is None or runs_on_python(request) or not isinstance(request.get("model"), str):
return await self._python.call(op, request)
if not hasattr(request.get(_METADATA_FIELDS.get(op, "metadata"), {}), "update"):
# `_update_kwargs_before_fallbacks` raises on metadata it cannot update, before any attempt.
return await self._python.call(op, request)
fields: Final = {
**{key: request[key] for key in _REQUEST_FIELDS if request.get(key) is not None},
**{name: request[key] for key, name in _RENAMED_REQUEST_FIELDS.items() if request.get(key) is not None},
}
team_id: Final = request_team_id(request)
try:
call: Final = handle.start(
op, _plain_document({"model": request["model"], **fields, **({"team_id": team_id} if team_id else {})})
)
except ValueError:
return await self._python.call(op, request)
return await _Driver(call, self._python, op, request).run()
# A fallback chain list in the shape the Python core's fallback walk reads it.
_Chains: TypeAlias = Sequence[Mapping[str, list[str]]] | None
# The request field `_update_kwargs_before_fallbacks` writes the model group into, per op.
_METADATA_FIELDS: Final = MappingProxyType(
{"aanthropic_messages": "litellm_metadata", "aresponses": "litellm_metadata"}
)
def _ttl(step: Mapping[str, object]) -> float | None:
ttl: Final = step.get("ttl")
return ttl if isinstance(ttl, int | float) else None
def _store_for(python: PythonCore, key: str) -> "DualCache":
"""Cooldowns live in the router's cooldown store; every other routing key in its main DualCache."""
health: Final = python.health()
return health.cooldowns.cooldown_store if key.endswith(":cooldown") else health.store
async def answer_cache_step(python: PythonCore, step: Mapping[str, object]) -> object:
"""Run one cache op the Rust core asked for against the router's own DualCaches (memory plus Redis)."""
key: Final = str(step.get("key"))
ttl: Final = _ttl(step)
match step.get("op"):
case "get":
return _plain(await _store_for(python, key).async_get_cache(key)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case "batch_get":
keys: Final = [str(k) for k in _as_list(step.get("keys"))]
values: Final[list[object]] = [ # mutable-ok: built here, then made plain
await _store_for(python, k).async_get_cache(k) # pyright: ignore[reportUnknownMemberType] # DualCache's reads are untyped
for k in keys
]
return _plain(values)
case "set":
store: Final = _store_for(python, key)
if ttl is None:
await store.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
else:
await store.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
return None
case "increment":
amount: Final = step.get("value")
by: Final = amount if isinstance(amount, int | float) and not isinstance(amount, bool) else 1
counter: Final = _store_for(python, key)
total: Final = (
await counter.async_increment_cache(key, by) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
if ttl is None
else await counter.async_increment_cache(key, by, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
)
return total if isinstance(total, int | float) else by
case "delete":
await _store_for(python, key).async_delete_cache(key)
return None
case other:
raise RuntimeError(f"router-core asked for an unknown cache op {other!r}")
@dataclass(slots=True)
class _Hop:
"""One target of a fallback walk and what its level ended on."""
group: str
pending: Mapping[str, object] | None
error: BaseException | None = None
failed: Mapping[str, object] | None = None
children: "list[_Hop]" = field(default_factory=list) # mutable-ok: the fallbacks tried from this target
class _Driver:
"""Answers the steps one Rust call loop blocks on until it is done.
Each attempt runs on the Python core with retries and fallbacks off, so the driver puts back
what Python's own retry and fallback layers add to the outcome: the retry stamp on an error
whose retry loop ran out, the retry and fallback headers on a response, and the explanation
the fallback layer appends to the error it finally raises.
"""
def __init__(self, call: "RouterCall", python: PythonCore, op: str, request: Mapping[str, object]) -> None:
self._call: Final = call
self._python: Final = python
self._op: Final = op
self._request: Final = request
self._last_response: object = None
self._last_error: BaseException | None = None
# The last failure at each fallback level: a level raises its own last error.
self._level_errors: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call loop runs
self._attempted: bool = False
self._walk: Final = AttemptedFallbackTargets()
# Fallback targets in the order Rust walked them, and the groups an attempt was made in.
self._hop_targets: Final[list[str]] = [] # mutable-ok: filled as the call loop runs
self._attempted_groups: Final[set[str]] = set() # mutable-ok: filled as the call loop runs
# The walk in order: ("hop", (level, group, order)), ("attempt", level), ("error", (level,
# error, request)), ("sleep", None) before each retry, and ("exhausted", group) when a retry
# loop ran out.
self._timeline: Final[list[tuple[str, object]]] = [] # mutable-ok: filled as the call loop runs
async def run(self) -> object:
outcome: object = None # rebind-ok: each step's outcome answers it
while True:
step = _STEP.validate_python(self._call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome against the step it answers
self._take_events(step.get("events"))
match step["kind"]:
case "attempt":
outcome = await self._attempt(step)
case "sleep":
self._timeline.append(("sleep", None))
await self._python._host.sleep(float(str(step["seconds"]))) # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
outcome = None
case "allow_fallback":
outcome = await self._allow_fallback(str(step["target"]))
case "cache":
outcome = await answer_cache_step(self._python, step)
case "select":
outcome = step["candidates"]
case "done":
return self._with_headers(step)
case "failed":
primary = self._level_errors.get(0)
if primary is not None:
if self._walks_order_tiers(primary, self._request.get("model")):
# Python's order-tier walk raises what its last fallback raised, as is.
raise await self._walk_error(self._walk_tree(), primary)
await self._fill_unattempted_level()
self._explain(primary)
raise primary
if not self._attempted:
# No deployment was tried (unknown model, nothing healthy): Python raises its own error.
return await self._python.call(self._op, self._request)
raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}")
case other:
raise RuntimeError(f"router-core returned an unknown step {other!r}")
def _take_events(self, events: object) -> None:
for raw_event in _as_list(events):
event = as_str_mapping(raw_event)
if event is not None and event.get("kind") == "fallback":
self._timeline.append(("hop", (event.get("level"), event.get("to"), event.get("order"))))
if event is None or event.get("kind") != "retries_exhausted":
continue
self._timeline.append(("exhausted", event.get("model_group")))
error: BaseException | None = self._last_error
# `async_function_with_retries` stamps the error its exhausted retry loop raises.
if error is not None and type(error) in self._python._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
error.max_retries = event.get("max") # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes
error.num_retries = event.get("attempted") # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes
def _walks_order_tiers(self, error: BaseException, group: object) -> bool:
"""`async_function_with_fallbacks_common_utils` hands a group with more than one order level
to `run_async_fallback` outside its try, so the walk's own error is raised with no suffix."""
import litellm.exceptions as litellm_exceptions
python: Final = self._python
request: Final = self._request
if (
request.get("disable_fallbacks") is True
or not isinstance(group, str)
or is_guardrail_intervention(error) # pyright: ignore[reportArgumentType] # every attempt error is an Exception
or is_bad_request_without_fallback(error, python.fallback_on_unsupported_params) # pyright: ignore[reportArgumentType] # every attempt error is an Exception
or isinstance(
error,
litellm_exceptions.ContextWindowExceededError | litellm_exceptions.ContentPolicyViolationError,
)
):
return False
metadata: Final = as_str_mapping(request.get("metadata"))
team: Final = metadata.get("user_api_key_team_id") if metadata is not None else None
deployments: Final = (
python._core_get_model_list( # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
model_name=group, team_id=team if isinstance(team, str) else None
)
or ()
)
orders: Final = {
order
for deployment in deployments
if (order := (deployment.get("litellm_params") or {}).get("order")) is not None # pyright: ignore[reportUnknownMemberType] # a deployment's untyped litellm_params
}
return len(orders) > 1
def _walk_tree(self) -> _Hop:
"""The fallback walk as Rust ran it: each hop under the level it was tried from."""
root: Final = _Hop(str(self._request.get("model")), None)
current: Final[dict[int, _Hop]] = {0: root} # mutable-ok: the newest hop at each level
recent: _Hop | None = None # rebind-ok: the hop whose attempt failed last
for kind, value in self._timeline:
if kind == "hop" and isinstance(value, tuple):
level, group, order = value # pyright: ignore[reportUnknownVariableType] # the timeline's tuples are untyped
hop = _Hop(str(group), {**self._request, "model": str(group)}) # pyright: ignore[reportUnknownArgumentType] # the timeline's tuples are untyped
if isinstance(order, int):
hop.pending = {**self._request, "model": str(group), "_target_order": order} # pyright: ignore[reportUnknownArgumentType] # the timeline's tuples are untyped
parent = current.get(level - 1 if isinstance(level, int) else -1)
if parent is not None and isinstance(level, int):
parent.children.append(hop)
current[level] = hop
recent = None
elif kind == "attempt" and isinstance(value, int) and value in current:
current[value].pending = None
elif kind == "error" and isinstance(value, tuple) and value[0] in current:
recent = current[value[0]]
recent.error, recent.failed = value[1], value[2]
elif kind == "sleep" and recent is not None:
recent.pending = recent.failed # pyright: ignore[reportUnknownMemberType] # the timeline's tuples are untyped
elif kind == "exhausted" and recent is not None and recent.pending is recent.failed: # pyright: ignore[reportUnknownMemberType] # the timeline's tuples are untyped
recent.pending = None
return root
async def _walk_error(self, hop: _Hop, primary: BaseException) -> BaseException:
"""The error a hop's level raises. Python raises its own error where Rust made no attempt:
for a target with no deployment, and for a retry that found no deployment left to try. A
group with more than one order level raises its walk's last error instead
(`run_async_fallback`), whatever that walk tried."""
own: Final = await self._own_error(hop, primary)
if not hop.children or not self._walks_order_tiers(own, hop.group):
return own
last: BaseException = own # rebind-ok: each fallback tried replaces it
for child in hop.children:
last = await self._walk_error(child, primary)
return last
async def _own_error(self, hop: _Hop, primary: BaseException) -> BaseException:
"""The error the hop's own target raised: Python's, replayed, when Rust made no attempt."""
if hop.pending is not None:
try:
await self._python.call(self._op, {**hop.pending, "num_retries": 0, "disable_fallbacks": True})
except Exception as error: # noqa: BLE001 # the error Python's walk would end on
return error
return hop.error or primary
async def _fill_unattempted_level(self) -> None:
"""A first fallback that resolved no deployment made no attempt, so Rust holds no error for
it: ask Python, which raises its own no-deployment error without calling a provider."""
if 1 in self._level_errors:
return
target: Final = next((t for t in self._hop_targets if t not in self._attempted_groups), None)
if target is None:
return
try:
await self._python.call(
self._op, {**self._request, "model": target, "num_retries": 0, "disable_fallbacks": True}
)
except Exception as error: # noqa: BLE001 # the error is only quoted in the primary error's message
self._level_errors[1] = error
def _with_headers(self, step: Mapping[str, object]) -> object:
response: object = self._last_response # rebind-ok: the response gains its headers in turn
retries: Final = step.get("attempted_retries")
if isinstance(retries, int) and retries > 0:
max_retries: Final = step.get("max_retries")
response = add_retry_headers_to_response( # rebind-ok: the response gains its headers in turn
response=response,
attempted_retries=retries,
max_retries=max_retries if isinstance(max_retries, int) else None,
)
fallbacks: Final = step.get("attempted_fallbacks")
return add_fallback_headers_to_response(
response=response, attempted_fallbacks=fallbacks if isinstance(fallbacks, int) else 0
)
def _explain(self, error: BaseException) -> None:
"""Append what `async_function_with_fallbacks_common_utils` appends to the error it raises."""
python: Final = self._python
request: Final = self._request
if (
request.get("disable_fallbacks") is True
or is_guardrail_intervention(error) # pyright: ignore[reportArgumentType] # every attempt error is an Exception
or is_bad_request_without_fallback(error, python.fallback_on_unsupported_params) # pyright: ignore[reportArgumentType] # every attempt error is an Exception
or not hasattr(error, "message")
or not python._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
):
return
model_group: Final = request.get("model")
model_name: Final = model_group if isinstance(model_group, str) else None
try:
self._explain_fallbacks(error, model_name)
except Exception: # noqa: BLE001 # Python looks the chains up inside its fallback try: a bad chain means none was attempted
error.message += format_fallback_outcome_message( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix
model_name, None, "", skipped=self._walk.skipped
)
def _explain_fallbacks(self, error: BaseException, model_name: str | None) -> None:
import litellm.exceptions as litellm_exceptions
python: Final = self._python
request: Final = self._request
fallbacks: Final[_Chains] = python._core_fallbacks # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType, reportAssignmentType] # RustCore drives its own PythonCore, whose chains are untyped lists
context_window_fallbacks: Final[_Chains] = python._core_context_window_fallbacks # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType, reportAssignmentType] # RustCore drives its own PythonCore, whose chains are untyped lists
content_policy_fallbacks: Final[_Chains] = python._core_content_policy_fallbacks # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType, reportAssignmentType] # RustCore drives its own PythonCore, whose chains are untyped lists
lookup_groups: Final = fallback_lookup_groups(request, model_name)
context_window: Final = isinstance(error, litellm_exceptions.ContextWindowExceededError)
content_policy: Final = isinstance(error, litellm_exceptions.ContentPolicyViolationError)
attempted_group: Sequence[object] | None = None # rebind-ok: set below
orders: Final = sorted(
{
order
for deployment in python._core_get_model_list(model_name=model_name) or () # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
if (order := (deployment.get("litellm_params") or {}).get("order")) is not None # pyright: ignore[reportUnknownMemberType] # a deployment's untyped litellm_params
}
)
if len(orders) > 1 and not (context_window or content_policy):
external, generic = (
get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=lookup_groups)
if fallbacks is not None and lookup_groups
else (None, None)
)
if external is None and generic is not None and fallbacks is not None:
external = fallbacks[generic]["*"] # rebind-ok: the generic chain when no specific one matched
combined: Final = [{"model": model_name, "_target_order": o} for o in orders if o > orders[0]] + list(
external or ()
)
if combined:
attempted_group = combined # rebind-ok: settled in the order Python's walk does
if attempted_group is None:
specific: Final = (
context_window_fallbacks if context_window else content_policy_fallbacks if content_policy else None
)
if (context_window or content_policy) and specific is not None:
attempted_group, _ = ( # rebind-ok: settled in the order Python's walk does
get_fallback_model_group_for_lookup_groups(fallbacks=specific, lookup_groups=lookup_groups)
)
if attempted_group is None:
# Python raises before its generic fallbacks: "No fallback was attempted".
error.message += format_fallback_outcome_message( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix
model_name, None, "", skipped=self._walk.skipped
)
return
elif context_window:
error.message += ( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix
f"\nmodel={model_name}. context_window_fallbacks={mask_sensitive_structure(context_window_fallbacks)}."
f" fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'context_window_fallback' -"
" https://docs.litellm.ai/docs/routing#fallbacks"
)
elif content_policy:
error.message += ( # rebind-ok: the error Python raises carries the same suffix
f"\n{format_content_policy_fallback_hint(model_name, content_policy_fallbacks, fallbacks)}" # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above
)
if attempted_group is None and fallbacks is not None and lookup_groups:
attempted_group, generic_index = ( # rebind-ok: settled in the order Python's walk does
get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=lookup_groups)
)
if attempted_group is None and generic_index is not None:
attempted_group = fallbacks[generic_index][ # rebind-ok: settled in the order Python's walk does
"*"
]
if attempted_group is None:
error.message += format_no_fallback_group_message(lookup_groups, fallbacks) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix
return
detail: Final = (
truncate_fallback_error_detail(redact_string(str(self._level_errors.get(1) or error)))
if attempted_group
else ""
)
error.message += format_fallback_outcome_message( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the error Python raises carries the same suffix
model_name, attempted_group, detail, skipped=self._walk.skipped
)
def _pinned_request(self, step: Mapping[str, object]) -> Mapping[str, object]:
"""The request as the Python router's own fallback walk would send it, narrowed to the one
deployment Rust picked: the name Rust resolved, with every other deployment it resolved to
skipped."""
deployment_id: Final = str(step["deployment_id"])
name: Final = str(step.get("model") or step["model_group"])
candidates: Final = step.get("candidate_ids")
others: Final = tuple(str(i) for i in _as_list(candidates) if str(i) != deployment_id)
pinned: Final[dict[str, object]] = { # mutable-ok: built here, then handed to the Python core as its request
**self._request,
"model": name,
}
if others:
pinned["_retry_skipped_deployment_ids"] = others
target_order: Final = step.get("target_order")
if isinstance(target_order, int):
pinned["_target_order"] = target_order
return pinned
async def _attempt(self, step: Mapping[str, object]) -> Mapping[str, object]:
self._attempted = True
self._attempted_groups.add(str(step["model_group"]))
self._timeline.append(("attempt", step.get("fallback_level")))
request: Final = self._pinned_request(step)
try:
response: Final = await self._python.call(
self._op, {**request, "num_retries": 0, "disable_fallbacks": True}
)
except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop
self._last_error = error
level: Final = step.get("fallback_level")
self._level_errors[level if isinstance(level, int) else 0] = error
self._timeline.append(("error", (level if isinstance(level, int) else 0, error, request)))
return attempt_failure(error)
self._last_response = response
return {"ok": True, "total_tokens": total_tokens(response)}
async def _allow_fallback(self, target: str) -> bool:
"""`_is_fallback_target_authorized` then `_is_fallback_target_within_budget`."""
ctx: Final = self._python._core_ctx # pyright: ignore[reportPrivateUsage] # RustCore drives its own PythonCore
settings: Final = ctx.settings()
request: Final = dict(self._request)
original: Final = self._request.get("model")
if target == original:
return True
if settings.fallback_access_check is not None and not await ctx.fallback_access_allowed(
model=target, request_kwargs=request
):
self._walk.record_skip(target, FallbackSkipReason.NOT_AUTHORIZED)
return False
if settings.fallback_budget_check is not None and not await ctx.fallback_budget_allowed(
model=target, request_kwargs=request
):
self._walk.record_skip(target, FallbackSkipReason.OVER_BUDGET)
return False
self._hop_targets.append(target)
return True

View file

@ -1,6 +1,6 @@
from asyncio import Future
from collections.abc import Mapping, Sequence
from typing import Never, final
from typing import Literal, Never, NotRequired, TypedDict, final
import httpx
from pydantic import JsonValue
@ -218,6 +218,8 @@ __all__ = [
"NativeDiagnosticProcessor",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RouterCall",
"RouterCoreHandle",
"RustBridgeDeclined",
"RustUpstreamError",
"TokenCounter",
@ -325,3 +327,115 @@ class NativeCacheHandle:
def ping(self) -> Future[bool]: ...
def disconnect(self) -> Future[None]: ...
def delete(self, keys: Sequence[str]) -> Future[None]: ...
# The FailureKind names an attempt outcome's `kind` takes (router-core `FailureKind`).
_FailureKind = Literal[
"RateLimitError",
"ContextWindowExceededError",
"ContentPolicyViolationError",
"AuthenticationError",
"PermissionDeniedError",
"NotFoundError",
"BadRequestError",
"Timeout",
"APIConnectionError",
"InternalServerError",
"ServiceUnavailableError",
"GuardrailIntervention",
"Declined",
"CallerCredentialError",
"APIError",
]
class _AttemptStep(TypedDict):
kind: Literal["attempt"]
deployment_id: str
model_group: str
retry: int
fallback_depth: int
timeout: float | None
class _SleepStep(TypedDict):
kind: Literal["sleep"]
seconds: float
class _AllowFallbackStep(TypedDict):
kind: Literal["allow_fallback"]
target: str
class _SelectStep(TypedDict):
kind: Literal["select"]
stage: Literal["plugins", "strategy"]
strategy: str | None
model_group: str
candidates: list[str]
# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`,
# which sets `keys`; `value` is the JSON to store (`set`) or the amount (`increment`); `ttl` is
# in seconds. Outcomes: get -> the JSON value or None; batch_get -> one JSON value or None per
# key, in order; increment -> the new int or float; set and delete -> None.
class _CacheStep(TypedDict):
kind: Literal["cache"]
op: Literal["get", "set", "increment", "batch_get", "delete"]
key: str | None
keys: list[str] | None
value: JsonValue
ttl: float | None
class _DoneStep(TypedDict):
kind: Literal["done"]
deployment_id: str
model_group: str
attempted_retries: int
attempted_fallbacks: int
total_tokens: int
class _FailedStep(TypedDict):
kind: Literal["failed"]
error: str
error_kind: Literal["config", "unknown_model", "no_healthy_deployments", "attempt", "cache", "retired"]
attempt_failed: bool
status: int | None
failure_kind: _FailureKind | None
# The end of a `RouterCoreHandle.select`.
class _SelectedStep(TypedDict):
kind: Literal["done"]
deployment_id: str
model_group: str
_RouterStep = _AttemptStep | _SleepStep | _AllowFallbackStep | _SelectStep | _CacheStep | _DoneStep | _FailedStep
class _AttemptSucceeded(TypedDict):
ok: Literal[True]
total_tokens: NotRequired[int]
class _AttemptFailed(TypedDict):
ok: Literal[False]
kind: _FailureKind
status: NotRequired[int | None]
retry_after: NotRequired[float | None]
num_retries: NotRequired[int | None]
detail: NotRequired[str | None]
# Each document is a mapping or its JSON text.
_Document = Mapping[str, object] | str
@final
class RouterCoreHandle:
def __new__(cls, config: _Document, dynamic: _Document) -> RouterCoreHandle: ...
def start(self, op: str, request: _Document) -> RouterCall: ...
# Driven like `start`; its `done` step is a `_SelectedStep`.
def select(self, model: str, request: _Document) -> RouterCall: ...
def apply(self, changes: _Document) -> int: ...
def registry_version(self) -> int: ...
def retire(self) -> None: ...
@final
class RouterCall:
# attempt: _AttemptSucceeded | _AttemptFailed; sleep: None; allow_fallback: bool;
# select: the remaining candidate ids, pick first; cache: see _CacheStep; first call: None.
# A call from `select` steps through select and cache only, and ends in _SelectedStep or failed.
def next(
self, outcome: _AttemptSucceeded | _AttemptFailed | bool | Sequence[str] | JsonValue | None
) -> _RouterStep | _SelectedStep: ...

View file

@ -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"})),
)

View file

@ -21,6 +21,7 @@ CHECKED = (
"litellm/router_core.py",
"litellm/router_deprecated.py",
"litellm/router_python_core.py",
"litellm/router_rust_core.py",
)

View file

@ -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",
}

View file

@ -452,25 +452,36 @@ def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwarg
return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds
CoreBuilder = Callable[..., RouterCore]
"""``builder(host, **router_arguments)``: a router core on ``host``."""
CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core}
"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a
configuration raises ``CoreDeclined``, and the test is skipped for that core."""
class CoreDeclined(Exception):
"""The core does not serve this configuration (the router would stay on the Python core)."""
CoreBuilder = Callable[..., RouterCore]
"""``builder(host, **router_arguments)``: a router core on ``host``."""
def rust_core(host: ScriptedHost, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments
"""A ``RustCore`` on ``host``, built the way ``make_core`` builds one when the Rust router is on."""
from litellm.router_rust_core import RustCore
try:
return RustCore.build(build_router_config(arguments), host)
except Exception as declined: # noqa: BLE001 # make_core keeps such a router on the Python core
raise CoreDeclined(str(declined)) from declined
CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core, "rust": rust_core}
"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a
configuration raises ``CoreDeclined``, and the test is skipped for that core."""
__all__ = [
"Call",
"HostCallbacks",
"INFERENCE_OPS",
"CORE_BUILDERS",
"INFERENCE_OPS",
"Call",
"CoreBuilder",
"CoreDeclined",
"HostCallbacks",
"ScriptedHost",
"auth_error",
"bad_request",
@ -482,6 +493,7 @@ __all__ = [
"embedding_response",
"python_core",
"rate_limit",
"rust_core",
"server_error",
"timeout",
]

View file

@ -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)