mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
feat(router): add the Rust router core with the host as its only boundary (#49513)
Adds the litellm-router-core crate, its step-driven PyO3 bridge (RouterCoreHandle / RouterCall) and RustCore (litellm/router_rust/), served behind LITELLM_RUST through the ROUTER route rule. The host call-out is the only boundary: every attempt goes through Host::invoke, there is no native executor path, and RustCore never builds, wraps or forwards to PythonCore, so router_python_core.py stays deletable. Attempt preparation lives in litellm/router_attempt.py so it outlives router_python_core.py. RustCore matches PythonCore's behavior, including its known bugs (fixed in follow-up PRs). Routing decisions (strategy picks, routing groups, PTU, pre-routing, pre-call checks, fallback walk, health routing, access groups, weights, cooldowns, allowed fails, usage) live in the crate; work that needs live Python objects (request shaping, endpoints, streaming, model info facts, callbacks, alerting, prompt management) lives in the host-side litellm/router_rust/ package (steps/, invoke/, surface/) driven by Rust steps. The crate owns the registry. It runs the deployment build loop (secret references, ids, organization fan-out, install order) and asks the host for what needs Python, one deployment at a time so errors surface where Python raised them: secrets, validate, provider and added steps. It holds the model list, the index maps and the pattern routers with the Python core's history rules (host_index.rs), and answers the registry lookups (lookup.rs). The Python views of that state in litellm/router_rust/surface/ (model_list, index(), names(), patterns()) read the handle, and every write to them becomes a changeset. Model info, deployment model info, configured limits and model group info are merged in the crate (model_info.rs, group_info.rs) from catalog rows and facts the host answers on a model_info step asked at query time, so price reloads and register_model stay visible. The usage cache key format is defined once, in the crate. The crate records each attempt's outcome itself rather than through litellm callbacks. A streamed attempt is deferred: it counts its request when the stream opens, and the host's stream wrappers report the success with the final usage, or a failure after the first chunk, when the stream ends, as the Python core's callbacks do. tests/unit/test_router/core characterizes PythonCore's public behavior and runs every test on both cores, with nothing pending; only the names of the Python core's own outcome callbacks are asserted on Python alone. litellm.Router is now final: subclassing it raises TypeError at the class statement, and the machinery that kept subclass overrides working is gone. The Rust core never honoured subclass overrides, so they would have silently stopped running there. Custom routing goes through register_plugin(), settings through apply() and outcomes through observe(). This is a breaking change for anyone subclassing Router. Fuzz parity fixes on top of the rebase: the Rust core keeps num_retries, max_fallbacks, allowed_fails, retry_after, timeout, the tag settings, routing_strategy_args and the fallback lists as given (the crate reads them as numbers, update_settings casts as Python does); a deprecated `model_list = [...]` write reaches the Rust core; a failed set_model_list is all or nothing on both cores (the deployments served before the call are put back and the error raised); config_deployments() is a generator and `del` of a plain attribute works on the Rust core. Least-busy pick reservations from main now work on the Rust core: its selectors share one reservation store, routed calls reserve and free their unstarted picks, and select() and batch cancel do not reserve. On the Python core a dropped upsert takes its name back off deployment_names, and get_deployment_order reads every numeric spelling as one int so mixed orders no longer raise TypeError. Rulings on the remaining differences, the Python core changed where it was the odd one out: sync picks and picks by deployment id apply the team filter as async picks do, so another team's deployment is refused; the fallback lists and default_litellm_params in RouterSettings are copies, so an in-place edit is ignored and settings change only through apply(); one health record with an unreadable timestamp drops out alone; an empty fallback entry is skipped; a bool cooldown_time is refused at construction, a bool or string cooldown_time on a deployment or a call falls back to the router default, and a failure status that no HTTP status can be is recorded as 500; an infinite or NaN cost and a whole number wider than 64 bits (ptu_shares) are refused on both cores with the same error (litellm/router_utils/deployment_numbers.py); and a /v1/messages stream whose provider iterator raises cools the deployment and retries like a non-stream call, with the Rust core waiting for that cooldown before its retry picks. Differences the proxy parity fuzz and its regression tests found on the Rust core: a file upload or a batch on an uploaded file falls back only within its own model group, as on the Python core, since the file lives with one provider (the host sends same_group_fallbacks_only and the walk skips targets in other groups); Decisions calls drop deployments whose provider has no Decisions config, and deployments that refuse the audio or videos a call carries, unless that would leave none, as the Python core does, while the retry count still reads the whole group; a file upload that fails on one order level is created on the next level's deployments rather than the first level again; and a response that cannot be converted (a text_completion body without a message) reports only its own traceback, so a rate limit handled by the attempt before it no longer turns the error into a 429, and both cores return 500.
This commit is contained in:
parent
c4faa55249
commit
9752869d84
195 changed files with 69945 additions and 772 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -419,6 +419,7 @@ jobs:
|
|||
|
||||
- shard: router
|
||||
test-path: >-
|
||||
tests/unit/router_rust
|
||||
tests/unit/test_router
|
||||
python-version: "3.12"
|
||||
workers: 4
|
||||
|
|
|
|||
20
litellm-rust/Cargo.lock
generated
20
litellm-rust/Cargo.lock
generated
|
|
@ -4288,6 +4288,7 @@ dependencies = [
|
|||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"aws-sdk-secretsmanager",
|
||||
"bytes",
|
||||
"criterion",
|
||||
|
|
@ -4312,6 +4313,7 @@ dependencies = [
|
|||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-model-catalog",
|
||||
"litellm-router-core",
|
||||
"litellm-router-types",
|
||||
"litellm-secret-value",
|
||||
"litellm-secrets",
|
||||
|
|
@ -4393,6 +4395,24 @@ dependencies = [
|
|||
[[package]]
|
||||
name = "litellm-router-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"async-trait",
|
||||
"jiff",
|
||||
"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"
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
@ -76,6 +77,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"] }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -47,6 +48,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 }
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
@ -115,6 +118,8 @@ mod tests {
|
|||
"bundled_model_cost_map",
|
||||
"TokenCounter",
|
||||
"Tokenizer",
|
||||
"RouterCoreHandle",
|
||||
"RouterCall",
|
||||
"gil_stats",
|
||||
"process_state_started",
|
||||
"reserve_process_for_forking",
|
||||
|
|
|
|||
408
litellm-rust/crates/python-bridge/src/router/config.rs
Normal file
408
litellm-rust/crates/python-bridge/src/router/config.rs
Normal file
|
|
@ -0,0 +1,408 @@
|
|||
//! 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::select::PassThrough;
|
||||
use litellm_router_core::{
|
||||
ChangeSet, ConfigError, DeploymentSpec, OutcomeCallbacks, RouterConfig, RouterDynamicConfig,
|
||||
RouterRequest,
|
||||
config::{Fallbacks, RetryPolicy},
|
||||
host::UncountedRetries,
|
||||
mock::MockFailure,
|
||||
registry::IndexOverrides,
|
||||
};
|
||||
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",
|
||||
"routing_group_members",
|
||||
"plugins",
|
||||
"cooldown_time",
|
||||
"disable_cooldowns",
|
||||
"allowed_fails_policy",
|
||||
"retry_policy",
|
||||
"model_group_retry_policy",
|
||||
"retry_after",
|
||||
"timeout",
|
||||
"request_timeout",
|
||||
"drop_params",
|
||||
"enable_pre_call_checks",
|
||||
"routing_plugins",
|
||||
"callback_filters",
|
||||
"enable_health_check_routing",
|
||||
"health_check_staleness_threshold",
|
||||
"background_health_check_model_groups",
|
||||
"enable_weighted_failover",
|
||||
"fallback_on_unsupported_params",
|
||||
"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",
|
||||
"allowed_fails_is_default",
|
||||
"fallback_keys",
|
||||
"default_deployment",
|
||||
"default_deployment_raises",
|
||||
];
|
||||
|
||||
/// 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("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::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 parse_router_config(value: Value) -> Result<RouterConfig, Error> {
|
||||
let object = routing_keys(Document::Static, value, STATIC_KEYS)?;
|
||||
shaped(Document::Static, object)
|
||||
}
|
||||
|
||||
pub(crate) fn parse_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>,
|
||||
listing_team_id: Option<String>,
|
||||
proxy_admin: Option<bool>,
|
||||
allowed_models: Option<Vec<String>>,
|
||||
admin_model: Option<String>,
|
||||
specific_deployment: Option<bool>,
|
||||
num_retries: Option<u32>,
|
||||
uncounted_retries: Option<UncountedRetries>,
|
||||
disable_fallbacks: Option<bool>,
|
||||
router_weights: Option<Value>,
|
||||
pinned_deployment_id: Option<String>,
|
||||
routing_strategy: Option<String>,
|
||||
filter_tags: Option<bool>,
|
||||
/// `listed` or `healthy`: which pass-through pick the request makes.
|
||||
pass_through: Option<String>,
|
||||
target_order: Option<i64>,
|
||||
excluded_ids: Option<Vec<String>>,
|
||||
retry_skipped_ids: Option<Vec<String>>,
|
||||
pre_routes: Option<bool>,
|
||||
has_image: Option<bool>,
|
||||
endpoint: Option<String>,
|
||||
web_search: Option<bool>,
|
||||
route_handler: Option<String>,
|
||||
decisions_media: Option<bool>,
|
||||
wildcard_request: Option<String>,
|
||||
unrouted: Option<bool>,
|
||||
passthrough_on_no_deployment: Option<bool>,
|
||||
skip_weighted_failover: Option<bool>,
|
||||
same_group_fallbacks_only: Option<bool>,
|
||||
model_group_retry_policy: Option<BTreeMap<String, RetryPolicy>>,
|
||||
retry_policy_raises: Option<Vec<String>>,
|
||||
#[serde(default, deserialize_with = "present")]
|
||||
fallbacks: Option<Option<Fallbacks>>,
|
||||
#[serde(default, deserialize_with = "present")]
|
||||
context_window_fallbacks: Option<Option<Fallbacks>>,
|
||||
#[serde(default, deserialize_with = "present")]
|
||||
content_policy_fallbacks: Option<Option<Fallbacks>>,
|
||||
max_fallbacks: Option<u32>,
|
||||
mock_failure: Option<MockFailure>,
|
||||
cooldown_time: Option<f64>,
|
||||
health_check_probe: Option<bool>,
|
||||
outcome_callbacks: Option<OutcomeCallbacks>,
|
||||
body: Option<Value>,
|
||||
}
|
||||
|
||||
/// A key the request passed, even as `None`: `Some(None)`.
|
||||
fn present<'de, D, T>(deserializer: D) -> Result<Option<Option<T>>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
T: Deserialize<'de>,
|
||||
{
|
||||
Option::<T>::deserialize(deserializer).map(Some)
|
||||
}
|
||||
|
||||
/// Read a request; `model` replaces the request's own when given.
|
||||
pub(crate) fn parse_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,
|
||||
listing_team_id: spec.listing_team_id,
|
||||
proxy_admin: spec.proxy_admin.unwrap_or(false),
|
||||
allowed_models: spec.allowed_models.unwrap_or_default(),
|
||||
admin_model: spec.admin_model,
|
||||
specific_deployment: spec.specific_deployment.unwrap_or(false),
|
||||
num_retries: spec.num_retries,
|
||||
uncounted_retries: spec.uncounted_retries,
|
||||
disable_fallbacks: spec.disable_fallbacks.unwrap_or(false),
|
||||
router_weights: spec.router_weights,
|
||||
pinned_deployment_id: spec.pinned_deployment_id,
|
||||
routing_strategy: spec.routing_strategy,
|
||||
filter_tags: spec.filter_tags.unwrap_or(false),
|
||||
pass_through: match spec.pass_through.as_deref() {
|
||||
None => None,
|
||||
Some("healthy") => Some(PassThrough::Healthy),
|
||||
Some(_) => Some(PassThrough::Listed),
|
||||
},
|
||||
target_order: spec.target_order,
|
||||
excluded_ids: spec.excluded_ids.unwrap_or_default(),
|
||||
retry_skipped_ids: spec.retry_skipped_ids.unwrap_or_default(),
|
||||
pre_routes: spec.pre_routes.unwrap_or(false),
|
||||
has_image: spec.has_image.unwrap_or(false),
|
||||
endpoint: spec.endpoint,
|
||||
web_search: spec.web_search.unwrap_or(false),
|
||||
route_handler: spec.route_handler,
|
||||
decisions_media: spec.decisions_media.unwrap_or(false),
|
||||
wildcard_request: spec.wildcard_request,
|
||||
unrouted: spec.unrouted.unwrap_or(false),
|
||||
passthrough_on_no_deployment: spec.passthrough_on_no_deployment.unwrap_or(false),
|
||||
skip_weighted_failover: spec.skip_weighted_failover.unwrap_or(false),
|
||||
same_group_fallbacks_only: spec.same_group_fallbacks_only.unwrap_or(false),
|
||||
model_group_retry_policy: spec.model_group_retry_policy,
|
||||
retry_policy_raises: spec.retry_policy_raises.unwrap_or_default(),
|
||||
fallbacks: spec.fallbacks,
|
||||
context_window_fallbacks: spec.context_window_fallbacks,
|
||||
content_policy_fallbacks: spec.content_policy_fallbacks,
|
||||
max_fallbacks: spec.max_fallbacks,
|
||||
mock_failure: spec.mock_failure,
|
||||
cooldown_time: spec.cooldown_time,
|
||||
health_check_probe: spec.health_check_probe.unwrap_or(false),
|
||||
outcome_callbacks: spec.outcome_callbacks.unwrap_or_default(),
|
||||
payload: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct ChangesSpec {
|
||||
/// Every deployment, replacing the current list before the upserts and deletes.
|
||||
#[serde(default)]
|
||||
replace: Option<Vec<DeploymentSpec>>,
|
||||
/// Deployments added at the end, beside any that share their id.
|
||||
#[serde(default)]
|
||||
append: Vec<DeploymentSpec>,
|
||||
#[serde(default)]
|
||||
upsert: Vec<DeploymentSpec>,
|
||||
#[serde(default)]
|
||||
delete: Vec<String>,
|
||||
/// A new dynamic document (the same keys `RouterCoreHandle(static, dynamic)` takes).
|
||||
#[serde(default)]
|
||||
settings: Option<Value>,
|
||||
/// The registry indexes the host wrote (`IndexOverrides`); `{}` drops them.
|
||||
#[serde(default)]
|
||||
index: Option<IndexOverrides>,
|
||||
/// The deployment at the id map's position for this id, removed with every index shifted.
|
||||
#[serde(default)]
|
||||
remove: Option<String>,
|
||||
/// A write to the indexes the host hands out (`HostIndexes`), field by field.
|
||||
#[serde(default)]
|
||||
host_index: Option<litellm_router_core::host_index::IndexWrite>,
|
||||
}
|
||||
|
||||
pub(crate) fn parse_changes(value: Value) -> Result<ChangeSet, Error> {
|
||||
let spec: ChangesSpec = shaped(Document::Changes, object(Document::Changes, value)?)?;
|
||||
Ok(ChangeSet {
|
||||
replace: spec.replace,
|
||||
appends: spec.append,
|
||||
upserts: spec.upsert,
|
||||
deletes: spec.delete,
|
||||
settings: spec.settings.map(parse_dynamic_config).transpose()?,
|
||||
index: spec.index,
|
||||
remove: spec.remove,
|
||||
host_index: spec.host_index,
|
||||
})
|
||||
}
|
||||
|
||||
#[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 = parse_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}))]
|
||||
#[case::pre_call_checks_on(json!({"enable_pre_call_checks": true}))]
|
||||
fn unset_and_ignored_keys_are_accepted(#[case] json: Value) {
|
||||
parse_router_config(json).unwrap();
|
||||
}
|
||||
|
||||
#[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!(parse_router_config(json).unwrap_err().declined(), declined);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn python_router_defaults_read_as_unset() {
|
||||
let dynamic = parse_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!(parse_request(json, model).unwrap().model, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::negative("negative", UncountedRetries::Negative)]
|
||||
#[case::not_positive("not_positive", UncountedRetries::NotPositive)]
|
||||
#[case::raises("raises", UncountedRetries::Raises)]
|
||||
fn a_request_reads_a_num_retries_that_is_no_count(
|
||||
#[case] spelled: &str,
|
||||
#[case] expected: UncountedRetries,
|
||||
) {
|
||||
let request =
|
||||
parse_request(json!({"model": "gpt", "uncounted_retries": spelled}), None).unwrap();
|
||||
assert_eq!(request.uncounted_retries, Some(expected));
|
||||
assert_eq!(request.num_retries, None);
|
||||
}
|
||||
|
||||
#[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!(!parse_request(json, None).unwrap_err().declined());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_change_set_reads_upserts_and_deletes() {
|
||||
let changes = parse_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"]);
|
||||
}
|
||||
}
|
||||
1569
litellm-rust/crates/python-bridge/src/router/host.rs
Normal file
1569
litellm-rust/crates/python-bridge/src/router/host.rs
Normal file
File diff suppressed because it is too large
Load diff
2051
litellm-rust/crates/python-bridge/src/router/mod.rs
Normal file
2051
litellm-rust/crates/python-bridge/src/router/mod.rs
Normal file
File diff suppressed because it is too large
Load diff
55
litellm-rust/crates/python-bridge/src/router/retries.rs
Normal file
55
litellm-rust/crates/python-bridge/src/router/retries.rs
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
//! The retry decisions `RouterCoreHandle` exposes as members: `get_num_retries_from_retry_policy`
|
||||
//! and `should_retry_this_error`, decided by router-core for a failure Python describes the way it
|
||||
//! describes a failed attempt.
|
||||
|
||||
use litellm_router_core::{
|
||||
AttemptError, Router,
|
||||
retry::{RetryContext, resolve_retry_policy, should_retry as core_should_retry},
|
||||
};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
|
||||
|
||||
fn failure(outcome: &Bound<'_, PyDict>) -> PyResult<AttemptError> {
|
||||
match super::attempt_outcome(outcome)? {
|
||||
Err(failure) => Ok(failure),
|
||||
Ok(_) => Err(PyValueError::new_err("a failure needs `ok: False`")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn policy_retries(
|
||||
router: &Router,
|
||||
outcome: &Bound<'_, PyDict>,
|
||||
model_group: Option<&str>,
|
||||
) -> PyResult<Option<u32>> {
|
||||
let failure = failure(outcome)?;
|
||||
let settings = router.settings();
|
||||
let policy = match model_group {
|
||||
Some(group) => resolve_retry_policy(
|
||||
settings.retry_policy.as_ref(),
|
||||
group,
|
||||
Some(&settings.model_group_retry_policy),
|
||||
),
|
||||
None => settings.retry_policy.as_ref(),
|
||||
};
|
||||
Ok(policy.and_then(|policy| policy.retries_for(&failure)))
|
||||
}
|
||||
|
||||
pub(crate) fn allowed_fails(router: &Router, outcome: &Bound<'_, PyDict>) -> PyResult<Option<u32>> {
|
||||
Ok(router.settings().allowed_fails_for(&failure(outcome)?))
|
||||
}
|
||||
|
||||
pub(crate) fn should_retry(
|
||||
outcome: &Bound<'_, PyDict>,
|
||||
(healthy, all): (usize, usize),
|
||||
(context_window, content_policy, regular): (bool, bool, bool),
|
||||
) -> PyResult<bool> {
|
||||
Ok(core_should_retry(
|
||||
&failure(outcome)?,
|
||||
RetryContext {
|
||||
healthy,
|
||||
all,
|
||||
has_context_window_fallbacks: context_window,
|
||||
has_content_policy_fallbacks: content_policy,
|
||||
has_regular_fallbacks: regular,
|
||||
},
|
||||
))
|
||||
}
|
||||
94
litellm-rust/crates/python-bridge/src/router/stream.rs
Normal file
94
litellm-rust/crates/python-bridge/src/router/stream.rs
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
//! The stream-recovery entry points of `RouterCoreHandle`: what may still take over a stream that
|
||||
//! failed before content, the same-group stream retry, and the fallback walk re-entered from a
|
||||
//! stream. The host's stream wrappers call them; the calls are driven like `start`.
|
||||
|
||||
use litellm_router_core::{
|
||||
FailureKind, Op, Router,
|
||||
call::{StreamOrigin, StreamPolicy},
|
||||
};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
|
||||
|
||||
use super::{RouterCall, attempt_outcome, config, config_error_to_pyerr, document, host::Finish};
|
||||
|
||||
/// A stream recovery a host wrapper asks for.
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum Recovery {
|
||||
/// `_astream_retry_same_group`.
|
||||
Retry,
|
||||
/// The fallback walk re-entered from the stream wrapper.
|
||||
Fallback,
|
||||
}
|
||||
|
||||
fn origin(value: &Bound<'_, PyAny>) -> PyResult<StreamOrigin> {
|
||||
serde_json::from_value(document(value)?)
|
||||
.map_err(|error| PyValueError::new_err(format!("a stream origin: {error}")))
|
||||
}
|
||||
|
||||
fn kind(kind: Option<&str>) -> PyResult<Option<FailureKind>> {
|
||||
kind.map(|kind| {
|
||||
kind.parse()
|
||||
.map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind")))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// `{"can_retry", "can_fall_back", "refusal_fallback", "retry_budget": int | None,
|
||||
/// "policy_applies": bool}`.
|
||||
pub(super) fn policy<'py>(
|
||||
py: Python<'py>,
|
||||
router: &Router,
|
||||
request: &Bound<'py, PyAny>,
|
||||
stream_origin: &Bound<'py, PyAny>,
|
||||
failure_kind: Option<&str>,
|
||||
) -> PyResult<Bound<'py, PyDict>> {
|
||||
let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?;
|
||||
let StreamPolicy {
|
||||
can_retry,
|
||||
can_fall_back,
|
||||
refusal_fallback,
|
||||
retry_budget,
|
||||
} = router.stream_policy(&request, &origin(stream_origin)?, kind(failure_kind)?);
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("can_retry", can_retry)?;
|
||||
dict.set_item("can_fall_back", can_fall_back)?;
|
||||
dict.set_item("refusal_fallback", refusal_fallback)?;
|
||||
dict.set_item("retry_budget", retry_budget.map(|(budget, _)| budget))?;
|
||||
dict.set_item(
|
||||
"policy_applies",
|
||||
retry_budget.is_some_and(|(_, applies)| applies),
|
||||
)?;
|
||||
Ok(dict)
|
||||
}
|
||||
|
||||
/// A recovery call for a stream from `stream_origin` that failed with `failure`.
|
||||
pub(super) fn recover(
|
||||
router: &Router,
|
||||
recovery: Recovery,
|
||||
op: &str,
|
||||
request: &Bound<'_, PyAny>,
|
||||
stream_origin: &Bound<'_, PyAny>,
|
||||
failure: &Bound<'_, PyDict>,
|
||||
) -> PyResult<RouterCall> {
|
||||
let op = Op::new(op);
|
||||
let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?;
|
||||
let stream_origin = origin(stream_origin)?;
|
||||
let Err(failure) = attempt_outcome(failure)? else {
|
||||
return Err(PyValueError::new_err("a stream failure needs `ok: False`"));
|
||||
};
|
||||
let router = router.clone();
|
||||
Ok(RouterCall::new(async move {
|
||||
let response = match recovery {
|
||||
Recovery::Retry => {
|
||||
router
|
||||
.retry_stream(op, request, stream_origin, failure)
|
||||
.await
|
||||
}
|
||||
Recovery::Fallback => {
|
||||
router
|
||||
.fallback_stream(op, request, stream_origin, failure)
|
||||
.await
|
||||
}
|
||||
};
|
||||
response.map(Finish::Call)
|
||||
}))
|
||||
}
|
||||
|
|
@ -7,3 +7,21 @@ 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
|
||||
jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-zoneinfo", "tzdb-bundle-always"] }
|
||||
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
|
||||
|
|
|
|||
236
litellm-rust/crates/router-core/src/access.rs
Normal file
236
litellm-rust/crates/router-core/src/access.rs
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
//! Who may reach a deployment: model access windows (`litellm_core_utils/access_windows.py`,
|
||||
//! a deployment reserved for some teams during a daily wall-clock window) and model access
|
||||
//! groups (`_filter_deployments_by_model_access_groups`, a caller allowed only through access
|
||||
//! groups is kept to its groups' deployments).
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc};
|
||||
|
||||
use jiff::{Timestamp, tz::TimeZone};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::registry::Deployment;
|
||||
|
||||
/// One `model_info.access_windows` entry (`ModelAccessWindow`), validated when the deployment
|
||||
/// was added.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct AccessWindow {
|
||||
/// Microseconds after midnight.
|
||||
pub start: u64,
|
||||
pub end: u64,
|
||||
pub timezone: String,
|
||||
pub team_ids: Vec<String>,
|
||||
}
|
||||
|
||||
const MICROS_PER_SECOND: u64 = 1_000_000;
|
||||
|
||||
/// A `datetime.time` as pydantic reads it: `HH:MM[:SS[.ffffff]]`, or seconds after midnight.
|
||||
fn wall_clock(value: &Value) -> Option<u64> {
|
||||
match value {
|
||||
Value::Number(seconds) => {
|
||||
let seconds = seconds.as_f64()?;
|
||||
(0.0..86_400.0)
|
||||
.contains(&seconds)
|
||||
.then(|| (seconds * MICROS_PER_SECOND as f64).round() as u64)
|
||||
}
|
||||
Value::String(text) => {
|
||||
let mut parts = text.splitn(3, ':');
|
||||
let hour: u64 = parts.next()?.parse().ok()?;
|
||||
let minute: u64 = parts.next()?.parse().ok()?;
|
||||
let (second, micros) = match parts.next() {
|
||||
None => (0, 0),
|
||||
Some(rest) => {
|
||||
let (whole, fraction) = rest.split_once('.').unwrap_or((rest, ""));
|
||||
let digits: String = fraction.chars().take(6).collect();
|
||||
let micros = if digits.is_empty() {
|
||||
0
|
||||
} else {
|
||||
format!("{digits:0<6}").parse().ok()?
|
||||
};
|
||||
(whole.parse().ok()?, micros)
|
||||
}
|
||||
};
|
||||
(hour < 24 && minute < 60 && second < 60)
|
||||
.then(|| ((hour * 60 + minute) * 60 + second) * MICROS_PER_SECOND + micros)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
impl AccessWindow {
|
||||
fn parse(value: &Value) -> Option<Self> {
|
||||
let window = value.as_object()?;
|
||||
Some(Self {
|
||||
start: wall_clock(window.get("start")?)?,
|
||||
end: wall_clock(window.get("end")?)?,
|
||||
timezone: window.get("timezone")?.as_str()?.to_owned(),
|
||||
team_ids: window
|
||||
.get("team_ids")?
|
||||
.as_array()?
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::to_owned)
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
|
||||
/// `is_window_active`: the window's wall clock contains `now`, wrapping past midnight when
|
||||
/// it ends before it starts.
|
||||
pub fn is_active(&self, now: Timestamp) -> bool {
|
||||
let Ok(zone) = TimeZone::get(&self.timezone) else {
|
||||
return false;
|
||||
};
|
||||
let local = now.to_zoned(zone).time();
|
||||
let micros = ((u64::from(local.hour().unsigned_abs()) * 60
|
||||
+ u64::from(local.minute().unsigned_abs()))
|
||||
* 60
|
||||
+ u64::from(local.second().unsigned_abs()))
|
||||
* MICROS_PER_SECOND
|
||||
+ u64::from(local.subsec_nanosecond().unsigned_abs()) / 1_000;
|
||||
if self.start < self.end {
|
||||
self.start <= micros && micros < self.end
|
||||
} else {
|
||||
micros >= self.start || micros < self.end
|
||||
}
|
||||
}
|
||||
|
||||
/// The window's end as `{end:%H:%M}`.
|
||||
pub fn end_hhmm(&self) -> String {
|
||||
let minutes = self.end / MICROS_PER_SECOND / 60;
|
||||
format!("{:02}:{:02}", minutes / 60, minutes % 60)
|
||||
}
|
||||
}
|
||||
|
||||
/// A deployment's access windows.
|
||||
pub fn access_windows(deployment: &Deployment) -> Vec<AccessWindow> {
|
||||
deployment
|
||||
.model_info
|
||||
.raw
|
||||
.get("access_windows")
|
||||
.and_then(Value::as_array)
|
||||
.map(|windows| windows.iter().filter_map(AccessWindow::parse).collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// `_reservation_blocking_window`: the first active window, unless the caller's team is named
|
||||
/// by one of the active windows.
|
||||
fn blocking_window(
|
||||
deployment: &Deployment,
|
||||
team_id: Option<&str>,
|
||||
now: Timestamp,
|
||||
) -> Option<AccessWindow> {
|
||||
let active: Vec<AccessWindow> = access_windows(deployment)
|
||||
.into_iter()
|
||||
.filter(|window| window.is_active(now))
|
||||
.collect();
|
||||
if let Some(team) = team_id
|
||||
&& active
|
||||
.iter()
|
||||
.any(|window| window.team_ids.iter().any(|id| id == team))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
active.into_iter().next()
|
||||
}
|
||||
|
||||
/// `filter_reserved_deployments`: the deployments no active window withholds from the caller,
|
||||
/// and the first window that withheld one.
|
||||
pub fn filter_reserved(
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
team_id: Option<&str>,
|
||||
now: Timestamp,
|
||||
) -> (Vec<Arc<Deployment>>, Option<AccessWindow>) {
|
||||
let mut blocking = None;
|
||||
let kept = deployments
|
||||
.into_iter()
|
||||
.filter(
|
||||
|deployment| match blocking_window(deployment, team_id, now) {
|
||||
None => true,
|
||||
Some(window) => {
|
||||
blocking.get_or_insert(window);
|
||||
false
|
||||
}
|
||||
},
|
||||
)
|
||||
.collect();
|
||||
(kept, blocking)
|
||||
}
|
||||
|
||||
/// `_filter_deployments_by_model_access_groups`: a caller whose allowed models (`models` plus
|
||||
/// `team_models`) reach `model` only through access groups is kept to the deployments in those
|
||||
/// groups. `model_groups` are the access groups the rows listed for `model` carry.
|
||||
pub fn filter_by_access_groups(
|
||||
model: &str,
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
allowed_models: &[String],
|
||||
model_groups: impl FnOnce() -> BTreeSet<String>,
|
||||
) -> Vec<Arc<Deployment>> {
|
||||
if deployments.is_empty() || allowed_models.is_empty() {
|
||||
return deployments;
|
||||
}
|
||||
let allowed = |name: &str| allowed_models.iter().any(|m| m == name);
|
||||
if allowed(model) || allowed("*") || allowed("all-proxy-models") {
|
||||
return deployments;
|
||||
}
|
||||
let groups: BTreeSet<String> = model_groups()
|
||||
.into_iter()
|
||||
.filter(|group| allowed(group))
|
||||
.collect();
|
||||
if groups.is_empty() {
|
||||
return deployments;
|
||||
}
|
||||
deployments
|
||||
.into_iter()
|
||||
.filter(|d| {
|
||||
d.model_info
|
||||
.access_groups
|
||||
.iter()
|
||||
.any(|g| groups.contains(g))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::minutes(json!("09:30"), Some(34_200_000_000))]
|
||||
#[case::seconds(json!("09:30:15"), Some(34_215_000_000))]
|
||||
#[case::fraction(json!("00:00:00.5"), Some(500_000))]
|
||||
#[case::number(json!(60), Some(60_000_000))]
|
||||
#[case::bad(json!("25:00"), None)]
|
||||
fn wall_clock_reads_pydantic_times(#[case] raw: Value, #[case] expected: Option<u64>) {
|
||||
assert_eq!(wall_clock(&raw), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::inside("2026-01-01T10:00:00Z", true)]
|
||||
#[case::at_end("2026-01-01T17:00:00Z", false)]
|
||||
#[case::before("2026-01-01T08:59:59Z", false)]
|
||||
fn a_window_is_active_inside_its_wall_clock(#[case] now: &str, #[case] expected: bool) {
|
||||
let window = AccessWindow {
|
||||
start: wall_clock(&json!("09:00")).unwrap(),
|
||||
end: wall_clock(&json!("17:00")).unwrap(),
|
||||
timezone: "UTC".to_owned(),
|
||||
team_ids: vec!["t".to_owned()],
|
||||
};
|
||||
assert_eq!(window.is_active(now.parse().unwrap()), expected);
|
||||
assert_eq!(window.end_hhmm(), "17:00");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_overnight_window_wraps_past_midnight() {
|
||||
let window = AccessWindow {
|
||||
start: wall_clock(&json!("22:00")).unwrap(),
|
||||
end: wall_clock(&json!("02:00")).unwrap(),
|
||||
timezone: "UTC".to_owned(),
|
||||
team_ids: vec!["t".to_owned()],
|
||||
};
|
||||
assert!(window.is_active("2026-01-01T23:00:00Z".parse().unwrap()));
|
||||
assert!(window.is_active("2026-01-01T01:00:00Z".parse().unwrap()));
|
||||
assert!(!window.is_active("2026-01-01T12:00:00Z".parse().unwrap()));
|
||||
}
|
||||
}
|
||||
386
litellm-rust/crates/router-core/src/build.rs
Normal file
386
litellm-rust/crates/router-core/src/build.rs
Normal file
|
|
@ -0,0 +1,386 @@
|
|||
//! The deployment build: how configured entries become the registry's deployments.
|
||||
//!
|
||||
//! The core owns the loop (`set_model_list`, `_create_deployment` and `_add_deployment`): which
|
||||
//! values are secret references, the id of a deployment that declares none, one deployment per
|
||||
//! `organization`, and the order everything happens in. What needs the embedding process is a
|
||||
//! host step, asked for one deployment at a time so an error surfaces exactly where the Python
|
||||
//! router raises it:
|
||||
//! - [`Host::secrets`]: the `os.environ/` references the entry holds;
|
||||
//! - [`Host::validate`]: the host's checks, and the entry as it normalizes it;
|
||||
//! - [`Host::provider`]: the provider that serves it;
|
||||
//! - [`Host::added`]: the host's side effects, and the entry as the registry stores it.
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
error::{ConfigError, Error},
|
||||
host::{BuildAsk, BuildPath, Built, Host, Provided, SecretAt, SecretRef},
|
||||
registry::{DeploymentSpec, deployment_id},
|
||||
};
|
||||
|
||||
/// The `litellm_params` keys whose `os.environ/` value is a secret pointer the identity source
|
||||
/// dereferences at use time, so the build leaves it as written. The router's copy of
|
||||
/// `_WIF_POINTER_FIELDS` (litellm/types/router.py), which the proxy's config loader also reads.
|
||||
pub const SECRET_POINTER_PARAMS: &[&str] = &[
|
||||
"anthropic_issuer_signing_key_ref",
|
||||
"anthropic_keycloak_client_secret_ref",
|
||||
];
|
||||
|
||||
const SECRET_PREFIX: &str = "os.environ/";
|
||||
|
||||
/// How a build's deployments join the registry.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum Install {
|
||||
/// They become every deployment (`set_model_list`).
|
||||
Replace,
|
||||
/// They are added at the end (a deployment retried once its provider resolves).
|
||||
Append,
|
||||
/// Each replaces the deployment with its id, if any, and is added at the end
|
||||
/// (`add_deployment`, an upsert).
|
||||
Upsert,
|
||||
}
|
||||
|
||||
/// One deployment the build accepted, as the registry stores it.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct BuiltDeployment {
|
||||
/// The entry index in the list the build was given.
|
||||
pub index: usize,
|
||||
pub spec: DeploymentSpec,
|
||||
}
|
||||
|
||||
/// Build `entries` (model_list entries) through `host`, in order. `resolve_secrets` resolves
|
||||
/// top-level `os.environ/` references first, as `set_model_list` does; `add_deployment` and
|
||||
/// upserts take their params as given. `install` takes each deployment as soon as it is
|
||||
/// accepted, so it joins the registry before the next one is built, as the Python router adds
|
||||
/// each in turn.
|
||||
pub async fn build(
|
||||
host: &dyn Host,
|
||||
path: BuildPath,
|
||||
entries: Vec<Map<String, Value>>,
|
||||
resolve_secrets: bool,
|
||||
mut install: impl FnMut(&BuiltDeployment) -> Result<(), Error> + Send,
|
||||
) -> Result<Vec<BuiltDeployment>, Error> {
|
||||
let declared: Vec<Option<String>> = entries.iter().map(declared_id).collect();
|
||||
let mut built = Vec::new();
|
||||
for (index, mut entry) in entries.into_iter().enumerate() {
|
||||
let model_name = text(entry.get("model_name"));
|
||||
let mut params = match entry.remove("litellm_params") {
|
||||
Some(Value::Object(params)) => params,
|
||||
_ => Map::new(),
|
||||
};
|
||||
if resolve_secrets {
|
||||
let refs = param_refs(¶ms);
|
||||
if !refs.is_empty() {
|
||||
let values = host.secrets(index, &refs).await?;
|
||||
for (secret, value) in refs.into_iter().zip(values) {
|
||||
if let SecretAt::Param(key) = secret.at {
|
||||
params.insert(key, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut info = match entry.remove("model_info") {
|
||||
Some(Value::Object(info)) => info,
|
||||
_ => Map::new(),
|
||||
};
|
||||
let declared_id = declared[index].clone();
|
||||
if !info.contains_key("id") {
|
||||
info.insert(
|
||||
"id".to_owned(),
|
||||
Value::String(deployment_id(&model_name, ¶ms)),
|
||||
);
|
||||
}
|
||||
// `set_model_list` adds one deployment per organization of a list, each with the
|
||||
// organization set on the same params (so every one keeps the list's id).
|
||||
let organizations = match (path, params.get("organization")) {
|
||||
(BuildPath::Create, Some(Value::Array(organizations))) => organizations.clone(),
|
||||
_ => vec![Value::Null],
|
||||
};
|
||||
let duplicate_id = path == BuildPath::Create
|
||||
&& declared_id.as_ref().is_some_and(|id| {
|
||||
declared
|
||||
.iter()
|
||||
.filter(|other| other.as_ref() == Some(id))
|
||||
.count()
|
||||
> 1
|
||||
});
|
||||
for organization in organizations {
|
||||
if !organization.is_null() {
|
||||
params.insert("organization".to_owned(), organization);
|
||||
}
|
||||
let mut deployment = entry.clone();
|
||||
deployment.insert("model_name".to_owned(), Value::String(model_name.clone()));
|
||||
deployment.insert("litellm_params".to_owned(), Value::Object(params.clone()));
|
||||
deployment.insert("model_info".to_owned(), Value::Object(info.clone()));
|
||||
if let Some(spec) = build_one(
|
||||
host,
|
||||
path,
|
||||
index,
|
||||
&deployment,
|
||||
declared_id.as_deref(),
|
||||
duplicate_id,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
let deployment = BuiltDeployment { index, spec };
|
||||
install(&deployment)?;
|
||||
built.push(deployment);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(built)
|
||||
}
|
||||
|
||||
/// One deployment's steps: validate, provider, the data-source references, added.
|
||||
async fn build_one(
|
||||
host: &dyn Host,
|
||||
path: BuildPath,
|
||||
index: usize,
|
||||
entry: &Map<String, Value>,
|
||||
declared_id: Option<&str>,
|
||||
duplicate_id: bool,
|
||||
) -> Result<Option<DeploymentSpec>, Error> {
|
||||
let ask = |entry| BuildAsk {
|
||||
path,
|
||||
index,
|
||||
entry,
|
||||
declared_id,
|
||||
duplicate_id,
|
||||
};
|
||||
let validated = match host.validate(&ask(entry)).await? {
|
||||
Built::Entry(validated) => validated,
|
||||
Built::Skip | Built::Invalid => return Ok(None),
|
||||
};
|
||||
let provider = match host.provider(&ask(&validated)).await? {
|
||||
Provided::Provider(provider) => provider,
|
||||
Provided::Invalid => return Ok(None),
|
||||
};
|
||||
let refs = data_source_refs(&validated);
|
||||
if !refs.is_empty() {
|
||||
let _ = host.secrets(index, &refs).await?;
|
||||
}
|
||||
let stored = match host.added(&ask(&validated), provider.as_deref()).await? {
|
||||
Built::Entry(stored) => stored,
|
||||
Built::Skip | Built::Invalid => return Ok(None),
|
||||
};
|
||||
let model_name = text(stored.get("model_name"));
|
||||
let mut spec: DeploymentSpec = serde_json::from_value(Value::Object(stored))
|
||||
.map_err(|source| ConfigError::InvalidDeployment { model_name, source })?;
|
||||
spec.provider = provider;
|
||||
Ok(Some(spec))
|
||||
}
|
||||
|
||||
/// The top-level `litellm_params` values `set_model_list` resolves.
|
||||
fn param_refs(params: &Map<String, Value>) -> Vec<SecretRef> {
|
||||
params
|
||||
.iter()
|
||||
.filter(|(key, _)| !SECRET_POINTER_PARAMS.contains(&key.as_str()))
|
||||
.filter_map(|(key, value)| {
|
||||
let name = value
|
||||
.as_str()
|
||||
.filter(|name| name.starts_with(SECRET_PREFIX))?;
|
||||
Some(SecretRef {
|
||||
name: name.to_owned(),
|
||||
at: SecretAt::Param(key.clone()),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The `dataSources[].parameters.endpoint` / `key` references `_add_deployment` reads from the
|
||||
/// environment.
|
||||
fn data_source_refs(entry: &Map<String, Value>) -> Vec<SecretRef> {
|
||||
let sources = entry
|
||||
.get("litellm_params")
|
||||
.and_then(|params| params.get("dataSources"))
|
||||
.and_then(Value::as_array);
|
||||
let mut refs = Vec::new();
|
||||
for (source, data_source) in sources.into_iter().flatten().enumerate() {
|
||||
let Some(parameters) = data_source.get("parameters").and_then(Value::as_object) else {
|
||||
continue;
|
||||
};
|
||||
for key in ["endpoint", "key"] {
|
||||
if let Some(name) = parameters
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|name| name.starts_with(SECRET_PREFIX))
|
||||
{
|
||||
refs.push(SecretRef {
|
||||
name: name.to_owned(),
|
||||
at: SecretAt::DataSource {
|
||||
source,
|
||||
key: key.to_owned(),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
refs
|
||||
}
|
||||
|
||||
/// The `model_info.id` an entry declares, as Python's `str()` spells it.
|
||||
fn declared_id(entry: &Map<String, Value>) -> Option<String> {
|
||||
match entry.get("model_info")?.get("id")? {
|
||||
Value::Null => None,
|
||||
Value::String(id) => Some(id.clone()),
|
||||
other => Some(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// A name as Python's `str()` spells it.
|
||||
fn text(value: Option<&Value>) -> String {
|
||||
match value {
|
||||
Some(Value::String(text)) => text.clone(),
|
||||
Some(Value::Null) | None => "None".to_owned(),
|
||||
Some(other) => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
cache::{InMemoryRouterCache, RouterCache},
|
||||
error::AttemptError,
|
||||
host::{Attempt, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage},
|
||||
registry::Deployment,
|
||||
};
|
||||
|
||||
/// Records every build step and resolves `os.environ/X` to `secret-X`.
|
||||
#[derive(Default)]
|
||||
struct Recording {
|
||||
steps: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Host for Recording {
|
||||
async fn invoke(&self, _attempt: &Attempt) -> Result<Output, AttemptError> {
|
||||
unreachable!("the build runs no attempt")
|
||||
}
|
||||
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,
|
||||
) -> Result<bool, crate::error::AttemptError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn sleep(&self, _duration: std::time::Duration) {}
|
||||
fn observe(&self, _event: &RouterEvent) {}
|
||||
fn cache(&self) -> Arc<dyn RouterCache> {
|
||||
Arc::new(InMemoryRouterCache::default())
|
||||
}
|
||||
async fn secrets(&self, _index: usize, refs: &[SecretRef]) -> Result<Vec<Value>, Error> {
|
||||
self.steps.lock().unwrap().push(format!("secrets {refs:?}"));
|
||||
Ok(refs
|
||||
.iter()
|
||||
.map(|r| json!(format!("secret-{}", &r.name[SECRET_PREFIX.len()..])))
|
||||
.collect())
|
||||
}
|
||||
async fn validate(&self, ask: &BuildAsk<'_>) -> Result<Built, Error> {
|
||||
self.steps.lock().unwrap().push(format!(
|
||||
"validate {} {}",
|
||||
ask.index, ask.entry["litellm_params"]
|
||||
));
|
||||
Ok(Built::Entry(ask.entry.clone()))
|
||||
}
|
||||
async fn provider(&self, ask: &BuildAsk<'_>) -> Result<Provided, Error> {
|
||||
self.steps
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("provider {}", ask.index));
|
||||
Ok(Provided::Provider(Some("azure".to_owned())))
|
||||
}
|
||||
async fn added(&self, ask: &BuildAsk<'_>, provider: Option<&str>) -> Result<Built, Error> {
|
||||
self.steps
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(format!("added {} {provider:?}", ask.index));
|
||||
Ok(Built::Entry(ask.entry.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
fn entries(value: Value) -> Vec<Map<String, Value>> {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_resolves_secrets_then_builds_one_deployment_per_organization() {
|
||||
let host = Recording::default();
|
||||
let built = build(
|
||||
&host,
|
||||
BuildPath::Create,
|
||||
entries(json!([{
|
||||
"model_name": "g",
|
||||
"litellm_params": {
|
||||
"model": "azure/x",
|
||||
"api_key": "os.environ/KEY",
|
||||
"anthropic_issuer_signing_key_ref": "os.environ/POINTER",
|
||||
"organization": ["a", "b"],
|
||||
},
|
||||
}])),
|
||||
true,
|
||||
|_| Ok(()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let steps = host.steps.lock().unwrap().clone();
|
||||
assert_eq!(steps.len(), 7, "{steps:?}");
|
||||
assert!(steps[0].starts_with("secrets") && steps[0].contains("KEY"));
|
||||
assert!(!steps[0].contains("POINTER"));
|
||||
assert!(steps[1].contains(r#""organization":"a""#));
|
||||
assert!(steps[1].contains("secret-KEY"));
|
||||
assert_eq!(steps[3], "added 0 Some(\"azure\")");
|
||||
assert!(steps[4].contains(r#""organization":"b""#));
|
||||
assert_eq!(built.len(), 2);
|
||||
assert_eq!(built[0].spec.provider.as_deref(), Some("azure"));
|
||||
assert_eq!(
|
||||
built[0].spec.model_info["id"],
|
||||
built[1].spec.model_info["id"]
|
||||
);
|
||||
let deployment = Deployment::from_spec(built[0].spec.clone()).unwrap();
|
||||
assert_eq!(deployment.provider(), "azure");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn add_leaves_params_as_given_and_asks_for_data_source_references() {
|
||||
let host = Recording::default();
|
||||
let built = build(
|
||||
&host,
|
||||
BuildPath::Add,
|
||||
entries(json!([{
|
||||
"model_name": "g",
|
||||
"litellm_params": {
|
||||
"model": "azure/x",
|
||||
"api_key": "os.environ/KEY",
|
||||
"organization": ["a", "b"],
|
||||
"dataSources": [{"parameters": {"endpoint": "os.environ/END", "key": "plain"}}],
|
||||
},
|
||||
"model_info": {"id": "given"},
|
||||
}])),
|
||||
false,
|
||||
|_| Ok(()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let steps = host.steps.lock().unwrap().clone();
|
||||
assert_eq!(steps.len(), 4, "{steps:?}");
|
||||
assert!(steps[0].contains("os.environ/KEY"));
|
||||
assert!(steps[2].contains("DataSource { source: 0, key: \"endpoint\" }"));
|
||||
assert_eq!(built.len(), 1);
|
||||
assert_eq!(built[0].spec.model_info["id"], "given");
|
||||
}
|
||||
}
|
||||
159
litellm-rust/crates/router-core/src/cache.rs
Normal file
159
litellm-rust/crates/router-core/src/cache.rs
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
//! [`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>;
|
||||
|
||||
/// Read this worker's in-memory tier only (`local_only=True`).
|
||||
async fn async_get_local(&self, key: &str) -> Result<Option<Value>, Error> {
|
||||
self.async_get_cache(key).await
|
||||
}
|
||||
|
||||
/// Add to a counter in this worker's in-memory tier only (`local_only=True`).
|
||||
async fn async_increment_local(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: f64,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<f64, Error> {
|
||||
self.async_increment(key, amount, ttl).await
|
||||
}
|
||||
|
||||
/// `_increment_deployment_usage`'s write: add each `(key, amount)` in one pipeline. With a
|
||||
/// shared tier, this worker's in-memory tier then holds the shared totals, so its reads see
|
||||
/// what other workers counted up to now.
|
||||
async fn async_increment_usage(
|
||||
&self,
|
||||
increments: &[(String, f64)],
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
for (key, amount) in increments {
|
||||
self.async_increment(key, *amount, ttl).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[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(())
|
||||
}
|
||||
}
|
||||
1555
litellm-rust/crates/router-core/src/call.rs
Normal file
1555
litellm-rust/crates/router-core/src/call.rs
Normal file
File diff suppressed because it is too large
Load diff
456
litellm-rust/crates/router-core/src/call/stream.rs
Normal file
456
litellm-rust/crates/router-core/src/call/stream.rs
Normal file
|
|
@ -0,0 +1,456 @@
|
|||
//! Recovering a stream that failed after its attempt returned it. The host wraps every streamed
|
||||
//! response; when the stream fails before the caller received content, the wrapper asks the core
|
||||
//! what may still take over and hands the failure back here. Mirrors the decisions of the Python
|
||||
//! router's stream wrappers: `_stream_can_fall_back`, `_refusal_fallback_available`,
|
||||
//! `_anthropic_messages_stream_can_retry`, `_stream_retry_budget`, `_astream_retry_same_group`,
|
||||
//! and the `async_function_with_fallbacks_common_utils` re-entry of the mid-stream fallback.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::{Failed, FallbackChain, FallbackHop, Lists, Target, Tried, group_members};
|
||||
use crate::{
|
||||
Router, RouterStateSnapshot,
|
||||
config::Op,
|
||||
error::{AttemptError, Error, FailureKind},
|
||||
fallback::FallbackTarget,
|
||||
health::should_retry_status,
|
||||
host::{RouterEvent, RouterRequest, UncountedRetries},
|
||||
select::{Candidates, SelectInputs},
|
||||
};
|
||||
|
||||
/// Where a streamed response came from: the attempt that returned it.
|
||||
#[derive(Clone, Debug, Default, Deserialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct StreamOrigin {
|
||||
/// The name the attempt served: the requested group, or the fallback target.
|
||||
pub group: String,
|
||||
/// The `order` level an order fallback targeted.
|
||||
pub order: Option<i64>,
|
||||
/// `fallback_depth` of the attempt.
|
||||
pub depth: u32,
|
||||
pub level: u32,
|
||||
/// The deployment the attempt ran on.
|
||||
pub deployment_id: Option<String>,
|
||||
/// Same-group stream retries this request already spent (`attempted_retries_for_request`).
|
||||
pub attempted_retries: u32,
|
||||
/// The budget the first stream retry committed to (`committed_retry_budget_for_request`).
|
||||
pub committed_budget: Option<u32>,
|
||||
}
|
||||
|
||||
/// What may still take over a stream that fails before content.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub struct StreamPolicy {
|
||||
/// `_anthropic_messages_stream_can_retry`.
|
||||
pub can_retry: bool,
|
||||
/// `_stream_can_fall_back`.
|
||||
pub can_fall_back: bool,
|
||||
/// `_refusal_fallback_available`.
|
||||
pub refusal_fallback: bool,
|
||||
/// `_stream_retry_budget` for the failure kind asked about: the budget, and whether a retry
|
||||
/// policy governs the retry.
|
||||
pub retry_budget: Option<(u32, bool)>,
|
||||
}
|
||||
|
||||
impl StreamOrigin {
|
||||
fn target(&self) -> Target {
|
||||
Target {
|
||||
group: self.group.clone(),
|
||||
order: self.order,
|
||||
params: None,
|
||||
excluded: Arc::from([]),
|
||||
failover_excluded: Arc::from([]),
|
||||
}
|
||||
}
|
||||
|
||||
/// `fallback_lookup_groups`: the served group, then the group the request was for.
|
||||
fn lookup<'a>(&'a self, request: &'a RouterRequest) -> Vec<&'a str> {
|
||||
let mut lookup = vec![self.group.as_str()];
|
||||
if request.model != self.group {
|
||||
lookup.push(request.model.as_str());
|
||||
}
|
||||
lookup
|
||||
}
|
||||
|
||||
/// `has_unattempted_fallback_target`. The first attempt carries no attempted targets, so any
|
||||
/// resolved chain counts; a fallback hop has already tried its own group and the original one.
|
||||
fn has_unattempted(&self, request: &RouterRequest, chain: Option<Vec<String>>) -> bool {
|
||||
let Some(chain) = chain else {
|
||||
return false;
|
||||
};
|
||||
if self.depth == 0 {
|
||||
return true;
|
||||
}
|
||||
chain
|
||||
.iter()
|
||||
.any(|target| *target != self.group && *target != request.model)
|
||||
}
|
||||
}
|
||||
|
||||
/// The group names a chain targets (a dict target counts by its `model`).
|
||||
fn names(chain: Option<Vec<FallbackTarget>>) -> Option<Vec<String>> {
|
||||
chain.map(|targets| {
|
||||
targets
|
||||
.iter()
|
||||
.filter_map(|target| target.model().map(str::to_owned))
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
|
||||
/// The `content_policy_fallbacks` chain for the lookup groups: the request's list where it
|
||||
/// passed one, else the router's (`_has_content_policy_fallback`).
|
||||
fn content_policy_targets(lists: &Lists<'_>, lookup: &[&str]) -> Option<Option<Vec<String>>> {
|
||||
lists
|
||||
.content_policy
|
||||
.map(|list| names(list.lookup_groups(lookup, lists.keys)))
|
||||
}
|
||||
|
||||
/// The router's `default_fallbacks`, which Python folds into its `fallbacks` as `{"*": [...]}`: a
|
||||
/// request's own list replaces them.
|
||||
fn folded_defaults<'a>(snapshot: &'a RouterStateSnapshot, request: &RouterRequest) -> &'a [String] {
|
||||
match request.fallbacks {
|
||||
Some(_) => &[],
|
||||
None => &snapshot.settings.dynamic.default_fallbacks,
|
||||
}
|
||||
}
|
||||
|
||||
/// The regular fallback targets for the lookup groups, then the folded `default_fallbacks`.
|
||||
fn regular_targets(
|
||||
snapshot: &RouterStateSnapshot,
|
||||
request: &RouterRequest,
|
||||
lists: &Lists<'_>,
|
||||
lookup: &[&str],
|
||||
) -> Option<Vec<String>> {
|
||||
let defaults = folded_defaults(snapshot, request);
|
||||
names(
|
||||
lists
|
||||
.fallbacks
|
||||
.and_then(|list| list.lookup_groups(lookup, lists.keys)),
|
||||
)
|
||||
.or_else(|| (!defaults.is_empty()).then(|| defaults.to_vec()))
|
||||
}
|
||||
|
||||
/// `_has_default_fallbacks`: the router's list, whatever the request passed.
|
||||
fn has_default_fallbacks(snapshot: &RouterStateSnapshot) -> bool {
|
||||
let dynamic = &snapshot.settings.dynamic;
|
||||
dynamic.fallbacks.has_star() || !dynamic.default_fallbacks.is_empty()
|
||||
}
|
||||
|
||||
impl Router {
|
||||
/// What may still take over a stream from `origin` that fails before content, and the retry
|
||||
/// budget for a failure of `kind`.
|
||||
pub fn stream_policy(
|
||||
&self,
|
||||
request: &RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
kind: Option<FailureKind>,
|
||||
) -> StreamPolicy {
|
||||
let snapshot = self.core.snapshot.load_full();
|
||||
StreamPolicy {
|
||||
can_retry: stream_can_retry(&snapshot, request, origin),
|
||||
can_fall_back: stream_can_fall_back(&snapshot, request, origin),
|
||||
refusal_fallback: refusal_fallback_available(&snapshot, request, origin),
|
||||
retry_budget: kind.map(|kind| stream_retry_budget(&snapshot, request, origin, kind)),
|
||||
}
|
||||
}
|
||||
|
||||
/// The fallback chain for a stream from `origin` that failed with `failure` before content:
|
||||
/// `async_function_with_fallbacks_common_utils` re-entered from the stream wrapper, so the
|
||||
/// group's own retry loop does not run again.
|
||||
pub async fn fallback_stream(
|
||||
&self,
|
||||
op: Op,
|
||||
request: RouterRequest,
|
||||
origin: StreamOrigin,
|
||||
failure: AttemptError,
|
||||
) -> Result<super::RouterResponse, Error> {
|
||||
self.live()?;
|
||||
let snapshot = self.core.snapshot.load_full();
|
||||
let request = Arc::new(request);
|
||||
let mut chain = FallbackChain::default();
|
||||
if request.model != origin.group {
|
||||
chain
|
||||
.attempted
|
||||
.push(Tried::Target(FallbackTarget::Name(request.model.clone())));
|
||||
}
|
||||
self.core
|
||||
.fall_back(
|
||||
&snapshot,
|
||||
&op,
|
||||
&request,
|
||||
origin.target(),
|
||||
FallbackHop {
|
||||
depth: origin.depth,
|
||||
level: origin.level,
|
||||
scope: 0,
|
||||
},
|
||||
&mut chain,
|
||||
Failed::from(Error::Attempt(failure)),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// `_astream_retry_same_group`: re-run the attempt in the stream's own group until a retry
|
||||
/// returns a response or the budget runs out. The failed result carries the newest failure.
|
||||
pub async fn retry_stream(
|
||||
&self,
|
||||
op: Op,
|
||||
request: RouterRequest,
|
||||
origin: StreamOrigin,
|
||||
failure: AttemptError,
|
||||
) -> Result<super::RouterResponse, Error> {
|
||||
self.live()?;
|
||||
let snapshot = self.core.snapshot.load_full();
|
||||
let request = Arc::new(request);
|
||||
let target = origin.target();
|
||||
let at = FallbackHop {
|
||||
depth: origin.depth,
|
||||
level: origin.level,
|
||||
scope: 0,
|
||||
};
|
||||
let counting = SelectInputs {
|
||||
count_only: true,
|
||||
..SelectInputs::default()
|
||||
};
|
||||
let counted = |c: &Candidates, n: usize| if c.kind.is_single() { 0 } else { n };
|
||||
let (healthy, all) = self
|
||||
.core
|
||||
.candidates(&snapshot, &op, &request, &target.group, &counting)
|
||||
.await
|
||||
.map(|c| (counted(&c, c.healthy.len()), counted(&c, c.all.len())))
|
||||
.unwrap_or((0, 0));
|
||||
let (budget, policy_applies) =
|
||||
stream_retry_budget(&snapshot, &request, &origin, failure.kind);
|
||||
let inputs = SelectInputs {
|
||||
target_order: target.order,
|
||||
requested: Some(target.group.clone()),
|
||||
fallback_hop: at.level > 0,
|
||||
..SelectInputs::default()
|
||||
};
|
||||
let mut last = failure;
|
||||
for attempt in origin.attempted_retries..budget {
|
||||
if last.agentic_followup {
|
||||
break;
|
||||
}
|
||||
if !policy_applies
|
||||
&& !self
|
||||
.core
|
||||
.should_retry(&snapshot, &request, &last, healthy, all)
|
||||
{
|
||||
break;
|
||||
}
|
||||
let delay = self.core.time_to_sleep(
|
||||
&snapshot,
|
||||
last.retry_after,
|
||||
budget - attempt,
|
||||
budget,
|
||||
healthy,
|
||||
all,
|
||||
);
|
||||
self.core.host.sleep(delay).await;
|
||||
self.core.host.observe(&RouterEvent::Retry {
|
||||
model_group: target.group.clone(),
|
||||
retry: attempt + 1,
|
||||
delay,
|
||||
});
|
||||
match self
|
||||
.core
|
||||
.attempt(
|
||||
&snapshot,
|
||||
&op,
|
||||
&request,
|
||||
&target,
|
||||
&inputs,
|
||||
(attempt + 1, budget),
|
||||
(at, &[]),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(mut response) => {
|
||||
response.attempted_retries = attempt + 1;
|
||||
response.max_retries = Some(budget);
|
||||
return Ok(response);
|
||||
}
|
||||
Err(failed) => match *failed.error {
|
||||
// `_stream_fallback_error_for_raised`: a status the pre-stream retry rule
|
||||
// would not retry ends the retries with that failure.
|
||||
Error::Attempt(failure) => {
|
||||
let stop = failure.status.is_some_and(|s| !should_retry_status(s));
|
||||
last = failure;
|
||||
if stop {
|
||||
break;
|
||||
}
|
||||
}
|
||||
other => return Err(other),
|
||||
},
|
||||
}
|
||||
}
|
||||
Err(Error::Attempt(last))
|
||||
}
|
||||
}
|
||||
|
||||
/// The request's `num_retries`, then the routed deployment's, then the router's
|
||||
/// (`_anthropic_messages_plain_retry_budget`).
|
||||
fn plain_retry_budget(
|
||||
snapshot: &RouterStateSnapshot,
|
||||
request: &RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
) -> u32 {
|
||||
if let Some(retries) = request.num_retries {
|
||||
return retries;
|
||||
}
|
||||
if request.uncounted_retries == Some(UncountedRetries::Negative) {
|
||||
return 0;
|
||||
}
|
||||
let deployment = origin
|
||||
.deployment_id
|
||||
.as_deref()
|
||||
.and_then(|id| snapshot.registry.deployment(id));
|
||||
let configured = deployment.and_then(|d| match d.raw_params.get("num_retries") {
|
||||
Some(serde_json::Value::Number(n)) => n.as_u64().and_then(|n| u32::try_from(n).ok()),
|
||||
Some(serde_json::Value::String(s)) if s.chars().all(|c| c.is_ascii_digit()) => {
|
||||
s.parse().ok()
|
||||
}
|
||||
_ => None,
|
||||
});
|
||||
configured.unwrap_or(snapshot.settings.num_retries)
|
||||
}
|
||||
|
||||
/// The retry policy for the request's model group (`_anthropic_messages_resolved_retry_policy`),
|
||||
/// unless the request opted out with `num_retries=0`. A policy that does not read as one governs
|
||||
/// nothing here, with no fall through to the router's `retry_policy`.
|
||||
fn retry_policy<'a>(
|
||||
snapshot: &'a RouterStateSnapshot,
|
||||
request: &'a RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
) -> Option<&'a crate::config::RetryPolicy> {
|
||||
if request.num_retries == Some(0) || request.retry_policy_raises.contains(&origin.group) {
|
||||
return None;
|
||||
}
|
||||
let settings = &snapshot.settings;
|
||||
crate::retry::resolve_retry_policy(
|
||||
settings.retry_policy.as_ref(),
|
||||
&origin.group,
|
||||
Some(
|
||||
request
|
||||
.model_group_retry_policy
|
||||
.as_ref()
|
||||
.unwrap_or(&settings.model_group_retry_policy),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
fn stream_retry_budget(
|
||||
snapshot: &RouterStateSnapshot,
|
||||
request: &RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
kind: FailureKind,
|
||||
) -> (u32, bool) {
|
||||
let failure = AttemptError::new(kind, None);
|
||||
let policy_retries =
|
||||
retry_policy(snapshot, request, origin).and_then(|p| p.retries_for(&failure));
|
||||
if origin.attempted_retries > 0
|
||||
&& let Some(committed) = origin.committed_budget
|
||||
{
|
||||
return (committed, policy_retries.is_some());
|
||||
}
|
||||
match policy_retries {
|
||||
Some(retries) => (retries, true),
|
||||
None => (plain_retry_budget(snapshot, request, origin), false),
|
||||
}
|
||||
}
|
||||
|
||||
fn stream_can_retry(
|
||||
snapshot: &RouterStateSnapshot,
|
||||
request: &RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
) -> bool {
|
||||
let attempted = origin.attempted_retries;
|
||||
if attempted > 0
|
||||
&& let Some(committed) = origin.committed_budget
|
||||
{
|
||||
return committed > attempted;
|
||||
}
|
||||
let plain = plain_retry_budget(snapshot, request, origin);
|
||||
// `_retry_policy_ceiling`: every field the policy dumps.
|
||||
let ceiling = retry_policy(snapshot, request, origin).map_or(plain, |policy| {
|
||||
let granted = [
|
||||
policy.BadRequestErrorRetries,
|
||||
policy.AuthenticationErrorRetries,
|
||||
policy.TimeoutErrorRetries,
|
||||
policy.RateLimitErrorRetries,
|
||||
policy.ContentPolicyViolationErrorRetries,
|
||||
policy.InternalServerErrorRetries,
|
||||
policy.ServiceUnavailableErrorRetries,
|
||||
policy.NotFoundErrorRetries,
|
||||
policy.DefaultRetries,
|
||||
];
|
||||
plain.max(granted.into_iter().flatten().max().unwrap_or(0))
|
||||
});
|
||||
ceiling > attempted
|
||||
}
|
||||
|
||||
fn stream_can_fall_back(
|
||||
snapshot: &RouterStateSnapshot,
|
||||
request: &RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
) -> bool {
|
||||
if request.disable_fallbacks {
|
||||
return false;
|
||||
}
|
||||
if snapshot.settings.enable_weighted_failover {
|
||||
return true;
|
||||
}
|
||||
// `_stream_order_levels`, the levels the fallback walk sees for the group.
|
||||
let listing_team = request.listing_team_id.as_deref();
|
||||
let mut orders: Vec<i64> = group_members(snapshot, &origin.group, listing_team)
|
||||
.iter()
|
||||
.filter(|d| d.listed_for_team(listing_team))
|
||||
.filter_map(|d| d.order())
|
||||
.collect();
|
||||
orders.sort_unstable();
|
||||
orders.dedup();
|
||||
if orders.len() > 1 {
|
||||
let current = origin.order.unwrap_or(orders[0]);
|
||||
if orders.iter().any(|order| *order > current) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
let lists = Lists::of(snapshot, request);
|
||||
let lookup = origin.lookup(request);
|
||||
if let Some(chain) = content_policy_targets(&lists, &lookup)
|
||||
&& origin.has_unattempted(request, chain)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
if !lists.has_regular() && folded_defaults(snapshot, request).is_empty() {
|
||||
return false;
|
||||
}
|
||||
if lists
|
||||
.fallbacks
|
||||
.is_some_and(|list| list.is_client_side(lists.keys))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
origin.has_unattempted(request, regular_targets(snapshot, request, &lists, &lookup))
|
||||
}
|
||||
|
||||
fn refusal_fallback_available(
|
||||
snapshot: &RouterStateSnapshot,
|
||||
request: &RouterRequest,
|
||||
origin: &StreamOrigin,
|
||||
) -> bool {
|
||||
if request.disable_fallbacks {
|
||||
return false;
|
||||
}
|
||||
let lists = Lists::of(snapshot, request);
|
||||
let lookup = origin.lookup(request);
|
||||
if let Some(chain) = content_policy_targets(&lists, &lookup) {
|
||||
return origin.has_unattempted(request, chain);
|
||||
}
|
||||
if has_default_fallbacks(snapshot) {
|
||||
return true;
|
||||
}
|
||||
origin.has_unattempted(request, regular_targets(snapshot, request, &lists, &lookup))
|
||||
}
|
||||
355
litellm-rust/crates/router-core/src/config.rs
Normal file
355
litellm-rust/crates/router-core/src/config.rs
Normal file
|
|
@ -0,0 +1,355 @@
|
|||
//! 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::AttemptError,
|
||||
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: the Python router's endpoint method name (`acompletion`,
|
||||
/// `aembedding`, ...). The core runs every op alike; what differs per endpoint is the host's.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct Op(String);
|
||||
|
||||
impl Op {
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self(name.into())
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for Op {
|
||||
fn from(name: &str) -> Self {
|
||||
Self::new(name)
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Op {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// `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,
|
||||
/// `routing_strategy=None` passed explicitly (sent as `"None"`): stored as None, so no
|
||||
/// strategy arm matches and every pick finds no deployment.
|
||||
#[strum(serialize = "None")]
|
||||
Unset,
|
||||
#[strum(default, transparent)]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
impl RoutingStrategy {
|
||||
pub fn is_native(&self) -> bool {
|
||||
matches!(self, Self::SimpleShuffle | Self::Unset)
|
||||
}
|
||||
}
|
||||
|
||||
/// `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>,
|
||||
pub ServiceUnavailableErrorRetries: Option<u32>,
|
||||
pub NotFoundErrorRetries: Option<u32>,
|
||||
pub DefaultRetries: Option<u32>,
|
||||
}
|
||||
|
||||
/// `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>,
|
||||
pub NotFoundErrorAllowedFails: Option<u32>,
|
||||
}
|
||||
|
||||
impl AllowedFailsPolicy {
|
||||
/// `_core_get_allowed_fails_from_policy`: the first set field whose exception class the
|
||||
/// failure is an instance of, in Python's order.
|
||||
pub fn allowed_fails_for(&self, failure: &AttemptError) -> Option<u32> {
|
||||
[
|
||||
("AuthenticationError", self.AuthenticationErrorAllowedFails),
|
||||
("Timeout", self.TimeoutErrorAllowedFails),
|
||||
("RateLimitError", self.RateLimitErrorAllowedFails),
|
||||
(
|
||||
"ContentPolicyViolationError",
|
||||
self.ContentPolicyViolationErrorAllowedFails,
|
||||
),
|
||||
("BadRequestError", self.BadRequestErrorAllowedFails),
|
||||
("InternalServerError", self.InternalServerErrorAllowedFails),
|
||||
(
|
||||
"ServiceUnavailableError",
|
||||
self.ServiceUnavailableErrorAllowedFails,
|
||||
),
|
||||
("BadGatewayError", self.BadGatewayErrorAllowedFails),
|
||||
("NotFoundError", self.NotFoundErrorAllowedFails),
|
||||
]
|
||||
.into_iter()
|
||||
.find_map(|(class, allowed)| allowed.filter(|_| failure.is_a(class)))
|
||||
}
|
||||
}
|
||||
|
||||
pub use crate::fallback::{FallbackKeys, Fallbacks};
|
||||
|
||||
/// 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>,
|
||||
/// `_model_to_group` as Python last built it (member model group to routing group), set
|
||||
/// only when a raw `routing_groups` write replaced the groups since: member calls keep
|
||||
/// routing by it, as the Python router's stale map does.
|
||||
pub routing_group_members: Option<BTreeMap<String, String>>,
|
||||
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>,
|
||||
/// `model_group_retry_policy`: a group's own policy, which wins over `retry_policy`.
|
||||
pub model_group_retry_policy: BTreeMap<String, 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,
|
||||
/// `Router(plugins=[...])` is set: the host runs the routing-plugin pipeline.
|
||||
pub routing_plugins: bool,
|
||||
/// The host's `CustomLogger` deployment filters run on every async candidate list
|
||||
/// (`async_callback_filter_deployments`).
|
||||
pub callback_filters: bool,
|
||||
pub enable_health_check_routing: bool,
|
||||
/// How old (seconds) a background health check result may be and still count
|
||||
/// (`health_check_staleness_threshold`, else the default interval times its multiplier).
|
||||
pub health_check_staleness_threshold: Option<f64>,
|
||||
/// `background_health_check_model_groups`: when set, only these groups are health-filtered.
|
||||
pub background_health_check_model_groups: Option<Vec<String>>,
|
||||
/// `enable_weighted_failover`: a failed deployment's group is re-picked without it before
|
||||
/// any cross-group fallback (simple-shuffle only).
|
||||
pub enable_weighted_failover: bool,
|
||||
/// `fallback_on_unsupported_params`; `None` is Python's default, `True`.
|
||||
pub fallback_on_unsupported_params: Option<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: `Host::supports` must accept each one 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`, as the router copied it when it was built.
|
||||
pub allowed_fails: Option<u32>,
|
||||
/// `allowed_fails` is only the process default: it was not passed to the router and still
|
||||
/// equals `litellm.allowed_fails`. Anything else switches cooldowns from the failure-ratio
|
||||
/// rule to the allowed-fails counter (`_is_allowed_fails_set_on_router`).
|
||||
pub allowed_fails_is_default: bool,
|
||||
/// Other spellings a fallback key matches.
|
||||
pub fallback_keys: FallbackKeys,
|
||||
/// `default_deployment`: what a name no deployment, alias, routing group, team route or
|
||||
/// wildcard serves is routed to, as a copy whose `litellm_params.model` is the name
|
||||
/// (`_try_early_resolve_deployments_for_model_not_in_names`).
|
||||
pub default_deployment: Option<DeploymentSpec>,
|
||||
/// The `default_deployment` is a value Python cannot copy for a name: resolving to it
|
||||
/// raises the copy's error ([`crate::Error::DefaultDeploymentRaises`]).
|
||||
pub default_deployment_raises: bool,
|
||||
}
|
||||
|
||||
/// `model_group_alias` as `resolve_model_group_alias` reads it: an alias maps to a non-empty
|
||||
/// group name, or to `{"model": <non-empty group name>, ...}`; any other value rewrites nothing.
|
||||
fn alias_targets<'de, D>(deserializer: D) -> Result<BTreeMap<String, String>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let aliases = BTreeMap::<String, serde_json::Value>::deserialize(deserializer)?;
|
||||
Ok(aliases
|
||||
.into_iter()
|
||||
.filter_map(|(alias, target)| {
|
||||
let group = match &target {
|
||||
serde_json::Value::Object(entry) => entry.get("model")?,
|
||||
other => other,
|
||||
};
|
||||
group
|
||||
.as_str()
|
||||
.filter(|group| !group.is_empty())
|
||||
.map(|group| (alias, group.to_owned()))
|
||||
})
|
||||
.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,
|
||||
allowed_fails_is_default: false,
|
||||
fallback_keys: FallbackKeys::default(),
|
||||
default_deployment: None,
|
||||
default_deployment_raises: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl RouterDynamicConfig {
|
||||
/// The `default_deployment` as resolution serves it.
|
||||
pub fn default_ref(&self) -> Option<crate::resolve::DefaultDeployment<'_>> {
|
||||
if self.default_deployment_raises {
|
||||
return Some(crate::resolve::DefaultDeployment::Raises);
|
||||
}
|
||||
self.default_deployment
|
||||
.as_ref()
|
||||
.map(crate::resolve::DefaultDeployment::Spec)
|
||||
}
|
||||
|
||||
/// The regular fallback targets for the lookup groups, then `default_fallbacks`, which
|
||||
/// Python folds into `fallbacks` as `{"*": [...]}`: what `_has_fallback_group` reads.
|
||||
pub fn fallback_targets(&self, groups: &[&str]) -> Option<Vec<String>> {
|
||||
if groups
|
||||
.iter()
|
||||
.any(|group| self.fallbacks.has_keyed_chain(group, &self.fallback_keys))
|
||||
{
|
||||
return Some(Vec::new());
|
||||
}
|
||||
(!self.default_fallbacks.is_empty()).then(|| self.default_fallbacks.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[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);
|
||||
}
|
||||
}
|
||||
240
litellm-rust/crates/router-core/src/eligibility.rs
Normal file
240
litellm-rust/crates/router-core/src/eligibility.rs
Normal file
|
|
@ -0,0 +1,240 @@
|
|||
//! Which deployments may serve a request's endpoint: a port of
|
||||
//! `litellm/router_utils/endpoint_eligibility.py`. A deployment that sets
|
||||
//! `enforce_supported_endpoints` and lists none of the chat-shaped endpoints serves none of them;
|
||||
//! the host names the endpoint for the calls `endpoint_selection` wraps.
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc};
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{error::Error, registry::Deployment, resolve::Caller};
|
||||
|
||||
/// `CHAT_ENDPOINTS`: `_endpoint_path` of each value of `CHAT_ENDPOINT_BY_CALL`.
|
||||
const CHAT_ENDPOINTS: [&str; 8] = [
|
||||
"/chat/completions",
|
||||
"/completions",
|
||||
"/messages",
|
||||
"/responses",
|
||||
"/generatecontent",
|
||||
"/streamgeneratecontent",
|
||||
"/interactions",
|
||||
"/responses/compact",
|
||||
];
|
||||
|
||||
/// `_endpoint_path`.
|
||||
fn endpoint_path(endpoint: &str) -> String {
|
||||
let rooted = format!("/{}", endpoint.trim().trim_matches('/').to_lowercase());
|
||||
let path = rooted.strip_prefix("/openai").unwrap_or(&rooted);
|
||||
let path = path.strip_prefix("/v1beta").unwrap_or(path);
|
||||
let path = path.strip_prefix("/v1").unwrap_or(path);
|
||||
match path.rsplit_once(':') {
|
||||
Some((_, action)) => format!("/{action}"),
|
||||
None => path.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The flag as pydantic's lax `bool | None` reads it: `None` for a value it refuses, which fails
|
||||
/// the whole declaration.
|
||||
fn enforce_flag(value: Option<&Value>) -> Option<bool> {
|
||||
match value {
|
||||
None | Some(Value::Null) => Some(false),
|
||||
Some(Value::Bool(flag)) => Some(*flag),
|
||||
Some(Value::Number(number)) => match number.as_f64() {
|
||||
Some(0.0) => Some(false),
|
||||
Some(1.0) => Some(true),
|
||||
_ => None,
|
||||
},
|
||||
Some(Value::String(text)) => match text.to_lowercase().as_str() {
|
||||
"0" | "off" | "f" | "false" | "n" | "no" => Some(false),
|
||||
"1" | "on" | "t" | "true" | "y" | "yes" => Some(true),
|
||||
_ => None,
|
||||
},
|
||||
Some(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `_enforced_endpoints`: the declared endpoints of a deployment that enforces them. One endpoint
|
||||
/// may be written as a string; a declaration that does not validate enforces nothing.
|
||||
fn enforced_endpoints(model_info: &Map<String, Value>) -> Option<Vec<String>> {
|
||||
if !enforce_flag(model_info.get("enforce_supported_endpoints"))? {
|
||||
return None;
|
||||
}
|
||||
match model_info.get("supported_endpoints")? {
|
||||
Value::String(endpoint) => Some(vec![endpoint.clone()]),
|
||||
Value::Array(endpoints) => endpoints
|
||||
.iter()
|
||||
.map(|endpoint| endpoint.as_str().map(str::to_owned))
|
||||
.collect(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `non_chat_endpoints`: what a deployment serves when it serves no chat-shaped endpoint, else
|
||||
/// `None`.
|
||||
fn non_chat_endpoints(deployment: &Deployment) -> Option<BTreeSet<String>> {
|
||||
let declared = enforced_endpoints(&deployment.model_info.raw)?;
|
||||
let chat = declared
|
||||
.iter()
|
||||
.any(|endpoint| CHAT_ENDPOINTS.contains(&endpoint_path(endpoint).as_str()));
|
||||
(!chat).then(|| declared.into_iter().collect())
|
||||
}
|
||||
|
||||
/// `eligible_deployments`: the deployments that may serve the caller's endpoint. When none may,
|
||||
/// the refusal names what they do serve.
|
||||
pub(crate) fn eligible_deployments(
|
||||
model: &str,
|
||||
caller: Caller<'_>,
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
let Some(endpoint) = caller.endpoint.filter(|_| !deployments.is_empty()) else {
|
||||
return Ok(deployments);
|
||||
};
|
||||
let declared: Vec<Option<BTreeSet<String>>> = deployments
|
||||
.iter()
|
||||
.map(|deployment| non_chat_endpoints(deployment))
|
||||
.collect();
|
||||
let eligible: Vec<Arc<Deployment>> = deployments
|
||||
.into_iter()
|
||||
.zip(&declared)
|
||||
.filter(|(_, endpoints)| endpoints.is_none())
|
||||
.map(|(deployment, _)| deployment)
|
||||
.collect();
|
||||
if !eligible.is_empty() {
|
||||
return Ok(eligible);
|
||||
}
|
||||
let supported: BTreeSet<String> = declared.into_iter().flatten().flatten().collect();
|
||||
let supported = if supported.is_empty() {
|
||||
"none".to_owned()
|
||||
} else {
|
||||
Vec::from_iter(supported).join(", ")
|
||||
};
|
||||
let requested = if caller.requested.is_empty() {
|
||||
model
|
||||
} else {
|
||||
caller.requested
|
||||
};
|
||||
Err(Error::EndpointNotSupported {
|
||||
model: model.to_owned(),
|
||||
message: format!(
|
||||
"{endpoint} is not supported by model={requested}. Supported endpoints: {supported}"
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::registry::DeploymentSpec;
|
||||
|
||||
fn deployment(id: &str, info: Value) -> Arc<Deployment> {
|
||||
let mut model_info = info;
|
||||
model_info["id"] = json!(id);
|
||||
let spec: DeploymentSpec = serde_json::from_value(json!({
|
||||
"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": model_info,
|
||||
}))
|
||||
.unwrap();
|
||||
Arc::new(Deployment::from_spec(spec).unwrap())
|
||||
}
|
||||
|
||||
fn chat<'a>() -> Caller<'a> {
|
||||
Caller {
|
||||
endpoint: Some("/v1/chat/completions"),
|
||||
requested: "g",
|
||||
..Caller::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::v1_prefix("/v1/chat/completions", "/chat/completions")]
|
||||
#[case::openai_prefix("/openai/v1/responses/", "/responses")]
|
||||
#[case::bare_gemini_action("generateContent", "/generatecontent")]
|
||||
#[case::beta_prefix("/v1beta/interactions", "/interactions")]
|
||||
#[case::gemini_model_path(
|
||||
"/v1beta/models/gemini:streamGenerateContent",
|
||||
"/streamgeneratecontent"
|
||||
)]
|
||||
fn endpoint_paths_normalize(#[case] endpoint: &str, #[case] expected: &str) {
|
||||
assert_eq!(endpoint_path(endpoint), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::embeddings_only(
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": ["/v1/embeddings"]}),
|
||||
false
|
||||
)]
|
||||
#[case::one_endpoint_as_a_string(
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": "/v1/embeddings"}),
|
||||
false
|
||||
)]
|
||||
#[case::string_flag(
|
||||
json!({"enforce_supported_endpoints": "yes", "supported_endpoints": ["/v1/embeddings"]}),
|
||||
false
|
||||
)]
|
||||
#[case::an_empty_list_serves_nothing(
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": []}),
|
||||
false
|
||||
)]
|
||||
#[case::any_chat_shaped_endpoint_serves_every_one(
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": ["/v1/responses"]}),
|
||||
true
|
||||
)]
|
||||
#[case::not_enforced(
|
||||
json!({"enforce_supported_endpoints": false, "supported_endpoints": ["/v1/embeddings"]}),
|
||||
true
|
||||
)]
|
||||
#[case::unset_flag(json!({"supported_endpoints": ["/v1/embeddings"]}), true)]
|
||||
#[case::enforced_without_a_list(json!({"enforce_supported_endpoints": true}), true)]
|
||||
#[case::an_invalid_flag_enforces_nothing(
|
||||
json!({"enforce_supported_endpoints": "maybe", "supported_endpoints": ["/v1/embeddings"]}),
|
||||
true
|
||||
)]
|
||||
#[case::an_invalid_list_enforces_nothing(
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": [1]}),
|
||||
true
|
||||
)]
|
||||
fn a_declaration_keeps_or_drops_a_deployment(#[case] info: Value, #[case] kept: bool) {
|
||||
let open = deployment("open", json!({}));
|
||||
let declared = deployment("declared", info);
|
||||
let eligible = eligible_deployments("g", chat(), vec![declared, open]).unwrap();
|
||||
assert_eq!(eligible.len(), if kept { 2 } else { 1 });
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_endpoint_keeps_everything() {
|
||||
let declared = deployment(
|
||||
"d",
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": []}),
|
||||
);
|
||||
let eligible = eligible_deployments("g", Caller::default(), vec![declared]).unwrap();
|
||||
assert_eq!(eligible.len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lists_what_they_serve(
|
||||
json!(["/v1/rerank", "/v1/embeddings"]),
|
||||
"/v1/chat/completions is not supported by model=g. \
|
||||
Supported endpoints: /v1/embeddings, /v1/rerank"
|
||||
)]
|
||||
#[case::nothing_declared(
|
||||
json!([]),
|
||||
"/v1/chat/completions is not supported by model=g. Supported endpoints: none"
|
||||
)]
|
||||
fn every_deployment_refusing_names_the_endpoint(
|
||||
#[case] endpoints: Value,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let declared = deployment(
|
||||
"d",
|
||||
json!({"enforce_supported_endpoints": true, "supported_endpoints": endpoints}),
|
||||
);
|
||||
let refused = eligible_deployments("resolved", chat(), vec![declared]);
|
||||
assert!(matches!(
|
||||
refused,
|
||||
Err(Error::EndpointNotSupported { model, message })
|
||||
if model == "resolved" && message == expected
|
||||
));
|
||||
}
|
||||
}
|
||||
473
litellm-rust/crates/router-core/src/error.rs
Normal file
473
litellm-rust/crates/router-core/src/error.rs
Normal file
|
|
@ -0,0 +1,473 @@
|
|||
//! The router's errors: [`ConfigError`] when a router is built or changed, [`Error`] when a call
|
||||
//! or a selection fails, and [`AttemptError`], the typed facts of one failed attempt.
|
||||
|
||||
use std::{any::Any, fmt, sync::Arc, time::Duration};
|
||||
|
||||
use crate::config::{Op, RoutingStrategy};
|
||||
|
||||
/// What a host cannot serve, found when the router is built rather than partway through a call.
|
||||
#[derive(Clone, Debug, PartialEq, thiserror::Error)]
|
||||
pub enum Unsupported {
|
||||
#[error("routing strategy `{0}` needs a host that can select")]
|
||||
Strategy(RoutingStrategy),
|
||||
#[error("routing plugin `{0}` needs a host that can select")]
|
||||
Plugin(String),
|
||||
#[error("the host cannot run operation `{op}` on provider `{provider}`")]
|
||||
Op { op: Op, provider: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ConfigError {
|
||||
#[error("the router config asks for something this host cannot serve")]
|
||||
Unsupported(#[from] Unsupported),
|
||||
#[error("deployment `{model_name}` has invalid litellm_params")]
|
||||
InvalidDeployment {
|
||||
model_name: String,
|
||||
#[source]
|
||||
source: serde_json::Error,
|
||||
},
|
||||
#[error("deployment id `{id}` is configured more than once")]
|
||||
DuplicateId { id: String },
|
||||
/// A removal the host's indexes could not finish ([`crate::host_index::RemovalError`]).
|
||||
#[error("{0}")]
|
||||
Removal(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,
|
||||
/// A blocked deployment named by its id (`model` is its litellm model).
|
||||
#[error("deployment model `{model}` is paused")]
|
||||
Paused { model: String },
|
||||
/// A pass-through pick named a deployment id without `use_in_pass_through`.
|
||||
#[error("deployment `{id}` does not support pass-through")]
|
||||
NotPassThrough { id: String },
|
||||
/// A pass-through pick of a group with healthy deployments, none of them pass-through.
|
||||
#[error("model `{model}` has no pass-through deployment")]
|
||||
NoPassThroughDeployments { model: String },
|
||||
/// Every deployment `model` resolved to is a pre-routing strategy marker (`auto_router/...`).
|
||||
#[error("every deployment of `{model}` is a strategy router marker")]
|
||||
OnlyMarkers { model: String },
|
||||
/// Every deployment `model` resolved to refuses the request's endpoint
|
||||
/// (`eligible_deployments`); `message` names what they serve.
|
||||
#[error("{message}")]
|
||||
EndpointNotSupported { model: String, message: String },
|
||||
/// Every deployment `model` resolved to is split into PTU shares the caller holds none of.
|
||||
#[error("deployment `{model}` is reserved for the teams holding a PTU share of it")]
|
||||
PtuReserved { model: String },
|
||||
/// A proxy admin without a team named a public model name several teams use.
|
||||
#[error("model name `{model}` matches deployments from multiple teams")]
|
||||
MultipleTeams { model: String },
|
||||
/// Every deployment `model` resolved to is reserved for other teams by an active access
|
||||
/// window, the first of which ends at `until` (`HH:MM`) in `timezone`.
|
||||
#[error("deployment `{model}` is reserved for another team until {until} {timezone}")]
|
||||
Reserved {
|
||||
model: String,
|
||||
until: String,
|
||||
timezone: String,
|
||||
},
|
||||
/// The pre-call checks found no deployment whose context window fits the prompt; `detail`
|
||||
/// names each one dropped for it.
|
||||
#[error("no deployment of `{model}` has a context window large enough for this call")]
|
||||
ContextWindowExceeded { model: String, detail: String },
|
||||
/// The pre-call checks found no deployment of `model` that takes the requested `params`.
|
||||
#[error("no deployment of `{model}` supports {params:?}")]
|
||||
UnsupportedParams { model: String, params: Vec<String> },
|
||||
/// The routing plugins left none of the healthy deployments' models.
|
||||
#[error("no deployments left after routing-plugin filtering: {models:?}")]
|
||||
RoutingPluginsLeftNothing { models: Vec<String> },
|
||||
/// The request's `_router_weights` failed `validate_router_weights` when simple-shuffle picked.
|
||||
#[error("the request's router weights are invalid")]
|
||||
InvalidRouterWeights,
|
||||
/// A deployment's `litellm_params` `metric` (weight, rpm or tpm) is not a number.
|
||||
#[error("deployment {metric} must be numeric")]
|
||||
NonNumericWeight { metric: &'static str },
|
||||
/// A router with routing plugins was asked for a pick on the sync path, which never runs them.
|
||||
#[error("routing plugins only run on the async selection path")]
|
||||
RoutingPluginsNeedAsync,
|
||||
/// The stale routing-group member map (`_model_to_group`, kept as it was when a raw
|
||||
/// `routing_groups` write replaced the groups) names a group that is no longer configured:
|
||||
/// Python's `_get_routing_context` raises `KeyError(group)` for every pick of the name.
|
||||
#[error("routing group `{group}` is no longer configured")]
|
||||
MissingRoutingGroup { group: String },
|
||||
/// The router's `default_deployment` is not a value Python can copy for `model`
|
||||
/// (`_try_early_resolve_deployments_for_model_not_in_names` raises the copy's error).
|
||||
#[error("the default deployment cannot be served for `{model}`")]
|
||||
DefaultDeploymentRaises { model: String },
|
||||
/// A written model-name or team index names a position past the end of the model list:
|
||||
/// Python's lookup of `model` raises `IndexError`.
|
||||
#[error("a written index of `{model}` names a position past the model list")]
|
||||
IndexOutOfRange { model: String },
|
||||
/// A model-info query raised where the Python router raises: `exception` names the Python
|
||||
/// exception class (`KeyError`, `TypeError`, `AttributeError`, `ValueError`) and `detail` its
|
||||
/// argument.
|
||||
#[error("the model-info query raised {exception}({detail:?})")]
|
||||
ModelInfo {
|
||||
exception: &'static str,
|
||||
detail: String,
|
||||
},
|
||||
}
|
||||
|
||||
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
|
||||
| Self::Paused { .. }
|
||||
| Self::NotPassThrough { .. }
|
||||
| Self::NoPassThroughDeployments { .. }
|
||||
| Self::OnlyMarkers { .. }
|
||||
| Self::EndpointNotSupported { .. }
|
||||
| Self::PtuReserved { .. }
|
||||
| Self::MultipleTeams { .. }
|
||||
| Self::Reserved { .. }
|
||||
| Self::ContextWindowExceeded { .. }
|
||||
| Self::UnsupportedParams { .. }
|
||||
| Self::RoutingPluginsLeftNothing { .. }
|
||||
| Self::InvalidRouterWeights
|
||||
| Self::NonNumericWeight { .. }
|
||||
| Self::RoutingPluginsNeedAsync
|
||||
| Self::MissingRoutingGroup { .. }
|
||||
| Self::DefaultDeploymentRaises { .. }
|
||||
| Self::IndexOutOfRange { .. }
|
||||
| Self::ModelInfo { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The Python exception a pick that failed with this error raises, as the retry loop
|
||||
/// decides by it (its class names, most specific first, and its status): `None` for an
|
||||
/// attempt's failure and for errors no pick raises.
|
||||
pub fn pick_failure(&self) -> Option<AttemptError> {
|
||||
const BAD_REQUEST: &[&str] = &["BadRequestError"];
|
||||
let (classes, status): (&[&str], Option<u16>) = match self {
|
||||
Self::UnknownModel { .. }
|
||||
| Self::OnlyMarkers { .. }
|
||||
| Self::PtuReserved { .. }
|
||||
| Self::NotPassThrough { .. }
|
||||
| Self::NoPassThroughDeployments { .. }
|
||||
| Self::MultipleTeams { .. }
|
||||
| Self::Reserved { .. } => (BAD_REQUEST, Some(400)),
|
||||
Self::EndpointNotSupported { .. } => (
|
||||
&["EndpointNotSupportedByModelError", "BadRequestError"],
|
||||
Some(400),
|
||||
),
|
||||
Self::ContextWindowExceeded { .. } => (
|
||||
&["ContextWindowExceededError", "BadRequestError"],
|
||||
Some(400),
|
||||
),
|
||||
Self::UnsupportedParams { .. } => {
|
||||
(&["UnsupportedParamsError", "BadRequestError"], Some(400))
|
||||
}
|
||||
Self::Paused { .. } => (&["ServiceUnavailableError"], Some(503)),
|
||||
Self::NoHealthyDeployments { .. } => (&["RouterRateLimitError", "ValueError"], None),
|
||||
Self::RoutingPluginsLeftNothing { .. } | Self::RoutingPluginsNeedAsync => {
|
||||
(&["ValueError"], None)
|
||||
}
|
||||
Self::InvalidRouterWeights => (&["ValidationError", "ValueError"], None),
|
||||
Self::NonNumericWeight { .. } => (&["TypeError"], None),
|
||||
Self::MissingRoutingGroup { .. } => (&["KeyError", "LookupError"], None),
|
||||
Self::IndexOutOfRange { .. } => (&["IndexError", "LookupError"], None),
|
||||
Self::DefaultDeploymentRaises { .. } => (&["Exception"], None),
|
||||
Self::Attempt(_)
|
||||
| Self::Config(_)
|
||||
| Self::Cache(_)
|
||||
| Self::Retired
|
||||
| Self::ModelInfo { .. } => return None,
|
||||
};
|
||||
let mut failure = AttemptError::host_raised(classes);
|
||||
failure.status = status;
|
||||
Some(failure)
|
||||
}
|
||||
|
||||
/// A copy of a pick's error, for the events that name it; `None` for any other error.
|
||||
pub fn pick_copy(&self) -> Option<Self> {
|
||||
Some(match self {
|
||||
Self::UnknownModel { model } => Self::UnknownModel {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::NoHealthyDeployments { model } => Self::NoHealthyDeployments {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::Paused { model } => Self::Paused {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::NotPassThrough { id } => Self::NotPassThrough { id: id.clone() },
|
||||
Self::NoPassThroughDeployments { model } => Self::NoPassThroughDeployments {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::OnlyMarkers { model } => Self::OnlyMarkers {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::EndpointNotSupported { model, message } => Self::EndpointNotSupported {
|
||||
model: model.clone(),
|
||||
message: message.clone(),
|
||||
},
|
||||
Self::PtuReserved { model } => Self::PtuReserved {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::MultipleTeams { model } => Self::MultipleTeams {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::Reserved {
|
||||
model,
|
||||
until,
|
||||
timezone,
|
||||
} => Self::Reserved {
|
||||
model: model.clone(),
|
||||
until: until.clone(),
|
||||
timezone: timezone.clone(),
|
||||
},
|
||||
Self::ContextWindowExceeded { model, detail } => Self::ContextWindowExceeded {
|
||||
model: model.clone(),
|
||||
detail: detail.clone(),
|
||||
},
|
||||
Self::UnsupportedParams { model, params } => Self::UnsupportedParams {
|
||||
model: model.clone(),
|
||||
params: params.clone(),
|
||||
},
|
||||
Self::RoutingPluginsLeftNothing { models } => Self::RoutingPluginsLeftNothing {
|
||||
models: models.clone(),
|
||||
},
|
||||
Self::InvalidRouterWeights => Self::InvalidRouterWeights,
|
||||
Self::NonNumericWeight { metric } => Self::NonNumericWeight { metric },
|
||||
Self::RoutingPluginsNeedAsync => Self::RoutingPluginsNeedAsync,
|
||||
Self::MissingRoutingGroup { group } => Self::MissingRoutingGroup {
|
||||
group: group.clone(),
|
||||
},
|
||||
Self::DefaultDeploymentRaises { model } => Self::DefaultDeploymentRaises {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::IndexOutOfRange { model } => Self::IndexOutOfRange {
|
||||
model: model.clone(),
|
||||
},
|
||||
Self::Attempt(_)
|
||||
| Self::Config(_)
|
||||
| Self::Cache(_)
|
||||
| Self::Retired
|
||||
| Self::ModelInfo { .. } => return 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 stored credential was rejected, not the deployment's
|
||||
/// (`CallerCredentialAuthenticationError`): never cools the deployment down.
|
||||
#[strum(serialize = "CallerCredentialAuthenticationError")]
|
||||
CallerCredentialAuthentication,
|
||||
/// A per-user credential exchange was rate limited (`CallerCredentialRateLimitError`): never
|
||||
/// cools the deployment down.
|
||||
#[strum(serialize = "CallerCredentialRateLimitError")]
|
||||
CallerCredentialRateLimit,
|
||||
#[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` header (`_get_retry_after_from_exception_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 deployment the exception itself names (`_deployment_ids_to_skip_on_retry` reads its
|
||||
/// `retry_skip_deployment_id`, then `failed_deployment_id`): an exception raised again keeps
|
||||
/// the id its first failure stamped, whichever deployment raised it now.
|
||||
pub failed_deployment_id: Option<String>,
|
||||
/// The provider's error text, kept only for the cooldown record Python writes
|
||||
/// (`exception_received`, masked past 50 characters).
|
||||
pub detail: Option<Arc<str>>,
|
||||
/// The exception's class names, most specific first (its `__mro__`), for the rules Python
|
||||
/// decides by `isinstance`. Empty when the host reports none; the kind then stands in.
|
||||
pub classes: Vec<Arc<str>>,
|
||||
/// The host already recorded what this failure does to the deployment's health (a routing
|
||||
/// pre-call check's failure, which Python's attempt handles itself: a rate limit cools the
|
||||
/// deployment down, anything else records nothing), or no failure callback runs for it (it
|
||||
/// was raised outside the litellm call, or a shared logging object already logged a failure),
|
||||
/// so the core records nothing.
|
||||
pub recorded: bool,
|
||||
/// The exception names its provider (`llm_provider` is set), which marks a provider's own
|
||||
/// 400 for `fallback_on_unsupported_params`.
|
||||
pub names_provider: bool,
|
||||
/// An agentic follow-up call failed after the first call was served
|
||||
/// (`is_agentic_followup_failure`): never retried, since a retry reruns every paid round. The
|
||||
/// failure keeps its own class, which decides its fallbacks and cooldown as for any failure.
|
||||
pub agentic_followup: bool,
|
||||
/// One of `deployment_callback_on_failure`'s early exits holds: the failure says nothing about
|
||||
/// the deployment, so it is neither counted nor cooled down (a failed request still counts
|
||||
/// against the minute's requests, which `async_deployment_callback_on_failure` does apart).
|
||||
pub outside_deployment: bool,
|
||||
/// One of `_should_run_cooldown_logic`'s per-caller exits holds (`_is_per_user_provider_request`
|
||||
/// or a per-user Copilot deployment): the failure is counted but never cools the deployment.
|
||||
pub caller_scoped: bool,
|
||||
/// The exception has no `status_code` (Python's status `""`): `_is_cooldown_required` holds
|
||||
/// no cooldown for it, unlike a status the host could not read, which cools down as a 500.
|
||||
pub status_absent: bool,
|
||||
/// The failure counts against no per-minute requests: `async_deployment_callback_on_failure`
|
||||
/// does not run for a sync call, and returns at once for a batch retrieve.
|
||||
pub uncounted_request: bool,
|
||||
pub payload: Option<Payload>,
|
||||
}
|
||||
|
||||
impl AttemptError {
|
||||
pub fn new(kind: FailureKind, status: Option<u16>) -> Self {
|
||||
Self {
|
||||
kind,
|
||||
status,
|
||||
retry_after: None,
|
||||
num_retries: None,
|
||||
failed_deployment_id: None,
|
||||
detail: None,
|
||||
classes: Vec::new(),
|
||||
recorded: false,
|
||||
names_provider: false,
|
||||
agentic_followup: false,
|
||||
outside_deployment: false,
|
||||
caller_scoped: false,
|
||||
status_absent: false,
|
||||
uncounted_request: false,
|
||||
payload: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The exception's class names, most specific first: the host's, else the kind's.
|
||||
pub fn class_names(&self) -> Vec<&str> {
|
||||
if self.classes.is_empty() {
|
||||
return crate::retry::kind_classes(self.kind).to_vec();
|
||||
}
|
||||
self.classes.iter().map(AsRef::as_ref).collect()
|
||||
}
|
||||
|
||||
/// `isinstance(error, <class named name>)`.
|
||||
pub fn is_a(&self, name: &str) -> bool {
|
||||
self.class_names().contains(&name)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_retry_after(self, retry_after: Duration) -> Self {
|
||||
Self {
|
||||
retry_after: Some(retry_after),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
/// A plain Python exception the host raised outside any attempt (a `KeyError` from a stale
|
||||
/// routing-group lookup, an error a retry policy lookup raised): no status, no provider,
|
||||
/// nothing recorded against a deployment. `classes` are its class names, most specific first.
|
||||
pub fn host_raised(classes: &[&str]) -> Self {
|
||||
Self {
|
||||
classes: classes.iter().map(|class| (*class).into()).collect(),
|
||||
recorded: true,
|
||||
..Self::new(FailureKind::Other, None)
|
||||
}
|
||||
}
|
||||
|
||||
/// 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_auth(
|
||||
"CallerCredentialAuthenticationError",
|
||||
FailureKind::CallerCredentialAuthentication
|
||||
)]
|
||||
#[case::caller_credential_rate_limit(
|
||||
"CallerCredentialRateLimitError",
|
||||
FailureKind::CallerCredentialRateLimit
|
||||
)]
|
||||
#[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());
|
||||
}
|
||||
}
|
||||
330
litellm-rust/crates/router-core/src/fallback.rs
Normal file
330
litellm-rust/crates/router-core/src/fallback.rs
Normal file
|
|
@ -0,0 +1,330 @@
|
|||
//! Fallback lists as Python spells them, and how a failed group's chain is found in one.
|
||||
//! Mirrors `get_fallback_model_group`, `get_fallback_model_group_for_lookup_groups`,
|
||||
//! `check_non_standard_fallback_format` and `fallback_attempt_key` in
|
||||
//! `router_utils/fallback_event_handlers.py`.
|
||||
//!
|
||||
//! A list holds keyed entries (`{"gpt-4": ["gpt-3.5", {"model": "x", "temperature": 0}]}`),
|
||||
//! bare names (`"claude"`, a chain for every group) and, in a client-side list, dict targets
|
||||
//! (`{"model": "x", "messages": [...]}`), which override request params for their attempt.
|
||||
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
/// One item of a fallback list, kept as Python holds it.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
enum Entry {
|
||||
Name(String),
|
||||
Dict(Map<String, Value>),
|
||||
/// Anything else: Python skips it when it looks a chain up.
|
||||
Other,
|
||||
}
|
||||
|
||||
/// A fallback list, in order. `"*"` keys the default chain.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
|
||||
#[serde(from = "Vec<Value>")]
|
||||
pub struct Fallbacks(Vec<Entry>);
|
||||
|
||||
impl From<Vec<Value>> for Fallbacks {
|
||||
fn from(entries: Vec<Value>) -> Self {
|
||||
Self(
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|entry| match entry {
|
||||
Value::String(name) => Entry::Name(name),
|
||||
Value::Object(dict) => Entry::Dict(dict),
|
||||
_ => Entry::Other,
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// One target of a chain: a group name, or a dict target that also overrides request params.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum FallbackTarget {
|
||||
Name(String),
|
||||
Params(Map<String, Value>),
|
||||
}
|
||||
|
||||
impl FallbackTarget {
|
||||
fn of(value: &Value) -> Option<Self> {
|
||||
match value {
|
||||
Value::String(name) => Some(Self::Name(name.clone())),
|
||||
Value::Object(dict) => Some(Self::Params(dict.clone())),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `_get_fallback_target_model_group`: the group the target names, if any.
|
||||
pub fn model(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Name(name) => Some(name),
|
||||
Self::Params(dict) => dict.get("model").and_then(Value::as_str),
|
||||
}
|
||||
}
|
||||
|
||||
/// The `model` of a dict target that sets a non-string one: the hop routes that value, which
|
||||
/// no group answers (`_get_fallback_target_model_group` finds no group to check, but
|
||||
/// `kwargs.update(target)` still rebinds the request's model).
|
||||
pub fn unroutable_model(&self) -> Option<&Value> {
|
||||
match self {
|
||||
Self::Params(dict) => dict.get("model").filter(|model| !model.is_string()),
|
||||
Self::Name(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `fallback_attempt_key`: a bare name and `{"model": name}` are the same attempt; a dict
|
||||
/// carrying anything else is its own.
|
||||
pub fn identity(&self) -> FallbackTarget {
|
||||
match self {
|
||||
Self::Params(dict) if dict.len() == 1 => match dict.get("model") {
|
||||
Some(Value::String(name)) => Self::Name(name.clone()),
|
||||
_ => self.clone(),
|
||||
},
|
||||
_ => self.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The params a dict target merges into the request (`kwargs.update(target)`).
|
||||
pub fn params(&self) -> Option<&Map<String, Value>> {
|
||||
match self {
|
||||
Self::Params(dict) => Some(dict),
|
||||
Self::Name(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What fallback keys match besides the exact group name: Python's `provider_list`, for
|
||||
/// `"openai/gpt-4o"` to use a `"gpt-4o"` entry, the provider-prefixed spelling of each bare
|
||||
/// group name (`_provider_prefixed_model_group`), for `"gpt-4o"` to use an `"openai/gpt-4o"`
|
||||
/// entry, and the `LiteLLMParamsTypedDict` keys that mark a client-side dict target.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct FallbackKeys {
|
||||
pub providers: Vec<String>,
|
||||
pub provider_prefixed: BTreeMap<String, String>,
|
||||
pub param_keys: Vec<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
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A chain found for a group, and whether it is the `"*"` chain itself.
|
||||
struct Found {
|
||||
targets: Vec<FallbackTarget>,
|
||||
generic: bool,
|
||||
}
|
||||
|
||||
/// The targets of a keyed entry's value: a list of names and dicts. A bare string is walked
|
||||
/// character by character, as Python iterates it.
|
||||
fn chain_of(value: &Value) -> Vec<FallbackTarget> {
|
||||
match value {
|
||||
Value::Array(items) => items.iter().filter_map(FallbackTarget::of).collect(),
|
||||
Value::String(text) => text
|
||||
.chars()
|
||||
.map(|c| FallbackTarget::Name(c.to_string()))
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
impl Fallbacks {
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.0.is_empty()
|
||||
}
|
||||
|
||||
/// `check_non_standard_fallback_format`: a client-side list of names, or of dicts that set
|
||||
/// a litellm param, is one chain for every group.
|
||||
pub fn is_client_side(&self, keys: &FallbackKeys) -> bool {
|
||||
if self.0.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if self.0.iter().all(|entry| matches!(entry, Entry::Name(_))) {
|
||||
return true;
|
||||
}
|
||||
let dicts: Vec<_> = self
|
||||
.0
|
||||
.iter()
|
||||
.filter_map(|entry| match entry {
|
||||
Entry::Dict(dict) => Some(dict),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
dicts.len() == self.0.len()
|
||||
&& dicts.iter().any(|dict| {
|
||||
keys.param_keys
|
||||
.iter()
|
||||
.any(|key| dict.get(key).is_some_and(|value| !value.is_array()))
|
||||
})
|
||||
}
|
||||
|
||||
/// The whole list as one chain (a client-side list).
|
||||
pub fn as_chain(&self) -> Vec<FallbackTarget> {
|
||||
self.0
|
||||
.iter()
|
||||
.filter_map(|entry| match entry {
|
||||
Entry::Name(name) => Some(FallbackTarget::Name(name.clone())),
|
||||
Entry::Dict(dict) => Some(FallbackTarget::Params(dict.clone())),
|
||||
Entry::Other => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `_get_first_default_fallback`: the first target of the first `"*"` entry.
|
||||
pub fn default_target(&self) -> Option<&str> {
|
||||
self.0.iter().find_map(|entry| match entry {
|
||||
Entry::Dict(dict) => match dict.get("*") {
|
||||
Some(Value::Array(targets)) => targets.first().and_then(Value::as_str),
|
||||
_ => None,
|
||||
},
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
/// `_has_default_fallbacks`: some entry keys `"*"`.
|
||||
pub fn has_star(&self) -> bool {
|
||||
self.0
|
||||
.iter()
|
||||
.any(|entry| matches!(entry, Entry::Dict(dict) if dict.contains_key("*")))
|
||||
}
|
||||
|
||||
/// `get_fallback_model_group`.
|
||||
fn lookup(&self, group: &str, keys: &FallbackKeys) -> Option<Found> {
|
||||
let prefixed = keys.provider_prefixed.get(group).map(String::as_str);
|
||||
let mut generic = None;
|
||||
let mut stripped = None;
|
||||
let mut found = None;
|
||||
for entry in &self.0 {
|
||||
match entry {
|
||||
Entry::Dict(dict) => {
|
||||
let Some((key, value)) = dict.iter().next() else {
|
||||
continue;
|
||||
};
|
||||
if key == group {
|
||||
found = Some(chain_of(value));
|
||||
break;
|
||||
} else if Some(key.as_str()) == prefixed || keys.is_stripped(group, key) {
|
||||
stripped = Some(chain_of(value));
|
||||
} else if key == "*" {
|
||||
generic = Some(value);
|
||||
}
|
||||
}
|
||||
Entry::Name(name) => found = Some(vec![FallbackTarget::Name(name.clone())]),
|
||||
Entry::Other => {}
|
||||
}
|
||||
}
|
||||
found
|
||||
.or(stripped)
|
||||
.map(|targets| Found {
|
||||
targets,
|
||||
generic: false,
|
||||
})
|
||||
.or_else(|| {
|
||||
generic.map(|value| Found {
|
||||
targets: chain_of(value),
|
||||
generic: 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<FallbackTarget>> {
|
||||
let results: Vec<Found> = groups
|
||||
.iter()
|
||||
.filter_map(|group| self.lookup(group, keys))
|
||||
.collect();
|
||||
let index = results
|
||||
.iter()
|
||||
.position(|found| !found.generic)
|
||||
.or_else(|| (!results.is_empty()).then_some(0))?;
|
||||
results.into_iter().nth(index).map(|found| found.targets)
|
||||
}
|
||||
|
||||
/// `_has_fallback_group`'s view: only keyed entries whose targets are names or dicts count.
|
||||
pub fn has_keyed_chain(&self, group: &str, keys: &FallbackKeys) -> bool {
|
||||
let keyed = Fallbacks(
|
||||
self.0
|
||||
.iter()
|
||||
.filter(|entry| matches!(entry, Entry::Dict(_)))
|
||||
.cloned()
|
||||
.collect(),
|
||||
);
|
||||
keyed
|
||||
.lookup(group, keys)
|
||||
.is_some_and(|found| !found.targets.is_empty())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::an_int(json!({"model": 5}), true)]
|
||||
#[case::none(json!({"model": null}), true)]
|
||||
#[case::a_string(json!({"model": "b"}), false)]
|
||||
#[case::no_model(json!({"temperature": 0}), false)]
|
||||
fn only_a_set_non_string_model_is_unroutable(#[case] target: Value, #[case] unroutable: bool) {
|
||||
let target = FallbackTarget::of(&target).unwrap();
|
||||
assert_eq!(target.unroutable_model().is_some(), unroutable);
|
||||
}
|
||||
|
||||
fn names(targets: Option<Vec<FallbackTarget>>) -> Option<Vec<String>> {
|
||||
targets.map(|targets| {
|
||||
targets
|
||||
.iter()
|
||||
.map(|t| t.model().unwrap_or_default().to_owned())
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_entry(json!([{"a": ["b"]}, {"*": ["d"]}]), "a", Some(vec!["b"]))]
|
||||
#[case::star_entry(json!([{"a": ["b"]}, {"*": ["d"]}]), "x", Some(vec!["d"]))]
|
||||
#[case::string_is_every_groups_chain(json!(["b", {"x": ["c"]}]), "a", Some(vec!["b"]))]
|
||||
#[case::keyed_entry_wins(json!(["c", {"a": ["b"]}]), "a", Some(vec!["b"]))]
|
||||
#[case::string_value_is_iterated(json!([{"a": "bc"}]), "a", Some(vec!["b", "c"]))]
|
||||
#[case::none(json!([{"x": ["c"]}]), "a", None)]
|
||||
fn chains_follow_python(
|
||||
#[case] list: Value,
|
||||
#[case] group: &str,
|
||||
#[case] expected: Option<Vec<&str>>,
|
||||
) {
|
||||
let fallbacks: Fallbacks = serde_json::from_value(list).unwrap();
|
||||
assert_eq!(
|
||||
names(fallbacks.lookup_groups(&[group], &FallbackKeys::default())),
|
||||
expected.map(|e| e.into_iter().map(str::to_owned).collect())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::names(json!(["a", "b"]), true)]
|
||||
#[case::param_dicts(json!([{"model": "a", "temperature": 1}]), true)]
|
||||
#[case::keyed(json!([{"a": ["b"]}]), false)]
|
||||
#[case::mixed(json!(["a", {"a": ["b"]}]), false)]
|
||||
fn client_side_lists_follow_python(#[case] list: Value, #[case] expected: bool) {
|
||||
let fallbacks: Fallbacks = serde_json::from_value(list).unwrap();
|
||||
let keys = FallbackKeys {
|
||||
param_keys: vec!["model".to_owned(), "temperature".to_owned()],
|
||||
..FallbackKeys::default()
|
||||
};
|
||||
assert_eq!(fallbacks.is_client_side(&keys), expected);
|
||||
}
|
||||
}
|
||||
646
litellm-rust/crates/router-core/src/group_info.rs
Normal file
646
litellm-rust/crates/router-core/src/group_info.rs
Normal file
|
|
@ -0,0 +1,646 @@
|
|||
//! Model group info and the configured-value reads: what the router reports a model group, or a
|
||||
//! concrete deployment behind a listed name, can do.
|
||||
//!
|
||||
//! The rules are the Python router core's (`_set_model_group_info`, `get_model_listing_info`,
|
||||
//! `get_configured_token_limits`, `get_configured_mode`, `get_configured_display_name`,
|
||||
//! `get_configured_service_tiers`). The host picks the deployments (the registry lookups) and
|
||||
//! answers what needs its objects as [`Host::model_info`] lookups asked while the query runs: a
|
||||
//! deployment's validated params, its merged model info and provider, its capabilities, and the
|
||||
//! group's first entry as its model validates it. The core decides which deployments count and
|
||||
//! how their values combine.
|
||||
|
||||
use serde_json::{Map, Number, Value};
|
||||
|
||||
use crate::{
|
||||
error::Error,
|
||||
host::{Host, ModelInfoLookup},
|
||||
model_info::{
|
||||
discovered, is_opaque, lookup, malformed, python_float, python_int, raised, truthy,
|
||||
without_digit_underscores,
|
||||
},
|
||||
};
|
||||
|
||||
/// One `_set_model_group_info` query.
|
||||
#[derive(Clone, Debug, Default, serde::Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct GroupInfoQuery {
|
||||
/// The model group the deployments serve (an alias already resolved).
|
||||
pub model_group: String,
|
||||
/// The name the group is reported as.
|
||||
pub user_facing_model_group_name: String,
|
||||
/// The group is a team's public model name the team's deployments serve.
|
||||
pub serves_team_deployments: bool,
|
||||
/// A wildcard pattern routes `model_group` (asked only when an entry needs it).
|
||||
pub pattern_routes: bool,
|
||||
/// The first string `model_info.description` among the deployments.
|
||||
pub description: Value,
|
||||
/// The deployments the group routes to, as stored.
|
||||
pub deployments: Vec<Value>,
|
||||
}
|
||||
|
||||
/// The group info: `None` when no deployment matched, else the fields the merge assigned after
|
||||
/// the host built the group from its first deployment ([`ModelInfoLookup::GroupSeed`]).
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct GroupInfo {
|
||||
pub assignments: Map<String, Value>,
|
||||
}
|
||||
|
||||
/// The type name Python reports for a value, in its error messages.
|
||||
pub(crate) fn type_name(value: &Value) -> &'static str {
|
||||
match value {
|
||||
Value::Null => "NoneType",
|
||||
Value::Bool(_) => "bool",
|
||||
Value::Number(number) if number.is_f64() => "float",
|
||||
Value::Number(_) => "int",
|
||||
Value::String(_) => "str",
|
||||
Value::Array(_) => "list",
|
||||
Value::Object(_) if is_opaque(value) => "object",
|
||||
Value::Object(_) => "dict",
|
||||
}
|
||||
}
|
||||
|
||||
/// A number as Python's arithmetic reads it (a bool is 0 or 1), or `None`.
|
||||
fn arithmetic(value: &Value) -> Option<Number> {
|
||||
match value {
|
||||
Value::Bool(flag) => Some(Number::from(u8::from(*flag))),
|
||||
Value::Number(number) => Some(number.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `total += value`, with `total` starting at `0`; adding a non-number raises `TypeError`.
|
||||
fn add(total: Option<Value>, value: &Value) -> Result<Value, Error> {
|
||||
let total = total.unwrap_or_else(|| Value::Number(Number::from(0)));
|
||||
let (Some(left), Some(right)) = (arithmetic(&total), arithmetic(value)) else {
|
||||
return Err(raised(
|
||||
"TypeError",
|
||||
format!(
|
||||
"unsupported operand type(s) for +=: '{}' and '{}'",
|
||||
type_name(&total),
|
||||
type_name(value)
|
||||
),
|
||||
));
|
||||
};
|
||||
if let (Some(left), Some(right)) = (left.as_i64(), right.as_i64())
|
||||
&& let Some(sum) = left.checked_add(right)
|
||||
{
|
||||
return Ok(Value::Number(Number::from(sum)));
|
||||
}
|
||||
let sum = left.as_f64().unwrap_or(0.0) + right.as_f64().unwrap_or(0.0);
|
||||
Ok(Number::from_f64(sum).map_or(Value::Null, Value::Number))
|
||||
}
|
||||
|
||||
/// `left > right`, for numbers or strings; any other pair raises `TypeError`.
|
||||
fn greater(left: &Value, right: &Value) -> Result<bool, Error> {
|
||||
if let (Some(left), Some(right)) = (arithmetic(left), arithmetic(right)) {
|
||||
return Ok(left.as_f64().unwrap_or(0.0) > right.as_f64().unwrap_or(0.0));
|
||||
}
|
||||
if let (Value::String(left), Value::String(right)) = (left, right) {
|
||||
return Ok(left > right);
|
||||
}
|
||||
Err(raised(
|
||||
"TypeError",
|
||||
format!(
|
||||
"'>' not supported between instances of '{}' and '{}'",
|
||||
type_name(left),
|
||||
type_name(right)
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
/// `_cost_value_as_float`: `float(value)`, or `None` when it is unset or not a number.
|
||||
fn cost_as_float(value: Option<&Value>) -> Option<f64> {
|
||||
match value? {
|
||||
Value::Bool(flag) => Some(f64::from(u8::from(*flag))),
|
||||
Value::Number(number) => number.as_f64(),
|
||||
Value::String(text) => {
|
||||
without_digit_underscores(text.trim()).and_then(|digits| python_float(&digits))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn float(value: Option<f64>) -> Value {
|
||||
value
|
||||
.and_then(Number::from_f64)
|
||||
.map_or(Value::Null, Value::Number)
|
||||
}
|
||||
|
||||
/// `value.get(key)` on a dict; a value that is no dict raises `AttributeError`, as `.get` does.
|
||||
fn get<'a>(value: &'a Value, key: &str) -> Result<Option<&'a Value>, Error> {
|
||||
match value {
|
||||
Value::Object(map) if !is_opaque(value) => Ok(map.get(key)),
|
||||
other => Err(raised(
|
||||
"AttributeError",
|
||||
format!("'{}' object has no attribute 'get'", type_name(other)),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `model.get(key)`, then `litellm_params.get(key)`, then `model_info.get(key)`: the first that is
|
||||
/// not `None`.
|
||||
fn first_set(
|
||||
model: &Value,
|
||||
params: &Value,
|
||||
info: &Value,
|
||||
key: &str,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
for layer in [model, params, info] {
|
||||
if let Some(value) = get(layer, key)?.filter(|value| !value.is_null()) {
|
||||
return Ok(Some(value.clone()));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn object(answer: Value, kind: &str) -> Result<Map<String, Value>, Error> {
|
||||
match answer {
|
||||
Value::Object(map) => Ok(map),
|
||||
_ => Err(malformed(kind)),
|
||||
}
|
||||
}
|
||||
|
||||
/// The group being merged: the host's validated first entry, and what the merge assigned since.
|
||||
struct Merged {
|
||||
current: Map<String, Value>,
|
||||
assignments: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl Merged {
|
||||
fn get(&self, field: &str) -> &Value {
|
||||
self.current.get(field).unwrap_or(&Value::Null)
|
||||
}
|
||||
|
||||
fn set(&mut self, field: &str, value: Value) {
|
||||
self.current.insert(field.to_owned(), value.clone());
|
||||
self.assignments.insert(field.to_owned(), value);
|
||||
}
|
||||
|
||||
/// `if info.get(field) is not None and info[field] > group.field (or group.field is None)`.
|
||||
fn widen(&mut self, field: &str, info: &Map<String, Value>) -> Result<(), Error> {
|
||||
let Some(value) = info.get(field).filter(|value| !value.is_null()) else {
|
||||
return Ok(());
|
||||
};
|
||||
let current = self.get(field);
|
||||
if current.is_null() || greater(value, current)? {
|
||||
self.set(field, value.clone());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The cost `field` widened to `float(info[field])` when that exceeds `group.field or 0.0`.
|
||||
fn widen_cost(&mut self, field: &str, info: &Map<String, Value>) {
|
||||
let Some(cost) = cost_as_float(info.get(field)) else {
|
||||
return;
|
||||
};
|
||||
let current = self.get(field);
|
||||
let floor = if truthy(current) {
|
||||
current.as_f64().unwrap_or(0.0)
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
if current.is_null() || cost > floor {
|
||||
self.set(field, float(Some(cost)));
|
||||
}
|
||||
}
|
||||
|
||||
/// `if info.get(field) is True: group.field = True`.
|
||||
fn any(&mut self, field: &str, info: &Map<String, Value>) {
|
||||
if info.get(field) == Some(&Value::Bool(true)) {
|
||||
self.set(field, Value::Bool(true));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `tuple(dict.fromkeys((*current, *added))) or None`.
|
||||
fn union(current: &Value, added: &[Value]) -> Value {
|
||||
let mut kept: Vec<Value> = Vec::new();
|
||||
let existing = current.as_array().map(Vec::as_slice).unwrap_or_default();
|
||||
for value in existing.iter().chain(added) {
|
||||
if !kept.contains(value) {
|
||||
kept.push(value.clone());
|
||||
}
|
||||
}
|
||||
if kept.is_empty() {
|
||||
Value::Null
|
||||
} else {
|
||||
Value::Array(kept)
|
||||
}
|
||||
}
|
||||
|
||||
/// `_set_model_group_info`: the group's info merged across the deployments that serve it.
|
||||
pub async fn model_group_info(
|
||||
host: &dyn Host,
|
||||
query: &GroupInfoQuery,
|
||||
) -> Result<Option<GroupInfo>, Error> {
|
||||
let mut merged: Option<Merged> = None;
|
||||
let mut totals: [(&str, Option<Value>); 4] =
|
||||
[("tpm", None), ("rpm", None), ("itpm", None), ("otpm", None)];
|
||||
let mut clientside_auth_params = Value::Null;
|
||||
let mut efforts_initialized = false;
|
||||
let mut efforts_unknown = false;
|
||||
let empty = Value::Object(Map::new());
|
||||
|
||||
for (index, model) in query.deployments.iter().enumerate() {
|
||||
let named = get(model, "model_name")?;
|
||||
let is_match = query.serves_team_deployments
|
||||
|| named.is_some_and(|name| name.as_str() == Some(query.model_group.as_str()))
|
||||
|| (named.is_some() && query.pattern_routes);
|
||||
if !is_match {
|
||||
continue;
|
||||
}
|
||||
let params = object(
|
||||
lookup(host, ModelInfoLookup::GroupParams { index }).await?,
|
||||
"group_params",
|
||||
)?;
|
||||
clientside_auth_params = params
|
||||
.get("configurable_clientside_auth_params")
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null);
|
||||
let model_params = get(model, "litellm_params")?.unwrap_or(&empty);
|
||||
let model_info_dict = get(model, "model_info")?.unwrap_or(&empty);
|
||||
let mut deployment = [None, None, None, None];
|
||||
for (slot, (metric, _)) in deployment.iter_mut().zip(totals.iter()) {
|
||||
*slot = first_set(model, model_params, model_info_dict, metric)?;
|
||||
}
|
||||
|
||||
let facts = object(
|
||||
lookup(
|
||||
host,
|
||||
ModelInfoLookup::GroupDeployment {
|
||||
index,
|
||||
model: params.get("model").cloned().unwrap_or(Value::Null),
|
||||
},
|
||||
)
|
||||
.await?,
|
||||
"group_deployment",
|
||||
)?;
|
||||
let mapped = facts.get("mapped") == Some(&Value::Bool(true));
|
||||
let litellm_model = facts.get("litellm_model").cloned().unwrap_or(Value::Null);
|
||||
let llm_provider = facts.get("llm_provider").cloned().unwrap_or(Value::Null);
|
||||
let model_info = match facts.get("model_info") {
|
||||
Some(Value::Object(info)) if !is_opaque(&Value::Object(info.clone())) => info.clone(),
|
||||
Some(Value::Null) | None => {
|
||||
// No model info: what the deployment's own (database) model_info says.
|
||||
let db_info = get(model, "model_info")?.unwrap_or(&empty);
|
||||
let mut info = Map::new();
|
||||
info.insert("key".to_owned(), Value::String(query.model_group.clone()));
|
||||
for key in ["max_tokens", "max_input_tokens", "max_output_tokens"] {
|
||||
info.insert(key.to_owned(), Value::Null);
|
||||
}
|
||||
for key in ["input_cost_per_token", "output_cost_per_token"] {
|
||||
info.insert(key.to_owned(), float(cost_as_float(get(db_info, key)?)));
|
||||
}
|
||||
info.insert("litellm_provider".to_owned(), llm_provider.clone());
|
||||
info.insert(
|
||||
"mode".to_owned(),
|
||||
match get(db_info, "mode")? {
|
||||
Some(mode) => mode.clone(),
|
||||
None => Value::String("chat".to_owned()),
|
||||
},
|
||||
);
|
||||
info.insert(
|
||||
"supported_openai_params".to_owned(),
|
||||
match facts.get("supported_openai_params") {
|
||||
Some(Value::Null) | None => Value::Array(Vec::new()),
|
||||
Some(params) => params.clone(),
|
||||
},
|
||||
);
|
||||
info.insert("supports_system_messages".to_owned(), Value::Null);
|
||||
info
|
||||
}
|
||||
Some(_) => return Err(malformed("group_deployment")),
|
||||
};
|
||||
|
||||
match merged.as_mut() {
|
||||
None => {
|
||||
let mut seed = Map::new();
|
||||
seed.insert(
|
||||
"model_group".to_owned(),
|
||||
Value::String(query.user_facing_model_group_name.clone()),
|
||||
);
|
||||
seed.insert(
|
||||
"providers".to_owned(),
|
||||
Value::Array(vec![llm_provider.clone()]),
|
||||
);
|
||||
seed.insert("description".to_owned(), query.description.clone());
|
||||
seed.extend(model_info.clone());
|
||||
seed.insert(
|
||||
"supports_continuation".to_owned(),
|
||||
Value::Bool(
|
||||
model_info.get("supports_continuation") == Some(&Value::Bool(true)),
|
||||
),
|
||||
);
|
||||
seed.insert("supports_fast_mode".to_owned(), Value::Bool(true));
|
||||
seed.insert("supported_reasoning_efforts".to_owned(), Value::Null);
|
||||
seed.insert("supported_endpoints".to_owned(), Value::Null);
|
||||
let current = object(
|
||||
lookup(host, ModelInfoLookup::GroupSeed { info: seed }).await?,
|
||||
"group_seed",
|
||||
)?;
|
||||
merged = Some(Merged {
|
||||
current,
|
||||
assignments: Map::new(),
|
||||
});
|
||||
}
|
||||
Some(group) => {
|
||||
let providers = group
|
||||
.get("providers")
|
||||
.as_array()
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
if !providers.contains(&llm_provider) {
|
||||
let mut providers = providers;
|
||||
providers.push(llm_provider.clone());
|
||||
group.set("providers", Value::Array(providers));
|
||||
}
|
||||
group.widen("max_input_tokens", &model_info)?;
|
||||
group.widen("max_output_tokens", &model_info)?;
|
||||
group.widen_cost("input_cost_per_token", &model_info);
|
||||
group.widen_cost("output_cost_per_token", &model_info);
|
||||
for field in [
|
||||
"supports_parallel_function_calling",
|
||||
"supports_vision",
|
||||
"supports_function_calling",
|
||||
"supports_web_search",
|
||||
"supports_url_context",
|
||||
"supports_continuation",
|
||||
"supports_reasoning",
|
||||
] {
|
||||
group.any(field, &model_info);
|
||||
}
|
||||
if let Some(params) = model_info
|
||||
.get("supported_openai_params")
|
||||
.filter(|params| !params.is_null())
|
||||
{
|
||||
group.set("supported_openai_params", params.clone());
|
||||
}
|
||||
for (slot, metric) in deployment.iter_mut().zip(["tpm", "rpm"]) {
|
||||
if slot.is_none()
|
||||
&& let Some(value) = model_info.get(metric).filter(|v| !v.is_null())
|
||||
{
|
||||
*slot = Some(value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let Some(group) = merged.as_mut() else {
|
||||
unreachable!("the group was just built");
|
||||
};
|
||||
|
||||
let capabilities = object(
|
||||
lookup(
|
||||
host,
|
||||
ModelInfoLookup::GroupCapabilities {
|
||||
index,
|
||||
model_info: Value::Object(model_info),
|
||||
litellm_model,
|
||||
llm_provider,
|
||||
mapped,
|
||||
},
|
||||
)
|
||||
.await?,
|
||||
"group_capabilities",
|
||||
)?;
|
||||
let endpoints = capabilities
|
||||
.get("supported_endpoints")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let endpoints = union(group.get("supported_endpoints"), &endpoints);
|
||||
group.set("supported_endpoints", endpoints);
|
||||
let fast_mode = truthy(group.get("supports_fast_mode"))
|
||||
&& capabilities.get("supports_fast_mode") == Some(&Value::Bool(true));
|
||||
group.set("supports_fast_mode", Value::Bool(fast_mode));
|
||||
match capabilities
|
||||
.get("reasoning_efforts")
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null)
|
||||
{
|
||||
Value::Null => {
|
||||
efforts_unknown = true;
|
||||
group.set("supported_reasoning_efforts", Value::Null);
|
||||
}
|
||||
resolved if !efforts_initialized => {
|
||||
efforts_initialized = true;
|
||||
if !efforts_unknown {
|
||||
group.set("supported_reasoning_efforts", resolved);
|
||||
}
|
||||
}
|
||||
resolved if !efforts_unknown => {
|
||||
let intersected = lookup(
|
||||
host,
|
||||
ModelInfoLookup::ReasoningIntersect {
|
||||
current: group.get("supported_reasoning_efforts").clone(),
|
||||
resolved,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
group.set("supported_reasoning_efforts", intersected);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
for ((_, total), value) in totals.iter_mut().zip(deployment) {
|
||||
if let Some(value) = value {
|
||||
*total = Some(add(total.take(), &value)?);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let Some(mut group) = merged else {
|
||||
return Ok(None);
|
||||
};
|
||||
for (metric, total) in totals {
|
||||
if let Some(total) = total {
|
||||
group.set(metric, total);
|
||||
}
|
||||
}
|
||||
if !clientside_auth_params.is_null() {
|
||||
group.set(
|
||||
"configurable_clientside_auth_params",
|
||||
clientside_auth_params,
|
||||
);
|
||||
}
|
||||
Ok(Some(GroupInfo {
|
||||
assignments: group.assignments,
|
||||
}))
|
||||
}
|
||||
|
||||
/// `coerce_token_limit`: a token limit as an int, `None` when it is unset or not a usable number
|
||||
/// (a bool never is).
|
||||
pub fn coerce_token_limit(value: Option<&Value>) -> Option<i64> {
|
||||
match value? {
|
||||
Value::Number(number) => number.as_i64().or_else(|| {
|
||||
number
|
||||
.as_f64()
|
||||
.filter(|n| n.is_finite())
|
||||
.map(|n| n.trunc() as i64)
|
||||
}),
|
||||
Value::String(text) => without_digit_underscores(text.trim())
|
||||
.and_then(|digits| python_int(&digits))
|
||||
.and_then(|number| number.as_i64()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The deployment's discovered limits under its configured `model_info` values that are set.
|
||||
async fn effective_model_info(
|
||||
host: &dyn Host,
|
||||
id: Value,
|
||||
configured: &Map<String, Value>,
|
||||
) -> Result<Map<String, Value>, Error> {
|
||||
let mut info = discovered(host, id).await?;
|
||||
info.extend(
|
||||
configured
|
||||
.iter()
|
||||
.filter(|(_, value)| !value.is_null())
|
||||
.map(|(key, value)| (key.clone(), value.clone())),
|
||||
);
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
/// `get_configured_token_limits`: the (max input, max output) tokens a concrete deployment
|
||||
/// configures, or its refresh discovered, from its `model_info` (`exclude_none` dump) and id.
|
||||
pub async fn configured_token_limits(
|
||||
host: &dyn Host,
|
||||
id: Value,
|
||||
configured: &Map<String, Value>,
|
||||
) -> Result<(Option<i64>, Option<i64>), Error> {
|
||||
let info = effective_model_info(host, id, configured).await?;
|
||||
Ok((
|
||||
coerce_token_limit(info.get("max_input_tokens")),
|
||||
coerce_token_limit(info.get("max_output_tokens")),
|
||||
))
|
||||
}
|
||||
|
||||
/// `get_configured_mode` / `get_configured_display_name`: a configured string that is not blank.
|
||||
pub fn configured_text(value: Option<&Value>) -> Option<String> {
|
||||
value
|
||||
.and_then(Value::as_str)
|
||||
.filter(|text| !text.trim().is_empty())
|
||||
.map(str::to_owned)
|
||||
}
|
||||
|
||||
/// `get_configured_service_tiers`: each deployment's `model_info.service_tiers`, unvalidated.
|
||||
pub fn configured_service_tiers(model_infos: &[Map<String, Value>]) -> Vec<Value> {
|
||||
model_infos
|
||||
.iter()
|
||||
.map(|info| info.get("service_tiers").cloned().unwrap_or(Value::Null))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// What a listed name's deployments contribute to its `/v1/models` entry.
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct ListingInfo {
|
||||
pub cost_map_keys: Vec<String>,
|
||||
pub max_input_tokens: Option<i64>,
|
||||
pub max_output_tokens: Option<i64>,
|
||||
}
|
||||
|
||||
fn named_text(value: Option<&Value>) -> Option<&str> {
|
||||
value
|
||||
.and_then(Value::as_str)
|
||||
.filter(|text| !text.is_empty())
|
||||
}
|
||||
|
||||
/// `deployment_pricing_models(deployment)[0]`: `model_info.base_model`, then
|
||||
/// `litellm_params.base_model`, then `litellm_params.model`; malformed blocks count as unset.
|
||||
fn pricing_key(deployment: &Value) -> Option<String> {
|
||||
let block = |key: &str| {
|
||||
deployment
|
||||
.get(key)
|
||||
.and_then(Value::as_object)
|
||||
.filter(|_| !is_opaque(deployment.get(key).unwrap_or(&Value::Null)))
|
||||
};
|
||||
let params = block("litellm_params");
|
||||
let info = block("model_info");
|
||||
let base_model = info
|
||||
.and_then(|info| info.get("base_model"))
|
||||
.filter(|base| truthy(base))
|
||||
.or_else(|| params.and_then(|params| params.get("base_model")));
|
||||
named_text(base_model)
|
||||
.or_else(|| named_text(params.and_then(|params| params.get("model"))))
|
||||
.map(str::to_owned)
|
||||
}
|
||||
|
||||
/// `get_model_listing_info` over the deployments a listed name names: the cost-map keys of their
|
||||
/// models, deduplicated, and the widest token limits they configure or discovered.
|
||||
pub async fn model_listing_info(
|
||||
host: &dyn Host,
|
||||
deployments: &[Value],
|
||||
) -> Result<ListingInfo, Error> {
|
||||
let mut listing = ListingInfo::default();
|
||||
let empty = Value::Object(Map::new());
|
||||
for deployment in deployments {
|
||||
let info = match get(deployment, "model_info")? {
|
||||
Some(info) if truthy(info) => info,
|
||||
_ => &empty,
|
||||
};
|
||||
let id = get(info, "id")?.cloned().unwrap_or(Value::Null);
|
||||
let configured = match info {
|
||||
Value::Object(map) => map.clone(),
|
||||
_ => Map::new(),
|
||||
};
|
||||
let effective = effective_model_info(host, id, &configured).await?;
|
||||
for (field, widest) in [
|
||||
("max_input_tokens", &mut listing.max_input_tokens),
|
||||
("max_output_tokens", &mut listing.max_output_tokens),
|
||||
] {
|
||||
if let Some(limit) = coerce_token_limit(effective.get(field)) {
|
||||
*widest = Some(widest.map_or(limit, |current| current.max(limit)));
|
||||
}
|
||||
}
|
||||
}
|
||||
for key in deployments.iter().filter_map(pricing_key) {
|
||||
if !listing.cost_map_keys.contains(&key) {
|
||||
listing.cost_map_keys.push(key);
|
||||
}
|
||||
}
|
||||
Ok(listing)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn coerce_token_limit_reads_python_ints() {
|
||||
assert_eq!(coerce_token_limit(Some(&json!(128))), Some(128));
|
||||
assert_eq!(coerce_token_limit(Some(&json!(12.9))), Some(12));
|
||||
assert_eq!(coerce_token_limit(Some(&json!(" 1_000 "))), Some(1000));
|
||||
assert_eq!(coerce_token_limit(Some(&json!("128,000"))), None);
|
||||
assert_eq!(coerce_token_limit(Some(&json!(true))), None);
|
||||
assert_eq!(coerce_token_limit(None), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn totals_add_like_python() {
|
||||
assert_eq!(add(None, &json!(5)).unwrap(), json!(5));
|
||||
assert_eq!(add(Some(json!(5)), &json!(true)).unwrap(), json!(6));
|
||||
assert_eq!(add(Some(json!(5)), &json!(0.5)).unwrap(), json!(5.5));
|
||||
assert!(add(Some(json!(5)), &json!("1")).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn union_keeps_first_seen_order_and_empty_reads_as_none() {
|
||||
assert_eq!(
|
||||
union(&json!(["a", "b"]), &[json!("b"), json!("c")]),
|
||||
json!(["a", "b", "c"])
|
||||
);
|
||||
assert_eq!(union(&Value::Null, &[]), Value::Null);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pricing_key_prefers_the_base_model() {
|
||||
let deployment = json!({
|
||||
"litellm_params": {"model": "azure/x", "base_model": "params-base"},
|
||||
"model_info": {"base_model": "info-base"},
|
||||
});
|
||||
assert_eq!(pricing_key(&deployment).as_deref(), Some("info-base"));
|
||||
let deployment = json!({"litellm_params": {"model": "azure/x"}, "model_info": "junk"});
|
||||
assert_eq!(pricing_key(&deployment).as_deref(), Some("azure/x"));
|
||||
}
|
||||
}
|
||||
960
litellm-rust/crates/router-core/src/health.rs
Normal file
960
litellm-rust/crates/router-core/src/health.rs
Normal file
|
|
@ -0,0 +1,960 @@
|
|||
//! [`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::RouterCache,
|
||||
error::{AttemptError, Error, FailureKind},
|
||||
outcome::{OutcomeCallbacks, OutcomeWrite},
|
||||
registry::{Deployment, Registry},
|
||||
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, strum::EnumString)]
|
||||
pub enum UsageMetric {
|
||||
#[strum(serialize = "tpm")]
|
||||
Tpm,
|
||||
#[strum(serialize = "rpm")]
|
||||
Rpm,
|
||||
/// Input tokens per minute, which the io-token rate limit check counts.
|
||||
#[strum(serialize = "itpm")]
|
||||
Itpm,
|
||||
/// Output tokens per minute.
|
||||
#[strum(serialize = "otpm")]
|
||||
Otpm,
|
||||
}
|
||||
|
||||
/// `RouterCacheEnum.TPM` / `RPM` / `ITPM` / `OTPM`: `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 or an empty status never does.
|
||||
pub fn is_cooldown_required(failure: &AttemptError) -> bool {
|
||||
if failure.kind == FailureKind::ApiConnection || failure.status_absent {
|
||||
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
|
||||
}
|
||||
|
||||
/// `_EXCEPTION_POLICY_FIELDS`: a deployment policy's fields in the order Python matches them
|
||||
/// (a content-policy violation is a bad request, so it comes first).
|
||||
const DEPLOYMENT_POLICY_FIELDS: [(&str, &str); 9] = [
|
||||
(
|
||||
"ContentPolicyViolationError",
|
||||
"ContentPolicyViolationErrorAllowedFails",
|
||||
),
|
||||
("BadRequestError", "BadRequestErrorAllowedFails"),
|
||||
("AuthenticationError", "AuthenticationErrorAllowedFails"),
|
||||
("Timeout", "TimeoutErrorAllowedFails"),
|
||||
("RateLimitError", "RateLimitErrorAllowedFails"),
|
||||
("InternalServerError", "InternalServerErrorAllowedFails"),
|
||||
(
|
||||
"ServiceUnavailableError",
|
||||
"ServiceUnavailableErrorAllowedFails",
|
||||
),
|
||||
("BadGatewayError", "BadGatewayErrorAllowedFails"),
|
||||
("NotFoundError", "NotFoundErrorAllowedFails"),
|
||||
];
|
||||
|
||||
/// `_resolve_allowed_fails_from_policy` over a deployment's `model_info.allowed_fails_policy`:
|
||||
/// the first field set for a class the failure is an instance of.
|
||||
pub fn deployment_allowed_fails(
|
||||
policy: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
failure: &AttemptError,
|
||||
) -> Option<f64> {
|
||||
let policy = policy?;
|
||||
DEPLOYMENT_POLICY_FIELDS
|
||||
.iter()
|
||||
.filter(|(class, _)| failure.is_a(class))
|
||||
.find_map(|(_, field)| policy.get(*field).and_then(serde_json::Value::as_f64))
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
/// `_core_team_model_has_alternatives`: another routable deployment serves the deployment's team
|
||||
/// public name.
|
||||
pub(crate) fn team_model_has_alternatives(registry: &Registry, deployment: &Deployment) -> bool {
|
||||
let info = &deployment.model_info;
|
||||
let (Some(team_id), Some(public_name)) = (&info.team_id, &info.team_public_model_name) else {
|
||||
return false;
|
||||
};
|
||||
registry
|
||||
.team_named(team_id, public_name)
|
||||
.is_some_and(|siblings| {
|
||||
siblings
|
||||
.iter()
|
||||
.filter(|sibling| !sibling.model_info.blocked)
|
||||
.count()
|
||||
> 1
|
||||
})
|
||||
}
|
||||
|
||||
/// `_core_routing_group_has_alternatives`: `model_group` names a callable routing group whose
|
||||
/// members span more than one deployment.
|
||||
pub(crate) fn routing_group_has_alternatives(registry: &Registry, model_group: &str) -> bool {
|
||||
let resolved = registry.alias_of(model_group).unwrap_or(model_group);
|
||||
registry.routing_group_named(resolved).is_some_and(|group| {
|
||||
group
|
||||
.models
|
||||
.iter()
|
||||
.map(|member| registry.named(member).map_or(0, |members| members.len()))
|
||||
.sum::<usize>()
|
||||
> 1
|
||||
})
|
||||
}
|
||||
|
||||
/// 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>,
|
||||
clock: Arc<dyn Clock>,
|
||||
}
|
||||
|
||||
impl Health {
|
||||
pub fn new(cache: Arc<dyn RouterCache>, clock: Arc<dyn Clock>) -> Self {
|
||||
Self { cache, clock }
|
||||
}
|
||||
|
||||
pub fn cache(&self) -> &Arc<dyn RouterCache> {
|
||||
&self.cache
|
||||
}
|
||||
|
||||
pub(crate) fn now_time(&self) -> SystemTime {
|
||||
self.clock.now()
|
||||
}
|
||||
|
||||
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(())
|
||||
}
|
||||
|
||||
async fn local_count(&self, key: &str) -> Result<u64, Error> {
|
||||
Ok(self
|
||||
.cache
|
||||
.async_get_local(key)
|
||||
.await?
|
||||
.and_then(|value| value.as_f64())
|
||||
.map_or(0, |count| count as u64))
|
||||
}
|
||||
|
||||
/// The per-minute (successes, fails) this worker saw (`{id}:successes`, `{id}:fails`, kept
|
||||
/// `local_only`).
|
||||
pub async fn outcomes_this_minute(&self, id: &str) -> Result<(u64, u64), Error> {
|
||||
Ok((
|
||||
self.local_count(&format!("{id}:successes")).await?,
|
||||
self.local_count(&format!("{id}:fails")).await?,
|
||||
))
|
||||
}
|
||||
|
||||
/// `increment_deployment_successes_for_current_minute`: returns the key it counted under.
|
||||
pub async fn count_success(&self, id: &str) -> Result<String, Error> {
|
||||
let key = format!("{id}:successes");
|
||||
self.cache
|
||||
.async_increment_local(&key, 1.0, Some(USAGE_TTL))
|
||||
.await?;
|
||||
Ok(key)
|
||||
}
|
||||
|
||||
/// `increment_deployment_failures_for_current_minute`.
|
||||
pub async fn count_failure(&self, id: &str) -> Result<(), Error> {
|
||||
self.cache
|
||||
.async_increment_local(&format!("{id}:fails"), 1.0, Some(USAGE_TTL))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 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::cooldown_for_call(settings, deployment, None, failure)
|
||||
}
|
||||
|
||||
/// `cooldown_duration` for a call whose own params carry `cooldown_time` (a request's
|
||||
/// value wins over the deployment's `litellm_params`, not over its `model_info`).
|
||||
pub fn cooldown_for_call(
|
||||
settings: &Settings,
|
||||
deployment: &Deployment,
|
||||
call_cooldown: Option<f64>,
|
||||
failure: &AttemptError,
|
||||
) -> Duration {
|
||||
let configured = deployment
|
||||
.model_info
|
||||
.cooldown_time
|
||||
.or(call_cooldown)
|
||||
.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());
|
||||
configured
|
||||
.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,
|
||||
call_cooldown: Option<f64>,
|
||||
failure: &AttemptError,
|
||||
) -> Result<Option<Duration>, Error> {
|
||||
if failure.outside_deployment {
|
||||
return Ok(None);
|
||||
}
|
||||
self.count_failure(&deployment.id).await?;
|
||||
let cooldown = Self::cooldown_for_call(settings, deployment, call_cooldown, failure);
|
||||
self.cool_down_after(settings, deployment, shape, failure, cooldown)
|
||||
.await
|
||||
}
|
||||
|
||||
/// `set_cooldown_deployments` for a failure already counted: cool the deployment down for
|
||||
/// `cooldown` when the rules say so. Returns the cooldown applied, if any.
|
||||
pub async fn cool_down_after(
|
||||
&self,
|
||||
settings: &Settings,
|
||||
deployment: &Deployment,
|
||||
shape: GroupShape,
|
||||
failure: &AttemptError,
|
||||
cooldown: Duration,
|
||||
) -> Result<Option<Duration>, Error> {
|
||||
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::CallerCredentialAuthentication
|
||||
| FailureKind::CallerCredentialRateLimit
|
||||
) && !failure.caller_scoped
|
||||
&& !cooldown.is_zero()
|
||||
&& !settings.disable_cooldowns
|
||||
&& (is_cooldown_required(failure)
|
||||
|| deployment_allowed_fails(
|
||||
deployment.model_info.allowed_fails_policy.as_ref(),
|
||||
failure,
|
||||
)
|
||||
.is_some())
|
||||
&& !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);
|
||||
let lone = shape.single_deployment && !shape.has_fallback_group;
|
||||
let deployment_policy = deployment.model_info.allowed_fails_policy.as_ref();
|
||||
let deployment_fails = deployment_allowed_fails(deployment_policy, failure);
|
||||
if shape.single_deployment
|
||||
&& status == 402
|
||||
&& deployment_fails.is_none()
|
||||
&& policy_fails.is_none()
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
let deployment_allowed = deployment.model_info.allowed_fails;
|
||||
if deployment_policy.is_some() || deployment_allowed.is_some() {
|
||||
// `_should_cooldown_based_on_deployment_policy`: a deployment-wide `allowed_fails`
|
||||
// keeps the lone-deployment safety net; a policy entry naming the error does not.
|
||||
if deployment_fails.is_none()
|
||||
&& lone
|
||||
&& (deployment_allowed.is_some() || policy_fails.is_none())
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
let (allowed, suffix) = match deployment_fails {
|
||||
Some(allowed) => (
|
||||
Some(allowed),
|
||||
failure.class_names().first().map(|name| (*name).to_owned()),
|
||||
),
|
||||
None => (
|
||||
deployment_allowed.map(f64::from),
|
||||
deployment_allowed.map(|_| "generic".to_owned()),
|
||||
),
|
||||
};
|
||||
let allowed = allowed
|
||||
.or_else(|| policy_fails.map(f64::from))
|
||||
.unwrap_or_else(|| f64::from(settings.effective_allowed_fails()));
|
||||
let ttl =
|
||||
Self::configured_cooldown(deployment).unwrap_or_else(|| settings.fails_window());
|
||||
return self
|
||||
.exceeds_allowed_fails(&deployment.id, suffix.as_deref(), allowed, ttl)
|
||||
.await;
|
||||
}
|
||||
if settings.allowed_fails_policy.is_none() && !settings.allowed_fails_set() {
|
||||
let (successes, fails) = self.outcomes_this_minute(&deployment.id).await?;
|
||||
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,
|
||||
f64::from(allowed),
|
||||
settings.fails_window(),
|
||||
)
|
||||
.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: f64,
|
||||
ttl: Duration,
|
||||
) -> Result<bool, Error> {
|
||||
let fails = self
|
||||
.cache
|
||||
.async_increment(&allowed_fails_key(id, suffix), 1.0, Some(ttl))
|
||||
.await?;
|
||||
Ok(fails > allowed)
|
||||
}
|
||||
|
||||
/// Count a success, and the tokens and request against the deployment's per-minute usage
|
||||
/// when `tracks` says it has limits to track. The count is the success callback's; the usage
|
||||
/// is `make_call`'s, which runs whatever the callbacks do.
|
||||
pub async fn record_success(
|
||||
&self,
|
||||
callbacks: OutcomeCallbacks,
|
||||
deployment: &Deployment,
|
||||
total_tokens: u64,
|
||||
tracks: bool,
|
||||
) -> Result<(), Error> {
|
||||
callbacks
|
||||
.record(OutcomeWrite::Success, self.count_success(&deployment.id))
|
||||
.await?;
|
||||
if !tracks {
|
||||
return Ok(());
|
||||
}
|
||||
callbacks
|
||||
.record(
|
||||
OutcomeWrite::AttemptUsage,
|
||||
self.add_usage(deployment, None, total_tokens as f64, 1.0),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests
|
||||
/// under `model` (metadata's `deployment`, which a wildcard route fills; else the
|
||||
/// deployment's own), unless there is nothing to add.
|
||||
pub async fn add_usage(
|
||||
&self,
|
||||
deployment: &Deployment,
|
||||
model: Option<&str>,
|
||||
total_tokens: f64,
|
||||
requests: f64,
|
||||
) -> Result<(), Error> {
|
||||
if total_tokens <= 0.0 && requests <= 0.0 {
|
||||
return Ok(());
|
||||
}
|
||||
let minute = current_minute(self.clock.now());
|
||||
let model = model.unwrap_or(&deployment.params.model);
|
||||
self.cache
|
||||
.async_increment_usage(
|
||||
&[
|
||||
(
|
||||
usage_key(&deployment.id, model, UsageMetric::Tpm, &minute),
|
||||
total_tokens,
|
||||
),
|
||||
(
|
||||
usage_key(&deployment.id, model, UsageMetric::Rpm, &minute),
|
||||
requests,
|
||||
),
|
||||
],
|
||||
Some(USAGE_TTL),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `async_deployment_callback_on_failure`: a failed request counts against the deployment's
|
||||
/// per-minute requests under `model` (see [`Self::add_usage`]), whether or not it tracks usage.
|
||||
pub async fn count_failed_request(
|
||||
&self,
|
||||
deployment: &Deployment,
|
||||
model: Option<&str>,
|
||||
) -> Result<(), Error> {
|
||||
let minute = current_minute(self.clock.now());
|
||||
self.cache
|
||||
.async_increment(
|
||||
&usage_key(
|
||||
&deployment.id,
|
||||
model.unwrap_or(&deployment.params.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 std::collections::BTreeMap;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
cache::InMemoryRouterCache,
|
||||
config::{RouterConfig, RouterDynamicConfig},
|
||||
registry::{DeploymentSpec, RoutingGroupSpec},
|
||||
};
|
||||
|
||||
#[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::unreadable_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
|
||||
);
|
||||
}
|
||||
|
||||
/// An exception with no `status_code` is Python's status `""`, which no rule cools down.
|
||||
#[rstest]
|
||||
fn an_empty_status_requires_no_cooldown() {
|
||||
let mut failure = AttemptError::new(FailureKind::Other, None);
|
||||
failure.status_absent = true;
|
||||
assert!(!is_cooldown_required(&failure));
|
||||
}
|
||||
|
||||
#[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]
|
||||
#[case::content_policy_first(&["ContentPolicyViolationError", "BadRequestError"], Some(1.0))]
|
||||
#[case::bad_request(&["BadRequestError"], Some(2.0))]
|
||||
#[case::unnamed(&["Timeout"], None)]
|
||||
fn a_deployment_policy_matches_in_python_order(
|
||||
#[case] classes: &[&str],
|
||||
#[case] expected: Option<f64>,
|
||||
) {
|
||||
let policy = serde_json::json!({
|
||||
"BadRequestErrorAllowedFails": 2,
|
||||
"ContentPolicyViolationErrorAllowedFails": 1,
|
||||
});
|
||||
let mut failure = AttemptError::new(FailureKind::BadRequest, Some(400));
|
||||
failure.classes = classes.iter().map(|class| (*class).into()).collect();
|
||||
assert_eq!(
|
||||
deployment_allowed_fails(policy.as_object(), &failure),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
fn deployment(id: &str, info: serde_json::Value) -> DeploymentSpec {
|
||||
let mut model_info = serde_json::json!({"id": id});
|
||||
model_info
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.extend(info.as_object().unwrap().clone());
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"model_name": id,
|
||||
"litellm_params": {"model": format!("openai/{id}")},
|
||||
"model_info": model_info,
|
||||
}))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn health() -> Health {
|
||||
Health::new(
|
||||
Arc::new(InMemoryRouterCache::default()),
|
||||
Arc::new(SystemClock),
|
||||
)
|
||||
}
|
||||
|
||||
fn settings(dynamic: RouterDynamicConfig) -> Settings {
|
||||
Settings::new(&RouterConfig::default(), Arc::new(dynamic))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::held(false, false, 1, true)]
|
||||
#[case::caller_scoped(false, true, 1, false)]
|
||||
#[case::outside_the_deployment(true, false, 0, false)]
|
||||
#[tokio::test]
|
||||
async fn the_failure_exits_decide_what_is_counted_and_cooled(
|
||||
#[case] outside_deployment: bool,
|
||||
#[case] caller_scoped: bool,
|
||||
#[case] counted: u64,
|
||||
#[case] cooled: bool,
|
||||
) {
|
||||
let health = health();
|
||||
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
|
||||
let mut failure = AttemptError::new(FailureKind::RateLimit, Some(429));
|
||||
failure.outside_deployment = outside_deployment;
|
||||
failure.caller_scoped = caller_scoped;
|
||||
|
||||
let cooldown = health
|
||||
.record_failure(
|
||||
&settings(RouterDynamicConfig::default()),
|
||||
&deployment,
|
||||
GroupShape::default(),
|
||||
None,
|
||||
&failure,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cooldown.is_some(), cooled);
|
||||
assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, counted);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::authentication(
|
||||
FailureKind::CallerCredentialAuthentication,
|
||||
401,
|
||||
"AuthenticationError"
|
||||
)]
|
||||
#[case::rate_limit(FailureKind::CallerCredentialRateLimit, 429, "RateLimitError")]
|
||||
#[tokio::test]
|
||||
async fn a_caller_credential_failure_counts_without_cooling(
|
||||
#[case] kind: FailureKind,
|
||||
#[case] status: u16,
|
||||
#[case] base_class: &str,
|
||||
) {
|
||||
let health = health();
|
||||
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
|
||||
let failure = AttemptError::new(kind, Some(status));
|
||||
|
||||
let cooldown = health
|
||||
.record_failure(
|
||||
&settings(RouterDynamicConfig::default()),
|
||||
&deployment,
|
||||
GroupShape::default(),
|
||||
None,
|
||||
&failure,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cooldown, None);
|
||||
assert!(failure.is_a(base_class));
|
||||
assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1);
|
||||
}
|
||||
|
||||
/// `litellm.allowed_fails` written after the build: the copied 3 starts counting instead of
|
||||
/// the failure-ratio rule, which cools a 429 down at once.
|
||||
#[rstest]
|
||||
#[case::still_the_process_default(true, true)]
|
||||
#[case::the_process_default_moved(false, false)]
|
||||
#[tokio::test]
|
||||
async fn the_copied_allowed_fails_counts_once_it_differs_from_the_process_default(
|
||||
#[case] allowed_fails_is_default: bool,
|
||||
#[case] cooled: bool,
|
||||
) {
|
||||
let settings = settings(RouterDynamicConfig {
|
||||
allowed_fails: Some(3),
|
||||
allowed_fails_is_default,
|
||||
..RouterDynamicConfig::default()
|
||||
});
|
||||
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
|
||||
let failure = AttemptError::new(FailureKind::RateLimit, Some(429));
|
||||
|
||||
let cooldown = health()
|
||||
.record_failure(
|
||||
&settings,
|
||||
&deployment,
|
||||
GroupShape::default(),
|
||||
None,
|
||||
&failure,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cooldown.is_some(), cooled);
|
||||
}
|
||||
|
||||
fn team(id: &str, team_id: &str, blocked: bool) -> DeploymentSpec {
|
||||
deployment(
|
||||
id,
|
||||
serde_json::json!({"team_id": team_id, "team_public_model_name": "tp", "blocked": blocked}),
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::a_sibling(team("b", "t1", false), true)]
|
||||
#[case::a_blocked_sibling(team("b", "t1", true), false)]
|
||||
#[case::another_teams_deployment(team("b", "t2", false), false)]
|
||||
#[case::no_team(deployment("b", serde_json::json!({})), false)]
|
||||
fn a_team_deployment_has_alternatives_under_its_public_name(
|
||||
#[case] other: DeploymentSpec,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
let registry = Registry::build(
|
||||
vec![team("a", "t1", false), other],
|
||||
BTreeMap::new(),
|
||||
Vec::new(),
|
||||
)
|
||||
.unwrap();
|
||||
let first = registry.deployment("a").unwrap();
|
||||
assert_eq!(team_model_has_alternatives(®istry, &first), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::spans_two_deployments(&["a", "b"], "rg", true)]
|
||||
#[case::one_deployment(&["a"], "rg", false)]
|
||||
#[case::through_its_alias(&["a", "b"], "alias", true)]
|
||||
#[case::a_member_called_directly(&["a", "b"], "a", false)]
|
||||
fn a_routing_group_has_alternatives_past_one_deployment(
|
||||
#[case] members: &[&str],
|
||||
#[case] requested: &str,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
let group: RoutingGroupSpec = serde_json::from_value(serde_json::json!({
|
||||
"group_name": "rg", "models": members, "routing_strategy": "simple-shuffle",
|
||||
}))
|
||||
.unwrap();
|
||||
let registry = Registry::build(
|
||||
vec![
|
||||
deployment("a", serde_json::json!({})),
|
||||
deployment("b", serde_json::json!({})),
|
||||
],
|
||||
BTreeMap::from([("alias".to_owned(), "rg".to_owned())]),
|
||||
vec![group],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
routing_group_has_alternatives(®istry, requested),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
/// An empty status is counted and cools nothing down, unless the deployment's own policy
|
||||
/// names the error.
|
||||
#[rstest]
|
||||
#[case::no_policy(serde_json::json!({}), false)]
|
||||
#[case::policy_names_the_error(
|
||||
serde_json::json!({"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}),
|
||||
true
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn an_empty_status_on_the_attempt_path_needs_the_deployments_policy_to_cool(
|
||||
#[case] info: serde_json::Value,
|
||||
#[case] cooled: bool,
|
||||
) {
|
||||
let health = health();
|
||||
let deployment = Deployment::from_spec(deployment("a", info)).unwrap();
|
||||
let mut failure = AttemptError::new(FailureKind::BadRequest, None);
|
||||
failure.status_absent = true;
|
||||
|
||||
let cooldown = health
|
||||
.record_failure(
|
||||
&settings(RouterDynamicConfig::default()),
|
||||
&deployment,
|
||||
GroupShape::default(),
|
||||
None,
|
||||
&failure,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(cooldown.is_some(), cooled);
|
||||
assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::the_deployments_own_model(None, "openai/a")]
|
||||
#[case::a_filled_wildcard(Some("openai/gpt-4o"), "openai/gpt-4o")]
|
||||
#[tokio::test]
|
||||
async fn usage_is_keyed_by_the_model_served(#[case] model: Option<&str>, #[case] keyed: &str) {
|
||||
let cache = InMemoryRouterCache::default();
|
||||
let clock = Arc::new(SystemClock);
|
||||
let health = Health::new(Arc::new(cache.clone()), clock.clone());
|
||||
let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap();
|
||||
|
||||
health
|
||||
.add_usage(&deployment, model, 15.0, 1.0)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let minute = current_minute(clock.now());
|
||||
assert_eq!(
|
||||
(
|
||||
cache.get(&usage_key("a", keyed, UsageMetric::Tpm, &minute)),
|
||||
cache.get(&usage_key("a", keyed, UsageMetric::Rpm, &minute)),
|
||||
),
|
||||
(Some(serde_json::json!(15)), Some(serde_json::json!(1))),
|
||||
);
|
||||
}
|
||||
|
||||
#[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(),
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
126
litellm-rust/crates/router-core/src/health_routing.rs
Normal file
126
litellm-rust/crates/router-core/src/health_routing.rs
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
//! Health-check-driven routing: the background health checks' results, kept in the shared
|
||||
//! [`RouterCache`](crate::cache::RouterCache) under one key as Python's `DeploymentHealthCache`
|
||||
//! writes them, and the filter `_async_filter_health_check_unhealthy_deployments` /
|
||||
//! `_filter_health_check_unhealthy_deployments` run before the cooldown filter.
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc, time::UNIX_EPOCH};
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{health::Health, registry::Deployment, settings::Settings};
|
||||
|
||||
/// `DeploymentHealthCache.CACHE_KEY`.
|
||||
pub const HEALTH_STATE_KEY: &str = "litellm:health_check:deployment_health_state";
|
||||
|
||||
/// `_extract_unhealthy_ids`: the ids marked unhealthy whose result is younger than `staleness`.
|
||||
pub fn unhealthy_ids(raw: Option<&Value>, now: f64, staleness: f64) -> BTreeSet<String> {
|
||||
let Some(Value::Object(states)) = raw else {
|
||||
return BTreeSet::new();
|
||||
};
|
||||
states
|
||||
.iter()
|
||||
.filter(|(_, state)| {
|
||||
let Value::Object(state) = state else {
|
||||
return false;
|
||||
};
|
||||
// `not state.get("is_healthy", True)`: any falsy value is unhealthy.
|
||||
let unhealthy = match state.get("is_healthy") {
|
||||
None => false,
|
||||
Some(Value::Null) => true,
|
||||
Some(Value::Bool(value)) => !value,
|
||||
Some(Value::Number(value)) => value.as_f64() == Some(0.0),
|
||||
Some(Value::String(value)) => value.is_empty(),
|
||||
Some(Value::Array(value)) => value.is_empty(),
|
||||
Some(Value::Object(value)) => value.is_empty(),
|
||||
};
|
||||
let timestamp = state
|
||||
.get("timestamp")
|
||||
.and_then(Value::as_f64)
|
||||
.unwrap_or(0.0);
|
||||
unhealthy && now - timestamp < staleness
|
||||
})
|
||||
.map(|(id, _)| id.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// What the filter does when it would leave nothing.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum WhenAllUnhealthy {
|
||||
/// The safety net: every candidate is kept.
|
||||
KeepAll,
|
||||
/// A health check probe sees no deployment.
|
||||
Nothing,
|
||||
}
|
||||
|
||||
impl Health {
|
||||
/// The deployments the background health checks have not marked unhealthy, as the
|
||||
/// router's health-check routing settings say.
|
||||
pub async fn health_check_filter(
|
||||
&self,
|
||||
settings: &Settings,
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
when_all_unhealthy: WhenAllUnhealthy,
|
||||
) -> Vec<Arc<Deployment>> {
|
||||
if !settings.enable_health_check_routing {
|
||||
return deployments;
|
||||
}
|
||||
let scoped = settings.background_health_check_model_groups.as_deref();
|
||||
// With an allowed_fails_policy cooldowns are the only exclusion, unless the filter is
|
||||
// scoped to listed groups.
|
||||
if settings.allowed_fails_policy.is_some() && scoped.is_none() {
|
||||
return deployments;
|
||||
}
|
||||
// A read that fails counts as no health state.
|
||||
let raw = self
|
||||
.cache()
|
||||
.async_get_cache(HEALTH_STATE_KEY)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
let now = self
|
||||
.now_time()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|elapsed| elapsed.as_secs_f64())
|
||||
.unwrap_or_default();
|
||||
let unhealthy = unhealthy_ids(
|
||||
raw.as_ref(),
|
||||
now,
|
||||
settings.health_check_staleness.as_secs_f64(),
|
||||
);
|
||||
if unhealthy.is_empty() {
|
||||
return deployments;
|
||||
}
|
||||
let filtered: Vec<Arc<Deployment>> = deployments
|
||||
.iter()
|
||||
.filter(|d| {
|
||||
!unhealthy.contains(&d.id)
|
||||
|| scoped.is_some_and(|groups| !groups.contains(&d.model_name))
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
match (filtered.is_empty(), when_all_unhealthy) {
|
||||
(false, _) => filtered,
|
||||
(true, WhenAllUnhealthy::KeepAll) => deployments,
|
||||
(true, WhenAllUnhealthy::Nothing) => Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::fresh_unhealthy(json!({"a": {"is_healthy": false, "timestamp": 95.0}}), vec!["a"])]
|
||||
#[case::stale(json!({"a": {"is_healthy": false, "timestamp": 50.0}}), vec![])]
|
||||
#[case::healthy(json!({"a": {"is_healthy": true, "timestamp": 95.0}}), vec![])]
|
||||
#[case::no_flag(json!({"a": {"timestamp": 95.0}}), vec![])]
|
||||
#[case::not_a_dict(json!({"a": 1}), vec![])]
|
||||
fn unhealthy_ids_follow_python(#[case] raw: Value, #[case] expected: Vec<&str>) {
|
||||
let ids = unhealthy_ids(Some(&raw), 100.0, 10.0);
|
||||
assert_eq!(ids.iter().map(String::as_str).collect::<Vec<_>>(), expected);
|
||||
}
|
||||
}
|
||||
726
litellm-rust/crates/router-core/src/host.rs
Normal file
726
litellm-rust/crates/router-core/src/host.rs
Normal file
|
|
@ -0,0 +1,726 @@
|
|||
//! [`Host`]: the router core's only boundary. Everything the core needs from the embedding
|
||||
//! process goes through it, and every attempt runs through [`Host::invoke`]: the core never runs
|
||||
//! an operation itself. `PythonHost` serves it in a Python process (attempts call `litellm.<op>`,
|
||||
//! and the deployment build asks it for secrets, validation, providers and side effects).
|
||||
|
||||
use std::{collections::BTreeMap, sync::Arc, time::Duration};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
cache::RouterCache,
|
||||
config::{Op, RetryPolicy, RoutingStrategy},
|
||||
error::{AttemptError, Error, FailureKind, Payload},
|
||||
fallback::Fallbacks,
|
||||
health::Health,
|
||||
mock::MockFailure,
|
||||
outcome::OutcomeCallbacks,
|
||||
precall::PreCallFacts,
|
||||
prerouting::{PreRoute, PreRouted},
|
||||
registry::{Deployment, Registry},
|
||||
settings::Settings,
|
||||
strategy::StrategyScope,
|
||||
};
|
||||
|
||||
/// 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 the plugins and host steps that read them.
|
||||
pub body: Value,
|
||||
/// `user_api_key_team_id`, from the first metadata bucket that holds one
|
||||
/// (`get_request_team_id`).
|
||||
pub team_id: Option<String>,
|
||||
/// The `metadata` bucket's `user_api_key_team_id` alone: the team a failed level lists the
|
||||
/// group's deployments for (`async_function_with_fallbacks_common_utils`,
|
||||
/// `_stream_order_levels`). `None` lists every team's.
|
||||
pub listing_team_id: Option<String>,
|
||||
/// The caller's `user_role` is `proxy_admin`.
|
||||
pub proxy_admin: bool,
|
||||
/// The caller's `models` plus `team_models` (`user_api_key_auth`), which narrow a caller
|
||||
/// allowed only through access groups to those groups' deployments.
|
||||
pub allowed_models: Vec<String>,
|
||||
/// For a proxy admin, the model the request names (`model`, else the metadata
|
||||
/// `model_group`), which keeps one team's deployments listed under it (`filter_team_based_models`).
|
||||
pub admin_model: Option<String>,
|
||||
pub specific_deployment: bool,
|
||||
/// A per-request `num_retries`, which wins over the router's.
|
||||
pub num_retries: Option<u32>,
|
||||
/// A per-request `num_retries` that is no count, which replaces the router's as well.
|
||||
pub uncounted_retries: Option<UncountedRetries>,
|
||||
pub disable_fallbacks: bool,
|
||||
/// `_router_weights` as the request passed it: model group to deployment id to weight,
|
||||
/// validated when simple-shuffle picks ([`crate::weights::validate`]).
|
||||
pub router_weights: Option<Value>,
|
||||
/// The deployment a call on an existing resource (a video id) names: picked whenever the
|
||||
/// filters keep it (`_narrow_to_video_deployment`).
|
||||
pub pinned_deployment_id: Option<String>,
|
||||
/// A per-request `routing_strategy` override, as Python normalizes it.
|
||||
pub routing_strategy: Option<String>,
|
||||
/// Whether tag routing may apply to this request (router, request or model setting), so the
|
||||
/// host's tag stage runs.
|
||||
pub filter_tags: bool,
|
||||
/// A pass-through pick: only `use_in_pass_through` deployments may serve it.
|
||||
pub pass_through: Option<crate::select::PassThrough>,
|
||||
/// `_target_order`: the order level to pick from instead of the lowest.
|
||||
pub target_order: Option<i64>,
|
||||
/// `_excluded_deployment_ids`: never picked, even when that leaves nothing.
|
||||
pub excluded_ids: Vec<String>,
|
||||
/// `_retry_skipped_deployment_ids`: skipped unless that leaves nothing.
|
||||
pub retry_skipped_ids: Vec<String>,
|
||||
/// The endpoint `endpoint_selection` names around the call, which deployments that declare
|
||||
/// `enforce_supported_endpoints` may refuse ([`crate::eligibility`]).
|
||||
pub endpoint: Option<String>,
|
||||
/// A web search request on an async pick: `filter_web_search_deployments` drops the
|
||||
/// deployments that declare no support. The sync pick never runs that filter.
|
||||
pub web_search: bool,
|
||||
/// The handler name of a call only some providers serve (`_ROUTE_CAPABILITY_RESOLVERS`):
|
||||
/// the deployments whose provider cannot serve it are dropped unless that leaves nothing.
|
||||
pub route_handler: Option<String>,
|
||||
/// A Decisions call carrying audio or videos: the deployments that refuse them are skipped
|
||||
/// unless that leaves nothing (`_without_decisions_deployments_refusing_media`).
|
||||
pub decisions_media: bool,
|
||||
/// `f"{provider}/{model}"` for a lookup that cannot ask the host
|
||||
/// (`_core_get_routable_deployments`); a pick asks [`Host::provider_request`].
|
||||
pub wildcard_request: Option<String>,
|
||||
/// The request carries an image (`_request_carries_image`), so deployments that declare no
|
||||
/// vision are skipped while any other remains.
|
||||
pub has_image: bool,
|
||||
/// An async pick: the host's pre-routing step runs before each pick resolves its name.
|
||||
pub pre_routes: bool,
|
||||
/// Which of the router's outcome callbacks run for this request.
|
||||
pub outcome_callbacks: OutcomeCallbacks,
|
||||
/// The requested name is no model group (a search tool name): nothing is selected, every
|
||||
/// attempt is handed to the host as is, and no outcome is recorded for a deployment
|
||||
/// (`_asearch_with_fallbacks`, whose helper picks among the router's search tools).
|
||||
pub unrouted: bool,
|
||||
/// `passthrough_on_no_deployment`: a pick that fails hands the attempt to the host under the
|
||||
/// requested name, inside the level's retries and fallbacks, with no outcome recorded for a
|
||||
/// deployment (`_ageneric_api_call_with_fallbacks_helper`).
|
||||
pub passthrough_on_no_deployment: bool,
|
||||
/// The endpoint's attempt stamps no `failed_deployment_id` on its error (file and batch
|
||||
/// creation, whose result belongs to the deployment that served it), so weighted failover
|
||||
/// never runs for it.
|
||||
pub skip_weighted_failover: bool,
|
||||
/// The request creates or names a resource only one provider holds (a file upload, a batch
|
||||
/// on an uploaded file): it falls back within its own model group only
|
||||
/// (`run_async_fallback`'s `same_model_group_only`).
|
||||
pub same_group_fallbacks_only: bool,
|
||||
/// Fallback lists the request passed (`Some(None)` when it passed `None`), which replace
|
||||
/// the router's.
|
||||
pub fallbacks: Option<Option<Fallbacks>>,
|
||||
pub context_window_fallbacks: Option<Option<Fallbacks>>,
|
||||
pub content_policy_fallbacks: Option<Option<Fallbacks>>,
|
||||
/// A per-request `max_fallbacks`, which replaces the router's.
|
||||
pub max_fallbacks: Option<u32>,
|
||||
/// A per-request `cooldown_time`: the call's own params carry it, so it cools a failed
|
||||
/// deployment down for that long unless its `model_info` sets one.
|
||||
pub cooldown_time: Option<f64>,
|
||||
/// A `mock_testing_*` flag the request set.
|
||||
pub mock_failure: Option<MockFailure>,
|
||||
/// A per-request `model_group_retry_policy`, which replaces the router's.
|
||||
pub model_group_retry_policy: Option<BTreeMap<String, RetryPolicy>>,
|
||||
/// The model groups whose `model_group_retry_policy` lookup raises (the policy is no mapping
|
||||
/// of policies, or the group's entry is not one): the first failure of a hop serving one ends
|
||||
/// the hop with the host's error instead ([`RouterEvent::RetryPolicyRaised`]), no retry.
|
||||
pub retry_policy_raises: Vec<String>,
|
||||
/// `health_check_probe`: a health check's own pick, which sees no deployment when every
|
||||
/// one is marked unhealthy (the safety net keeps them all for any other request).
|
||||
pub health_check_probe: bool,
|
||||
/// What the host needs to run the call itself, such as the Python kwargs.
|
||||
pub payload: Option<Payload>,
|
||||
}
|
||||
|
||||
/// A per-request `num_retries` that is no count of retries. `async_function_with_retries` only
|
||||
/// reads it once a hop's first failure was not taken over by a retry policy: a call that succeeds
|
||||
/// first time never does.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum UncountedRetries {
|
||||
/// A negative int: `num_retries > 0` is false, so the failure stands, and a stream's plain
|
||||
/// budget is spent (`isinstance(num_retries, int)`).
|
||||
Negative,
|
||||
/// Any other value `num_retries > 0` is false for (a negative float, a NaN): the failure
|
||||
/// stands, and a stream falls through to the deployment's budget, then the router's.
|
||||
NotPositive,
|
||||
/// `num_retries > 0` or `range(num_retries)` raises (a string, a positive float): the hop ends
|
||||
/// with that error ([`RouterEvent::RetryCountRaised`]).
|
||||
Raises,
|
||||
}
|
||||
|
||||
/// One `os.environ/` reference the deployment build found ([`crate::build`]).
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SecretRef {
|
||||
/// The value as written, `os.environ/<name>`.
|
||||
pub name: String,
|
||||
/// Where the value sits in the deployment.
|
||||
pub at: SecretAt,
|
||||
}
|
||||
|
||||
/// Where a [`SecretRef`] sits in a deployment.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum SecretAt {
|
||||
/// A top-level `litellm_params` value: resolved through the host's secret manager.
|
||||
Param(String),
|
||||
/// `litellm_params.dataSources[source].parameters[key]` (`endpoint` or `key`): read from the
|
||||
/// environment alone, empty when unset.
|
||||
DataSource { source: usize, key: String },
|
||||
}
|
||||
|
||||
/// Which of the router's adds a build runs.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum BuildPath {
|
||||
/// `set_model_list` (and the constructor's list): config entries, one deployment per
|
||||
/// organization, ignored when invalid if the router ignores invalid deployments.
|
||||
Create,
|
||||
/// `add_deployment` (and upserts): one validated deployment.
|
||||
Add,
|
||||
}
|
||||
|
||||
/// What the build asks the host about one deployment.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct BuildAsk<'a> {
|
||||
pub path: BuildPath,
|
||||
/// The entry's position in the list the build was given.
|
||||
pub index: usize,
|
||||
/// The entry as the build holds it: `model_name`, `litellm_params` (secrets resolved, the
|
||||
/// organization of this deployment set), `model_info` (with its `id`) and any other keys.
|
||||
pub entry: &'a Map<String, Value>,
|
||||
/// The `model_info.id` the entry declared, before the build assigned one.
|
||||
pub declared_id: Option<&'a str>,
|
||||
/// Another entry of the same list declares the same id.
|
||||
pub duplicate_id: bool,
|
||||
}
|
||||
|
||||
/// The host's answer about one deployment.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Built {
|
||||
Entry(Map<String, Value>),
|
||||
/// Left out without an error (`model_info.supported_environments` excludes this one, or an
|
||||
/// `add_deployment` id that already exists).
|
||||
Skip,
|
||||
/// The host failed the deployment and the router ignores invalid deployments.
|
||||
Invalid,
|
||||
}
|
||||
|
||||
/// The host's provider answer.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Provided {
|
||||
Provider(Option<String>),
|
||||
Invalid,
|
||||
}
|
||||
|
||||
/// What the model-info merge ([`crate::model_info`]) asks the host while it answers one query.
|
||||
/// Asked at query time, so a price reload or a registration the host made since the router was
|
||||
/// built is what the next query reads.
|
||||
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
|
||||
#[serde(tag = "lookup", rename_all = "snake_case")]
|
||||
pub enum ModelInfoLookup {
|
||||
/// The process cost map's row for `model` (`get_model_info(model=...)`), the host choosing
|
||||
/// the keys it tries. Answered with the row, or `None`. A lookup that raises propagates,
|
||||
/// unless `tolerate` is set: then the host answers `{"failed": True}`. `exact` also asks
|
||||
/// whether the row came from a fallback-generalization rule (`{"row": ..., "generalized":
|
||||
/// bool}`).
|
||||
ModelInfo {
|
||||
model: Value,
|
||||
tolerate: bool,
|
||||
exact: bool,
|
||||
},
|
||||
/// The raw cost-map entry under `key` (`model_cost.get(key)`): answered with
|
||||
/// `{"entry": value}`, or `None` when the key is absent. `process` reads the process global
|
||||
/// (`litellm.model_cost`) rather than the host's live settings, where the Python router does.
|
||||
CostEntry { key: String, process: bool },
|
||||
/// `get_llm_provider` for the queried deployment's `model` and validated params: answered
|
||||
/// with `[model, provider]`.
|
||||
Provider { model: String },
|
||||
/// The deployments a wildcard pattern routes `name` to: answered with `[{id, model}]`, or
|
||||
/// `None` when no pattern matches.
|
||||
PatternModels { name: String },
|
||||
/// The base model a bedrock model id names (`get_bedrock_base_model`).
|
||||
BedrockBase { model_id: String },
|
||||
/// The limits the background model-info refresh discovered for deployment `id`.
|
||||
Discovered { id: Value },
|
||||
/// The queried deployment's own `model_info`, read when the merge reaches it (a malformed one
|
||||
/// raises there).
|
||||
UserModelInfo,
|
||||
/// A log record, emitted where the Python router logs it.
|
||||
Log { level: LogLevel, message: String },
|
||||
/// Group info: deployment `index`'s validated params (`LiteLLM_Params(**litellm_params)`, which
|
||||
/// raises on bad params): answered with `{"model", "configurable_clientside_auth_params"}`.
|
||||
GroupParams { index: usize },
|
||||
/// Group info: what the host knows of deployment `index` serving `model`: answered with
|
||||
/// `{"model_info"}` (its merged model info, `None` when that fails), `"mapped"`
|
||||
/// (`deployment_is_catalog_mapped`), `"litellm_model"` and `"llm_provider"`, and, when it has
|
||||
/// no model info, `"supported_openai_params"`.
|
||||
GroupDeployment { index: usize, model: Value },
|
||||
/// Group info: deployment `index`'s capabilities under `model_info`: answered with
|
||||
/// `{"supported_endpoints"}` (empty when the team may not use it), `"supports_fast_mode"` and
|
||||
/// `"reasoning_efforts"` (`resolve_supported_reasoning_efforts`, or `None`).
|
||||
GroupCapabilities {
|
||||
index: usize,
|
||||
model_info: Value,
|
||||
litellm_model: Value,
|
||||
llm_provider: Value,
|
||||
mapped: bool,
|
||||
},
|
||||
/// Group info: the group's first deployment as the host builds it (`ModelGroupInfo(**info)`,
|
||||
/// whose validation may raise): answered with the fields the merge compares, as validated.
|
||||
GroupSeed { info: Map<String, Value> },
|
||||
/// `intersect_supported_reasoning_efforts(current, resolved)`, in the host's advertisement order.
|
||||
ReasoningIntersect { current: Value, resolved: Value },
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum LogLevel {
|
||||
Debug,
|
||||
Error,
|
||||
}
|
||||
|
||||
/// 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 chain is: 0 for the requested name.
|
||||
pub fallback_level: u32,
|
||||
pub timeout: Option<Duration>,
|
||||
/// Request params the fallback chain's dict targets merged in (`kwargs.update(target)`).
|
||||
pub params: Option<Arc<Map<String, Value>>>,
|
||||
/// The chain level whose request metadata this attempt carries (see [`RouterEvent::Fallback`]).
|
||||
pub scope: u32,
|
||||
/// The retry budget the level's loop stamps on the request before this attempt.
|
||||
pub max_retries: u32,
|
||||
/// How many deployments `get_model_list` lists for the level's group (`model_group_size`).
|
||||
pub model_group_size: usize,
|
||||
}
|
||||
|
||||
/// 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,
|
||||
/// A response the provider blocked, which the router may raise to fall back.
|
||||
pub blocked: Option<Blocked>,
|
||||
/// A stream the host reports the outcome of when it ends ([`crate::Router::report_success`],
|
||||
/// [`crate::Router::report_failure`]), or a realtime session, whose success the host counts
|
||||
/// only when it ended cleanly: the attempt counts its request and `total_tokens` now, and no
|
||||
/// success.
|
||||
pub deferred: bool,
|
||||
/// Whether the attempt counts against the deployment's per-minute usage, when the host knows
|
||||
/// better than the deployment's own limits ([`Deployment::tracks_usage`]): a limit only the
|
||||
/// model catalog sets (`_deployment_tracks_usage`), or a realtime session, whose `None`
|
||||
/// response counts nothing (`increment_deployment_usage_for_response`).
|
||||
pub tracks_usage: Option<bool>,
|
||||
}
|
||||
|
||||
/// How a provider blocked a response the call otherwise served.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum Blocked {
|
||||
/// A chat response whose first choice was content-filtered (or that has no choices).
|
||||
ContentFilter,
|
||||
/// An Anthropic /v1/messages safeguard refusal carrying `stop_details`.
|
||||
Refusal,
|
||||
}
|
||||
|
||||
/// 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,
|
||||
/// The name the request resolved from (Python's `model`), which strategies key state by.
|
||||
pub name: &'a str,
|
||||
/// Whose selector a strategy stage reads.
|
||||
pub scope: StrategyScope,
|
||||
/// The deployments the name resolved to for this caller, before health filters
|
||||
/// (`request_deployments`), for the tag stage.
|
||||
pub pool: &'a [Arc<Deployment>],
|
||||
/// `_excluded_deployment_ids` a weighted failover hop put on the request, which the callback
|
||||
/// filters read (`deployment_ids_failed_in_request`).
|
||||
pub excluded_ids: &'a [String],
|
||||
/// `_retry_skipped_deployment_ids` the retry loop put on the request, as `excluded_ids`.
|
||||
pub skipped_ids: &'a [String],
|
||||
/// `_target_order` as the filters read it: the level an order fallback targets, else the one
|
||||
/// the request named.
|
||||
pub target_order: Option<i64>,
|
||||
}
|
||||
|
||||
/// 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),
|
||||
/// Tag routing (`get_deployments_for_tag`), not yet ported: return the deployments that remain.
|
||||
Tags,
|
||||
/// `_deployment_declares_no_vision`: return the candidates that declare no vision support.
|
||||
TextOnly,
|
||||
/// `_deployment_serves_route`: return the candidates whose provider cannot serve the call
|
||||
/// `handler` names (a Decisions call on a provider with no Decisions config).
|
||||
RouteIncapable { handler: String },
|
||||
/// `_decisions_deployment_refuses_media`: return the candidates that refuse the audio or
|
||||
/// videos a Decisions call carries.
|
||||
MediaRefusing,
|
||||
}
|
||||
|
||||
/// Routing events, delivered through [`Host::observe`].
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum RouterEvent {
|
||||
Selected {
|
||||
model_group: String,
|
||||
deployment_id: String,
|
||||
/// How long the pick took (`_acompletion`'s `async_get_available_deployment` service event).
|
||||
pick: Duration,
|
||||
},
|
||||
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,
|
||||
/// The new level's request metadata scope, copied from `parent`'s at the hop.
|
||||
scope: u32,
|
||||
parent: u32,
|
||||
/// The `max_fallbacks` the chain reads, stamped on the hop's request.
|
||||
max_fallbacks: u32,
|
||||
/// The non-string `model` of the dict target this hop routes
|
||||
/// ([`crate::fallback::FallbackTarget::unroutable_model`]): the hop fails on it before any
|
||||
/// attempt, and the host raises what Python's level raises for that value.
|
||||
unroutable: Option<Value>,
|
||||
},
|
||||
/// A hop of the fallback chain returned: `log_success_fallback_event` or
|
||||
/// `log_failure_fallback_event` for the level at `parent`.
|
||||
FallbackDone {
|
||||
from: String,
|
||||
to: String,
|
||||
level: u32,
|
||||
scope: u32,
|
||||
parent: u32,
|
||||
succeeded: bool,
|
||||
},
|
||||
/// The router raised a `ContentPolicyViolationError` for a blocked response, since a
|
||||
/// fallback would serve it (`_should_raise_content_policy_error`,
|
||||
/// `_should_raise_anthropic_refusal_error`).
|
||||
Blocked {
|
||||
level: u32,
|
||||
model: String,
|
||||
blocked: Blocked,
|
||||
},
|
||||
/// `log_retry`: the router recorded a failure in `scope`'s request metadata
|
||||
/// (`previous_models`, `request_retry_count`) before retrying or falling back.
|
||||
RetryLogged {
|
||||
scope: u32,
|
||||
level: u32,
|
||||
/// The group the failed level served (`kwargs["model"]`).
|
||||
model: String,
|
||||
/// The retry number stamped before the failed try; `None` keeps the scope's own.
|
||||
attempted_retries: Option<u32>,
|
||||
cause: RetryCause,
|
||||
},
|
||||
/// `async_function_with_retries` opened a level's retry loop, before any pick: `scope`'s
|
||||
/// request metadata tracks no retry yet, `max_retries` of them, and the group's size.
|
||||
RetriesOpened {
|
||||
scope: u32,
|
||||
max_retries: u32,
|
||||
model_group_size: usize,
|
||||
},
|
||||
/// The level's first failure settled its overrides: the budget `scope` then tracks.
|
||||
RetryBudget {
|
||||
scope: u32,
|
||||
max_retries: u32,
|
||||
},
|
||||
/// A pick of the async selection or the sync one failed on `error`: the async one's `except`
|
||||
/// logs every refusal to the request's logging object (`_async_get_available_deployment`).
|
||||
PickFailed {
|
||||
level: u32,
|
||||
model: String,
|
||||
error: PickError,
|
||||
},
|
||||
/// A level of the fallback chain serving `model` ended on an error a pick raised (no
|
||||
/// deployment for the caller, none available, a stale routing group, ...): Python raises
|
||||
/// that error for the level, whatever its earlier attempts raised.
|
||||
LevelFailed {
|
||||
level: u32,
|
||||
model: String,
|
||||
error: PickError,
|
||||
/// Raised by the level's fallback handling, while the level's own error was handled
|
||||
/// (so it chains to that error), rather than by the level's retry loop.
|
||||
handling: bool,
|
||||
},
|
||||
/// The router raised a `mock_testing_*` failure for `model` instead of calling it.
|
||||
MockFailed {
|
||||
mock: MockFailure,
|
||||
level: u32,
|
||||
model: String,
|
||||
num_retries: Option<u32>,
|
||||
},
|
||||
/// The first failure of a hop serving `model` looked up a `model_group_retry_policy` that
|
||||
/// raises for it: the hop ends with the error the lookup raised (the host raises it, from the
|
||||
/// failure `cause` names), without retrying, and falls back as from any plain exception.
|
||||
RetryPolicyRaised {
|
||||
level: u32,
|
||||
model: String,
|
||||
cause: RetryCause,
|
||||
},
|
||||
/// The first failure of a hop serving `model` read a request `num_retries` that raises
|
||||
/// ([`UncountedRetries::Raises`]): the hop ends with that error (the host raises it, from the
|
||||
/// failure `cause` names), and falls back as from any plain exception.
|
||||
RetryCountRaised {
|
||||
level: u32,
|
||||
model: String,
|
||||
cause: RetryCause,
|
||||
},
|
||||
/// The level serving `model` reached the fallback handling that appends the context window
|
||||
/// and content policy hints to its error, before any hop runs
|
||||
/// (`async_function_with_fallbacks_common_utils`). The order path never gets there.
|
||||
FallingBack {
|
||||
level: u32,
|
||||
model: String,
|
||||
},
|
||||
/// The order fallback chain of the level at `level` served nothing, so the level raises from
|
||||
/// `async_function_with_fallbacks_common_utils`' order path, with no outcome message: what the
|
||||
/// chain's last hop raised when `last_hop`, else its own error (a context window error, or a
|
||||
/// chain no hop of which ran). Weighted failover raising what it does not catch
|
||||
/// (`_maybe_run_weighted_failover`) ends the level the same way.
|
||||
OrderChainFailed {
|
||||
level: u32,
|
||||
last_hop: bool,
|
||||
},
|
||||
/// The regular fallback handling of the level at `level` served nothing, where
|
||||
/// `async_function_with_fallbacks_common_utils` reaches its `except` and logs the failure.
|
||||
/// `chain` when the level had targets to walk: the error it logs is what the last hop raised
|
||||
/// when `last_hop`, else its own. With no `chain`, the host decides: only the host holds
|
||||
/// whether a level with no match had a fallback list at all, which Python does not log.
|
||||
/// The order path and weighted failover are outside that `try` and never report this.
|
||||
FallbacksFailed {
|
||||
level: u32,
|
||||
chain: bool,
|
||||
last_hop: bool,
|
||||
},
|
||||
/// 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,
|
||||
/// The level the retry loop served.
|
||||
level: u32,
|
||||
/// The last failure was a pick's error, which the level then raises
|
||||
/// ([`RouterEvent::LevelFailed`]), rather than an attempt's.
|
||||
picked: bool,
|
||||
},
|
||||
}
|
||||
|
||||
/// A pick's error, as the events that name it carry it.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PickError(pub Arc<Error>);
|
||||
|
||||
impl PickError {
|
||||
/// The copy of `error` events carry, when a pick raised it.
|
||||
pub fn of(error: &Error) -> Option<Self> {
|
||||
error.pick_copy().map(|copy| Self(Arc::new(copy)))
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for PickError {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
format!("{:?}", self.0) == format!("{:?}", other.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Which error a [`RouterEvent::RetryLogged`] records.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum RetryCause {
|
||||
/// The last attempt's failure.
|
||||
Attempt,
|
||||
/// The error the level ended on (`run_async_fallback`'s `original_exception`).
|
||||
Level,
|
||||
/// A pick raised this error: the group was unknown or had none available, the stale
|
||||
/// routing-group member map named a group no longer configured, ...
|
||||
Pick(PickError),
|
||||
/// Any other error no attempt raised.
|
||||
Other,
|
||||
}
|
||||
|
||||
impl RetryCause {
|
||||
/// The cause that names `error`.
|
||||
pub fn of(error: &Error) -> Self {
|
||||
match error {
|
||||
Error::Attempt(_) => Self::Attempt,
|
||||
_ => PickError::of(error).map_or(Self::Other, Self::Pick),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait Host: Send + Sync {
|
||||
/// Run one attempt against the picked deployment. The core records its outcome (cooldowns,
|
||||
/// allowed-fails counters, usage) itself; the host only runs it.
|
||||
async fn invoke(&self, attempt: &Attempt) -> Result<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>;
|
||||
|
||||
/// Pre-routing before a pick resolves `ask.name`: the model group to resolve instead, a
|
||||
/// custom strategy's deployment to serve, or `None` to resolve the name itself. Asked only
|
||||
/// for picks that pre-route (`pre_routes`).
|
||||
async fn pre_route(&self, ask: &PreRoute<'_>) -> Result<Option<PreRouted>, Error> {
|
||||
let _ = ask;
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Access and budget checks before a cross-group fallback target.
|
||||
/// `Err` when a check raised: the chain ends there, outside any target's `try`.
|
||||
async fn allow_fallback(
|
||||
&self,
|
||||
request: &RouterRequest,
|
||||
target: &str,
|
||||
) -> Result<bool, AttemptError>;
|
||||
|
||||
/// Whether `run_async_fallback` refuses a chain for one of its dict `targets`, a server-owned
|
||||
/// federation field or an `oidc/` api key. The chain is walked by `level`.
|
||||
async fn refuses_fallbacks(&self, _level: u32, _targets: &[Map<String, Value>]) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
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 `op` on `provider`. Checked for every served op and
|
||||
/// deployment when the router is built or changed, so an unsupported pair fails the config,
|
||||
/// never a request. `PythonHost` runs every op.
|
||||
fn supports(&self, _op: &Op, _provider: &str) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Whether [`Self::select`] can run strategies and plugins.
|
||||
fn can_select(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Run the routing plugins for `ctx.name` when `run` (the async pick's pre-routing hook), and
|
||||
/// report the candidate models they left on the request, if any.
|
||||
async fn routing_plugin_models(
|
||||
&self,
|
||||
_ctx: &RoutingContext<'_>,
|
||||
_run: bool,
|
||||
) -> Result<Option<Vec<String>>, Error> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// The facts the `enable_pre_call_checks` filter decides by (token count, context windows,
|
||||
/// param support), for `candidates` and the deployments in `ctx.pool`. A host without them
|
||||
/// reports no input, so nothing is checked.
|
||||
async fn pre_call_facts(
|
||||
&self,
|
||||
_ctx: &RoutingContext<'_>,
|
||||
_candidates: &[Arc<Deployment>],
|
||||
) -> Result<PreCallFacts, Error> {
|
||||
Ok(PreCallFacts::default())
|
||||
}
|
||||
|
||||
/// `f"{provider}/{model}"` for `ctx.name`, which `PatternMatchRouter.get_pattern` tries when
|
||||
/// no pattern matches the name itself: only the host can infer a provider. `None` from a
|
||||
/// host that infers none.
|
||||
async fn provider_request(&self, _ctx: &RoutingContext<'_>) -> Result<Option<String>, Error> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in
|
||||
/// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and
|
||||
/// falls back from like any attempt failure. A host without filters keeps every candidate.
|
||||
async fn callback_filters(
|
||||
&self,
|
||||
_ctx: &RoutingContext<'_>,
|
||||
candidates: Vec<Arc<Deployment>>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
Ok(candidates)
|
||||
}
|
||||
/// The values of `refs`, one per ref, in order: what the deployment build substitutes for
|
||||
/// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per
|
||||
/// deployment, as the build reaches it. A host without secrets leaves every reference as
|
||||
/// written.
|
||||
async fn secrets(&self, index: usize, refs: &[SecretRef]) -> Result<Vec<Value>, Error> {
|
||||
let _ = index;
|
||||
Ok(refs.iter().map(|r| Value::String(r.name.clone())).collect())
|
||||
}
|
||||
|
||||
/// Check one deployment before it is added, and answer the entry as the host normalized it,
|
||||
/// a skip (the deployment is left out), or [`Built::Invalid`] (it failed and the router
|
||||
/// ignores invalid deployments). A host without checks accepts the entry as given.
|
||||
async fn validate(&self, ask: &BuildAsk<'_>) -> Result<Built, Error> {
|
||||
Ok(Built::Entry(ask.entry.clone()))
|
||||
}
|
||||
|
||||
/// The provider that serves one validated deployment (`get_llm_provider`), `None` for one
|
||||
/// served without a provider (a prompt-management model), or [`Provided::Invalid`]. A host
|
||||
/// without provider resolution answers `None`, so [`Deployment::provider`] reads the model.
|
||||
async fn provider(&self, ask: &BuildAsk<'_>) -> Result<Provided, Error> {
|
||||
let _ = ask;
|
||||
Ok(Provided::Provider(None))
|
||||
}
|
||||
|
||||
/// Run the host's side effects for an accepted deployment and answer the entry as the
|
||||
/// registry stores it. A host without side effects stores the entry it validated.
|
||||
async fn added(&self, ask: &BuildAsk<'_>, provider: Option<&str>) -> Result<Built, Error> {
|
||||
let _ = provider;
|
||||
Ok(Built::Entry(ask.entry.clone()))
|
||||
}
|
||||
|
||||
/// Answer one model-info lookup ([`ModelInfoLookup`]). A host without a cost map answers
|
||||
/// `None` to everything, so no deployment has model info.
|
||||
async fn model_info(&self, lookup: &ModelInfoLookup) -> Result<Value, Error> {
|
||||
let _ = lookup;
|
||||
Ok(Value::Null)
|
||||
}
|
||||
}
|
||||
727
litellm-rust/crates/router-core/src/host_index.rs
Normal file
727
litellm-rust/crates/router-core/src/host_index.rs
Normal file
|
|
@ -0,0 +1,727 @@
|
|||
//! [`HostIndexes`]: the registry indexes the Python router hands out, kept the way it keeps them.
|
||||
//!
|
||||
//! The Python router keeps its indexes by history rather than deriving them from the deployment
|
||||
//! list: `model_names` is the set of names `set_model_list` was given (invalid entries included),
|
||||
//! `deployment_names` only ever grows, assigning `model_list` reindexes nothing, a removal shifts
|
||||
//! the positions every index holds (a caller-written one included), and the wildcard pattern
|
||||
//! routers keep the copy of a deployment they were given. The registry keeps one of these per
|
||||
//! version and applies those rules on its own writes; the host reads them back as the live
|
||||
//! objects `registry().index()`, `names()` and `patterns()` return, and hands a caller's write to
|
||||
//! any of them back as a changeset ([`IndexWrite`]).
|
||||
//!
|
||||
//! Values are JSON so whatever a caller wrote reads back as written, junk included.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// One wildcard pattern router (`PatternMatchRouter`): regex to the deployments it routes to,
|
||||
/// most specific pattern first.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct PatternRouter {
|
||||
pub patterns: Vec<(String, Vec<Value>)>,
|
||||
}
|
||||
|
||||
/// The indexes, as the host reads and writes them.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct HostIndexes {
|
||||
/// `model_id_to_deployment_index_map`: id to position, in insertion order.
|
||||
pub id_map: Map<String, Value>,
|
||||
/// `model_name_to_deployment_indices`: name to its positions.
|
||||
pub model_name: Map<String, Value>,
|
||||
/// `team_model_to_deployment_indices`: `[team_id, public name, positions]`.
|
||||
pub team_model: Vec<(String, String, Value)>,
|
||||
/// `model_names`: a set, its members in the order they joined.
|
||||
pub model_names: Vec<Value>,
|
||||
pub deployment_names: Vec<Value>,
|
||||
pub team_public_model_names: Vec<String>,
|
||||
pub provider_default_ids: Vec<String>,
|
||||
/// `pattern_router`.
|
||||
pub patterns: PatternRouter,
|
||||
/// `team_pattern_routers`: team id to its router.
|
||||
pub team_patterns: Vec<(String, PatternRouter)>,
|
||||
}
|
||||
|
||||
/// A host write to the indexes: each field it names replaces the one held.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct IndexWrite {
|
||||
pub id_map: Option<Map<String, Value>>,
|
||||
pub model_name: Option<Map<String, Value>>,
|
||||
pub team_model: Option<Vec<(String, String, Value)>>,
|
||||
pub model_names: Option<Vec<Value>>,
|
||||
pub deployment_names: Option<Vec<Value>>,
|
||||
pub team_public_model_names: Option<Vec<String>>,
|
||||
pub provider_default_ids: Option<Vec<String>>,
|
||||
pub patterns: Option<PatternRouter>,
|
||||
pub team_patterns: Option<Vec<(String, PatternRouter)>>,
|
||||
}
|
||||
|
||||
/// A removal the Python router could not finish: an index held a position it cannot compare
|
||||
/// (`_without_indices` raised `TypeError`). The indexes keep what it shifted before that.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
#[error("'>' not supported between instances of '{found}' and 'int'")]
|
||||
pub struct RemovalError {
|
||||
pub found: &'static str,
|
||||
}
|
||||
|
||||
impl HostIndexes {
|
||||
/// Apply a host write.
|
||||
pub fn write(&mut self, write: IndexWrite) {
|
||||
let IndexWrite {
|
||||
id_map,
|
||||
model_name,
|
||||
team_model,
|
||||
model_names,
|
||||
deployment_names,
|
||||
team_public_model_names,
|
||||
provider_default_ids,
|
||||
patterns,
|
||||
team_patterns,
|
||||
} = write;
|
||||
if let Some(value) = id_map {
|
||||
self.id_map = value;
|
||||
}
|
||||
if let Some(value) = model_name {
|
||||
self.model_name = value;
|
||||
}
|
||||
if let Some(value) = team_model {
|
||||
self.team_model = value;
|
||||
}
|
||||
if let Some(value) = model_names {
|
||||
self.model_names = unique(value);
|
||||
}
|
||||
if let Some(value) = deployment_names {
|
||||
self.deployment_names = value;
|
||||
}
|
||||
if let Some(value) = team_public_model_names {
|
||||
self.team_public_model_names = value;
|
||||
}
|
||||
if let Some(value) = provider_default_ids {
|
||||
self.provider_default_ids = value;
|
||||
}
|
||||
if let Some(value) = patterns {
|
||||
self.patterns = value;
|
||||
}
|
||||
if let Some(value) = team_patterns {
|
||||
self.team_patterns = value;
|
||||
}
|
||||
}
|
||||
|
||||
/// `set_model_list` starting over: the id, name and team indexes and the team routers go;
|
||||
/// the deployment names, the wildcard router and the provider default ids stay.
|
||||
pub fn reset(&mut self) {
|
||||
self.id_map.clear();
|
||||
self.model_name.clear();
|
||||
self.team_model.clear();
|
||||
self.team_patterns.clear();
|
||||
self.team_public_model_names.clear();
|
||||
}
|
||||
|
||||
/// `set_model_list`'s `model_names`: the names it was given.
|
||||
pub fn set_model_names(&mut self, names: Vec<String>) {
|
||||
self.model_names = unique(names.into_iter().map(Value::String).collect());
|
||||
}
|
||||
|
||||
/// A deployment joined the list at `position` (`_add_deployment` then
|
||||
/// `_add_model_to_list_and_index_map`); `add` is `add_deployment`, which also names its group.
|
||||
pub fn added(&mut self, entry: &Map<String, Value>, position: usize, add: bool) {
|
||||
let params = entry.get("litellm_params").and_then(Value::as_object);
|
||||
let info = entry.get("model_info").and_then(Value::as_object);
|
||||
let model_name = entry.get("model_name").and_then(Value::as_str);
|
||||
let id = info
|
||||
.and_then(|info| info.get("id"))
|
||||
.filter(|id| !id.is_null());
|
||||
self.deployment_names.push(
|
||||
params
|
||||
.and_then(|params| params.get("model"))
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
if let Some(name) = model_name.filter(|name| name.contains('*')) {
|
||||
self.patterns.add(name, Value::Object(entry.clone()));
|
||||
if let Some(id) = id.and_then(Value::as_str).filter(|id| !id.is_empty()) {
|
||||
self.provider_default_ids.push(id.to_owned());
|
||||
}
|
||||
}
|
||||
let team_id = info
|
||||
.and_then(|info| info.get("team_id"))
|
||||
.filter(|v| !v.is_null());
|
||||
let public = info
|
||||
.and_then(|info| info.get("team_public_model_name"))
|
||||
.filter(|v| !v.is_null());
|
||||
if let (Some(team_id), Some(Value::String(public))) = (team_id, public)
|
||||
&& public.contains('*')
|
||||
{
|
||||
let team = python_str(team_id);
|
||||
let router = match self.team_patterns.iter_mut().find(|(t, _)| *t == team) {
|
||||
Some((_, router)) => router,
|
||||
None => {
|
||||
self.team_patterns.push((team, PatternRouter::default()));
|
||||
&mut self.team_patterns.last_mut().expect("just pushed").1
|
||||
}
|
||||
};
|
||||
router.add(public, Value::Object(entry.clone()));
|
||||
}
|
||||
let position = Value::from(position);
|
||||
if let Some(id) = id {
|
||||
self.id_map.insert(python_str(id), position.clone());
|
||||
}
|
||||
if let Some(name) = model_name.filter(|name| !name.is_empty()) {
|
||||
push_position(
|
||||
self.model_name
|
||||
.entry(name.to_owned())
|
||||
.or_insert_with(|| Value::Array(Vec::new())),
|
||||
&position,
|
||||
);
|
||||
}
|
||||
if add {
|
||||
let name = Value::String(model_name.unwrap_or_default().to_owned());
|
||||
if !self.model_names.contains(&name) {
|
||||
self.model_names.push(name);
|
||||
}
|
||||
}
|
||||
// `_update_team_model_index`: a truthy team id and public name.
|
||||
if let (Some(team_id), Some(public)) = (team_id, public)
|
||||
&& truthy(team_id)
|
||||
&& truthy(public)
|
||||
{
|
||||
let (team, public) = (python_str(team_id), python_str(public));
|
||||
if !self.team_public_model_names.contains(&public) {
|
||||
self.team_public_model_names.push(public.clone());
|
||||
}
|
||||
let held = match self
|
||||
.team_model
|
||||
.iter_mut()
|
||||
.find(|(t, p, _)| *t == team && *p == public)
|
||||
{
|
||||
Some((_, _, held)) => held,
|
||||
None => {
|
||||
self.team_model
|
||||
.push((team, public, Value::Array(Vec::new())));
|
||||
&mut self.team_model.last_mut().expect("just pushed").2
|
||||
}
|
||||
};
|
||||
if let Value::Array(positions) = held
|
||||
&& !positions.contains(&position)
|
||||
{
|
||||
positions.push(position);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The deployment `id` left the list from `position` (`_update_deployment_indices_after_removal`).
|
||||
pub fn removed(&mut self, id: &str, position: usize) -> Result<(), RemovalError> {
|
||||
let mut id_map = Map::new();
|
||||
for (key, held) in std::mem::take(&mut self.id_map) {
|
||||
if key == id {
|
||||
continue;
|
||||
}
|
||||
let shifted = shift(&held, position)?.unwrap_or(held);
|
||||
id_map.insert(key, shifted);
|
||||
}
|
||||
self.id_map = id_map;
|
||||
let names: Vec<String> = self.model_name.keys().cloned().collect();
|
||||
for name in names {
|
||||
let kept = without(&self.model_name[&name], position)?;
|
||||
if kept.is_empty() {
|
||||
self.model_name.remove(&name);
|
||||
self.model_names
|
||||
.retain(|held| held.as_str() != Some(name.as_str()));
|
||||
} else {
|
||||
self.model_name.insert(name, Value::Array(kept));
|
||||
}
|
||||
}
|
||||
let mut team_model = Vec::new();
|
||||
let held = std::mem::take(&mut self.team_model);
|
||||
let mut rest = held.into_iter();
|
||||
for (team, public, positions) in rest.by_ref() {
|
||||
match without(&positions, position) {
|
||||
Ok(kept) if kept.is_empty() => {}
|
||||
Ok(kept) => team_model.push((team, public, Value::Array(kept))),
|
||||
Err(error) => {
|
||||
team_model.push((team, public, positions));
|
||||
team_model.extend(rest);
|
||||
self.team_model = team_model;
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
self.team_model = team_model;
|
||||
let mut publics: Vec<String> = Vec::new();
|
||||
for (_, public, _) in &self.team_model {
|
||||
if !publics.contains(public) {
|
||||
publics.push(public.clone());
|
||||
}
|
||||
}
|
||||
self.team_public_model_names = publics;
|
||||
self.patterns.remove(id);
|
||||
for (_, router) in &mut self.team_patterns {
|
||||
router.remove(id);
|
||||
}
|
||||
self.team_patterns
|
||||
.retain(|(_, router)| !router.patterns.is_empty());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Whether the id map holds `id` (`_has_model_id`).
|
||||
pub fn has_id(&self, id: &str) -> bool {
|
||||
self.id_map.contains_key(id)
|
||||
}
|
||||
|
||||
/// What the id map holds for `id`, as written (`model_id_to_deployment_index_map.get(id)`).
|
||||
pub fn id_entry(&self, id: &str) -> Option<&Value> {
|
||||
self.id_map.get(id)
|
||||
}
|
||||
|
||||
/// What the name index holds for `name`, as written.
|
||||
pub fn named(&self, name: &str) -> Option<&Value> {
|
||||
self.model_name.get(name)
|
||||
}
|
||||
|
||||
/// What the team index holds for `(team, public)`, as written.
|
||||
pub fn team_named(&self, team: &str, public: &str) -> Option<&Value> {
|
||||
self.team_model
|
||||
.iter()
|
||||
.find(|(t, p, _)| t == team && p == public)
|
||||
.map(|(_, _, positions)| positions)
|
||||
}
|
||||
|
||||
/// `team_public_model_names_of`: the public names the team index holds for `team`.
|
||||
pub fn team_public_names_of(&self, team: &str) -> Vec<String> {
|
||||
let mut names: Vec<String> = Vec::new();
|
||||
for (t, public, _) in &self.team_model {
|
||||
if t == team && !names.contains(public) {
|
||||
names.push(public.clone());
|
||||
}
|
||||
}
|
||||
names
|
||||
}
|
||||
|
||||
/// The position the id map holds for `id`, when it is a position.
|
||||
pub fn position_of(&self, id: &str) -> Option<usize> {
|
||||
self.id_map
|
||||
.get(id)
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|position| usize::try_from(position).ok())
|
||||
}
|
||||
}
|
||||
|
||||
impl PatternRouter {
|
||||
/// `add_pattern`: a new regex joins in specificity order (a stable sort, most specific
|
||||
/// first); a known one gains the deployment.
|
||||
pub fn add(&mut self, pattern: &str, deployment: Value) {
|
||||
let regex = pattern_to_regex(pattern);
|
||||
if let Some((_, deployments)) = self.patterns.iter_mut().find(|(held, _)| *held == regex) {
|
||||
deployments.push(deployment);
|
||||
return;
|
||||
}
|
||||
self.patterns.push((regex, vec![deployment]));
|
||||
self.patterns
|
||||
.sort_by_key(|(regex, _)| std::cmp::Reverse(regex_specificity(regex)));
|
||||
}
|
||||
|
||||
/// `remove_deployment`: every deployment with this id goes, and so does a pattern left empty.
|
||||
pub fn remove(&mut self, id: &str) {
|
||||
for (_, deployments) in &mut self.patterns {
|
||||
deployments.retain(|deployment| {
|
||||
deployment
|
||||
.get("model_info")
|
||||
.and_then(|info| info.get("id"))
|
||||
.and_then(Value::as_str)
|
||||
!= Some(id)
|
||||
});
|
||||
}
|
||||
self.patterns
|
||||
.retain(|(_, deployments)| !deployments.is_empty());
|
||||
}
|
||||
|
||||
/// `route`: the deployments of the first pattern `request` matches, each a copy whose
|
||||
/// `litellm_params.model` takes the parts the wildcards matched; `filtered` keeps to the
|
||||
/// patterns of those names.
|
||||
pub fn route(&self, request: &str, filtered: Option<&[String]>) -> Option<Vec<Value>> {
|
||||
let filtered: Option<Vec<String>> =
|
||||
filtered.map(|names| names.iter().map(|name| pattern_to_regex(name)).collect());
|
||||
for (regex, deployments) in &self.patterns {
|
||||
if filtered
|
||||
.as_ref()
|
||||
.is_some_and(|names| !names.contains(regex))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(groups) = regex_match(regex, request) else {
|
||||
continue;
|
||||
};
|
||||
let mut routed = Vec::with_capacity(deployments.len());
|
||||
for deployment in deployments {
|
||||
let mut copy = deployment.clone();
|
||||
let model = deployment
|
||||
.get("litellm_params")
|
||||
.and_then(|params| params.get("model"));
|
||||
let Some(Value::String(model)) = model else {
|
||||
// `deployment["litellm_params"]["model"]` raising ends the route with None.
|
||||
return None;
|
||||
};
|
||||
let named = deployment_model_name(&groups, request, model);
|
||||
if model.contains('*')
|
||||
&& let Some(params) =
|
||||
deployment.get("litellm_params").and_then(Value::as_object)
|
||||
&& let Some(info) = copy.get_mut("model_info").and_then(Value::as_object_mut)
|
||||
{
|
||||
for flag in cache_flags_left_to_the_catalog(params) {
|
||||
info.remove(*flag);
|
||||
}
|
||||
}
|
||||
if let Some(params) = copy
|
||||
.get_mut("litellm_params")
|
||||
.and_then(Value::as_object_mut)
|
||||
{
|
||||
params.insert("model".to_owned(), Value::String(named));
|
||||
}
|
||||
routed.push(copy);
|
||||
}
|
||||
return Some(routed);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// `get_models_by_pattern`'s pick: the first (request, pattern) pair that matches, trying
|
||||
/// `model` against every pattern, then `provider/model`.
|
||||
pub fn models_by_pattern(&self, requests: &[String]) -> Option<Vec<String>> {
|
||||
for request in requests {
|
||||
for (regex, deployments) in &self.patterns {
|
||||
if let Some(groups) = regex_match(regex, request) {
|
||||
return Some(
|
||||
deployments
|
||||
.iter()
|
||||
.filter_map(|deployment| {
|
||||
deployment
|
||||
.get("litellm_params")
|
||||
.and_then(|params| params.get("model"))
|
||||
.and_then(Value::as_str)
|
||||
.map(|model| deployment_model_name(&groups, request, model))
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// `_cache_flags_left_to_the_catalog`: the `model_info` flags a wildcard deployment does not speak
|
||||
/// for on the models it matches. An azure_ai wildcard's `supports_prompt_cache_breakpoint` does.
|
||||
/// `inferred_provider` names `azure_ai` only for a model carrying that prefix.
|
||||
pub(crate) fn cache_flags_left_to_the_catalog(
|
||||
params: &Map<String, Value>,
|
||||
) -> &'static [&'static str] {
|
||||
let provider = params
|
||||
.get("custom_llm_provider")
|
||||
.filter(|provider| truthy(provider))
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| {
|
||||
params
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|model| model.split_once('/'))
|
||||
.map(|(prefix, _)| prefix)
|
||||
});
|
||||
if provider == Some("azure_ai") {
|
||||
&["supports_prompt_caching"]
|
||||
} else {
|
||||
&[
|
||||
"supports_prompt_cache_breakpoint",
|
||||
"supports_prompt_caching",
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
/// `re.escape(pattern).replace(r"\*", "(.*)")`.
|
||||
pub fn pattern_to_regex(pattern: &str) -> String {
|
||||
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),
|
||||
}
|
||||
}
|
||||
regex
|
||||
}
|
||||
|
||||
/// `PatternUtils.calculate_pattern_specificity`: the regex's length, then how much regex syntax
|
||||
/// it has.
|
||||
pub fn regex_specificity(regex: &str) -> (usize, usize) {
|
||||
let complexity = regex
|
||||
.chars()
|
||||
.filter(|c| matches!(c, '*' | '+' | '?' | '\\' | '^' | '$' | '|' | '(' | ')'))
|
||||
.count();
|
||||
(regex.chars().count(), complexity)
|
||||
}
|
||||
|
||||
/// A pattern regex read back as its literal runs between wildcards, or None when it holds regex
|
||||
/// syntax `pattern_to_regex` never writes (a caller wrote it).
|
||||
fn literal_runs(regex: &str) -> Option<Vec<String>> {
|
||||
let mut runs = vec![String::new()];
|
||||
let mut chars = regex.chars().peekable();
|
||||
while let Some(c) = chars.next() {
|
||||
match c {
|
||||
'\\' => runs.last_mut()?.push(chars.next()?),
|
||||
'(' => {
|
||||
let rest: String = chars.by_ref().take(3).collect();
|
||||
if rest != ".*)" {
|
||||
return None;
|
||||
}
|
||||
runs.push(String::new());
|
||||
}
|
||||
'.' | '*' | '+' | '?' | '[' | ']' | '{' | '}' | '|' | '^' | '$' | ')' => return None,
|
||||
other => runs.last_mut()?.push(other),
|
||||
}
|
||||
}
|
||||
Some(runs)
|
||||
}
|
||||
|
||||
/// `re.match(regex, request)` for a pattern regex: anchored at the start only, each `(.*)`
|
||||
/// greedy. The groups it captured, or None when it does not match.
|
||||
pub fn regex_match(regex: &str, request: &str) -> Option<Vec<String>> {
|
||||
let runs = literal_runs(regex)?;
|
||||
let rest = request.strip_prefix(runs[0].as_str())?;
|
||||
let mut groups = Vec::new();
|
||||
greedy(&runs[1..], rest, &mut groups).then_some(groups)
|
||||
}
|
||||
|
||||
/// Match `runs` (each a literal that follows a group) against `rest`, longest group first.
|
||||
fn greedy(runs: &[String], rest: &str, groups: &mut Vec<String>) -> bool {
|
||||
let Some((literal, later)) = runs.split_first() else {
|
||||
return true;
|
||||
};
|
||||
let mut ends: Vec<usize> = rest.char_indices().map(|(at, _)| at).collect();
|
||||
ends.push(rest.len());
|
||||
for &end in ends.iter().rev() {
|
||||
let (group, after) = rest.split_at(end);
|
||||
let Some(after) = after.strip_prefix(literal.as_str()) else {
|
||||
continue;
|
||||
};
|
||||
groups.push(group.to_owned());
|
||||
if greedy(later, after, groups) {
|
||||
return true;
|
||||
}
|
||||
groups.pop();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// `set_deployment_model_name`: the deployment's model with each wildcard replaced by the group
|
||||
/// matched for it, the request itself when more groups matched than it has wildcards.
|
||||
pub fn deployment_model_name(groups: &[String], request: &str, model: &str) -> String {
|
||||
if !model.contains('*') {
|
||||
return model.to_owned();
|
||||
}
|
||||
if groups.len() > model.matches('*').count() {
|
||||
return request.to_owned();
|
||||
}
|
||||
let mut named = model.to_owned();
|
||||
for group in groups {
|
||||
named = named.replacen('*', group, 1);
|
||||
}
|
||||
named
|
||||
}
|
||||
|
||||
fn unique(values: Vec<Value>) -> Vec<Value> {
|
||||
let mut kept: Vec<Value> = Vec::with_capacity(values.len());
|
||||
for value in values {
|
||||
if !kept.contains(&value) {
|
||||
kept.push(value);
|
||||
}
|
||||
}
|
||||
kept
|
||||
}
|
||||
|
||||
fn push_position(held: &mut Value, position: &Value) {
|
||||
if let Value::Array(positions) = held {
|
||||
positions.push(position.clone());
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's truthiness for a JSON value.
|
||||
pub(crate) fn truthy(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Null => false,
|
||||
Value::Bool(flag) => *flag,
|
||||
Value::Number(number) => number.as_f64() != Some(0.0),
|
||||
Value::String(text) => !text.is_empty(),
|
||||
Value::Array(items) => !items.is_empty(),
|
||||
Value::Object(items) => !items.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `str()` for an id or team id.
|
||||
pub(crate) fn python_str(value: &Value) -> String {
|
||||
match value {
|
||||
Value::String(text) => text.clone(),
|
||||
Value::Bool(true) => "True".to_owned(),
|
||||
Value::Bool(false) => "False".to_owned(),
|
||||
Value::Null => "None".to_owned(),
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn type_name(value: &Value) -> &'static str {
|
||||
match value {
|
||||
Value::Null => "NoneType",
|
||||
Value::Bool(_) => "bool",
|
||||
Value::Number(number) if number.is_f64() => "float",
|
||||
Value::Number(_) => "int",
|
||||
Value::String(_) => "str",
|
||||
Value::Array(_) => "list",
|
||||
Value::Object(_) => "dict",
|
||||
}
|
||||
}
|
||||
|
||||
/// `idx - 1 if idx > removal_idx else idx`: None keeps `held` as is.
|
||||
fn shift(held: &Value, position: usize) -> Result<Option<Value>, RemovalError> {
|
||||
let removal = position as f64;
|
||||
let number = match held {
|
||||
Value::Bool(flag) => f64::from(u8::from(*flag)),
|
||||
Value::Number(number) => number.as_f64().unwrap_or_default(),
|
||||
other => {
|
||||
return Err(RemovalError {
|
||||
found: type_name(other),
|
||||
});
|
||||
}
|
||||
};
|
||||
if number <= removal {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(match held {
|
||||
Value::Number(n) if n.is_i64() || n.is_u64() => Value::from(n.as_i64().unwrap_or(0) - 1),
|
||||
_ => Value::from(number - 1.0),
|
||||
}))
|
||||
}
|
||||
|
||||
/// `_without_indices`: the positions without `position`, the later ones shifted down.
|
||||
fn without(held: &Value, position: usize) -> Result<Vec<Value>, RemovalError> {
|
||||
let Value::Array(positions) = held else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let removal = Value::from(position);
|
||||
let mut kept = Vec::with_capacity(positions.len());
|
||||
for index in positions {
|
||||
if *index == removal {
|
||||
continue;
|
||||
}
|
||||
kept.push(shift(index, position)?.unwrap_or_else(|| index.clone()));
|
||||
}
|
||||
Ok(kept)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn entry(name: &str, id: &str, model: &str) -> Map<String, Value> {
|
||||
serde_json::from_value(json!({
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": model},
|
||||
"model_info": {"id": id},
|
||||
}))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn removal_shifts_every_index_and_drops_emptied_names() {
|
||||
let mut indexes = HostIndexes::default();
|
||||
indexes.added(&entry("gpt", "a", "openai/a"), 0, false);
|
||||
indexes.added(&entry("gpt", "b", "openai/b"), 1, false);
|
||||
indexes.added(&entry("claude", "c", "anthropic/c"), 2, false);
|
||||
indexes.set_model_names(vec!["gpt".into(), "claude".into(), "broken".into()]);
|
||||
indexes.removed("a", 0).unwrap();
|
||||
assert_eq!(
|
||||
indexes.id_map,
|
||||
json!({"b": 0, "c": 1}).as_object().unwrap().clone()
|
||||
);
|
||||
assert_eq!(indexes.model_name["gpt"], json!([0]));
|
||||
indexes.removed("c", 1).unwrap();
|
||||
assert!(!indexes.model_name.contains_key("claude"));
|
||||
assert_eq!(indexes.model_names, vec![json!("gpt"), json!("broken")]);
|
||||
assert_eq!(indexes.deployment_names.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_pattern_routes_with_the_matched_parts() {
|
||||
let mut router = PatternRouter::default();
|
||||
router.add(
|
||||
"llmengine/*",
|
||||
json!({"litellm_params": {"model": "openai/*"}, "model_info": {"id": "w"}}),
|
||||
);
|
||||
router.add(
|
||||
"*meta.llama3*",
|
||||
json!({"litellm_params": {"model": "bedrock/meta.llama3*"}, "model_info": {"id": "m"}}),
|
||||
);
|
||||
let routed = router.route("llmengine/foo", None).unwrap();
|
||||
assert_eq!(routed[0]["litellm_params"]["model"], "openai/foo");
|
||||
let routed = router.route("hello-world-meta.llama3-70b", None).unwrap();
|
||||
assert_eq!(
|
||||
routed[0]["litellm_params"]["model"],
|
||||
"hello-world-meta.llama3-70b"
|
||||
);
|
||||
assert_eq!(
|
||||
regex_match(r"a(.*)b(.*)", "axbyb").unwrap(),
|
||||
vec!["xby", ""]
|
||||
);
|
||||
router.remove("w");
|
||||
assert!(router.route("llmengine/foo", None).is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::any_wildcard(
|
||||
json!({"model": "openai/*"}),
|
||||
json!({"id": "w", "keep": 1}),
|
||||
)]
|
||||
#[case::azure_ai_keeps_the_breakpoint_flag(
|
||||
json!({"model": "azure_ai/*"}),
|
||||
json!({"id": "w", "keep": 1, "supports_prompt_cache_breakpoint": true}),
|
||||
)]
|
||||
#[case::azure_ai_named_by_custom_llm_provider(
|
||||
json!({"model": "*", "custom_llm_provider": "azure_ai"}),
|
||||
json!({"id": "w", "keep": 1, "supports_prompt_cache_breakpoint": true}),
|
||||
)]
|
||||
#[case::a_fixed_model_speaks_for_itself(
|
||||
json!({"model": "openai/gpt-4o"}),
|
||||
json!({
|
||||
"id": "w",
|
||||
"keep": 1,
|
||||
"supports_prompt_cache_breakpoint": true,
|
||||
"supports_prompt_caching": true,
|
||||
}),
|
||||
)]
|
||||
fn a_wildcard_leaves_the_cache_flags_to_the_catalog(
|
||||
#[case] params: Value,
|
||||
#[case] expected: Value,
|
||||
) {
|
||||
let mut router = PatternRouter::default();
|
||||
router.add(
|
||||
"any/*",
|
||||
json!({
|
||||
"litellm_params": params,
|
||||
"model_info": {
|
||||
"id": "w",
|
||||
"keep": 1,
|
||||
"supports_prompt_cache_breakpoint": true,
|
||||
"supports_prompt_caching": true,
|
||||
},
|
||||
}),
|
||||
);
|
||||
let routed = router.route("any/gpt-4o", None).unwrap();
|
||||
assert_eq!(routed[0]["model_info"], expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -43,21 +43,706 @@
|
|||
//! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the
|
||||
//! policy is exhausted.
|
||||
//!
|
||||
//! # Execution
|
||||
//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs
|
||||
//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A
|
||||
//! started attempt is never replayed on a different executor.
|
||||
//! # Hosts: the only boundary
|
||||
//! The core reaches outside itself only through [`Host`], and every attempt runs through
|
||||
//! `Host::invoke`: the core picks the deployment, the host runs the call, the core records the
|
||||
//! outcome. There is no second, "native" attempt path in the core. Which code runs an attempt is
|
||||
//! the host's business:
|
||||
//! - `PythonHost` (inside a Python process) builds the request for the picked deployment and calls
|
||||
//! `litellm.<op>`. It never calls back into a Python router.
|
||||
//!
|
||||
//! # Hosts
|
||||
//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select`
|
||||
//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks,
|
||||
//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all
|
||||
//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost`
|
||||
//! serves them from the existing Python code when the router runs inside Python; `NativeHost`
|
||||
//! serves them in the Rust server, where an operation without a Rust implementation makes the
|
||||
//! config fail at build time, never partway through a request.
|
||||
//! An op or provider a host cannot run fails `Host::supports` when the router is built or
|
||||
//! changed, never partway through a request.
|
||||
//!
|
||||
//! The core also builds its deployments through the host ([`build`]): for each configured entry
|
||||
//! it finds the secret references and asks `secrets`, then `validate`, `provider` and `added`.
|
||||
//!
|
||||
//! Besides `invoke` and `supports`, a host provides `select` for strategies and plugins not yet
|
||||
//! ported, `allow_fallback` for access and budget checks, `refuses_fallbacks` for a chain's dict
|
||||
//! targets, `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that
|
||||
//! holds all routing state under the same keys, encodings and TTLs the Python router uses. The
|
||||
//! core writes cooldowns, allowed-fails counters and usage itself, through that cache.
|
||||
//!
|
||||
//! # Plugins
|
||||
//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive)
|
||||
//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and
|
||||
//! the current request), never the router itself.
|
||||
//!
|
||||
//! # Modules
|
||||
//! - [`registry`]: deployments and their indices, `deployment_id`, `ChangeSet`.
|
||||
//! - [`config`] / [`settings`]: static and dynamic config, and the frozen `Settings` view.
|
||||
//! - [`cache`] / [`health`]: the shared routing cache, cooldowns and usage.
|
||||
//! - [`select`]: resolve, filter chain, simple-shuffle.
|
||||
//! - [`call`]: retries and fallbacks.
|
||||
//! - [`eligibility`]: the endpoints a deployment refuses to serve.
|
||||
//! - [`precall`]: the `enable_pre_call_checks` filter.
|
||||
//! - [`model_info`]: model info merged from the cost map, asked of the host at query time.
|
||||
//! - [`host`]: the `Host` boundary.
|
||||
|
||||
pub mod access;
|
||||
pub mod build;
|
||||
pub mod cache;
|
||||
pub mod call;
|
||||
pub mod config;
|
||||
pub mod eligibility;
|
||||
pub mod error;
|
||||
pub mod fallback;
|
||||
pub mod group_info;
|
||||
pub mod health;
|
||||
pub mod health_routing;
|
||||
pub mod host;
|
||||
pub mod host_index;
|
||||
pub mod lookup;
|
||||
pub mod mock;
|
||||
pub mod model_info;
|
||||
pub mod outcome;
|
||||
pub mod precall;
|
||||
pub mod prerouting;
|
||||
pub mod ptu;
|
||||
pub mod registry;
|
||||
pub mod resolve;
|
||||
pub mod retry;
|
||||
pub mod routing_groups;
|
||||
pub mod routing_plugins;
|
||||
pub mod select;
|
||||
pub mod settings;
|
||||
pub mod strategy;
|
||||
pub mod weights;
|
||||
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use parking_lot::Mutex;
|
||||
|
||||
pub use crate::{
|
||||
cache::{InMemoryRouterCache, RouterCache},
|
||||
call::RouterResponse,
|
||||
config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy},
|
||||
error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported},
|
||||
health::{Clock, Health, Outcome, SystemClock},
|
||||
host::{
|
||||
Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause,
|
||||
RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage,
|
||||
},
|
||||
outcome::{OutcomeCallbacks, OutcomeWrite},
|
||||
prerouting::{PreRoute, PreRouted},
|
||||
registry::{ChangeSet, Deployment, DeploymentSpec, Registry, RegistryVersion},
|
||||
select::{Randomness, SeededRandomness, ThreadRandomness},
|
||||
settings::Settings,
|
||||
};
|
||||
use crate::{
|
||||
health::{GroupShape, routing_group_has_alternatives, team_model_has_alternatives},
|
||||
registry::DeploymentFilter,
|
||||
select::SelectInputs,
|
||||
};
|
||||
|
||||
/// The registry and settings one call reads, swapped atomically by `apply`.
|
||||
pub(crate) struct RouterStateSnapshot {
|
||||
pub registry: Arc<Registry>,
|
||||
pub settings: Arc<Settings>,
|
||||
}
|
||||
|
||||
impl RouterStateSnapshot {
|
||||
/// `_get_routing_context` for a pick of `name`: unless the request overrides the strategy,
|
||||
/// a stale member map naming a group no longer configured raises `KeyError(group)`.
|
||||
pub(crate) fn missing_routing_group(
|
||||
&self,
|
||||
request: &RouterRequest,
|
||||
name: &str,
|
||||
) -> Result<(), Error> {
|
||||
if strategy::request_override(request.routing_strategy.as_deref()).is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
match self.registry.missing_member_group(name) {
|
||||
Some(group) => Err(Error::MissingRoutingGroup {
|
||||
group: group.to_owned(),
|
||||
}),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct Core {
|
||||
host: Arc<dyn Host>,
|
||||
rng: Arc<dyn Randomness>,
|
||||
health: Health,
|
||||
config: RouterConfig,
|
||||
snapshot: ArcSwap<RouterStateSnapshot>,
|
||||
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()) {
|
||||
self.validate_deployment(&deployment)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reject a deployment whose provider this host cannot run a served op on.
|
||||
fn validate_deployment(&self, deployment: &Deployment) -> Result<(), ConfigError> {
|
||||
for op in &self.config.served_ops {
|
||||
if !self.host.supports(op, deployment.provider()) {
|
||||
return Err(Unsupported::Op {
|
||||
op: op.clone(),
|
||||
provider: deployment.provider().to_owned(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// How a reported failure is recorded.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct FailureReport {
|
||||
/// Cool down for the deployment's own `cooldown_time` (then the provider's retry-after),
|
||||
/// as the failure callback does, rather than the router default.
|
||||
pub deployment_cooldown: bool,
|
||||
/// The `cooldown_time` the failed call's own params carried (`model_info`, then
|
||||
/// `litellm_params`), which wins over the provider's retry-after for a callback report.
|
||||
pub cooldown_time: Option<f64>,
|
||||
/// Count the failure only: its status is one no cooldown rule reads (an empty status).
|
||||
pub count_only: bool,
|
||||
/// The model group the request was for, whose fallbacks keep a lone deployment cooling.
|
||||
pub requested_group: Option<String>,
|
||||
/// Cool down without counting the failure (`set_cooldown_deployments` alone).
|
||||
pub uncounted: bool,
|
||||
/// Also count the failed request against the deployment's per-minute requests
|
||||
/// (`async_deployment_callback_on_failure`).
|
||||
pub counts_request: bool,
|
||||
/// The model the request's usage is keyed by (metadata's `deployment`, which a wildcard route
|
||||
/// fills), when it is not the deployment's own.
|
||||
pub usage_model: Option<String>,
|
||||
/// The failure callbacks of a logged failure; a report no callback makes keeps the default.
|
||||
pub callbacks: OutcomeCallbacks,
|
||||
}
|
||||
|
||||
/// How a reported success is recorded.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct SuccessReport {
|
||||
/// See [`host::Output::tracks_usage`].
|
||||
pub tracks_usage: Option<bool>,
|
||||
/// See [`FailureReport::usage_model`].
|
||||
pub usage_model: Option<String>,
|
||||
/// See [`FailureReport::callbacks`].
|
||||
pub callbacks: OutcomeCallbacks,
|
||||
}
|
||||
|
||||
/// Builds a [`Router`], with test seams.
|
||||
pub struct RouterBuilder {
|
||||
config: RouterConfig,
|
||||
dynamic: RouterDynamicConfig,
|
||||
host: Arc<dyn Host>,
|
||||
rng: Arc<dyn Randomness>,
|
||||
clock: Arc<dyn Clock>,
|
||||
}
|
||||
|
||||
impl RouterBuilder {
|
||||
#[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(),
|
||||
)?
|
||||
.with_member_groups(self.config.routing_group_members.clone());
|
||||
let settings = Settings::new(&self.config, Arc::new(self.dynamic));
|
||||
let core = Core {
|
||||
health: Health::new(self.host.cache(), self.clock),
|
||||
host: self.host,
|
||||
rng: self.rng,
|
||||
config: self.config,
|
||||
snapshot: ArcSwap::from_pointee(RouterStateSnapshot {
|
||||
registry: Arc::new(Registry::default()),
|
||||
settings: Arc::new(settings.clone()),
|
||||
}),
|
||||
writer: Mutex::new(()),
|
||||
retired: AtomicBool::new(false),
|
||||
};
|
||||
core.validate(®istry, &settings)?;
|
||||
core.snapshot.store(Arc::new(RouterStateSnapshot {
|
||||
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 {
|
||||
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();
|
||||
routing_plugins::refuse_sync_pick(
|
||||
snapshot.settings.routing_plugins,
|
||||
&snapshot.settings.routing_strategy,
|
||||
request,
|
||||
)?;
|
||||
let routed = self
|
||||
.core
|
||||
.pre_routed(PreRoute {
|
||||
op,
|
||||
request,
|
||||
name: &request.model,
|
||||
fallback_depth: 0,
|
||||
retry: 0,
|
||||
})
|
||||
.await?;
|
||||
let name = match routed {
|
||||
PreRouted::Model(name) => name,
|
||||
PreRouted::Deployment(deployment) => return Ok(deployment),
|
||||
};
|
||||
snapshot.missing_routing_group(request, &name)?;
|
||||
let candidates = self
|
||||
.core
|
||||
.candidates(&snapshot, op, request, &name, &SelectInputs::default())
|
||||
.await?;
|
||||
self.core.pick(&snapshot, op, request, &candidates).await
|
||||
}
|
||||
|
||||
/// The deployments a pick for `request` would choose among: resolved, team- and
|
||||
/// cooldown-filtered, through the plugins and the lowest `order`.
|
||||
pub async fn healthy(
|
||||
&self,
|
||||
op: &Op,
|
||||
request: &RouterRequest,
|
||||
) -> Result<select::Candidates, Error> {
|
||||
self.live()?;
|
||||
let snapshot = self.core.snapshot.load_full();
|
||||
let candidates = self
|
||||
.core
|
||||
.candidates(
|
||||
&snapshot,
|
||||
op,
|
||||
request,
|
||||
&request.model,
|
||||
&SelectInputs::default(),
|
||||
)
|
||||
.await?;
|
||||
// `_core_async_get_healthy_deployments` raises when nothing is left.
|
||||
if candidates.healthy.is_empty() {
|
||||
return Err(Error::NoHealthyDeployments {
|
||||
model: candidates.failed_model(&snapshot.registry),
|
||||
});
|
||||
}
|
||||
Ok(candidates)
|
||||
}
|
||||
|
||||
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. Returns whether
|
||||
/// the deployment was cooled down.
|
||||
pub async fn report(&self, id: &str, outcome: Outcome) -> Result<bool, Error> {
|
||||
match outcome {
|
||||
Outcome::Success { total_tokens } => self
|
||||
.report_success(id, total_tokens as f64, 1.0, SuccessReport::default())
|
||||
.await
|
||||
.map(|_| false),
|
||||
Outcome::Failure(failure) => self
|
||||
.report_failure(id, &failure, FailureReport::default())
|
||||
.await
|
||||
.map(|cooldown| cooldown.is_some()),
|
||||
}
|
||||
}
|
||||
|
||||
/// A success a logging callback saw (`deployment_callback_on_success`): counted for a
|
||||
/// configured deployment, with its tokens and requests when it tracks usage. Returns the
|
||||
/// success counter's key, or `None` for an unknown deployment.
|
||||
pub async fn report_success(
|
||||
&self,
|
||||
id: &str,
|
||||
total_tokens: f64,
|
||||
requests: f64,
|
||||
report: SuccessReport,
|
||||
) -> Result<Option<String>, Error> {
|
||||
report
|
||||
.callbacks
|
||||
.record(
|
||||
OutcomeWrite::Success,
|
||||
self.reported_success(id, total_tokens, requests, &report),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn reported_success(
|
||||
&self,
|
||||
id: &str,
|
||||
total_tokens: f64,
|
||||
requests: f64,
|
||||
report: &SuccessReport,
|
||||
) -> Result<Option<String>, Error> {
|
||||
let snapshot = self.core.snapshot.load_full();
|
||||
let Some(deployment) = snapshot.registry.deployment(id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let key = self.core.health.count_success(id).await?;
|
||||
if report
|
||||
.tracks_usage
|
||||
.unwrap_or_else(|| deployment.tracks_usage())
|
||||
{
|
||||
self.core
|
||||
.health
|
||||
.add_usage(
|
||||
&deployment,
|
||||
report.usage_model.as_deref(),
|
||||
total_tokens,
|
||||
requests,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(Some(key))
|
||||
}
|
||||
|
||||
/// `sync_deployment_callback_on_success`: count a success against `id`, configured or not,
|
||||
/// and return the counter's key.
|
||||
pub async fn count_success(
|
||||
&self,
|
||||
id: &str,
|
||||
callbacks: OutcomeCallbacks,
|
||||
) -> Result<Option<String>, Error> {
|
||||
callbacks
|
||||
.record(OutcomeWrite::Success, async {
|
||||
self.core.health.count_success(id).await.map(Some)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// A failure the router did not see itself: counted against `id` (known or not), then
|
||||
/// cooled down by the rules (`record_failure`, `deployment_callback_on_failure`). Returns the
|
||||
/// cooldown applied, if any.
|
||||
pub async fn report_failure(
|
||||
&self,
|
||||
id: &str,
|
||||
failure: &AttemptError,
|
||||
report: FailureReport,
|
||||
) -> Result<Option<std::time::Duration>, Error> {
|
||||
report
|
||||
.callbacks
|
||||
.record(
|
||||
OutcomeWrite::Failure,
|
||||
self.reported_failure(id, failure, &report),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn reported_failure(
|
||||
&self,
|
||||
id: &str,
|
||||
failure: &AttemptError,
|
||||
report: &FailureReport,
|
||||
) -> Result<Option<std::time::Duration>, Error> {
|
||||
if !report.uncounted && !failure.outside_deployment {
|
||||
self.core.health.count_failure(id).await?;
|
||||
}
|
||||
let snapshot = self.core.snapshot.load_full();
|
||||
let Some(deployment) = snapshot.registry.deployment(id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if report.counts_request {
|
||||
self.core
|
||||
.health
|
||||
.count_failed_request(&deployment, report.usage_model.as_deref())
|
||||
.await?;
|
||||
}
|
||||
if failure.outside_deployment {
|
||||
return Ok(None);
|
||||
}
|
||||
// An empty status is no cooldown (`_is_cooldown_required`), unless the deployment's
|
||||
// own policy names the error.
|
||||
if report.count_only
|
||||
&& crate::health::deployment_allowed_fails(
|
||||
deployment.model_info.allowed_fails_policy.as_ref(),
|
||||
failure,
|
||||
)
|
||||
.is_none()
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let cooldown = if report.deployment_cooldown {
|
||||
report
|
||||
.cooldown_time
|
||||
.filter(|seconds| *seconds >= 0.0)
|
||||
.and_then(|seconds| std::time::Duration::try_from_secs_f64(seconds).ok())
|
||||
.or(failure.retry_after)
|
||||
.unwrap_or(snapshot.settings.cooldown_time)
|
||||
} else {
|
||||
snapshot.settings.cooldown_time
|
||||
};
|
||||
let shape = GroupShape {
|
||||
single_deployment: snapshot
|
||||
.registry
|
||||
.group(&deployment.model_name)
|
||||
.is_some_and(|group| group.len() == 1)
|
||||
&& !report
|
||||
.requested_group
|
||||
.as_deref()
|
||||
.is_some_and(|group| routing_group_has_alternatives(&snapshot.registry, group))
|
||||
&& !team_model_has_alternatives(&snapshot.registry, &deployment),
|
||||
has_fallback_group: snapshot
|
||||
.settings
|
||||
.dynamic
|
||||
.fallback_targets(&[report
|
||||
.requested_group
|
||||
.as_deref()
|
||||
.unwrap_or(&deployment.model_name)])
|
||||
.is_some(),
|
||||
};
|
||||
self.core
|
||||
.health
|
||||
.cool_down_after(&snapshot.settings, &deployment, shape, failure, cooldown)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 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(¤t.settings),
|
||||
};
|
||||
self.core.validate(®istry, &settings)?;
|
||||
let version = registry.version();
|
||||
self.core.snapshot.store(Arc::new(RouterStateSnapshot {
|
||||
registry: Arc::new(registry),
|
||||
settings: Arc::new(settings),
|
||||
}));
|
||||
Ok(version)
|
||||
}
|
||||
|
||||
/// `get_router_model_info` for one deployment, its cost-map reads asked of the host now
|
||||
/// ([`model_info::router_model_info`]).
|
||||
pub async fn router_model_info(
|
||||
&self,
|
||||
query: &model_info::RouterModelInfoQuery,
|
||||
) -> Result<Option<model_info::MergedModelInfo>, Error> {
|
||||
model_info::router_model_info(self.core.host.as_ref(), query).await
|
||||
}
|
||||
|
||||
/// `_set_model_group_info` for the deployments `query` names ([`group_info::model_group_info`]).
|
||||
pub async fn model_group_info(
|
||||
&self,
|
||||
query: &group_info::GroupInfoQuery,
|
||||
) -> Result<Option<group_info::GroupInfo>, Error> {
|
||||
group_info::model_group_info(self.core.host.as_ref(), query).await
|
||||
}
|
||||
|
||||
/// `get_configured_token_limits` for deployment `id` ([`group_info::configured_token_limits`]).
|
||||
pub async fn configured_token_limits(
|
||||
&self,
|
||||
id: serde_json::Value,
|
||||
configured: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Result<(Option<i64>, Option<i64>), Error> {
|
||||
group_info::configured_token_limits(self.core.host.as_ref(), id, configured).await
|
||||
}
|
||||
|
||||
/// `get_model_listing_info` over a listed name's deployments ([`group_info::model_listing_info`]).
|
||||
pub async fn model_listing_info(
|
||||
&self,
|
||||
deployments: &[serde_json::Value],
|
||||
) -> Result<group_info::ListingInfo, Error> {
|
||||
group_info::model_listing_info(self.core.host.as_ref(), deployments).await
|
||||
}
|
||||
|
||||
/// `get_deployment_model_info` for deployment `model_id` serving `model_name`
|
||||
/// ([`model_info::deployment_model_info`]).
|
||||
pub async fn deployment_model_info(
|
||||
&self,
|
||||
model_id: &str,
|
||||
model_name: &serde_json::Value,
|
||||
) -> Result<Option<serde_json::Map<String, serde_json::Value>>, Error> {
|
||||
model_info::deployment_model_info(self.core.host.as_ref(), model_id, model_name).await
|
||||
}
|
||||
|
||||
/// Build `entries` through the host ([`build::build`]), installing each deployment it
|
||||
/// accepts as soon as it does, as `install` says: a replace empties the registry first (and
|
||||
/// takes `set_model_list`'s `model_names` once every entry is built), an append or upsert
|
||||
/// adds at the end. When a step fails, what was installed before it stays, as the Python
|
||||
/// router keeps the deployments it added before the one that raised.
|
||||
///
|
||||
/// The build installs into a registry of its own, indexed in place, and serves it once it
|
||||
/// ends (or fails), so n deployments build in O(n). A write that lands while it runs (a step
|
||||
/// answer that applies a changeset) is kept: what the build installed is added over it.
|
||||
pub async fn build(
|
||||
&self,
|
||||
path: BuildPath,
|
||||
entries: Vec<serde_json::Map<String, serde_json::Value>>,
|
||||
resolve_secrets: bool,
|
||||
install: build::Install,
|
||||
) -> Result<Vec<build::BuiltDeployment>, Error> {
|
||||
let model_names: Vec<String> = entries
|
||||
.iter()
|
||||
.map(|entry| match entry.get("model_name") {
|
||||
Some(serde_json::Value::String(name)) => name.clone(),
|
||||
None | Some(serde_json::Value::Null) => "None".to_owned(),
|
||||
Some(other) => other.to_string(),
|
||||
})
|
||||
.collect();
|
||||
let replace = install == build::Install::Replace;
|
||||
let base = self.core.snapshot.load_full().registry.clone();
|
||||
let mut working = if replace {
|
||||
base.emptied()
|
||||
} else {
|
||||
Registry::clone(&base)
|
||||
};
|
||||
let mut installed: Vec<DeploymentSpec> = Vec::new();
|
||||
let add = path == BuildPath::Add;
|
||||
let core = &self.core;
|
||||
let built = build::build(
|
||||
core.host.as_ref(),
|
||||
path,
|
||||
entries,
|
||||
resolve_secrets,
|
||||
|deployment| {
|
||||
working.push_built(deployment.spec.clone(), add)?;
|
||||
installed.push(deployment.spec.clone());
|
||||
core.validate_deployment(working.last().expect("just installed"))
|
||||
.map_err(Error::from)
|
||||
},
|
||||
)
|
||||
.await;
|
||||
if built.is_ok() && replace && path == BuildPath::Create {
|
||||
working.set_model_names(model_names);
|
||||
}
|
||||
self.serve_built(&base, working, installed, add)?;
|
||||
built
|
||||
}
|
||||
|
||||
/// Serve the registry a build made from `base`. When another write replaced `base` while
|
||||
/// the build ran, the build's deployments are added over that write instead.
|
||||
fn serve_built(
|
||||
&self,
|
||||
base: &Arc<Registry>,
|
||||
built: Registry,
|
||||
installed: Vec<DeploymentSpec>,
|
||||
add: bool,
|
||||
) -> Result<(), ConfigError> {
|
||||
let _writer = self.core.writer.lock();
|
||||
let current = self.core.snapshot.load_full();
|
||||
let registry = if Arc::ptr_eq(¤t.registry, base) {
|
||||
built
|
||||
} else {
|
||||
let mut next = Registry::clone(¤t.registry);
|
||||
for spec in installed {
|
||||
next.push_built(spec, add)?;
|
||||
}
|
||||
next
|
||||
};
|
||||
self.core.snapshot.store(Arc::new(RouterStateSnapshot {
|
||||
registry: Arc::new(registry),
|
||||
settings: current.settings.clone(),
|
||||
}));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Keep the providers `from`'s host resolved for the deployments this router serves too (a
|
||||
/// router rebuilt over the same deployments for new static settings).
|
||||
pub fn adopt_providers(&self, from: &Self) {
|
||||
let _writer = self.core.writer.lock();
|
||||
let current = self.core.snapshot.load_full();
|
||||
let registry = current
|
||||
.registry
|
||||
.with_providers_of(&from.core.snapshot.load().registry);
|
||||
self.core.snapshot.store(Arc::new(RouterStateSnapshot {
|
||||
registry: Arc::new(registry),
|
||||
settings: current.settings.clone(),
|
||||
}));
|
||||
}
|
||||
|
||||
/// Stop serving: every later call and selection fails with [`Error::Retired`].
|
||||
pub fn retire(&self) {
|
||||
self.core.retired.store(true, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
883
litellm-rust/crates/router-core/src/lookup.rs
Normal file
883
litellm-rust/crates/router-core/src/lookup.rs
Normal file
|
|
@ -0,0 +1,883 @@
|
|||
//! The registry lookups the host serves as `registry()` queries, answered over the deployments and
|
||||
//! the host indexes ([`crate::host_index::HostIndexes`]) the way the Python router answers them.
|
||||
//!
|
||||
//! A lookup answers with [`Row`]s: positions in the deployment list (each with how the row renames
|
||||
//! the deployment), wildcard routes, or a position the Python list could not be read at. The host
|
||||
//! turns a position into the entry object it stores, so identity and in-place reads stay its own;
|
||||
//! a [`Row::Raw`] is handed back for the host to index its list by, which raises exactly where the
|
||||
//! Python router's own read raised (an index a caller wrote junk into).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
host_index::{HostIndexes, python_str, truthy},
|
||||
registry::{Deployment, Registry},
|
||||
};
|
||||
|
||||
/// One row of a lookup's answer.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Row {
|
||||
/// The deployment at `position`, under `name` when the row renames it (an alias or a routing
|
||||
/// group), as a routing group's member `member` of `group` when `group` is set (its row drops
|
||||
/// the member's `access_groups` and takes the group's priority as its `order`).
|
||||
At {
|
||||
position: usize,
|
||||
name: Option<String>,
|
||||
group: Option<(String, String)>,
|
||||
},
|
||||
/// A wildcard route: a copy of a pattern router's deployment, under `name`.
|
||||
Pattern { entry: Value, name: String },
|
||||
/// A position the list cannot be read at: the host indexes its list by it, and raises.
|
||||
Raw(Value),
|
||||
/// A position container that is not a list: the host iterates it as its list's positions.
|
||||
RawPositions(Value),
|
||||
}
|
||||
|
||||
impl Row {
|
||||
/// The row as the host reads it.
|
||||
pub fn wire(&self) -> Value {
|
||||
match self {
|
||||
Self::At {
|
||||
position,
|
||||
name,
|
||||
group,
|
||||
} => {
|
||||
let mut row = Map::new();
|
||||
row.insert("at".to_owned(), json!(position));
|
||||
if let Some(name) = name {
|
||||
row.insert("name".to_owned(), json!(name));
|
||||
}
|
||||
if let Some((group, member)) = group {
|
||||
row.insert("group".to_owned(), json!(group));
|
||||
row.insert("member".to_owned(), json!(member));
|
||||
}
|
||||
Value::Object(row)
|
||||
}
|
||||
Self::Pattern { entry, name } => json!({"pattern": entry, "name": name}),
|
||||
Self::Raw(value) => json!({"raw": value}),
|
||||
Self::RawPositions(value) => json!({"raw_positions": value}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// An alias `model_list` lists: its name and the group it targets (`hidden` aliases already left
|
||||
/// out by the host when the whole list is asked for).
|
||||
#[derive(Clone, Debug, PartialEq, serde::Deserialize)]
|
||||
pub struct AliasItem {
|
||||
pub alias: String,
|
||||
pub target: String,
|
||||
}
|
||||
|
||||
/// A routing group callable by its name (the host leaves out one a deployment or an alias
|
||||
/// shadows): its name and member model groups.
|
||||
#[derive(Clone, Debug, PartialEq, serde::Deserialize)]
|
||||
pub struct CallableGroup {
|
||||
pub name: String,
|
||||
pub members: Vec<String>,
|
||||
}
|
||||
|
||||
/// What `get_model_list` composes beyond the deployments: the aliases and callable routing groups
|
||||
/// it lists, and the wildcard requests a name no deployment serves is routed by.
|
||||
#[derive(Clone, Debug, Default, PartialEq, serde::Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct ListScope {
|
||||
pub aliases: Vec<AliasItem>,
|
||||
pub groups: Vec<CallableGroup>,
|
||||
/// `get_deployments_by_pattern`'s requests: the name, then the name under its provider.
|
||||
pub wildcard: Vec<String>,
|
||||
/// Only one part of the list: `"aliases"` (`_core_get_model_list_from_model_alias`) or
|
||||
/// `"groups"` (`_core_get_model_list_from_routing_groups`).
|
||||
pub only: Option<String>,
|
||||
}
|
||||
|
||||
/// A row read for a lookup's filters: the deployment's fields as the row shows them.
|
||||
struct View<'a> {
|
||||
name: &'a str,
|
||||
info: &'a Map<String, Value>,
|
||||
params: &'a Map<String, Value>,
|
||||
strips_access_groups: bool,
|
||||
}
|
||||
|
||||
/// The lookups over one registry version.
|
||||
pub struct Lookup<'a> {
|
||||
deployments: &'a [Arc<Deployment>],
|
||||
host: &'a HostIndexes,
|
||||
}
|
||||
|
||||
impl<'a> Lookup<'a> {
|
||||
pub fn new(registry: &'a Registry) -> Self {
|
||||
Self {
|
||||
deployments: registry.deployments_from(0),
|
||||
host: registry.host_indexes(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `list[value]` as Python reads it: an int (or bool) in range, negative from the end.
|
||||
fn position(&self, value: &Value) -> Option<usize> {
|
||||
let len = i64::try_from(self.deployments.len()).ok()?;
|
||||
let index = match value {
|
||||
Value::Bool(flag) => i64::from(*flag),
|
||||
Value::Number(number) => number.as_i64()?,
|
||||
_ => return None,
|
||||
};
|
||||
let index = if index < 0 { len + index } else { index };
|
||||
(0..len)
|
||||
.contains(&index)
|
||||
.then(|| usize::try_from(index).ok())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// Every position `held` (what an index holds) names, read before any is filtered, as the
|
||||
/// Python router reads them into a tuple first: `Err` with the row that raises.
|
||||
fn positions(&self, held: &Value) -> Result<Vec<usize>, Row> {
|
||||
let items = match held {
|
||||
Value::Array(items) => items,
|
||||
// Iterating an empty string or dict reads no position; anything else raises.
|
||||
Value::String(text) if text.is_empty() => return Ok(Vec::new()),
|
||||
Value::Object(map) if map.is_empty() => return Ok(Vec::new()),
|
||||
_ => return Err(Row::RawPositions(held.clone())),
|
||||
};
|
||||
items
|
||||
.iter()
|
||||
.map(|item| self.position(item).ok_or_else(|| Row::Raw(item.clone())))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn view(&self, position: usize) -> View<'a> {
|
||||
let deployment = &self.deployments[position];
|
||||
View {
|
||||
name: &deployment.model_name,
|
||||
info: &deployment.model_info.raw,
|
||||
params: &deployment.raw_params,
|
||||
strips_access_groups: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn row_view<'r>(&self, row: &'r Row) -> Option<View<'r>>
|
||||
where
|
||||
'a: 'r,
|
||||
{
|
||||
match row {
|
||||
Row::At {
|
||||
position,
|
||||
name,
|
||||
group,
|
||||
} => {
|
||||
let base = self.view(*position);
|
||||
Some(View {
|
||||
name: name.as_deref().unwrap_or(base.name),
|
||||
strips_access_groups: group.is_some(),
|
||||
..base
|
||||
})
|
||||
}
|
||||
Row::Pattern { entry, name } => Some(View {
|
||||
name,
|
||||
info: entry
|
||||
.get("model_info")
|
||||
.and_then(Value::as_object)
|
||||
.unwrap_or(empty()),
|
||||
params: entry
|
||||
.get("litellm_params")
|
||||
.and_then(Value::as_object)
|
||||
.unwrap_or(empty()),
|
||||
strips_access_groups: false,
|
||||
}),
|
||||
Row::Raw(_) | Row::RawPositions(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
// -- deployments by name -------------------------------------------------------------------
|
||||
|
||||
/// `_core_get_all_deployments`: the team's deployments under public name `model_name`, else
|
||||
/// the group's, else (for a team) every deployment the name includes; each under `rename`.
|
||||
pub fn all(&self, model_name: &str, team_id: Option<&str>, rename: Option<&str>) -> Vec<Row> {
|
||||
let rows = |positions: Vec<usize>| -> Vec<Row> {
|
||||
positions
|
||||
.into_iter()
|
||||
.filter(|position| should_include(model_name, &self.view(*position), team_id))
|
||||
.map(|position| Row::At {
|
||||
position,
|
||||
name: rename.map(str::to_owned),
|
||||
group: None,
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
if let Some(team) = team_id
|
||||
&& let Some(held) = self.host.team_named(team, model_name)
|
||||
{
|
||||
match self.positions(held) {
|
||||
Ok(positions) => {
|
||||
let found = rows(positions);
|
||||
if !found.is_empty() {
|
||||
return found;
|
||||
}
|
||||
}
|
||||
Err(raw) => return vec![raw],
|
||||
}
|
||||
}
|
||||
if let Some(held) = self.host.named(model_name) {
|
||||
return match self.positions(held) {
|
||||
Ok(positions) => rows(positions),
|
||||
Err(raw) => vec![raw],
|
||||
};
|
||||
}
|
||||
if team_id.is_some() {
|
||||
return rows((0..self.deployments.len()).collect());
|
||||
}
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// `_team_deployments_across_teams`: every team's deployments under public name `model`, and
|
||||
/// how many distinct teams own them.
|
||||
pub fn across_teams(&self, model: &str) -> (Vec<Row>, usize) {
|
||||
let mut rows = Vec::new();
|
||||
for (_, public, held) in &self.host.team_model {
|
||||
if public != model {
|
||||
continue;
|
||||
}
|
||||
match self.positions(held) {
|
||||
Ok(positions) => rows.extend(positions.into_iter().map(|position| Row::At {
|
||||
position,
|
||||
name: None,
|
||||
group: None,
|
||||
})),
|
||||
Err(raw) => return (vec![raw], 0),
|
||||
}
|
||||
}
|
||||
let mut teams: Vec<&Value> = Vec::new();
|
||||
for row in &rows {
|
||||
if let Some(view) = self.row_view(row)
|
||||
&& let Some(team) = view.info.get("team_id").filter(|team| !team.is_null())
|
||||
&& !teams.contains(&team)
|
||||
{
|
||||
teams.push(team);
|
||||
}
|
||||
}
|
||||
let count = teams.len();
|
||||
(rows, count)
|
||||
}
|
||||
|
||||
/// `_core_config_deployments`: the deployments a config declared (no `db_model: true`).
|
||||
pub fn config(&self) -> Vec<Row> {
|
||||
(0..self.deployments.len())
|
||||
.filter(|position| {
|
||||
self.view(*position).info.get("db_model") != Some(&Value::Bool(true))
|
||||
})
|
||||
.map(|position| Row::At {
|
||||
position,
|
||||
name: None,
|
||||
group: None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `_core_get_model_ids`: the ids of the group's deployments, or (no group) of the id map's
|
||||
/// entries, without team deployments when `exclude_team_models`. `Err` with the row that raises.
|
||||
pub fn model_ids(
|
||||
&self,
|
||||
model_name: Option<&str>,
|
||||
exclude_team_models: bool,
|
||||
) -> Result<Vec<String>, Row> {
|
||||
let kept = |view: &View<'_>| {
|
||||
view.info.contains_key("id")
|
||||
&& !(exclude_team_models && view.info.get("team_id").is_some_and(truthy))
|
||||
};
|
||||
if let Some(name) = model_name {
|
||||
let Some(held) = self.host.named(name) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
return Ok(self
|
||||
.positions(held)?
|
||||
.into_iter()
|
||||
.map(|position| self.view(position))
|
||||
.filter(kept)
|
||||
.map(|view| python_str(&view.info["id"]))
|
||||
.collect());
|
||||
}
|
||||
let mut ids = Vec::new();
|
||||
for (id, held) in &self.host.id_map {
|
||||
let position = self.position(held).ok_or_else(|| Row::Raw(held.clone()))?;
|
||||
if kept(&self.view(position)) {
|
||||
ids.push(id.clone());
|
||||
}
|
||||
}
|
||||
Ok(ids)
|
||||
}
|
||||
|
||||
// -- the model list --------------------------------------------------------------------------
|
||||
|
||||
/// `_core_get_model_list`: the name's deployments, its alias rows and routing group rows, its
|
||||
/// wildcard routes when nothing else serves it, or (no name) every alias and group row and
|
||||
/// then every deployment.
|
||||
pub fn model_list(
|
||||
&self,
|
||||
model_name: Option<&str>,
|
||||
team_id: Option<&str>,
|
||||
scope: &ListScope,
|
||||
) -> Vec<Row> {
|
||||
let only = scope.only.as_deref();
|
||||
let mut rows = Vec::new();
|
||||
if let Some(name) = model_name
|
||||
&& only.is_none()
|
||||
{
|
||||
rows.extend(self.all(name, team_id, None));
|
||||
}
|
||||
let aliases: &[AliasItem] = if only == Some("groups") {
|
||||
&[]
|
||||
} else {
|
||||
&scope.aliases
|
||||
};
|
||||
for alias in aliases {
|
||||
match scope.groups.iter().find(|group| group.name == alias.target) {
|
||||
Some(group) => {
|
||||
rows.extend(self.group_rows(group).into_iter().map(|row| match row {
|
||||
Row::At {
|
||||
position, group, ..
|
||||
} => Row::At {
|
||||
position,
|
||||
name: Some(alias.alias.clone()),
|
||||
group,
|
||||
},
|
||||
other => other,
|
||||
}))
|
||||
}
|
||||
None => rows.extend(self.all(&alias.target, None, Some(&alias.alias))),
|
||||
}
|
||||
}
|
||||
match (model_name, only) {
|
||||
(_, Some("aliases")) => {}
|
||||
(Some(name), _) => {
|
||||
if let Some(group) = scope.groups.iter().find(|group| group.name == name) {
|
||||
rows.extend(self.group_rows(group));
|
||||
}
|
||||
}
|
||||
(None, _) => {
|
||||
for group in &scope.groups {
|
||||
rows.extend(self.group_rows(group));
|
||||
}
|
||||
}
|
||||
}
|
||||
if only.is_some() {
|
||||
return rows;
|
||||
}
|
||||
if let Some(name) = model_name
|
||||
&& rows.is_empty()
|
||||
{
|
||||
rows.extend(self.wildcard(name, team_id, &scope.wildcard));
|
||||
}
|
||||
if model_name.is_none() {
|
||||
rows.extend((0..self.deployments.len()).map(|position| Row::At {
|
||||
position,
|
||||
name: None,
|
||||
group: None,
|
||||
}));
|
||||
}
|
||||
rows
|
||||
}
|
||||
|
||||
/// The whole list's rows, or the first row the host's read of the list raises at.
|
||||
fn readable(&self, rows: Vec<Row>) -> Result<Vec<Row>, Row> {
|
||||
match rows
|
||||
.iter()
|
||||
.find(|row| matches!(row, Row::Raw(_) | Row::RawPositions(_)))
|
||||
{
|
||||
Some(raw) => Err(raw.clone()),
|
||||
None => Ok(rows),
|
||||
}
|
||||
}
|
||||
|
||||
/// A callable routing group's rows: each member's deployments under the group's name.
|
||||
fn group_rows(&self, group: &CallableGroup) -> Vec<Row> {
|
||||
let mut rows = Vec::new();
|
||||
for member in &group.members {
|
||||
for row in self.all(member, None, Some(&group.name)) {
|
||||
rows.push(match row {
|
||||
Row::At { position, name, .. } => Row::At {
|
||||
position,
|
||||
name,
|
||||
group: Some((group.name.clone(), member.clone())),
|
||||
},
|
||||
other => other,
|
||||
});
|
||||
}
|
||||
}
|
||||
rows
|
||||
}
|
||||
|
||||
/// `_get_wildcard_deployments`: the global router's routes for the requests, then the team's.
|
||||
fn wildcard(&self, name: &str, team_id: Option<&str>, requests: &[String]) -> Vec<Row> {
|
||||
let route = |router: &crate::host_index::PatternRouter| -> Vec<Value> {
|
||||
requests
|
||||
.iter()
|
||||
.find_map(|request| {
|
||||
router
|
||||
.route(request, None)
|
||||
.filter(|routed| !routed.is_empty())
|
||||
})
|
||||
.unwrap_or_default()
|
||||
};
|
||||
let mut routed = route(&self.host.patterns);
|
||||
if let Some(team) = team_id
|
||||
&& let Some((_, router)) = self.host.team_patterns.iter().find(|(t, _)| t == team)
|
||||
{
|
||||
routed.extend(route(router));
|
||||
}
|
||||
routed
|
||||
.into_iter()
|
||||
.map(|entry| Row::Pattern {
|
||||
entry,
|
||||
name: name.to_owned(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The deployments `team_id`'s own wildcard pattern routes `name` to
|
||||
/// (`_core_team_pattern_routers`, keyed by `team_public_model_name`).
|
||||
pub fn team_wildcard(&self, name: &str, team_id: &str) -> Vec<Arc<Deployment>> {
|
||||
let Some((_, router)) = self.host.team_patterns.iter().find(|(t, _)| t == team_id) else {
|
||||
return Vec::new();
|
||||
};
|
||||
router
|
||||
.route(name, None)
|
||||
.unwrap_or_default()
|
||||
.iter()
|
||||
.filter_map(|entry| entry.get("model_info")?.get("id").map(python_str))
|
||||
.filter_map(|id| self.deployments.iter().find(|d| d.id == id).cloned())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `_core_get_model_names`: per row of the whole list, a team deployment's public name for
|
||||
/// its own team, else the row's name.
|
||||
pub fn model_names(
|
||||
&self,
|
||||
team_id: Option<&str>,
|
||||
scope: &ListScope,
|
||||
) -> Result<Vec<String>, Row> {
|
||||
let rows = self.readable(self.model_list(None, None, scope))?;
|
||||
let mut names = Vec::new();
|
||||
for row in &rows {
|
||||
let Some(view) = self.row_view(row) else {
|
||||
continue;
|
||||
};
|
||||
match view.info.get("team_id").filter(|team| truthy(team)) {
|
||||
Some(team) => {
|
||||
if team_id.is_some_and(|id| team == &Value::String(id.to_owned()))
|
||||
&& let Some(public) = view
|
||||
.info
|
||||
.get("team_public_model_name")
|
||||
.filter(|public| truthy(public))
|
||||
{
|
||||
names.push(python_str(public));
|
||||
}
|
||||
}
|
||||
None => names.push(view.name.to_owned()),
|
||||
}
|
||||
}
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
/// `_core_get_model_access_groups`: access group to the names of the rows that declare it,
|
||||
/// in the order they appear.
|
||||
pub fn access_groups(
|
||||
&self,
|
||||
model_name: Option<&str>,
|
||||
access_group: Option<&str>,
|
||||
team_id: Option<&str>,
|
||||
scope: &ListScope,
|
||||
) -> Result<Vec<(String, Vec<String>)>, Row> {
|
||||
let mut groups: Vec<(String, Vec<String>)> = Vec::new();
|
||||
for row in self.readable(self.model_list(model_name, team_id, scope))? {
|
||||
let Some(view) = self.row_view(&row) else {
|
||||
continue;
|
||||
};
|
||||
if view.info.is_empty() || view.strips_access_groups {
|
||||
continue;
|
||||
}
|
||||
// `access_groups or []`, iterated: a list's items, a string's characters.
|
||||
let declared: Vec<String> = match view.info.get("access_groups") {
|
||||
Some(Value::Array(items)) => items.iter().map(python_str).collect(),
|
||||
Some(Value::String(text)) => text.chars().map(String::from).collect(),
|
||||
_ => continue,
|
||||
};
|
||||
for group in declared {
|
||||
if access_group.is_some_and(|wanted| wanted != group) {
|
||||
continue;
|
||||
}
|
||||
match groups.iter_mut().find(|(held, _)| *held == group) {
|
||||
Some((_, names)) => names.push(view.name.to_owned()),
|
||||
None => groups.push((group, vec![view.name.to_owned()])),
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(groups)
|
||||
}
|
||||
|
||||
/// `_core_resolve_model_name_from_model_id`'s scan of the whole list: whether it lists any
|
||||
/// row, and the name of the first row whose `litellm_params.model` is `model_id`, or ends in
|
||||
/// `/model_id` or `:model_id`.
|
||||
pub fn name_serving(
|
||||
&self,
|
||||
model_id: &str,
|
||||
scope: &ListScope,
|
||||
) -> Result<(bool, Option<String>), Row> {
|
||||
let rows = self.readable(self.model_list(None, None, scope))?;
|
||||
let (slash, colon) = (format!("/{model_id}"), format!(":{model_id}"));
|
||||
let name = rows.iter().find_map(|row| {
|
||||
let view = self.row_view(row)?;
|
||||
let model = view.params.get("model").and_then(Value::as_str)?;
|
||||
(model == model_id || model.ends_with(&slash) || model.ends_with(&colon))
|
||||
.then_some(view.name)
|
||||
.filter(|name| !name.is_empty())
|
||||
.map(str::to_owned)
|
||||
});
|
||||
Ok((!rows.is_empty(), name))
|
||||
}
|
||||
|
||||
// -- teams and credentials -------------------------------------------------------------------
|
||||
|
||||
/// `_core_team_model_has_alternatives`: whether the team's public name of deployment `id` has
|
||||
/// another unpaused deployment of that team.
|
||||
pub fn team_model_has_alternatives(&self, id: &str) -> Result<bool, Row> {
|
||||
let Some(held) = self.host.id_entry(id) else {
|
||||
return Ok(false);
|
||||
};
|
||||
let position = self.position(held).ok_or_else(|| Row::Raw(held.clone()))?;
|
||||
let info = self.view(position).info;
|
||||
let (Some(Value::String(team)), Some(Value::String(public))) =
|
||||
(info.get("team_id"), info.get("team_public_model_name"))
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
let Some(siblings) = self.host.team_named(team, public) else {
|
||||
return Ok(false);
|
||||
};
|
||||
let unpaused = self
|
||||
.positions(siblings)?
|
||||
.into_iter()
|
||||
.filter(|position| self.view(*position).info.get("blocked") != Some(&Value::Bool(true)))
|
||||
.count();
|
||||
Ok(unpaused > 1)
|
||||
}
|
||||
|
||||
/// `_get_model_group_deployment_usable_by_team`: the first deployment of the group the team
|
||||
/// may use.
|
||||
pub fn group_deployment_usable_by_team(
|
||||
&self,
|
||||
model_group: &str,
|
||||
team_id: Option<&str>,
|
||||
) -> Result<Option<usize>, Row> {
|
||||
// `.get(name) or ()`: nothing held, or a falsy value, reads no position.
|
||||
let Some(held) = self.host.named(model_group).filter(|held| truthy(held)) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Value::Array(items) = held else {
|
||||
return Err(Row::RawPositions(held.clone()));
|
||||
};
|
||||
for item in items {
|
||||
let position = self.position(item).ok_or_else(|| Row::Raw(item.clone()))?;
|
||||
if usable_by_team(&self.view(position), team_id) {
|
||||
return Ok(Some(position));
|
||||
}
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// `_get_team_public_name_deployment`: the first deployment of the team's public name.
|
||||
pub fn team_public_name_deployment(
|
||||
&self,
|
||||
model_id: &str,
|
||||
team_id: &str,
|
||||
) -> Result<Option<usize>, Row> {
|
||||
match self.host.team_named(team_id, model_id) {
|
||||
Some(Value::Array(items)) => match items.first() {
|
||||
Some(item) => self
|
||||
.position(item)
|
||||
.map(Some)
|
||||
.ok_or_else(|| Row::Raw(item.clone())),
|
||||
None => Ok(None),
|
||||
},
|
||||
Some(held) if truthy(held) => Err(Row::RawPositions(held.clone())),
|
||||
Some(_) | None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
// -- routable deployments --------------------------------------------------------------------
|
||||
|
||||
/// `_routable_deployments`: the routable group's deployments for the team (those it may use
|
||||
/// when any), without paused ones.
|
||||
pub fn routable(&self, model_group: &str, team_id: Option<&str>) -> Vec<Row> {
|
||||
let named = self.all(model_group, team_id, None);
|
||||
if named
|
||||
.iter()
|
||||
.any(|row| matches!(row, Row::Raw(_) | Row::RawPositions(_)))
|
||||
{
|
||||
return named;
|
||||
}
|
||||
let usable: Vec<&Row> = named
|
||||
.iter()
|
||||
.filter(|row| {
|
||||
self.row_view(row)
|
||||
.is_some_and(|view| usable_by_team(&view, team_id))
|
||||
})
|
||||
.collect();
|
||||
let pool: Vec<&Row> = if usable.is_empty() {
|
||||
named.iter().collect()
|
||||
} else {
|
||||
usable
|
||||
};
|
||||
pool.into_iter()
|
||||
.filter(|row| {
|
||||
self.row_view(row)
|
||||
.is_some_and(|view| view.info.get("blocked") != Some(&Value::Bool(true)))
|
||||
})
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `_core_get_candidate_model_ids_for_route`'s ids of `rows`, in first-seen order.
|
||||
pub fn ids_of(&self, rows: &[Row]) -> Result<Vec<String>, Row> {
|
||||
let mut ids: Vec<String> = Vec::new();
|
||||
for row in rows {
|
||||
let view = match row {
|
||||
Row::Raw(_) | Row::RawPositions(_) => return Err(row.clone()),
|
||||
_ => self.row_view(row).expect("a readable row"),
|
||||
};
|
||||
if let Some(id) = view.info.get("id").filter(|id| !id.is_null()) {
|
||||
let id = python_str(id);
|
||||
if !ids.contains(&id) {
|
||||
ids.push(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(ids)
|
||||
}
|
||||
|
||||
/// `_core_get_candidate_model_ids_for_route` once the alias resolved `resolved`: the
|
||||
/// members' ids when it names a callable routing group (`members`), else, for a name that is
|
||||
/// no model group, the team's deployments under it or its wildcard routes, else the group's.
|
||||
pub fn candidate_ids(
|
||||
&self,
|
||||
resolved: &str,
|
||||
team_id: Option<&str>,
|
||||
members: Option<&[String]>,
|
||||
requests: &[String],
|
||||
) -> Result<Vec<String>, Row> {
|
||||
if let Some(members) = members {
|
||||
let mut ids: Vec<String> = Vec::new();
|
||||
for member in members {
|
||||
for id in self.ids_of(&self.all(member, team_id, None))? {
|
||||
if !ids.contains(&id) {
|
||||
ids.push(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
return Ok(ids);
|
||||
}
|
||||
let named = Value::String(resolved.to_owned());
|
||||
if !self.host.model_names.contains(&named) {
|
||||
if team_id.is_some() {
|
||||
let team_rows = self.all(resolved, team_id, None);
|
||||
if !team_rows.is_empty() {
|
||||
return self.ids_of(&team_rows);
|
||||
}
|
||||
}
|
||||
let wildcard = self.wildcard_route(resolved, team_id, requests);
|
||||
if !wildcard.is_empty() {
|
||||
return self.ids_of(&wildcard);
|
||||
}
|
||||
}
|
||||
self.ids_of(&self.all(resolved, team_id, None))
|
||||
}
|
||||
|
||||
/// `_core_wildcard_route_deployments`: the team router's routes for `model`, else the global
|
||||
/// router's.
|
||||
pub fn wildcard_route(
|
||||
&self,
|
||||
model: &str,
|
||||
team_id: Option<&str>,
|
||||
requests: &[String],
|
||||
) -> Vec<Row> {
|
||||
let route = |router: &crate::host_index::PatternRouter| -> Vec<Value> {
|
||||
requests
|
||||
.iter()
|
||||
.find_map(|request| {
|
||||
router
|
||||
.route(request, None)
|
||||
.filter(|routed| !routed.is_empty())
|
||||
})
|
||||
.unwrap_or_default()
|
||||
};
|
||||
let team_owned = team_id
|
||||
.and_then(|team| self.host.team_patterns.iter().find(|(t, _)| t == team))
|
||||
.map(|(_, router)| route(router))
|
||||
.unwrap_or_default();
|
||||
let routed = if team_owned.is_empty() {
|
||||
route(&self.host.patterns)
|
||||
} else {
|
||||
team_owned
|
||||
};
|
||||
routed
|
||||
.into_iter()
|
||||
.map(|entry| Row::Pattern {
|
||||
entry,
|
||||
name: model.to_owned(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// `_core_should_include_deployment`: the team's deployment under its public name, or a
|
||||
/// deployment of the group the team may use.
|
||||
fn should_include(model_name: &str, view: &View<'_>, team_id: Option<&str>) -> bool {
|
||||
let team_value = team_id.map(|team| Value::String(team.to_owned()));
|
||||
let model_team = view.info.get("team_id");
|
||||
if let Some(team) = &team_value
|
||||
&& model_team == Some(team)
|
||||
&& view.info.get("team_public_model_name") == Some(&Value::String(model_name.to_owned()))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
if view.name == model_name {
|
||||
return team_value.is_none()
|
||||
|| model_team.is_none_or(Value::is_null)
|
||||
|| model_team == team_value.as_ref();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// `deployment_usable_by_team`: a deployment no team owns, or one the team owns.
|
||||
fn usable_by_team(view: &View<'_>, team_id: Option<&str>) -> bool {
|
||||
match view.info.get("team_id").filter(|team| !team.is_null()) {
|
||||
None => true,
|
||||
Some(owner) => team_id.is_some_and(|team| owner == &Value::String(team.to_owned())),
|
||||
}
|
||||
}
|
||||
|
||||
fn empty() -> &'static Map<String, Value> {
|
||||
static EMPTY: std::sync::OnceLock<Map<String, Value>> = std::sync::OnceLock::new();
|
||||
EMPTY.get_or_init(Map::new)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::registry::DeploymentSpec;
|
||||
|
||||
fn registry(entries: Value) -> Registry {
|
||||
let mut registry = Registry::default();
|
||||
for entry in entries.as_array().unwrap() {
|
||||
let spec: DeploymentSpec = serde_json::from_value(entry.clone()).unwrap();
|
||||
registry.push_built(spec, false).unwrap();
|
||||
}
|
||||
registry
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_teams_wildcard_routes_only_the_names_its_public_pattern_matches() {
|
||||
let registry = registry(json!([
|
||||
{"model_name": "t1-w", "litellm_params": {"model": "openai/*"},
|
||||
"model_info": {"id": "tw1", "team_id": "T1", "team_public_model_name": "tw-*"}},
|
||||
{"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "gp"}},
|
||||
]));
|
||||
let lookup = Lookup::new(®istry);
|
||||
let ids = |name: &str, team: &str| -> Vec<String> {
|
||||
let routed = lookup.team_wildcard(name, team);
|
||||
routed.iter().map(|d| d.id.clone()).collect()
|
||||
};
|
||||
assert_eq!(ids("tw-gpt-4o", "T1"), ["tw1"]);
|
||||
assert!(ids("gpt-4o", "T1").is_empty());
|
||||
assert!(ids("tw-gpt-4o", "T2").is_empty());
|
||||
}
|
||||
|
||||
fn entry(name: &str, id: &str, info: Value) -> Value {
|
||||
let mut info = info.as_object().cloned().unwrap_or_default();
|
||||
info.insert("id".to_owned(), json!(id));
|
||||
json!({"model_name": name, "litellm_params": {"model": format!("openai/{id}")}, "model_info": info})
|
||||
}
|
||||
|
||||
fn at(position: usize) -> Row {
|
||||
Row::At {
|
||||
position,
|
||||
name: None,
|
||||
group: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_prefers_the_teams_public_name_then_the_group() {
|
||||
let registry = registry(json!([
|
||||
entry("gpt", "a", json!({})),
|
||||
entry(
|
||||
"team-gpt",
|
||||
"t",
|
||||
json!({"team_id": "T", "team_public_model_name": "gpt"})
|
||||
),
|
||||
entry("gpt", "b", json!({"team_id": "U"})),
|
||||
]));
|
||||
let lookup = Lookup::new(®istry);
|
||||
assert_eq!(lookup.all("gpt", Some("T"), None), vec![at(1)]);
|
||||
assert_eq!(lookup.all("gpt", Some("V"), None), vec![at(0)]);
|
||||
assert_eq!(lookup.all("gpt", None, None), vec![at(0), at(2)]);
|
||||
assert_eq!(lookup.all("missing", None, None), vec![]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_junk_position_is_handed_back_to_raise() {
|
||||
let mut registry = registry(json!([entry("gpt", "a", json!({}))]));
|
||||
let mut model_name = registry.host_indexes().model_name.clone();
|
||||
model_name.insert("gpt".to_owned(), json!([0, "x"]));
|
||||
let changes = crate::registry::ChangeSet {
|
||||
host_index: Some(crate::host_index::IndexWrite {
|
||||
model_name: Some(model_name),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
registry = registry.apply(&changes, None).unwrap();
|
||||
let lookup = Lookup::new(®istry);
|
||||
assert_eq!(lookup.all("gpt", None, None), vec![Row::Raw(json!("x"))]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_model_list_lists_aliases_and_groups_then_every_deployment() {
|
||||
let registry = registry(json!([
|
||||
entry("gpt", "a", json!({"access_groups": ["g1"]})),
|
||||
entry("claude", "b", json!({})),
|
||||
]));
|
||||
let lookup = Lookup::new(®istry);
|
||||
let scope = ListScope {
|
||||
aliases: vec![AliasItem {
|
||||
alias: "fast".to_owned(),
|
||||
target: "gpt".to_owned(),
|
||||
}],
|
||||
groups: vec![CallableGroup {
|
||||
name: "pool".to_owned(),
|
||||
members: vec!["gpt".to_owned(), "claude".to_owned()],
|
||||
}],
|
||||
wildcard: Vec::new(),
|
||||
only: None,
|
||||
};
|
||||
let rows = lookup.model_list(None, None, &scope);
|
||||
assert_eq!(rows.len(), 5);
|
||||
assert_eq!(
|
||||
rows[0],
|
||||
Row::At {
|
||||
position: 0,
|
||||
name: Some("fast".to_owned()),
|
||||
group: None
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
lookup.model_names(None, &scope).unwrap(),
|
||||
["fast", "pool", "pool", "gpt", "claude"]
|
||||
);
|
||||
assert_eq!(
|
||||
lookup.access_groups(None, None, None, &scope).unwrap(),
|
||||
vec![("g1".to_owned(), vec!["fast".to_owned(), "gpt".to_owned()])]
|
||||
);
|
||||
}
|
||||
}
|
||||
44
litellm-rust/crates/router-core/src/mock.rs
Normal file
44
litellm-rust/crates/router-core/src/mock.rs
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
//! The request's `mock_testing_*` flags: a failure the router raises itself, without calling a
|
||||
//! deployment (`_handle_mock_testing_fallbacks` and `_handle_mock_testing_rate_limit_error`).
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::{AttemptError, FailureKind};
|
||||
|
||||
/// Which mock failure a request asks for.
|
||||
#[derive(
|
||||
Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::IntoStaticStr, strum::EnumString,
|
||||
)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum MockFailure {
|
||||
/// `mock_testing_fallbacks`: an `InternalServerError` before the first level's retries.
|
||||
Fallbacks,
|
||||
/// `mock_testing_context_fallbacks`: a `ContextWindowExceededError`, likewise.
|
||||
ContextWindow,
|
||||
/// `mock_testing_content_policy_fallbacks`: a `ContentPolicyViolationError`, likewise.
|
||||
ContentPolicy,
|
||||
/// `mock_testing_rate_limit_error`: a `RateLimitError` in place of every level's first
|
||||
/// attempt, which the retry loop handles like a failed attempt.
|
||||
RateLimit,
|
||||
}
|
||||
|
||||
impl MockFailure {
|
||||
/// Whether the failure stands in for the first level's retry loop (raised before its retries).
|
||||
pub fn before_retries(self) -> bool {
|
||||
!matches!(self, Self::RateLimit)
|
||||
}
|
||||
|
||||
/// The failure as the rules see it.
|
||||
pub fn failure(self, num_retries: Option<u32>) -> AttemptError {
|
||||
let (kind, status) = match self {
|
||||
Self::Fallbacks => (FailureKind::InternalServer, 500),
|
||||
Self::ContextWindow => (FailureKind::ContextWindowExceeded, 400),
|
||||
Self::ContentPolicy => (FailureKind::ContentPolicyViolation, 400),
|
||||
Self::RateLimit => (FailureKind::RateLimit, 429),
|
||||
};
|
||||
let mut failure = AttemptError::new(kind, Some(status));
|
||||
failure.num_retries = num_retries;
|
||||
failure
|
||||
}
|
||||
}
|
||||
716
litellm-rust/crates/router-core/src/model_info.rs
Normal file
716
litellm-rust/crates/router-core/src/model_info.rs
Normal file
|
|
@ -0,0 +1,716 @@
|
|||
//! Model info: what the router reports a deployment can do and costs, merged from the process
|
||||
//! cost map, what the background refresh discovered, and the deployment's own `model_info`.
|
||||
//!
|
||||
//! The merge rules are the Python router core's (`get_router_model_info`,
|
||||
//! `get_deployment_model_info`); every cost-map read is a [`Host::model_info`] lookup asked while
|
||||
//! the query runs, so a price reload or a registration the host made since the router was built is
|
||||
//! what the next query reads. The host keeps choosing the keys a lookup tries (`get_model_info`);
|
||||
//! the core chooses which lookups a query makes.
|
||||
//!
|
||||
//! Values the host cannot spell as JSON travel as opaque tokens ([`is_opaque`]): the merge moves
|
||||
//! them between layers without reading them, and the host restores them.
|
||||
|
||||
use serde_json::{Map, Number, Value};
|
||||
|
||||
use crate::{
|
||||
error::Error,
|
||||
host::{Host, LogLevel, ModelInfoLookup},
|
||||
};
|
||||
|
||||
/// The key the host wraps a value it cannot spell as JSON in.
|
||||
pub const OPAQUE_KEY: &str = "__litellm_python_value__";
|
||||
|
||||
/// A value the host passed as an opaque token: the merge moves it, never reads it.
|
||||
pub fn is_opaque(value: &Value) -> bool {
|
||||
matches!(value, Value::Object(map) if map.len() == 1 && map.contains_key(OPAQUE_KEY))
|
||||
}
|
||||
|
||||
/// One `get_router_model_info` query, as the host reads it off the deployment.
|
||||
#[derive(Clone, Debug, Default, serde::Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct RouterModelInfoQuery {
|
||||
/// `model_info.id`.
|
||||
pub deployment_id: Value,
|
||||
/// `model_info.base_model`, then `litellm_params.base_model`: the first that is not `None`
|
||||
/// names the model whose limits and costs the deployment reports.
|
||||
pub model_info_base_model: Value,
|
||||
pub params_base_model: Value,
|
||||
/// The validated `litellm_params.model`.
|
||||
pub params_model: String,
|
||||
/// `base_model` among the validated params' extra fields.
|
||||
pub extra_base_model: Value,
|
||||
/// The name the caller asked for, which a wildcard deployment matches.
|
||||
pub received_model_name: String,
|
||||
}
|
||||
|
||||
/// A merged model info, and the keys whose values are the deployment's own (the last layer).
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct MergedModelInfo {
|
||||
pub info: Map<String, Value>,
|
||||
pub user_keys: Vec<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn raised(exception: &'static str, detail: impl Into<String>) -> Error {
|
||||
Error::ModelInfo {
|
||||
exception,
|
||||
detail: detail.into(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn malformed(lookup: &str) -> Error {
|
||||
raised(
|
||||
"ValueError",
|
||||
format!("the host answered a `{lookup}` lookup malformed"),
|
||||
)
|
||||
}
|
||||
|
||||
/// Python truthiness of a JSON value.
|
||||
pub(crate) fn truthy(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Null => false,
|
||||
Value::Bool(flag) => *flag,
|
||||
Value::Number(number) => number.as_f64().is_some_and(|n| n != 0.0),
|
||||
Value::String(text) => !text.is_empty(),
|
||||
Value::Array(items) => !items.is_empty(),
|
||||
Value::Object(map) => is_opaque(value) || !map.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `(value or 0) > 0`, for a value that is a number (anything else reads as not positive).
|
||||
fn positive_or_zero(value: Option<&Value>) -> bool {
|
||||
value
|
||||
.filter(|value| truthy(value))
|
||||
.and_then(Value::as_f64)
|
||||
.is_some_and(|n| n > 0.0)
|
||||
}
|
||||
|
||||
pub(crate) async fn lookup(host: &dyn Host, lookup: ModelInfoLookup) -> Result<Value, Error> {
|
||||
host.model_info(&lookup).await
|
||||
}
|
||||
|
||||
/// `get_model_info(model=...)`: the row, or `None` when the lookup tolerates a failure and failed.
|
||||
async fn catalog_row(
|
||||
host: &dyn Host,
|
||||
model: Value,
|
||||
tolerate: bool,
|
||||
) -> Result<Option<Map<String, Value>>, Error> {
|
||||
match lookup(
|
||||
host,
|
||||
ModelInfoLookup::ModelInfo {
|
||||
model,
|
||||
tolerate,
|
||||
exact: false,
|
||||
},
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Value::Null => Ok(None),
|
||||
Value::Object(row) if tolerate && row.get("failed") == Some(&Value::Bool(true)) => Ok(None),
|
||||
Value::Object(row) => Ok(Some(row)),
|
||||
_ => Err(malformed("model_info")),
|
||||
}
|
||||
}
|
||||
|
||||
/// `model_cost.get(key)`: `None` when the key is absent, else the entry as stored.
|
||||
async fn cost_entry(host: &dyn Host, key: &str, process: bool) -> Result<Option<Value>, Error> {
|
||||
match lookup(
|
||||
host,
|
||||
ModelInfoLookup::CostEntry {
|
||||
key: key.to_owned(),
|
||||
process,
|
||||
},
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Value::Null => Ok(None),
|
||||
Value::Object(mut answer) => Ok(Some(answer.remove("entry").unwrap_or(Value::Null))),
|
||||
_ => Err(malformed("cost_entry")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn discovered(host: &dyn Host, id: Value) -> Result<Map<String, Value>, Error> {
|
||||
match lookup(host, ModelInfoLookup::Discovered { id }).await? {
|
||||
Value::Null => Ok(Map::new()),
|
||||
Value::Object(limits) => Ok(limits),
|
||||
_ => Err(malformed("discovered")),
|
||||
}
|
||||
}
|
||||
|
||||
async fn log(host: &dyn Host, level: LogLevel, message: String) -> Result<(), Error> {
|
||||
lookup(host, ModelInfoLookup::Log { level, message })
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
/// `_provider_model_id`: the model id a cost-map key or base model names, bedrock base model ids
|
||||
/// normalized.
|
||||
async fn provider_model_id(host: &dyn Host, name: &str, provider: &str) -> Result<String, Error> {
|
||||
let model_id = name.rsplit('/').next().unwrap_or(name).to_owned();
|
||||
if !provider.starts_with("bedrock") {
|
||||
return Ok(model_id);
|
||||
}
|
||||
match lookup(host, ModelInfoLookup::BedrockBase { model_id }).await? {
|
||||
Value::String(base) => Ok(base),
|
||||
_ => Err(malformed("bedrock_base")),
|
||||
}
|
||||
}
|
||||
|
||||
/// `_exact_base_model_max_input_tokens`: the base model's own window, unless its row came from a
|
||||
/// fallback-generalization rule or it is unmapped.
|
||||
async fn exact_base_model_max_input_tokens(
|
||||
host: &dyn Host,
|
||||
base_model: &str,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
let answer = lookup(
|
||||
host,
|
||||
ModelInfoLookup::ModelInfo {
|
||||
model: Value::String(base_model.to_owned()),
|
||||
tolerate: true,
|
||||
exact: true,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let Value::Object(mut answer) = answer else {
|
||||
return Ok(None);
|
||||
};
|
||||
if answer.get("failed") == Some(&Value::Bool(true))
|
||||
|| answer.get("generalized") == Some(&Value::Bool(true))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(Value::Object(row)) = answer.remove("row") else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(row
|
||||
.get("max_input_tokens")
|
||||
.filter(|tokens| !tokens.is_null())
|
||||
.cloned())
|
||||
}
|
||||
|
||||
/// `_base_model_widened_window`: a base model naming the same model as the route can only widen
|
||||
/// its context window.
|
||||
async fn base_model_widened_window(
|
||||
host: &dyn Host,
|
||||
route: &Map<String, Value>,
|
||||
base_model: Option<&str>,
|
||||
provider: &str,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
let Some(route_max) = route
|
||||
.get("max_input_tokens")
|
||||
.filter(|tokens| !tokens.is_null())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(base_model) = base_model else {
|
||||
return Ok(None);
|
||||
};
|
||||
let route_key = match route.get("key") {
|
||||
Some(Value::String(key)) => key.clone(),
|
||||
Some(_) => return Err(raised("AttributeError", "key")),
|
||||
None => return Err(raised("KeyError", "key")),
|
||||
};
|
||||
let route_id = provider_model_id(host, &route_key, provider).await?;
|
||||
let base_id = provider_model_id(host, base_model, provider).await?;
|
||||
let unversioned = base_id
|
||||
.split_once('@')
|
||||
.map_or(base_id.as_str(), |(id, _)| id);
|
||||
if route_id != base_id && route_id != unversioned {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(base_max) = exact_base_model_max_input_tokens(host, base_model).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
match (base_max.as_f64(), route_max.as_f64()) {
|
||||
(Some(base), Some(route)) if base > route => Ok(Some(base_max)),
|
||||
(Some(_), Some(_)) => Ok(None),
|
||||
_ => Err(raised(
|
||||
"TypeError",
|
||||
"'<=' not supported between the base model's and the route's max_input_tokens",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `_declared_base_model`: a `base_model` that is a string.
|
||||
fn declared_base_model(value: &Value) -> Option<&str> {
|
||||
value.as_str()
|
||||
}
|
||||
|
||||
/// `get_router_model_info` after the host validated the deployment's params: the cost-map row of
|
||||
/// the model the deployment serves, widened by its base model, overlaid with what the background
|
||||
/// refresh discovered and then with every value the deployment sets itself. `None` when the
|
||||
/// lookup answers `None`; an unmapped model raises as the lookup raises.
|
||||
pub async fn router_model_info(
|
||||
host: &dyn Host,
|
||||
query: &RouterModelInfoQuery,
|
||||
) -> Result<Option<MergedModelInfo>, Error> {
|
||||
let base_model = if query.model_info_base_model.is_null() {
|
||||
query.params_base_model.clone()
|
||||
} else {
|
||||
query.model_info_base_model.clone()
|
||||
};
|
||||
let mut model = base_model.clone();
|
||||
|
||||
let (resolved_model, provider) = match lookup(
|
||||
host,
|
||||
ModelInfoLookup::Provider {
|
||||
model: query.params_model.clone(),
|
||||
},
|
||||
)
|
||||
.await?
|
||||
{
|
||||
Value::Array(pair) => match pair.as_slice() {
|
||||
[Value::String(model), Value::String(provider), ..] => {
|
||||
(model.clone(), provider.clone())
|
||||
}
|
||||
_ => return Err(malformed("provider")),
|
||||
},
|
||||
_ => return Err(malformed("provider")),
|
||||
};
|
||||
|
||||
if provider == "azure" && base_model.is_null() {
|
||||
// Router init registers every deployment name into the cost map as a zeroed stub, so
|
||||
// membership alone cannot tell a resolvable name apart; it needs usable limits or costs.
|
||||
let fallback_key = if resolved_model.starts_with("azure/") {
|
||||
resolved_model.clone()
|
||||
} else {
|
||||
format!("azure/{resolved_model}")
|
||||
};
|
||||
let resolves = match cost_entry(host, &fallback_key, false).await? {
|
||||
Some(Value::Object(entry)) if !is_opaque(&Value::Object(entry.clone())) => {
|
||||
positive_or_zero(entry.get("max_input_tokens"))
|
||||
|| positive_or_zero(entry.get("input_cost_per_token"))
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
if resolves {
|
||||
log(
|
||||
host,
|
||||
LogLevel::Debug,
|
||||
format!(
|
||||
"Azure deployment '{resolved_model}' has no base_model set; using '{fallback_key}' from the model cost map for max tokens, cost tracking, etc."
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
log(
|
||||
host,
|
||||
LogLevel::Error,
|
||||
format!(
|
||||
"Could not identify azure model '{resolved_model}'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models"
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
} else if provider != "azure" {
|
||||
model = Value::String(resolved_model.clone());
|
||||
if resolved_model.contains('*') {
|
||||
// Only a wildcard deployment asks which model its pattern routed the name to.
|
||||
if let Value::Array(routed) = lookup(
|
||||
host,
|
||||
ModelInfoLookup::PatternModels {
|
||||
name: query.received_model_name.clone(),
|
||||
},
|
||||
)
|
||||
.await?
|
||||
&& let Some(routed) = routed
|
||||
.iter()
|
||||
.find(|routed| routed.get("id").unwrap_or(&Value::Null) == &query.deployment_id)
|
||||
{
|
||||
model = routed.get("model").cloned().unwrap_or(Value::Null);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if model.is_null() {
|
||||
// An azure deployment without a base_model is looked up by its own model.
|
||||
model = Value::String(resolved_model);
|
||||
}
|
||||
let Value::String(model) = model else {
|
||||
return Err(raised(
|
||||
"AttributeError",
|
||||
format!(
|
||||
"'{}' object has no attribute 'startswith'",
|
||||
crate::group_info::type_name(&model)
|
||||
),
|
||||
));
|
||||
};
|
||||
let model_info_name = if model.starts_with(&format!("{provider}/")) {
|
||||
model
|
||||
} else {
|
||||
format!("{provider}/{model}")
|
||||
};
|
||||
|
||||
let Some(row) = catalog_row(host, Value::String(model_info_name), false).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let user_model_info = match lookup(host, ModelInfoLookup::UserModelInfo).await? {
|
||||
Value::Object(info) => info,
|
||||
Value::Null => Map::new(),
|
||||
_ => return Err(malformed("user_model_info")),
|
||||
};
|
||||
let declared = user_model_info
|
||||
.get("base_model")
|
||||
.and_then(declared_base_model)
|
||||
.or_else(|| declared_base_model(&query.extra_base_model));
|
||||
let widened = base_model_widened_window(host, &row, declared, &provider).await?;
|
||||
let discovered = discovered(host, query.deployment_id.clone()).await?;
|
||||
|
||||
let mut info = row;
|
||||
if let Some(max_input_tokens) = widened {
|
||||
info.insert("max_input_tokens".to_owned(), max_input_tokens);
|
||||
}
|
||||
info.extend(discovered);
|
||||
let mut user_keys = Vec::new();
|
||||
for (key, value) in &user_model_info {
|
||||
if !value.is_null() {
|
||||
info.insert(key.clone(), value.clone());
|
||||
user_keys.push(key.clone());
|
||||
}
|
||||
}
|
||||
Ok(Some(MergedModelInfo { info, user_keys }))
|
||||
}
|
||||
|
||||
/// `_convert_stringified_numbers`: a string that reads as a number becomes one (a float when it
|
||||
/// has an exponent or a point, else an int); any other string stays as written.
|
||||
fn convert_stringified_number(text: &str) -> Value {
|
||||
let trimmed = text.trim();
|
||||
let digits = without_digit_underscores(trimmed);
|
||||
let lower = text.to_lowercase();
|
||||
if lower.contains('e') || text.contains('.') {
|
||||
return digits
|
||||
.and_then(|digits| python_float(&digits))
|
||||
.and_then(Number::from_f64)
|
||||
.map_or_else(|| Value::String(text.to_owned()), Value::Number);
|
||||
}
|
||||
digits
|
||||
.and_then(|digits| python_int(&digits))
|
||||
.map_or_else(|| Value::String(text.to_owned()), Value::Number)
|
||||
}
|
||||
|
||||
/// The text with the underscores Python's `int()`/`float()` accept between digits removed, or
|
||||
/// `None` when an underscore sits anywhere else.
|
||||
pub(crate) fn without_digit_underscores(text: &str) -> Option<String> {
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut out = String::with_capacity(text.len());
|
||||
for (index, ch) in chars.iter().enumerate() {
|
||||
if *ch == '_' {
|
||||
let before = index.checked_sub(1).and_then(|i| chars.get(i));
|
||||
let after = chars.get(index + 1);
|
||||
if !(before.is_some_and(char::is_ascii_digit)
|
||||
&& after.is_some_and(char::is_ascii_digit))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
out.push(*ch);
|
||||
}
|
||||
Some(out)
|
||||
}
|
||||
|
||||
pub(crate) fn python_float(text: &str) -> Option<f64> {
|
||||
let lower = text.to_ascii_lowercase();
|
||||
let unsigned = lower.trim_start_matches(['+', '-']);
|
||||
if matches!(unsigned, "inf" | "infinity" | "nan") {
|
||||
// Not finite: JSON cannot carry it, so the value stays as written.
|
||||
return None;
|
||||
}
|
||||
let well_formed = text
|
||||
.chars()
|
||||
.all(|ch| ch.is_ascii_digit() || matches!(ch, '+' | '-' | '.' | 'e' | 'E'));
|
||||
if !well_formed {
|
||||
return None;
|
||||
}
|
||||
text.parse::<f64>().ok().filter(|n| n.is_finite())
|
||||
}
|
||||
|
||||
pub(crate) fn python_int(text: &str) -> Option<Number> {
|
||||
let unsigned = text.strip_prefix(['+', '-']).unwrap_or(text);
|
||||
if unsigned.is_empty() || !unsigned.chars().all(|ch| ch.is_ascii_digit()) {
|
||||
return None;
|
||||
}
|
||||
let text = text.strip_prefix('+').unwrap_or(text);
|
||||
text.parse::<i64>()
|
||||
.map(Number::from)
|
||||
.ok()
|
||||
.or_else(|| text.parse::<u64>().map(Number::from).ok())
|
||||
}
|
||||
|
||||
/// `update_dictionary`: `new`'s values over `existing`, skipping `None`s; a string that reads as a
|
||||
/// number becomes one, and a dict merges one level into the dict it replaces.
|
||||
pub fn update_dictionary(existing: &mut Map<String, Value>, new: &Map<String, Value>) {
|
||||
for (key, value) in new {
|
||||
match value {
|
||||
Value::Null => {}
|
||||
Value::String(text) => {
|
||||
existing.insert(key.clone(), convert_stringified_number(text));
|
||||
}
|
||||
Value::Object(nested) if !is_opaque(value) => {
|
||||
let merged = match existing.get(key) {
|
||||
Some(current @ Value::Object(current_map)) if !is_opaque(current) => {
|
||||
let mut merged = current_map.clone();
|
||||
merged.extend(nested.clone());
|
||||
merged
|
||||
}
|
||||
_ => nested.clone(),
|
||||
};
|
||||
existing.insert(key.clone(), Value::Object(merged));
|
||||
}
|
||||
_ => {
|
||||
existing.insert(key.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `cost_map_omits_token_price`: the cost-map entries under `keys` exist, and none carries a
|
||||
/// per-token price.
|
||||
async fn cost_map_omits_token_price(host: &dyn Host, keys: &[Option<&str>]) -> Result<bool, Error> {
|
||||
let mut entries = 0;
|
||||
for key in keys.iter().flatten() {
|
||||
if let Some(entry) = cost_entry(host, key, true).await?
|
||||
&& !is_opaque(&entry)
|
||||
&& let Value::Object(map) = entry
|
||||
{
|
||||
entries += 1;
|
||||
if map.contains_key("input_cost_per_token") || map.contains_key("output_cost_per_token")
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(entries > 0)
|
||||
}
|
||||
|
||||
/// `get_deployment_model_info`: deployment `model_id`'s own cost-map entry (with what the refresh
|
||||
/// discovered, over its base model's row), merged over the row of its litellm model `model_name`.
|
||||
/// A lookup that fails leaves its layer out, as the Python router ignores it.
|
||||
pub async fn deployment_model_info(
|
||||
host: &dyn Host,
|
||||
model_id: &str,
|
||||
model_name: &Value,
|
||||
) -> Result<Option<Map<String, Value>>, Error> {
|
||||
let mut custom: Option<Map<String, Value>> = None;
|
||||
if let Some(entry) = cost_entry(host, model_id, false).await? {
|
||||
let entry = if truthy(&entry) {
|
||||
match entry {
|
||||
Value::Object(map) if !is_opaque(&Value::Object(map.clone())) => Some(map),
|
||||
// Not a mapping: the Python router's TypeError, which it ignores.
|
||||
_ => None,
|
||||
}
|
||||
} else {
|
||||
Some(Map::new())
|
||||
};
|
||||
if let Some(mut entry) = entry {
|
||||
entry.extend(discovered(host, Value::String(model_id.to_owned())).await?);
|
||||
custom = Some(entry);
|
||||
}
|
||||
}
|
||||
|
||||
let builtin = catalog_row(host, model_name.clone(), true).await?;
|
||||
|
||||
let mut base_model_key: Option<String> = None;
|
||||
if let Some(current) = &custom
|
||||
&& let Some(base_model) = current.get("base_model").filter(|base| !base.is_null())
|
||||
&& let Some(mut base_row) = catalog_row(host, base_model.clone(), true).await?
|
||||
{
|
||||
base_model_key = base_row
|
||||
.get("key")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
// The base model provides defaults; the deployment's own entry overrides them.
|
||||
update_dictionary(&mut base_row, current);
|
||||
custom = Some(base_row);
|
||||
}
|
||||
|
||||
let builtin_key = builtin
|
||||
.as_ref()
|
||||
.and_then(|row| row.get("key"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let info = match (custom, builtin) {
|
||||
(Some(custom), Some(mut builtin)) => {
|
||||
update_dictionary(&mut builtin, &custom);
|
||||
builtin
|
||||
}
|
||||
(None, Some(builtin)) => builtin,
|
||||
(Some(custom), None) => custom,
|
||||
(None, None) => return Ok(None),
|
||||
};
|
||||
if cost_map_omits_token_price(
|
||||
host,
|
||||
&[
|
||||
Some(model_id),
|
||||
builtin_key.as_deref(),
|
||||
base_model_key.as_deref(),
|
||||
],
|
||||
)
|
||||
.await?
|
||||
{
|
||||
let mut info = info;
|
||||
info.insert("input_cost_per_token".to_owned(), Value::Null);
|
||||
info.insert("output_cost_per_token".to_owned(), Value::Null);
|
||||
return Ok(Some(info));
|
||||
}
|
||||
Ok(Some(info))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
cache::{InMemoryRouterCache, RouterCache},
|
||||
error::AttemptError,
|
||||
host::{Attempt, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage},
|
||||
registry::Deployment,
|
||||
};
|
||||
|
||||
/// A cost map and a provider, answering lookups the way `PythonHost` does.
|
||||
#[derive(Default)]
|
||||
struct Catalog {
|
||||
rows: BTreeMap<String, Value>,
|
||||
provider: (String, String),
|
||||
user_model_info: Value,
|
||||
lookups: Mutex<Vec<Value>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Host for Catalog {
|
||||
async fn invoke(&self, _attempt: &Attempt) -> Result<Output, AttemptError> {
|
||||
unreachable!("a model-info query runs no attempt")
|
||||
}
|
||||
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,
|
||||
) -> Result<bool, crate::error::AttemptError> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn sleep(&self, _duration: std::time::Duration) {}
|
||||
fn observe(&self, _event: &RouterEvent) {}
|
||||
fn cache(&self) -> Arc<dyn RouterCache> {
|
||||
Arc::new(InMemoryRouterCache::default())
|
||||
}
|
||||
async fn model_info(&self, lookup: &ModelInfoLookup) -> Result<Value, Error> {
|
||||
self.lookups
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(serde_json::to_value(lookup).unwrap());
|
||||
Ok(match lookup {
|
||||
ModelInfoLookup::ModelInfo {
|
||||
model, tolerate, ..
|
||||
} => match model.as_str().and_then(|model| self.rows.get(model)) {
|
||||
Some(row) => row.clone(),
|
||||
None if *tolerate => json!({"failed": true}),
|
||||
None => return Err(raised("ValueError", "unmapped")),
|
||||
},
|
||||
ModelInfoLookup::CostEntry { key, .. } => self
|
||||
.rows
|
||||
.get(key)
|
||||
.map_or(Value::Null, |entry| json!({"entry": entry})),
|
||||
ModelInfoLookup::Provider { .. } => {
|
||||
json!([self.provider.0, self.provider.1])
|
||||
}
|
||||
ModelInfoLookup::UserModelInfo => self.user_model_info.clone(),
|
||||
_ => Value::Null,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_dictionary_converts_numbers_and_merges_nested_dicts_one_level() {
|
||||
let mut existing = json!({"a": 1, "n": {"x": 1, "y": 2}, "k": "keep"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
let new = json!({"a": None::<i64>, "n": {"y": 3}, "i": "8_192", "f": "3e-07", "s": "abc"})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
update_dictionary(&mut existing, &new);
|
||||
assert_eq!(
|
||||
Value::Object(existing),
|
||||
json!({"a": 1, "n": {"x": 1, "y": 3}, "k": "keep", "i": 8192, "f": 3e-07, "s": "abc"})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn router_model_info_overlays_the_deployments_own_values_on_the_row() {
|
||||
let host = Catalog {
|
||||
rows: BTreeMap::from([(
|
||||
"openai/gpt-4o".to_owned(),
|
||||
json!({"key": "gpt-4o", "max_input_tokens": 100, "input_cost_per_token": 1.0}),
|
||||
)]),
|
||||
provider: ("gpt-4o".to_owned(), "openai".to_owned()),
|
||||
user_model_info: json!({"id": "d", "input_cost_per_token": 2.0, "mode": None::<String>}),
|
||||
..Catalog::default()
|
||||
};
|
||||
let query = RouterModelInfoQuery {
|
||||
deployment_id: json!("d"),
|
||||
params_model: "openai/gpt-4o".to_owned(),
|
||||
received_model_name: "g".to_owned(),
|
||||
..RouterModelInfoQuery::default()
|
||||
};
|
||||
let merged = router_model_info(&host, &query).await.unwrap().unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(merged.info),
|
||||
json!({"key": "gpt-4o", "max_input_tokens": 100, "input_cost_per_token": 2.0, "id": "d"})
|
||||
);
|
||||
assert_eq!(merged.user_keys, vec!["id", "input_cost_per_token"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deployment_model_info_merges_the_entry_over_the_builtin_row() {
|
||||
let host = Catalog {
|
||||
rows: BTreeMap::from([
|
||||
(
|
||||
"gpt-4o".to_owned(),
|
||||
json!({"key": "gpt-4o", "max_tokens": 10, "input_cost_per_token": 1.0}),
|
||||
),
|
||||
("d".to_owned(), json!({"max_tokens": "20"})),
|
||||
]),
|
||||
..Catalog::default()
|
||||
};
|
||||
let info = deployment_model_info(&host, "d", &json!("gpt-4o"))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(info),
|
||||
json!({"key": "gpt-4o", "max_tokens": 20, "input_cost_per_token": 1.0})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_unpriced_entry_reports_no_token_prices() {
|
||||
let host = Catalog {
|
||||
rows: BTreeMap::from([("d".to_owned(), json!({"max_tokens": 5}))]),
|
||||
..Catalog::default()
|
||||
};
|
||||
let info = deployment_model_info(&host, "d", &json!("unmapped"))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(info),
|
||||
json!({"max_tokens": 5, "input_cost_per_token": null, "output_cost_per_token": null})
|
||||
);
|
||||
}
|
||||
}
|
||||
125
litellm-rust/crates/router-core/src/outcome.rs
Normal file
125
litellm-rust/crates/router-core/src/outcome.rs
Normal file
|
|
@ -0,0 +1,125 @@
|
|||
//! Which of a request's outcomes reach the deployment's health. The Python router learns them from
|
||||
//! logging callbacks it registers on the process lists; `_core_serve_through_request_logging` and
|
||||
//! `_core_discard` take those off, and `_core_bind_to_request_logging` puts them back on one
|
||||
//! request's logging object. This core registers none, so the host answers, for each request, what
|
||||
//! would have run ([`OutcomeCallbacks`]), and every write a callback owns passes through
|
||||
//! [`OutcomeCallbacks::record`].
|
||||
//!
|
||||
//! What the Python router writes outside its callbacks is not gated, and reaches [`crate::Health`]
|
||||
//! with the default value: `make_call`'s usage, `async_routing_strategy_pre_call_checks`'
|
||||
//! cooldown, `_trigger_cooldown_for_failed_deployment`, `RouterHealth.record_failure` and a
|
||||
//! realtime session's success.
|
||||
|
||||
use std::future::Future;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::Error;
|
||||
|
||||
/// Which of the router's outcome callbacks litellm's logging runs for a request. The default is a
|
||||
/// router whose callbacks are on the process lists.
|
||||
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct OutcomeCallbacks {
|
||||
/// `deployment_callback_on_success`, `sync_deployment_callback_on_success` for a sync call.
|
||||
/// `_core_bind_to_request_logging` binds only the former.
|
||||
pub success: bool,
|
||||
/// `deployment_callback_on_failure` and `async_deployment_callback_on_failure`.
|
||||
pub failure: bool,
|
||||
/// The attempt runs through `make_call`, which counts its usage whatever the callbacks do.
|
||||
pub make_call: bool,
|
||||
}
|
||||
|
||||
impl Default for OutcomeCallbacks {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
success: true,
|
||||
failure: true,
|
||||
make_call: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What an outcome writes.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OutcomeWrite {
|
||||
/// The success counter, or the usage the success callback adds at the end of a stream.
|
||||
Success,
|
||||
/// The tokens and request of an attempt that returned: `make_call`'s write, or the success
|
||||
/// callback's for an endpoint that skips `make_call`.
|
||||
AttemptUsage,
|
||||
/// The failure counter, the allowed-fails counter, the cooldown and the failed request.
|
||||
Failure,
|
||||
}
|
||||
|
||||
impl OutcomeCallbacks {
|
||||
/// Run `record` when `write` happens for this request, else record nothing.
|
||||
pub(crate) async fn record<T: Default>(
|
||||
self,
|
||||
write: OutcomeWrite,
|
||||
record: impl Future<Output = Result<T, Error>>,
|
||||
) -> Result<T, Error> {
|
||||
let runs = match write {
|
||||
OutcomeWrite::Success => self.success,
|
||||
OutcomeWrite::AttemptUsage => self.success || self.make_call,
|
||||
OutcomeWrite::Failure => self.failure,
|
||||
};
|
||||
if runs { record.await } else { Ok(T::default()) }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
const REGISTERED: OutcomeCallbacks = OutcomeCallbacks {
|
||||
success: true,
|
||||
failure: true,
|
||||
make_call: true,
|
||||
};
|
||||
const UNBOUND: OutcomeCallbacks = OutcomeCallbacks {
|
||||
success: false,
|
||||
failure: false,
|
||||
make_call: true,
|
||||
};
|
||||
const UNBOUND_WITHOUT_MAKE_CALL: OutcomeCallbacks = OutcomeCallbacks {
|
||||
success: false,
|
||||
failure: false,
|
||||
make_call: false,
|
||||
};
|
||||
const BOUND_SYNC: OutcomeCallbacks = OutcomeCallbacks {
|
||||
success: false,
|
||||
failure: true,
|
||||
make_call: true,
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[case::registered_success(REGISTERED, OutcomeWrite::Success, true)]
|
||||
#[case::registered_usage(REGISTERED, OutcomeWrite::AttemptUsage, true)]
|
||||
#[case::registered_failure(REGISTERED, OutcomeWrite::Failure, true)]
|
||||
#[case::unbound_success(UNBOUND, OutcomeWrite::Success, false)]
|
||||
#[case::unbound_usage(UNBOUND, OutcomeWrite::AttemptUsage, true)]
|
||||
#[case::unbound_failure(UNBOUND, OutcomeWrite::Failure, false)]
|
||||
#[case::no_make_call_usage(UNBOUND_WITHOUT_MAKE_CALL, OutcomeWrite::AttemptUsage, false)]
|
||||
#[case::bound_sync_success(BOUND_SYNC, OutcomeWrite::Success, false)]
|
||||
#[case::bound_sync_failure(BOUND_SYNC, OutcomeWrite::Failure, true)]
|
||||
#[tokio::test]
|
||||
async fn a_write_runs_when_its_callback_does(
|
||||
#[case] callbacks: OutcomeCallbacks,
|
||||
#[case] write: OutcomeWrite,
|
||||
#[case] runs: bool,
|
||||
) {
|
||||
let recorded = callbacks.record(write, async { Ok(true) }).await;
|
||||
|
||||
assert_eq!(recorded.ok(), Some(runs));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_router_on_the_process_lists_is_the_default() {
|
||||
let parsed: OutcomeCallbacks = serde_json::from_str("{}").expect("an empty document");
|
||||
|
||||
assert_eq!(parsed, REGISTERED);
|
||||
}
|
||||
}
|
||||
300
litellm-rust/crates/router-core/src/precall.rs
Normal file
300
litellm-rust/crates/router-core/src/precall.rs
Normal file
|
|
@ -0,0 +1,300 @@
|
|||
//! `enable_pre_call_checks` (`_pre_call_checks`): before a pick, drop the deployments whose
|
||||
//! context window the prompt exceeds, those outside the request's `allowed_model_region`, and
|
||||
//! those that cannot take a requested param, then decide which error a group with nothing left
|
||||
//! raises. The facts that need live model data (the prompt's token count, each deployment's
|
||||
//! context window, which params it takes) come from [`Host::pre_call_facts`]; every decision is
|
||||
//! made here.
|
||||
//!
|
||||
//! [`Host::pre_call_facts`]: crate::host::Host::pre_call_facts
|
||||
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{error::Error, registry::Deployment};
|
||||
|
||||
/// What the host knows about one deployment for the pre-call checks.
|
||||
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
|
||||
#[serde(default)]
|
||||
pub struct DeploymentFacts {
|
||||
/// The deployment's context window (`get_router_model_info`'s `max_input_tokens`), when it
|
||||
/// declares one and resolves.
|
||||
pub max_input_tokens: Option<i64>,
|
||||
/// `_deployment_lacks_params`: it does not support one of the requested params.
|
||||
pub lacks_params: bool,
|
||||
/// `_deployment_rejects_params`: it lacks one the request and deployment neither drop nor allow.
|
||||
pub rejects_params: bool,
|
||||
}
|
||||
|
||||
/// What the host knows about a request for the pre-call checks.
|
||||
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
|
||||
#[serde(default)]
|
||||
pub struct PreCallFacts {
|
||||
/// The request carries `messages` or `input`: without either no check runs.
|
||||
pub has_input: bool,
|
||||
/// The prompt may be counted against a context window (no compaction is pending).
|
||||
pub countable: bool,
|
||||
/// The prompt's token count, when it was counted.
|
||||
pub input_tokens: Option<i64>,
|
||||
/// The request's `allowed_model_region`.
|
||||
pub allowed_model_region: Option<Value>,
|
||||
/// The requested params the check filters on, sorted (`_pre_call_check_requested_chat_params`).
|
||||
pub requested_params: Vec<String>,
|
||||
/// Facts by deployment id, for the candidates and the group's deployments.
|
||||
pub deployments: BTreeMap<String, DeploymentFacts>,
|
||||
}
|
||||
|
||||
impl PreCallFacts {
|
||||
fn of(&self, deployment: &Deployment) -> DeploymentFacts {
|
||||
self.deployments
|
||||
.get(&deployment.id)
|
||||
.cloned()
|
||||
.unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
/// The name a context-window error gives a deployment: its base model, else its model.
|
||||
fn deployment_model(deployment: &Deployment) -> String {
|
||||
let truthy = |value: Option<&Value>| match value {
|
||||
Some(Value::String(text)) if !text.is_empty() => Some(text.clone()),
|
||||
Some(Value::Null | Value::Bool(false)) | None => None,
|
||||
Some(Value::String(_)) => None,
|
||||
Some(other) => Some(python_str(other)),
|
||||
};
|
||||
truthy(deployment.model_info.raw.get("base_model"))
|
||||
.or_else(|| truthy(deployment.raw_params.get("base_model")))
|
||||
.unwrap_or_else(|| deployment.params.model.clone())
|
||||
}
|
||||
|
||||
/// `str(value)` for the scalar values a base model may hold.
|
||||
fn python_str(value: &Value) -> String {
|
||||
match value {
|
||||
Value::Bool(true) => "True".to_owned(),
|
||||
Value::Number(number) => number.to_string(),
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `is_region_allowed`: the deployment's `region_name` equals the allowed region.
|
||||
fn in_region(deployment: &Deployment, allowed: &Value) -> bool {
|
||||
deployment.raw_params.get("region_name") == Some(allowed)
|
||||
}
|
||||
|
||||
/// `_pre_call_checks` over `candidates`; `pool` is the group's deployments before health filters
|
||||
/// (`request_deployments`), which decide whether nothing left is an unsupported-params error.
|
||||
pub fn filter(
|
||||
model: &str,
|
||||
candidates: Vec<Arc<Deployment>>,
|
||||
pool: &[Arc<Deployment>],
|
||||
facts: &PreCallFacts,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
if !facts.has_input {
|
||||
return Ok(candidates);
|
||||
}
|
||||
let mut invalid = vec![false; candidates.len()];
|
||||
let mut context_window_error = false;
|
||||
let mut detail = String::new();
|
||||
for (index, deployment) in candidates.iter().enumerate() {
|
||||
let found = facts.of(deployment);
|
||||
if let (Some(limit), true, Some(tokens)) =
|
||||
(found.max_input_tokens, facts.countable, facts.input_tokens)
|
||||
&& tokens > limit
|
||||
{
|
||||
invalid[index] = true;
|
||||
context_window_error = true;
|
||||
detail.push_str(&format!(
|
||||
"Model={}, Max Input Tokens={limit}, Got={tokens}",
|
||||
deployment_model(deployment)
|
||||
));
|
||||
continue;
|
||||
}
|
||||
if let Some(allowed) = &facts.allowed_model_region
|
||||
&& !allowed.is_null()
|
||||
&& !in_region(deployment, allowed)
|
||||
{
|
||||
invalid[index] = true;
|
||||
}
|
||||
}
|
||||
let params_checked = !facts.requested_params.is_empty();
|
||||
let param_candidates: Vec<usize> = (0..candidates.len()).filter(|i| !invalid[*i]).collect();
|
||||
let lacking: Vec<bool> = param_candidates
|
||||
.iter()
|
||||
.map(|i| params_checked && facts.of(&candidates[*i]).lacks_params)
|
||||
.collect();
|
||||
let none_takes = !lacking.is_empty() && lacking.iter().all(|lacks| *lacks);
|
||||
let failing: Vec<bool> = if none_takes {
|
||||
param_candidates
|
||||
.iter()
|
||||
.map(|i| params_checked && facts.of(&candidates[*i]).rejects_params)
|
||||
.collect()
|
||||
} else {
|
||||
lacking
|
||||
};
|
||||
for (index, fails) in param_candidates.iter().zip(failing) {
|
||||
if fails {
|
||||
invalid[*index] = true;
|
||||
}
|
||||
}
|
||||
let excluded = invalid.iter().filter(|x| **x).count();
|
||||
if excluded == candidates.len() {
|
||||
if !context_window_error
|
||||
&& !pool.is_empty()
|
||||
&& pool
|
||||
.iter()
|
||||
.all(|d| params_checked && facts.of(d).rejects_params)
|
||||
{
|
||||
return Err(Error::UnsupportedParams {
|
||||
model: model.to_owned(),
|
||||
params: facts.requested_params.clone(),
|
||||
});
|
||||
}
|
||||
if context_window_error {
|
||||
return Err(Error::ContextWindowExceeded {
|
||||
model: model.to_owned(),
|
||||
detail,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(candidates
|
||||
.into_iter()
|
||||
.zip(invalid)
|
||||
.filter_map(|(deployment, invalid)| (!invalid).then_some(deployment))
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::registry::DeploymentSpec;
|
||||
|
||||
fn deployment(id: &str, params: Value, model_info: Value) -> Arc<Deployment> {
|
||||
let mut info = model_info;
|
||||
info["id"] = json!(id);
|
||||
let spec: DeploymentSpec = serde_json::from_value(json!({
|
||||
"model_name": "g", "litellm_params": params, "model_info": info,
|
||||
}))
|
||||
.unwrap();
|
||||
Arc::new(Deployment::from_spec(spec).unwrap())
|
||||
}
|
||||
|
||||
fn facts(entries: &[(&str, DeploymentFacts)]) -> PreCallFacts {
|
||||
PreCallFacts {
|
||||
has_input: true,
|
||||
countable: true,
|
||||
deployments: entries
|
||||
.iter()
|
||||
.map(|(id, f)| ((*id).to_owned(), f.clone()))
|
||||
.collect(),
|
||||
..PreCallFacts::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn window(limit: i64) -> DeploymentFacts {
|
||||
DeploymentFacts {
|
||||
max_input_tokens: Some(limit),
|
||||
..DeploymentFacts::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_prompt_over_a_context_window_drops_that_deployment() {
|
||||
let small = deployment("small", json!({"model": "openai/a"}), json!({}));
|
||||
let big = deployment("big", json!({"model": "openai/b"}), json!({}));
|
||||
let mut found = facts(&[("small", window(5)), ("big", window(10_000))]);
|
||||
found.input_tokens = Some(100);
|
||||
let kept = filter("g", vec![small, big], &[], &found).unwrap();
|
||||
assert_eq!(
|
||||
kept.iter().map(|d| d.id.as_str()).collect::<Vec<_>>(),
|
||||
["big"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn no_window_left_names_each_deployment_back_to_back() {
|
||||
let one = deployment("one", json!({"model": "openai/one"}), json!({}));
|
||||
let two = deployment(
|
||||
"two",
|
||||
json!({"model": "openai/two"}),
|
||||
json!({"base_model": "azure/base"}),
|
||||
);
|
||||
let mut found = facts(&[("one", window(5)), ("two", window(6))]);
|
||||
found.input_tokens = Some(9);
|
||||
let error = filter("g", vec![one, two], &[], &found).unwrap_err();
|
||||
let Error::ContextWindowExceeded { model, detail } = error else {
|
||||
panic!("expected a context window error, got {error:?}");
|
||||
};
|
||||
assert_eq!(model, "g");
|
||||
assert_eq!(
|
||||
detail,
|
||||
"Model=openai/one, Max Input Tokens=5, Got=9Model=azure/base, Max Input Tokens=6, Got=9"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn without_input_nothing_is_checked() {
|
||||
let small = deployment("small", json!({"model": "openai/a"}), json!({}));
|
||||
let mut found = facts(&[("small", window(5))]);
|
||||
found.input_tokens = Some(100);
|
||||
found.has_input = false;
|
||||
assert_eq!(filter("g", vec![small], &[], &found).unwrap().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::in_region("eu", 1)]
|
||||
#[case::out_of_region("ap", 0)]
|
||||
fn the_allowed_region_keeps_only_its_deployments(#[case] region: &str, #[case] kept: usize) {
|
||||
let eu = deployment(
|
||||
"eu",
|
||||
json!({"model": "openai/a", "region_name": "eu"}),
|
||||
json!({}),
|
||||
);
|
||||
let none = deployment("none", json!({"model": "openai/a"}), json!({}));
|
||||
let mut found = facts(&[]);
|
||||
found.allowed_model_region = Some(json!(region));
|
||||
assert_eq!(
|
||||
filter("g", vec![eu, none], &[], &found).unwrap().len(),
|
||||
kept
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_group_where_every_deployment_rejects_the_param_is_unsupported() {
|
||||
let old = deployment("old", json!({"model": "openai/old"}), json!({}));
|
||||
let mut found = facts(&[(
|
||||
"old",
|
||||
DeploymentFacts {
|
||||
lacks_params: true,
|
||||
rejects_params: true,
|
||||
..DeploymentFacts::default()
|
||||
},
|
||||
)]);
|
||||
found.requested_params = vec!["response_format".to_owned()];
|
||||
let error = filter("g", vec![old.clone()], &[old], &found).unwrap_err();
|
||||
assert!(
|
||||
matches!(error, Error::UnsupportedParams { params, .. } if params == ["response_format"])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_deployment_that_takes_the_param_wins_over_one_that_would_drop_it() {
|
||||
let old = deployment("old", json!({"model": "openai/old"}), json!({}));
|
||||
let new = deployment("new", json!({"model": "openai/new"}), json!({}));
|
||||
let mut found = facts(&[(
|
||||
"old",
|
||||
DeploymentFacts {
|
||||
lacks_params: true,
|
||||
..DeploymentFacts::default()
|
||||
},
|
||||
)]);
|
||||
found.requested_params = vec!["response_format".to_owned()];
|
||||
let kept = filter("g", vec![old, new], &[], &found).unwrap();
|
||||
assert_eq!(
|
||||
kept.iter().map(|d| d.id.as_str()).collect::<Vec<_>>(),
|
||||
["new"]
|
||||
);
|
||||
}
|
||||
}
|
||||
53
litellm-rust/crates/router-core/src/prerouting.rs
Normal file
53
litellm-rust/crates/router-core/src/prerouting.rs
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
//! Pre-routing: before an async pick resolves its name, the host may route the request to another
|
||||
//! model group (Python's `_core_async_pre_routing_hook`). The complexity, auto, adaptive and
|
||||
//! quality routers registered from `auto_router/...` marker deployments, and the Claude Code
|
||||
//! session binding, live host-side; the core asks once per pick, before resolution, and resolves
|
||||
//! the name the host answers. The attempt still serves the requested name (its `model_group`):
|
||||
//! only the deployments come from the routed group.
|
||||
//!
|
||||
//! A custom routing strategy (`register_plugin`, `set_custom_routing_strategy`, lar1) replaces the
|
||||
//! whole pick (`_core_async_get_available_deployment` and `get_available_deployment`): the host
|
||||
//! answers the same step with the deployment document the strategy returned, which may name a
|
||||
//! deployment the registry does not hold, and the core serves that deployment as is.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::{Core, config::Op, error::Error, host::RouterRequest, registry::Deployment};
|
||||
|
||||
/// What the host's pre-routing step sees.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct PreRoute<'a> {
|
||||
pub op: &'a Op,
|
||||
pub request: &'a RouterRequest,
|
||||
/// The name this pick resolves: the requested name or the fallback target.
|
||||
pub name: &'a str,
|
||||
/// How many fallbacks deep the chain is (Python's `fallback_depth`): only a first attempt
|
||||
/// binds a Claude Code session.
|
||||
pub fallback_depth: u32,
|
||||
/// 0 for the first attempt in this group, then the retry number.
|
||||
pub retry: u32,
|
||||
}
|
||||
|
||||
/// The host's answer to a pre-routing step.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum PreRouted {
|
||||
/// Resolve this model group instead of the requested name.
|
||||
Model(String),
|
||||
/// A custom strategy's pick: serve this deployment, without resolving or filtering.
|
||||
Deployment(Arc<Deployment>),
|
||||
}
|
||||
|
||||
impl Core {
|
||||
/// What a pick of `ask.name` serves: the host's pre-routed group or custom pick, else the
|
||||
/// name itself. Only picks that ask for it pre-route (`request.pre_routes`).
|
||||
pub(crate) async fn pre_routed(&self, ask: PreRoute<'_>) -> Result<PreRouted, Error> {
|
||||
if !ask.request.pre_routes {
|
||||
return Ok(PreRouted::Model(ask.name.to_owned()));
|
||||
}
|
||||
Ok(self
|
||||
.host
|
||||
.pre_route(&ask)
|
||||
.await?
|
||||
.unwrap_or_else(|| PreRouted::Model(ask.name.to_owned())))
|
||||
}
|
||||
}
|
||||
96
litellm-rust/crates/router-core/src/ptu.rs
Normal file
96
litellm-rust/crates/router-core/src/ptu.rs
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
//! PTU shares (`litellm_core_utils/ptu_shares.py`): a deployment whose `model_info.ptu_shares`
|
||||
//! splits it across teams serves only those teams, and the teams holding a share are kept to
|
||||
//! their share.
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc};
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::registry::Deployment;
|
||||
|
||||
/// `parsed_ptu_shares`: the teams of a non-empty `ptu_shares` map whose every entry is a team id
|
||||
/// and a positive whole number of PTUs (`2.0` reads as 2, a bool is no number), else None.
|
||||
pub fn share_teams(model_info: &Map<String, Value>) -> Option<BTreeSet<String>> {
|
||||
let shares = model_info.get("ptu_shares")?.as_object()?;
|
||||
if shares.is_empty() {
|
||||
return None;
|
||||
}
|
||||
shares
|
||||
.iter()
|
||||
.map(|(team, share)| (!team.is_empty() && whole_ptus(share)).then(|| team.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `_whole_ptus`: a positive whole number.
|
||||
fn whole_ptus(share: &Value) -> bool {
|
||||
match share {
|
||||
Value::Number(number) => number.as_i64().map_or_else(
|
||||
|| {
|
||||
number
|
||||
.as_u64()
|
||||
.map(|whole| whole > 0)
|
||||
.or_else(|| number.as_f64().map(|f| f.fract() == 0.0 && f > 0.0))
|
||||
.unwrap_or(false)
|
||||
},
|
||||
|whole| whole > 0,
|
||||
),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// `filter_ptu_shared_deployments`: the deployments `team_id` holds a share on when it holds one,
|
||||
/// else the ones split into no shares; a caller without a team holds none. Also whether any
|
||||
/// deployment was withheld.
|
||||
pub fn filter_shared(
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
team_id: Option<&str>,
|
||||
) -> (Vec<Arc<Deployment>>, bool) {
|
||||
let team = team_id.unwrap_or("");
|
||||
let checks: Vec<(Arc<Deployment>, Option<BTreeSet<String>>)> = deployments
|
||||
.into_iter()
|
||||
.map(|deployment| {
|
||||
let shares = share_teams(&deployment.model_info.raw);
|
||||
(deployment, shares)
|
||||
})
|
||||
.collect();
|
||||
let holds_share = checks
|
||||
.iter()
|
||||
.any(|(_, shares)| shares.as_ref().is_some_and(|teams| teams.contains(team)));
|
||||
let total = checks.len();
|
||||
let kept: Vec<Arc<Deployment>> = checks
|
||||
.into_iter()
|
||||
.filter(|(_, shares)| match shares {
|
||||
Some(teams) => teams.contains(team),
|
||||
None => !holds_share,
|
||||
})
|
||||
.map(|(deployment, _)| deployment)
|
||||
.collect();
|
||||
let withheld = kept.len() < total;
|
||||
(kept, withheld)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case(json!({"TA": 50}), Some(vec!["TA"]))]
|
||||
#[case(json!({"TA": 2.0, "TB": 1}), Some(vec!["TA", "TB"]))]
|
||||
#[case(json!({"TA": 2.5}), None)]
|
||||
#[case(json!({"TA": true}), None)]
|
||||
#[case(json!({"TA": 0}), None)]
|
||||
#[case(json!({}), None)]
|
||||
#[case(json!({"": 1}), None)]
|
||||
#[case(json!("TA"), None)]
|
||||
fn reads_the_share_teams(#[case] shares: Value, #[case] expected: Option<Vec<&str>>) {
|
||||
let mut model_info = Map::new();
|
||||
model_info.insert("ptu_shares".to_owned(), shares);
|
||||
assert_eq!(
|
||||
share_teams(&model_info),
|
||||
expected.map(|teams| teams.into_iter().map(str::to_owned).collect())
|
||||
);
|
||||
}
|
||||
}
|
||||
1548
litellm-rust/crates/router-core/src/registry.rs
Normal file
1548
litellm-rust/crates/router-core/src/registry.rs
Normal file
File diff suppressed because it is too large
Load diff
585
litellm-rust/crates/router-core/src/resolve.rs
Normal file
585
litellm-rust/crates/router-core/src/resolve.rs
Normal file
|
|
@ -0,0 +1,585 @@
|
|||
//! A requested name to the deployments that may serve this caller: a port of
|
||||
//! `_common_checks_available_deployment`. In order: a specific deployment, a deployment id, the
|
||||
//! alias, a routing group, then for a name no deployment has, a team's public name, every team's
|
||||
//! public name for a proxy admin, a wildcard; else the group (team-first), kept to the caller's
|
||||
//! access groups, then deployments whose litellm model is the name, then the first default
|
||||
//! fallback. Every arm drops strategy markers and deployments that refuse the request's endpoint,
|
||||
//! and withholds deployments reserved for other teams.
|
||||
|
||||
use std::{collections::BTreeSet, sync::Arc};
|
||||
|
||||
use jiff::Timestamp;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::{
|
||||
access::{filter_by_access_groups, filter_reserved},
|
||||
eligibility::eligible_deployments,
|
||||
error::Error,
|
||||
host_index::{pattern_to_regex, regex_match},
|
||||
registry::{Deployment, DeploymentSpec, Registry, Resolution, ResolveKind, routed_deployment},
|
||||
select::is_strategy_marker,
|
||||
};
|
||||
|
||||
/// Who is asking, as the request's auth metadata says.
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct Caller<'a> {
|
||||
/// `user_api_key_team_id`.
|
||||
pub team_id: Option<&'a str>,
|
||||
/// The caller's `user_role` is `proxy_admin`.
|
||||
pub proxy_admin: bool,
|
||||
/// The caller's `models` plus `team_models`.
|
||||
pub allowed_models: &'a [String],
|
||||
/// The endpoint the pick runs under (`endpoint_selection`), which deployments may refuse.
|
||||
pub endpoint: Option<&'a str>,
|
||||
/// The model `endpoint_selection` was given: the name an endpoint refusal reads.
|
||||
pub requested: &'a str,
|
||||
/// `f"{provider}/{model}"`, the spelling `PatternMatchRouter.get_pattern` tries when no
|
||||
/// pattern matches the name itself.
|
||||
pub wildcard_request: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// `_core_should_include_deployment`: the team's own deployment under its public name, or a
|
||||
/// deployment named `model_name` the team may use.
|
||||
fn included(deployment: &Deployment, model_name: &str, team_id: Option<&str>) -> bool {
|
||||
let info = &deployment.model_info;
|
||||
if team_id.is_some()
|
||||
&& info.team_id.as_deref() == team_id
|
||||
&& info.team_public_model_name.as_deref() == Some(model_name)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
deployment.model_name == model_name
|
||||
&& (team_id.is_none() || info.team_id.is_none() || info.team_id.as_deref() == team_id)
|
||||
}
|
||||
|
||||
/// `_core_get_all_deployments(model_name, team_id)`: the team's deployments under that public
|
||||
/// name when it has any, else the deployments named `model_name` the team may use; a name the
|
||||
/// model-name index does not hold is looked for in every deployment, for a team.
|
||||
pub fn all_deployments(
|
||||
registry: &Registry,
|
||||
model_name: &str,
|
||||
team_id: Option<&str>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
let keep = |deployments: Vec<Arc<Deployment>>| -> Vec<Arc<Deployment>> {
|
||||
deployments
|
||||
.into_iter()
|
||||
.filter(|d| included(d, model_name, team_id))
|
||||
.collect()
|
||||
};
|
||||
// A written index position past the model list raises `IndexError` where it is read.
|
||||
let out_of_range = || Error::IndexOutOfRange {
|
||||
model: model_name.to_owned(),
|
||||
};
|
||||
if let Some(team) = team_id {
|
||||
if registry.team_index_overflows(team, model_name) {
|
||||
return Err(out_of_range());
|
||||
}
|
||||
if let Some(own) = registry.team_named(team, model_name) {
|
||||
let own = keep(own);
|
||||
if !own.is_empty() {
|
||||
return Ok(own);
|
||||
}
|
||||
}
|
||||
}
|
||||
if registry.name_index_overflows(model_name) {
|
||||
return Err(out_of_range());
|
||||
}
|
||||
Ok(match registry.named(model_name) {
|
||||
Some(named) => keep(named),
|
||||
None if team_id.is_some() => keep(registry.all().to_vec()),
|
||||
None => Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// `_team_deployments_across_teams`: every team's deployments under public name `model`, for a
|
||||
/// proxy admin calling without a team; a name several teams share is refused.
|
||||
fn across_teams(registry: &Registry, model: &str) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
if registry.across_teams_index_overflows(model) {
|
||||
return Err(Error::IndexOutOfRange {
|
||||
model: model.to_owned(),
|
||||
});
|
||||
}
|
||||
let deployments = registry.team_named_across_teams(model);
|
||||
let teams: BTreeSet<&str> = deployments
|
||||
.iter()
|
||||
.filter_map(|d| d.model_info.team_id.as_deref())
|
||||
.collect();
|
||||
if teams.len() > 1 {
|
||||
return Err(Error::MultipleTeams {
|
||||
model: model.to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(deployments)
|
||||
}
|
||||
|
||||
/// The access groups of the rows `_core_get_model_list(model_name, team_id)` lists: the
|
||||
/// group's deployments, an alias's target deployments, a routing group's rows (which carry no
|
||||
/// access groups), and only when there are none of those, the wildcard's deployments.
|
||||
fn model_access_groups(
|
||||
registry: &Registry,
|
||||
model: &str,
|
||||
team_id: Option<&str>,
|
||||
) -> BTreeSet<String> {
|
||||
let mut rows = all_deployments(registry, model, team_id).unwrap_or_default();
|
||||
let mut stripped_rows = 0;
|
||||
if let Some(target) = registry.alias_of(model) {
|
||||
match registry.routing_group_named(target) {
|
||||
Some(group) => stripped_rows += routing_group_rows(registry, group, None).len(),
|
||||
None => rows.extend(all_deployments(registry, target, None).unwrap_or_default()),
|
||||
}
|
||||
}
|
||||
if let Some(group) = registry.routing_group_named(model) {
|
||||
stripped_rows += routing_group_rows(registry, group, None).len();
|
||||
}
|
||||
if rows.is_empty() && stripped_rows == 0 {
|
||||
rows = registry
|
||||
.wildcard_for(model)
|
||||
.map(|(_, deployments)| deployments)
|
||||
.unwrap_or_default();
|
||||
}
|
||||
rows.iter()
|
||||
.flat_map(|d| d.model_info.access_groups.iter().cloned())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn routing_group_rows(
|
||||
registry: &Registry,
|
||||
group: &crate::registry::RoutingGroupSpec,
|
||||
team_id: Option<&str>,
|
||||
) -> Vec<Arc<Deployment>> {
|
||||
crate::routing_groups::members(group, |member| {
|
||||
all_deployments(registry, member, team_id).unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
/// `_filter_reserved_deployments` for deployments resolved as `model`: markers dropped (only
|
||||
/// markers is refused), those that refuse the caller's endpoint dropped (`_drop_strategy_markers`),
|
||||
/// then the deployments an active access window withholds; when that leaves nothing, the window
|
||||
/// is named in the refusal.
|
||||
fn reserved(
|
||||
model: &str,
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
caller: Caller<'_>,
|
||||
now: Timestamp,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
let any = !deployments.is_empty();
|
||||
let selectable: Vec<Arc<Deployment>> = deployments
|
||||
.into_iter()
|
||||
.filter(|d| !is_strategy_marker(d))
|
||||
.collect();
|
||||
if any && selectable.is_empty() {
|
||||
return Err(Error::OnlyMarkers {
|
||||
model: model.to_owned(),
|
||||
});
|
||||
}
|
||||
let selectable = eligible_deployments(model, caller, selectable)?;
|
||||
let (kept, blocking) = filter_reserved(selectable, caller.team_id, now);
|
||||
match blocking {
|
||||
Some(window) if kept.is_empty() => Err(Error::Reserved {
|
||||
model: model.to_owned(),
|
||||
until: window.end_hhmm(),
|
||||
timezone: window.timezone,
|
||||
}),
|
||||
_ => Ok(kept),
|
||||
}
|
||||
}
|
||||
|
||||
/// `PatternMatchRouter.route`: copies of a pattern's deployments serving `request`.
|
||||
fn routed_wildcard(
|
||||
pattern: &str,
|
||||
request: &str,
|
||||
deployments: &[Arc<Deployment>],
|
||||
) -> Vec<Arc<Deployment>> {
|
||||
let groups = regex_match(&pattern_to_regex(pattern), request).unwrap_or_default();
|
||||
deployments
|
||||
.iter()
|
||||
.map(|deployment| routed_deployment(deployment, &groups, request))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `get_pattern` for `model`: the most specific pattern matching it as spelled, else its
|
||||
/// provider-prefixed spelling, with the pattern's deployments serving that request.
|
||||
fn wildcard_route(
|
||||
registry: &Registry,
|
||||
model: &str,
|
||||
wildcard_request: Option<&str>,
|
||||
) -> Option<(String, Vec<Arc<Deployment>>)> {
|
||||
let route = |request: &str| {
|
||||
registry
|
||||
.wildcard_for(request)
|
||||
.map(|(pattern, deployments)| {
|
||||
let routed = routed_wildcard(&pattern, request, &deployments);
|
||||
(pattern, routed)
|
||||
})
|
||||
};
|
||||
route(model).or_else(|| route(wildcard_request?))
|
||||
}
|
||||
|
||||
/// The name resolution would look for under its provider (`get_pattern`'s second spelling), which
|
||||
/// only the host can infer: one that reaches the pattern routers and that no pattern matches as
|
||||
/// spelled.
|
||||
pub(crate) fn wildcard_name<'a>(
|
||||
registry: &'a Registry,
|
||||
name: &'a str,
|
||||
specific_deployment: bool,
|
||||
) -> Option<&'a str> {
|
||||
let model = registry.alias_of(name).unwrap_or(name);
|
||||
(registry.has_wildcards()
|
||||
&& !specific_deployment
|
||||
&& !registry.has_name(model)
|
||||
&& (registry.has_name(name) || registry.deployment(name).is_none())
|
||||
&& registry.routing_group_named(model).is_none()
|
||||
&& registry.wildcard_for(model).is_none())
|
||||
.then_some(model)
|
||||
}
|
||||
|
||||
/// The first default fallback a name no deployment serves is routed to.
|
||||
pub fn default_target(dynamic: &crate::config::RouterDynamicConfig) -> Option<&str> {
|
||||
dynamic
|
||||
.fallbacks
|
||||
.default_target()
|
||||
.or_else(|| dynamic.default_fallbacks.first().map(String::as_str))
|
||||
}
|
||||
|
||||
/// The router's `default_deployment`, as resolution serves it.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub enum DefaultDeployment<'a> {
|
||||
Spec(&'a DeploymentSpec),
|
||||
/// A value Python cannot copy for a name: resolving to it raises.
|
||||
Raises,
|
||||
}
|
||||
|
||||
/// The default deployment as served for `model`: a copy whose `litellm_params.model` is the name.
|
||||
fn default_for(spec: &DeploymentSpec, model: &str) -> Option<Arc<Deployment>> {
|
||||
let mut spec = spec.clone();
|
||||
spec.litellm_params
|
||||
.insert("model".to_owned(), Value::String(model.to_owned()));
|
||||
Deployment::from_spec(spec).ok().map(Arc::new)
|
||||
}
|
||||
|
||||
/// `_core_get_routable_deployments`: the deployments resolution would consider for `request`,
|
||||
/// before any team, cooldown or strategy filter; nothing when it resolves to nothing or refuses.
|
||||
pub fn routable(
|
||||
registry: &Registry,
|
||||
dynamic: &crate::config::RouterDynamicConfig,
|
||||
request: &crate::host::RouterRequest,
|
||||
) -> Vec<Arc<Deployment>> {
|
||||
let caller = Caller {
|
||||
team_id: request.team_id.as_deref(),
|
||||
proxy_admin: request.proxy_admin,
|
||||
allowed_models: &request.allowed_models,
|
||||
endpoint: request.endpoint.as_deref(),
|
||||
requested: &request.model,
|
||||
wildcard_request: request.wildcard_request.as_deref(),
|
||||
};
|
||||
match resolve_request(
|
||||
registry,
|
||||
&request.model,
|
||||
false,
|
||||
caller,
|
||||
default_target(dynamic),
|
||||
dynamic.default_ref(),
|
||||
Timestamp::now(),
|
||||
) {
|
||||
Ok(Some(resolution)) => resolution.deployments,
|
||||
Ok(None) | Err(_) => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve `name` for `caller`. `None` when nothing serves it (`no_healthy_deployments`).
|
||||
pub fn resolve_request(
|
||||
registry: &Registry,
|
||||
name: &str,
|
||||
specific_deployment: bool,
|
||||
caller: Caller<'_>,
|
||||
default_target: Option<&str>,
|
||||
default_deployment: Option<DefaultDeployment<'_>>,
|
||||
now: Timestamp,
|
||||
) -> Result<Option<Resolution>, Error> {
|
||||
let team_id = caller.team_id;
|
||||
let resolved = |kind, group: &str, deployments| {
|
||||
Ok(Some(Resolution {
|
||||
kind,
|
||||
group: group.to_owned(),
|
||||
deployments,
|
||||
}))
|
||||
};
|
||||
if specific_deployment {
|
||||
let matching = registry
|
||||
.all()
|
||||
.iter()
|
||||
.filter(|d| d.params.model == name)
|
||||
.cloned()
|
||||
.collect();
|
||||
let kept = reserved(name, matching, caller, now)?;
|
||||
return resolved(ResolveKind::SpecificDeployment, name, kept);
|
||||
}
|
||||
if !registry.has_name(name)
|
||||
&& let Some(deployment) = registry.deployment(name)
|
||||
{
|
||||
let group = deployment.model_name.clone();
|
||||
let litellm_model = deployment.params.model.clone();
|
||||
let kept = reserved(&litellm_model, vec![deployment], caller, now)?;
|
||||
return resolved(ResolveKind::Id, &group, kept);
|
||||
}
|
||||
let aliased = registry.alias_of(name);
|
||||
let model = aliased.unwrap_or(name);
|
||||
let routing_group = registry
|
||||
.routing_group_named(model)
|
||||
.map(|group| routing_group_rows(registry, group, team_id));
|
||||
if routing_group.is_none() && !registry.has_name(model) {
|
||||
let early = match team_id {
|
||||
Some(team) => Some(all_deployments(registry, model, Some(team))?),
|
||||
None if caller.proxy_admin => Some(across_teams(registry, model)?),
|
||||
None => None,
|
||||
}
|
||||
.filter(|deployments| !deployments.is_empty());
|
||||
if let Some(deployments) = early {
|
||||
let kept = reserved(model, deployments, caller, now)?;
|
||||
return resolved(ResolveKind::Group, model, kept);
|
||||
}
|
||||
// `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one.
|
||||
let team_owned = team_id
|
||||
.map(|team| crate::lookup::Lookup::new(registry).team_wildcard(model, team))
|
||||
.filter(|deployments| !deployments.is_empty());
|
||||
if let Some(deployments) = team_owned {
|
||||
let routed = deployments.iter().flat_map(|deployment| {
|
||||
let public = deployment.model_info.team_public_model_name.as_deref();
|
||||
let pattern = public.unwrap_or(model);
|
||||
routed_wildcard(pattern, model, std::slice::from_ref(deployment))
|
||||
});
|
||||
let kept = reserved(model, routed.collect(), caller, now)?;
|
||||
return resolved(ResolveKind::Wildcard, model, kept);
|
||||
}
|
||||
if let Some((pattern, deployments)) =
|
||||
wildcard_route(registry, model, caller.wildcard_request)
|
||||
{
|
||||
let kept = reserved(model, deployments, caller, now)?;
|
||||
return resolved(ResolveKind::Wildcard, &pattern, kept);
|
||||
}
|
||||
match default_deployment {
|
||||
Some(DefaultDeployment::Raises) => {
|
||||
return Err(Error::DefaultDeploymentRaises {
|
||||
model: model.to_owned(),
|
||||
});
|
||||
}
|
||||
Some(DefaultDeployment::Spec(spec)) => {
|
||||
if let Some(deployment) = default_for(spec, model) {
|
||||
let kept = reserved(model, vec![deployment], caller, now)?;
|
||||
return resolved(ResolveKind::Default, model, kept);
|
||||
}
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
let (kind, listed) = match routing_group {
|
||||
Some(deployments) => (ResolveKind::RoutingGroup, deployments),
|
||||
None if aliased.is_some() => (
|
||||
ResolveKind::Alias,
|
||||
all_deployments(registry, model, team_id)?,
|
||||
),
|
||||
None => (
|
||||
ResolveKind::Group,
|
||||
all_deployments(registry, model, team_id)?,
|
||||
),
|
||||
};
|
||||
let narrow = |model: &str, deployments: Vec<Arc<Deployment>>| {
|
||||
let allowed = filter_by_access_groups(model, deployments, caller.allowed_models, || {
|
||||
model_access_groups(registry, model, team_id)
|
||||
});
|
||||
reserved(model, allowed, caller, now)
|
||||
};
|
||||
let listed_any = !listed.is_empty();
|
||||
let deployments = narrow(model, listed)?;
|
||||
if !deployments.is_empty() {
|
||||
return resolved(kind, model, deployments);
|
||||
}
|
||||
let mut access_emptied = listed_any;
|
||||
if !listed_any {
|
||||
// The caller may have named a deployment's litellm model instead.
|
||||
let by_model: Vec<Arc<Deployment>> = registry
|
||||
.all()
|
||||
.iter()
|
||||
.filter(|d| d.params.model == model)
|
||||
.cloned()
|
||||
.collect();
|
||||
let found = !by_model.is_empty();
|
||||
let deployments = narrow(model, by_model)?;
|
||||
if !deployments.is_empty() {
|
||||
return resolved(ResolveKind::Group, model, deployments);
|
||||
}
|
||||
access_emptied = found;
|
||||
}
|
||||
if let Some(target) = default_target.filter(|target| !access_emptied && !target.is_empty()) {
|
||||
let deployments = narrow(target, all_deployments(registry, target, team_id)?)?;
|
||||
if !deployments.is_empty() {
|
||||
return resolved(ResolveKind::Group, target, deployments);
|
||||
}
|
||||
// `_common_checks_available_deployment` re-assigned the name: its error names the target.
|
||||
return Err(Error::UnknownModel {
|
||||
model: target.to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
use crate::registry::DeploymentSpec;
|
||||
|
||||
fn registry(entries: Value) -> Registry {
|
||||
let specs: Vec<DeploymentSpec> = serde_json::from_value(entries).unwrap();
|
||||
Registry::build(specs, BTreeMap::new(), Vec::new()).unwrap()
|
||||
}
|
||||
|
||||
fn ids(resolution: Option<Resolution>) -> Vec<String> {
|
||||
resolution
|
||||
.map(|r| r.deployments.iter().map(|d| d.id.clone()).collect())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_access_group_only_caller_is_kept_to_its_groups() {
|
||||
let registry = registry(json!([
|
||||
{"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a", "access_groups": ["A"]}},
|
||||
{"model_name": "g", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "b", "access_groups": ["B"]}},
|
||||
]));
|
||||
let allowed = vec!["A".to_owned()];
|
||||
let caller = Caller {
|
||||
allowed_models: &allowed,
|
||||
..Caller::default()
|
||||
};
|
||||
let resolved = resolve_request(®istry, "g", false, caller, None, None, Timestamp::now());
|
||||
assert_eq!(ids(resolved.unwrap()), ["a"]);
|
||||
// The litellm model string is not narrowed (python-fixes-to-make.md).
|
||||
let by_model = resolve_request(
|
||||
®istry,
|
||||
"openai/b",
|
||||
false,
|
||||
caller,
|
||||
None,
|
||||
None,
|
||||
Timestamp::now(),
|
||||
);
|
||||
assert_eq!(ids(by_model.unwrap()), ["b"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_proxy_admin_is_refused_a_public_name_several_teams_share() {
|
||||
let registry = registry(json!([
|
||||
{"model_name": "x", "litellm_params": {"model": "openai/x"}, "model_info": {"id": "x", "team_id": "tx", "team_public_model_name": "pub"}},
|
||||
{"model_name": "y", "litellm_params": {"model": "openai/y"}, "model_info": {"id": "y", "team_id": "ty", "team_public_model_name": "pub"}},
|
||||
]));
|
||||
let caller = Caller {
|
||||
proxy_admin: true,
|
||||
..Caller::default()
|
||||
};
|
||||
let refused = resolve_request(
|
||||
®istry,
|
||||
"pub",
|
||||
false,
|
||||
caller,
|
||||
None,
|
||||
None,
|
||||
Timestamp::now(),
|
||||
);
|
||||
assert!(matches!(refused, Err(Error::MultipleTeams { model }) if model == "pub"));
|
||||
}
|
||||
|
||||
fn wildcard_group() -> Value {
|
||||
json!([{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*"},
|
||||
"model_info": {"id": "w", "supports_prompt_cache_breakpoint": true},
|
||||
}])
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_wildcard_serves_the_matched_model_without_its_cache_flags() {
|
||||
let registry = registry(wildcard_group());
|
||||
let resolved = resolve_request(
|
||||
®istry,
|
||||
"openai/gpt-4o",
|
||||
false,
|
||||
Caller::default(),
|
||||
None,
|
||||
None,
|
||||
Timestamp::now(),
|
||||
);
|
||||
let resolution = resolved.unwrap().unwrap();
|
||||
assert_eq!(resolution.kind, ResolveKind::Wildcard);
|
||||
assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o");
|
||||
assert_eq!(
|
||||
resolution.deployments[0].record()["model_info"],
|
||||
json!({"id": "w"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_alias_to_a_wildcard_fills_the_model_from_its_target() {
|
||||
let specs: Vec<DeploymentSpec> = serde_json::from_value(wildcard_group()).unwrap();
|
||||
let aliases = BTreeMap::from([("al".to_owned(), "openai/gpt-4o".to_owned())]);
|
||||
let registry = Registry::build(specs, aliases, Vec::new()).unwrap();
|
||||
let resolved = resolve_request(
|
||||
®istry,
|
||||
"al",
|
||||
false,
|
||||
Caller::default(),
|
||||
None,
|
||||
None,
|
||||
Timestamp::now(),
|
||||
);
|
||||
assert_eq!(
|
||||
resolved.unwrap().unwrap().deployments[0].params.model,
|
||||
"openai/gpt-4o"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_name_no_pattern_matches_is_tried_under_its_provider() {
|
||||
let registry = registry(wildcard_group());
|
||||
assert_eq!(wildcard_name(®istry, "gpt-4o", false), Some("gpt-4o"));
|
||||
assert_eq!(wildcard_name(®istry, "openai/gpt-4o", false), None);
|
||||
let caller = Caller {
|
||||
wildcard_request: Some("openai/gpt-4o"),
|
||||
..Caller::default()
|
||||
};
|
||||
let resolved = resolve_request(
|
||||
®istry,
|
||||
"gpt-4o",
|
||||
false,
|
||||
caller,
|
||||
None,
|
||||
None,
|
||||
Timestamp::now(),
|
||||
);
|
||||
let resolution = resolved.unwrap().unwrap();
|
||||
assert_eq!(resolution.group, "openai/*");
|
||||
assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_deployment_named_by_id_refuses_the_endpoint_like_any_other() {
|
||||
let registry = registry(json!([
|
||||
{"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": {
|
||||
"id": "a", "enforce_supported_endpoints": true, "supported_endpoints": ["/v1/embeddings"],
|
||||
}},
|
||||
]));
|
||||
let caller = Caller {
|
||||
endpoint: Some("/v1/chat/completions"),
|
||||
requested: "a",
|
||||
..Caller::default()
|
||||
};
|
||||
let refused = resolve_request(®istry, "a", false, caller, None, None, Timestamp::now());
|
||||
assert!(matches!(
|
||||
refused,
|
||||
Err(Error::EndpointNotSupported { model, .. }) if model == "openai/a"
|
||||
));
|
||||
}
|
||||
}
|
||||
219
litellm-rust/crates/router-core/src/retry.rs
Normal file
219
litellm-rust/crates/router-core/src/retry.rs
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
//! Retry decisions: how many retries a `RetryPolicy` grants for a failure, and whether a failure
|
||||
//! is retried in the same group at all. Mirrors `get_num_retries_from_retry_policy` and
|
||||
//! `resolve_retry_policy` in `litellm_core_utils/get_retry_from_policy.py`, and
|
||||
//! `Router.should_retry_this_error`.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::{
|
||||
config::RetryPolicy,
|
||||
error::{AttemptError, FailureKind},
|
||||
health::should_retry_status,
|
||||
};
|
||||
|
||||
/// `resolve_retry_policy`: the group's own policy when `model_group_retry_policy` has one, else
|
||||
/// the router's `retry_policy`.
|
||||
pub fn resolve_retry_policy<'a>(
|
||||
retry_policy: Option<&'a RetryPolicy>,
|
||||
model_group: &str,
|
||||
model_group_retry_policy: Option<&'a BTreeMap<String, RetryPolicy>>,
|
||||
) -> Option<&'a RetryPolicy> {
|
||||
model_group_retry_policy
|
||||
.and_then(|policies| policies.get(model_group))
|
||||
.or(retry_policy)
|
||||
}
|
||||
|
||||
impl RetryPolicy {
|
||||
/// `get_num_retries_from_retry_policy`: `DefaultRetries` for an `UnsupportedEndpointError`;
|
||||
/// else `NotFoundErrorRetries` for any 404 answer, then the most specific exception class
|
||||
/// the policy names, then `DefaultRetries`.
|
||||
pub fn retries_for(&self, failure: &AttemptError) -> Option<u32> {
|
||||
if failure.is_a("UnsupportedEndpointError") {
|
||||
return self.DefaultRetries;
|
||||
}
|
||||
let not_found = (failure.status == Some(404))
|
||||
.then_some(self.NotFoundErrorRetries)
|
||||
.flatten();
|
||||
not_found
|
||||
.or_else(|| {
|
||||
failure
|
||||
.class_names()
|
||||
.iter()
|
||||
.find_map(|class| self.field_for(class))
|
||||
})
|
||||
.or(self.DefaultRetries)
|
||||
}
|
||||
|
||||
/// `_RETRIES_BY_EXCEPTION_TYPE`: the policy field an exception class reads.
|
||||
fn field_for(&self, class: &str) -> Option<u32> {
|
||||
match class {
|
||||
"AuthenticationError" => self.AuthenticationErrorRetries,
|
||||
"Timeout" => self.TimeoutErrorRetries,
|
||||
"RateLimitError" => self.RateLimitErrorRetries,
|
||||
"ContentPolicyViolationError" => self.ContentPolicyViolationErrorRetries,
|
||||
"BadRequestError" => self.BadRequestErrorRetries,
|
||||
"ServiceUnavailableError" => self.ServiceUnavailableErrorRetries,
|
||||
"InternalServerError" => self.InternalServerErrorRetries,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What `should_retry_this_error` knows besides the failure.
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct RetryContext {
|
||||
pub healthy: usize,
|
||||
pub all: usize,
|
||||
/// `context_window_fallbacks is not None`.
|
||||
pub has_context_window_fallbacks: bool,
|
||||
/// `content_policy_fallbacks is not None`.
|
||||
pub has_content_policy_fallbacks: bool,
|
||||
/// `regular_fallbacks is not None and len(regular_fallbacks) > 0`.
|
||||
pub has_regular_fallbacks: bool,
|
||||
}
|
||||
|
||||
/// `should_retry_this_error`, as a predicate: `false` where Python raises the error.
|
||||
pub fn should_retry(failure: &AttemptError, ctx: RetryContext) -> bool {
|
||||
if failure.is_a("ContextWindowExceededError") && ctx.has_context_window_fallbacks {
|
||||
return false;
|
||||
}
|
||||
if failure.is_a("ContentPolicyViolationError") && ctx.has_content_policy_fallbacks {
|
||||
return false;
|
||||
}
|
||||
if let Some(status) = failure.status
|
||||
&& !should_retry_status(status)
|
||||
&& !matches!(status, 401 | 403)
|
||||
&& !failure.is_a("UnsupportedEndpointError")
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if failure.is_a("NotFoundError") {
|
||||
return false;
|
||||
}
|
||||
if failure.is_a("RateLimitError") && ctx.healthy == 0 && ctx.has_regular_fallbacks {
|
||||
return false;
|
||||
}
|
||||
if (failure.is_a("AuthenticationError") || failure.is_a("PermissionDeniedError"))
|
||||
&& ctx.all <= 1
|
||||
{
|
||||
return false;
|
||||
}
|
||||
ctx.healthy > 0
|
||||
}
|
||||
|
||||
/// `ContentPolicyViolationError`'s classes, most specific first.
|
||||
pub(crate) fn content_policy_classes() -> Vec<std::sync::Arc<str>> {
|
||||
[
|
||||
"ContentPolicyViolationError",
|
||||
"BadRequestError",
|
||||
"APIStatusError",
|
||||
"APIError",
|
||||
"OpenAIError",
|
||||
"Exception",
|
||||
"BaseException",
|
||||
"object",
|
||||
]
|
||||
.into_iter()
|
||||
.map(Into::into)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The exception classes a failure kind stands for, most specific first, for a host that
|
||||
/// reports no class names (`litellm/exceptions.py`).
|
||||
pub(crate) fn kind_classes(kind: FailureKind) -> &'static [&'static str] {
|
||||
match kind {
|
||||
FailureKind::RateLimit => &["RateLimitError"],
|
||||
FailureKind::ContextWindowExceeded => &["ContextWindowExceededError", "BadRequestError"],
|
||||
FailureKind::ContentPolicyViolation => &["ContentPolicyViolationError", "BadRequestError"],
|
||||
FailureKind::Authentication => &["AuthenticationError"],
|
||||
FailureKind::PermissionDenied => &["PermissionDeniedError"],
|
||||
FailureKind::NotFound => &["NotFoundError"],
|
||||
FailureKind::BadRequest => &["BadRequestError"],
|
||||
FailureKind::Timeout => &["Timeout", "APITimeoutError", "APIConnectionError"],
|
||||
FailureKind::ApiConnection => &["APIConnectionError"],
|
||||
FailureKind::InternalServer => &["InternalServerError"],
|
||||
FailureKind::ServiceUnavailable => &["ServiceUnavailableError"],
|
||||
FailureKind::BadGateway => &["BadGatewayError"],
|
||||
FailureKind::CallerCredentialAuthentication => {
|
||||
&["CallerCredentialAuthenticationError", "AuthenticationError"]
|
||||
}
|
||||
FailureKind::CallerCredentialRateLimit => {
|
||||
&["CallerCredentialRateLimitError", "RateLimitError"]
|
||||
}
|
||||
FailureKind::GuardrailIntervention | FailureKind::Declined | FailureKind::Other => &[],
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn failure(kind: FailureKind, status: Option<u16>, classes: &[&str]) -> AttemptError {
|
||||
let mut failure = AttemptError::new(kind, status);
|
||||
failure.classes = classes.iter().map(|class| (*class).into()).collect();
|
||||
failure
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_field(failure(FailureKind::InternalServer, Some(500), &[]), Some(4))]
|
||||
#[case::default(failure(FailureKind::BadGateway, Some(502), &[]), Some(7))]
|
||||
#[case::not_found_status(failure(FailureKind::Other, Some(404), &["APIError"]), Some(1))]
|
||||
#[case::unsupported_endpoint(
|
||||
failure(FailureKind::Declined, Some(400), &["UnsupportedEndpointError", "BadRequestError"]),
|
||||
Some(7)
|
||||
)]
|
||||
#[case::malformed_reads_bad_request(
|
||||
failure(FailureKind::Declined, Some(400), &["MalformedRequestFieldError", "BadRequestError"]),
|
||||
Some(3)
|
||||
)]
|
||||
fn retries_follow_the_most_specific_class(
|
||||
#[case] failure: AttemptError,
|
||||
#[case] expected: Option<u32>,
|
||||
) {
|
||||
let policy = RetryPolicy {
|
||||
InternalServerErrorRetries: Some(4),
|
||||
BadRequestErrorRetries: Some(3),
|
||||
NotFoundErrorRetries: Some(1),
|
||||
DefaultRetries: Some(7),
|
||||
..RetryPolicy::default()
|
||||
};
|
||||
assert_eq!(policy.retries_for(&failure), expected);
|
||||
}
|
||||
|
||||
#[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] has_regular_fallbacks: bool,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
let ctx = RetryContext {
|
||||
healthy,
|
||||
all,
|
||||
has_regular_fallbacks,
|
||||
..RetryContext::default()
|
||||
};
|
||||
assert_eq!(
|
||||
should_retry(&AttemptError::new(kind, status), ctx),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
48
litellm-rust/crates/router-core/src/routing_groups.rs
Normal file
48
litellm-rust/crates/router-core/src/routing_groups.rs
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
//! Routing groups: a callable name over the deployments of its member model groups
|
||||
//! (`_get_routing_group_deployments`). A priority group serves each member's priority as the
|
||||
//! member deployments' `order` (`apply_routing_group_priority`), so the order filter and order
|
||||
//! fallbacks walk the priorities.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::registry::{Deployment, RoutingGroupSpec};
|
||||
|
||||
/// The member deployments of `group`, in member order, with a priority group's priorities as
|
||||
/// their `order`.
|
||||
pub fn members(
|
||||
group: &RoutingGroupSpec,
|
||||
deployments_of: impl Fn(&str) -> Vec<Arc<Deployment>>,
|
||||
) -> Vec<Arc<Deployment>> {
|
||||
let priorities = group
|
||||
.model_priorities
|
||||
.as_ref()
|
||||
.filter(|_| is_priority(group));
|
||||
group
|
||||
.models
|
||||
.iter()
|
||||
.flat_map(|member| {
|
||||
let priority = priorities.and_then(|p| p.get(member)).copied();
|
||||
deployments_of(member)
|
||||
.into_iter()
|
||||
.map(move |deployment| match priority {
|
||||
Some(order) => with_order(&deployment, order),
|
||||
None => deployment,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether `group` is a priority group.
|
||||
pub fn is_priority(group: &RoutingGroupSpec) -> bool {
|
||||
matches!(&group.routing_strategy, crate::config::RoutingStrategy::Other(name) if name == "priority")
|
||||
}
|
||||
|
||||
fn with_order(deployment: &Arc<Deployment>, order: i64) -> Arc<Deployment> {
|
||||
let mut prioritized = Deployment::clone(deployment);
|
||||
prioritized
|
||||
.raw_params
|
||||
.insert("order".to_owned(), Value::from(order));
|
||||
Arc::new(prioritized)
|
||||
}
|
||||
114
litellm-rust/crates/router-core/src/routing_plugins.rs
Normal file
114
litellm-rust/crates/router-core/src/routing_plugins.rs
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
//! `Router(plugins=[...])`: the routing-plugin pipeline. The host runs the plugins (live Python
|
||||
//! objects) in the async pick's pre-routing hook and reports the candidate models they left on
|
||||
//! the request (`_routing_plugin_candidate_models`); router-core decides which deployments that
|
||||
//! keeps (`_filter_by_routing_plugin_candidates`) and refuses the paths that never run the
|
||||
//! pipeline (`_get_available_deployment_unresolved`).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::{config::RoutingStrategy, error::Error, host::RouterRequest, registry::Deployment};
|
||||
|
||||
/// Whether the router's own strategy has an async selector: the others pick on the sync path
|
||||
/// even for async callers (`_async_get_available_deployment`).
|
||||
pub(crate) fn picks_async(strategy: &RoutingStrategy) -> bool {
|
||||
matches!(
|
||||
strategy,
|
||||
RoutingStrategy::UsageBasedV2
|
||||
| RoutingStrategy::SimpleShuffle
|
||||
| RoutingStrategy::CostBased
|
||||
| RoutingStrategy::LatencyBased
|
||||
| RoutingStrategy::LeastBusy
|
||||
)
|
||||
}
|
||||
|
||||
/// The sync pick never runs the pipeline, so a router with plugins refuses it rather than
|
||||
/// letting a policy plugin be bypassed. `request` is a pick or a call, never a listing.
|
||||
pub(crate) fn refuse_sync_pick(
|
||||
has_plugins: bool,
|
||||
strategy: &RoutingStrategy,
|
||||
request: &RouterRequest,
|
||||
) -> Result<(), Error> {
|
||||
if has_plugins
|
||||
&& request.pass_through.is_none()
|
||||
&& (!request.pre_routes || !picks_async(strategy))
|
||||
{
|
||||
return Err(Error::RoutingPluginsNeedAsync);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `_filter_by_routing_plugin_candidates`: keep the deployments whose litellm model the plugins
|
||||
/// left; nothing left is an error, never a fallback to the unfiltered pool.
|
||||
pub(crate) fn narrow(
|
||||
healthy: Vec<Arc<Deployment>>,
|
||||
models: Option<&[String]>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
let Some(models) = models else {
|
||||
return Ok(healthy);
|
||||
};
|
||||
let kept: Vec<Arc<Deployment>> = healthy
|
||||
.into_iter()
|
||||
.filter(|d| models.contains(&d.params.model))
|
||||
.collect();
|
||||
if kept.is_empty() {
|
||||
return Err(Error::RoutingPluginsLeftNothing {
|
||||
models: models.to_vec(),
|
||||
});
|
||||
}
|
||||
Ok(kept)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::registry::DeploymentSpec;
|
||||
|
||||
fn deployment(id: &str, model: &str) -> Arc<Deployment> {
|
||||
let spec: DeploymentSpec = serde_json::from_value(json!({
|
||||
"model_name": "g", "litellm_params": {"model": model}, "model_info": {"id": id},
|
||||
}))
|
||||
.unwrap();
|
||||
Arc::new(Deployment::from_spec(spec).unwrap())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_plugins_models_narrow_the_pool() {
|
||||
let pool = vec![deployment("a", "openai/a"), deployment("b", "openai/b")];
|
||||
let kept = narrow(pool, Some(&["openai/b".to_owned()])).unwrap();
|
||||
assert_eq!(
|
||||
kept.iter().map(|d| d.id.as_str()).collect::<Vec<_>>(),
|
||||
["b"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn narrowing_to_nothing_is_an_error() {
|
||||
let pool = vec![deployment("a", "openai/a")];
|
||||
assert!(matches!(
|
||||
narrow(pool, Some(&[])),
|
||||
Err(Error::RoutingPluginsLeftNothing { models }) if models.is_empty()
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::async_pick(true, RoutingStrategy::SimpleShuffle, false)]
|
||||
#[case::sync_pick(false, RoutingStrategy::SimpleShuffle, true)]
|
||||
#[case::v1_strategy(true, RoutingStrategy::UsageBased, true)]
|
||||
fn only_the_async_pick_runs_plugins(
|
||||
#[case] pre_routes: bool,
|
||||
#[case] strategy: RoutingStrategy,
|
||||
#[case] refused: bool,
|
||||
) {
|
||||
let request = RouterRequest {
|
||||
pre_routes,
|
||||
..RouterRequest::default()
|
||||
};
|
||||
assert_eq!(
|
||||
refuse_sync_pick(true, &strategy, &request).is_err(),
|
||||
refused
|
||||
);
|
||||
}
|
||||
}
|
||||
795
litellm-rust/crates/router-core/src/select.rs
Normal file
795
litellm-rust/crates/router-core/src/select.rs
Normal file
|
|
@ -0,0 +1,795 @@
|
|||
//! Selection: resolve a name, filter its deployments in the order of
|
||||
//! `Router._core_async_get_healthy_deployments` (team, web search, cooldown, blocked, plugins,
|
||||
//! pre-call checks, tags, 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, RouterStateSnapshot,
|
||||
config::{Op, RoutingStrategy},
|
||||
error::Error,
|
||||
health_routing::WhenAllUnhealthy,
|
||||
host::{RouterRequest, RoutingContext, SelectStage},
|
||||
registry::{Deployment, Registry, ResolveKind},
|
||||
resolve::{Caller, resolve_request},
|
||||
strategy::{StrategyScope, routing_context},
|
||||
};
|
||||
|
||||
/// 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;
|
||||
|
||||
/// An index below `len`, as Python's `random.choice` draws one.
|
||||
fn index_below(&self, len: usize) -> usize {
|
||||
((self.next_f64() * len as f64) as usize).min(len - 1)
|
||||
}
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
/// Which pass-through pick a request makes.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum PassThrough {
|
||||
/// `get_available_deployment_for_pass_through` (and `select(kind="pass_through")`): the
|
||||
/// pass-through deployments are chosen first, then cooldown and blocked filters, and only a
|
||||
/// priority group filters by order.
|
||||
Listed,
|
||||
/// `async_get_available_deployment_for_pass_through`: the full healthy chain, then the
|
||||
/// pass-through filter.
|
||||
Healthy,
|
||||
}
|
||||
|
||||
/// A group's deployments before and after filtering.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Candidates {
|
||||
/// The name that was resolved (Python's `model`).
|
||||
pub name: String,
|
||||
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>>,
|
||||
/// The provider-prefixed spelling the name resolved under (`get_pattern`), when the host
|
||||
/// was asked for one.
|
||||
pub wildcard_request: Option<String>,
|
||||
}
|
||||
|
||||
impl Candidates {
|
||||
/// The `model` a no-deployments error names: what `_common_checks_available_deployment`
|
||||
/// returned, so the alias-resolved name a wildcard serves and not its pattern.
|
||||
pub(crate) fn failed_model(&self, registry: &Registry) -> String {
|
||||
let model = match self.kind {
|
||||
ResolveKind::Wildcard => registry.alias_of(&self.name).unwrap_or(&self.name),
|
||||
_ => self.model_group.as_str(),
|
||||
};
|
||||
model.to_owned()
|
||||
}
|
||||
}
|
||||
|
||||
/// 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>,
|
||||
/// A weighted failover hop's `_excluded_deployment_ids`: never picked.
|
||||
pub excluded_ids: Arc<[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,
|
||||
/// The level's name, which an endpoint refusal reads (`endpoint_selection(model=...)`): the
|
||||
/// name before pre-routing replaced it.
|
||||
pub requested: Option<String>,
|
||||
/// A fallback hop: `run_async_fallback` pops the request's `_target_order`, so only the
|
||||
/// hop's own order applies.
|
||||
pub fallback_hop: bool,
|
||||
}
|
||||
|
||||
/// `simple_shuffle`: weighted by `_router_weights[group]`, else `weight`, else `rpm`, else
|
||||
/// `tpm`, each normalized by its maximum, else uniform. A metric is read only when every set
|
||||
/// before it weighed nothing, and a non-numeric one fails the pick then (`_metric_weight`).
|
||||
pub fn simple_shuffle(
|
||||
deployments: &[Arc<Deployment>],
|
||||
model_group: &str,
|
||||
router_weights: Option<&BTreeMap<String, BTreeMap<String, f64>>>,
|
||||
rng: &dyn Randomness,
|
||||
) -> Result<Option<Arc<Deployment>>, Error> {
|
||||
if deployments.is_empty() {
|
||||
return Ok(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"].into_iter().map(|metric| {
|
||||
deployments
|
||||
.iter()
|
||||
.map(|d| d.metric_weight(metric))
|
||||
.collect::<Option<Vec<_>>>()
|
||||
.ok_or(Error::NonNumericWeight { metric })
|
||||
});
|
||||
for weights in scoped.into_iter().map(Ok).chain(metric_sets) {
|
||||
let weights = weights?;
|
||||
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 Ok(Some(deployments[index].clone()));
|
||||
}
|
||||
Ok(Some(
|
||||
deployments[rng.index_below(deployments.len())].clone(),
|
||||
))
|
||||
}
|
||||
|
||||
/// A pre-routing strategy marker (`classify_strategy_router_model`): its model is `auto_router/...`.
|
||||
pub fn is_strategy_marker(deployment: &Deployment) -> bool {
|
||||
deployment.params.model.starts_with("auto_router/")
|
||||
}
|
||||
|
||||
/// `litellm_params.use_in_pass_through`, read for truthiness as Python does.
|
||||
fn uses_pass_through(deployment: &Deployment) -> bool {
|
||||
match deployment.raw_params.get("use_in_pass_through") {
|
||||
None | Some(serde_json::Value::Null) => false,
|
||||
Some(serde_json::Value::Bool(value)) => *value,
|
||||
Some(serde_json::Value::Number(value)) => value.as_f64().is_some_and(|n| n != 0.0),
|
||||
Some(serde_json::Value::String(value)) => !value.is_empty(),
|
||||
Some(serde_json::Value::Array(value)) => !value.is_empty(),
|
||||
Some(serde_json::Value::Object(value)) => !value.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `filter_team_based_models`: a team-owned deployment serves only its team, except that a proxy
|
||||
/// admin without a team keeps deployments that are all team-owned and listed under the model the
|
||||
/// request names, unless they belong to several teams.
|
||||
fn filter_team_based(
|
||||
deployments: Vec<Arc<Deployment>>,
|
||||
request: &RouterRequest,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
let team_id = request.team_id.as_deref();
|
||||
if team_id.is_none()
|
||||
&& request.proxy_admin
|
||||
&& let Some(requested) = request.admin_model.as_deref()
|
||||
&& !deployments.is_empty()
|
||||
&& deployments.iter().all(|d| {
|
||||
d.model_info.team_id.is_some()
|
||||
&& (d.model_name == requested
|
||||
|| d.model_info.team_public_model_name.as_deref() == Some(requested))
|
||||
})
|
||||
{
|
||||
let teams: std::collections::BTreeSet<&str> = deployments
|
||||
.iter()
|
||||
.filter_map(|d| d.model_info.team_id.as_deref())
|
||||
.collect();
|
||||
if teams.len() > 1 {
|
||||
return Err(Error::MultipleTeams {
|
||||
model: requested.to_owned(),
|
||||
});
|
||||
}
|
||||
return Ok(deployments);
|
||||
}
|
||||
Ok(deployments
|
||||
.into_iter()
|
||||
.filter(|d| d.usable_by_team(team_id))
|
||||
.collect())
|
||||
}
|
||||
|
||||
impl Core {
|
||||
fn context<'a>(
|
||||
&'a self,
|
||||
snapshot: &'a RouterStateSnapshot,
|
||||
op: &'a Op,
|
||||
request: &'a RouterRequest,
|
||||
model_group: &'a str,
|
||||
name: &'a str,
|
||||
scope: StrategyScope,
|
||||
) -> RoutingContext<'a> {
|
||||
RoutingContext {
|
||||
registry: &snapshot.registry,
|
||||
settings: &snapshot.settings,
|
||||
health: &self.health,
|
||||
request,
|
||||
op,
|
||||
model_group,
|
||||
name,
|
||||
scope,
|
||||
pool: &[],
|
||||
excluded_ids: &[],
|
||||
skipped_ids: &[],
|
||||
target_order: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The eligible deployments for `name`, in filter-chain order.
|
||||
pub(crate) async fn candidates(
|
||||
&self,
|
||||
snapshot: &RouterStateSnapshot,
|
||||
op: &Op,
|
||||
request: &RouterRequest,
|
||||
name: &str,
|
||||
inputs: &SelectInputs,
|
||||
) -> Result<Candidates, Error> {
|
||||
let registry = &snapshot.registry;
|
||||
// The async pick's pre-routing hook runs the routing plugins before anything resolves;
|
||||
// every pick but a listed pass-through one reads what they left on the request.
|
||||
let plugin_models = if snapshot.settings.routing_plugins
|
||||
&& !inputs.count_only
|
||||
&& request.pass_through != Some(PassThrough::Listed)
|
||||
{
|
||||
let ctx = self.context(snapshot, op, request, name, name, StrategyScope::Default);
|
||||
self.host
|
||||
.routing_plugin_models(&ctx, request.pre_routes)
|
||||
.await?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// No deployment for the name: route to the first default fallback.
|
||||
let default_target = crate::resolve::default_target(&snapshot.settings.dynamic);
|
||||
// `get_pattern`: a name no pattern matches is tried again under its inferred provider.
|
||||
let wildcard_name =
|
||||
crate::resolve::wildcard_name(registry, name, request.specific_deployment);
|
||||
let wildcard_request = match wildcard_name {
|
||||
Some(model) => {
|
||||
let scope = StrategyScope::Default;
|
||||
let ctx = self.context(snapshot, op, request, model, model, scope);
|
||||
self.host.provider_request(&ctx).await?
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
// `endpoint_selection` wraps the attempt, not the count a retry decision reads.
|
||||
let endpoint = request.endpoint.as_deref().filter(|_| !inputs.count_only);
|
||||
let caller = Caller {
|
||||
team_id: request.team_id.as_deref(),
|
||||
proxy_admin: request.proxy_admin,
|
||||
allowed_models: &request.allowed_models,
|
||||
endpoint,
|
||||
requested: inputs.requested.as_deref().unwrap_or(name),
|
||||
wildcard_request: wildcard_request.as_deref(),
|
||||
};
|
||||
let resolution = resolve_request(
|
||||
registry,
|
||||
name,
|
||||
request.specific_deployment,
|
||||
caller,
|
||||
default_target,
|
||||
snapshot.settings.dynamic.default_ref(),
|
||||
jiff::Timestamp::now(),
|
||||
)?
|
||||
.ok_or_else(|| Error::UnknownModel {
|
||||
// The name `_common_checks_available_deployment` looked for, past the alias.
|
||||
model: registry.alias_of(name).unwrap_or(name).to_owned(),
|
||||
})?;
|
||||
let first_model = resolution
|
||||
.deployments
|
||||
.first()
|
||||
.map(|d| d.params.model.clone());
|
||||
let team_id = request.team_id.as_deref();
|
||||
let usable = filter_team_based(resolution.deployments.clone(), request)?;
|
||||
// `_filter_reserved_deployments`: PTU shares keep holders to their share and everyone
|
||||
// else to the deployments split into none.
|
||||
let (mut all, withheld) = crate::ptu::filter_shared(usable, team_id);
|
||||
if withheld && all.is_empty() {
|
||||
return Err(Error::PtuReserved {
|
||||
model: match resolution.kind {
|
||||
ResolveKind::Id => first_model.unwrap_or_else(|| name.to_owned()),
|
||||
ResolveKind::Wildcard | ResolveKind::SpecificDeployment => name.to_owned(),
|
||||
_ => resolution.group.clone(),
|
||||
},
|
||||
});
|
||||
}
|
||||
// `filter_web_search_deployments` runs before `request_deployments` is read, so every later
|
||||
// stage sees the group without the deployments it dropped.
|
||||
if request.web_search && !inputs.count_only && !resolution.kind.is_single() {
|
||||
all.retain(|d| d.supports_web_search());
|
||||
}
|
||||
// `_without_deployments_that_cannot_serve_route`, right after the web search filter: the
|
||||
// deployments whose provider cannot serve the call, unless that leaves nothing. Only a pick
|
||||
// runs it (`_route_capability_pin`): the retry count reads the whole group.
|
||||
if let Some(handler) = request.route_handler.as_deref()
|
||||
&& !inputs.count_only
|
||||
&& !resolution.kind.is_single()
|
||||
&& !all.is_empty()
|
||||
{
|
||||
let ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
let stage = SelectStage::RouteIncapable {
|
||||
handler: handler.to_owned(),
|
||||
};
|
||||
let incapable = self.host.select(&stage, &ctx, all.clone()).await?;
|
||||
let capable: Vec<Arc<Deployment>> = all
|
||||
.iter()
|
||||
.filter(|d| !incapable.iter().any(|i| i.id == d.id))
|
||||
.cloned()
|
||||
.collect();
|
||||
if !capable.is_empty() {
|
||||
all = capable;
|
||||
}
|
||||
}
|
||||
let listed =
|
||||
request.pass_through == Some(PassThrough::Listed) && !resolution.kind.is_single();
|
||||
if listed {
|
||||
all.retain(|d| uses_pass_through(d));
|
||||
if all.is_empty() {
|
||||
// The name `_common_checks_available_deployment` resolved: an alias's target, a
|
||||
// routing group's name, the alias-resolved name a wildcard serves.
|
||||
let model = match resolution.kind {
|
||||
ResolveKind::Wildcard => registry.alias_of(name).unwrap_or(name),
|
||||
ResolveKind::SpecificDeployment => name,
|
||||
_ => resolution.group.as_str(),
|
||||
};
|
||||
return Err(Error::NoPassThroughDeployments {
|
||||
model: model.to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// A deployment named by its id is served even while it cools down
|
||||
// (`get_deployment_by_model_id` returns before the cooldown filter).
|
||||
let settings = &snapshot.settings;
|
||||
let cooling = if resolution.kind.is_single() {
|
||||
Vec::new()
|
||||
} else {
|
||||
// `async_get_cooldown_deployments` reads every deployment's, as one throttled batch.
|
||||
self.health
|
||||
.active_cooldowns(&snapshot.registry.ids())
|
||||
.await?
|
||||
};
|
||||
// Health-check filtering runs before the cooldown filter, except for a deployment named
|
||||
// by its id and the healthy count a retry decision reads.
|
||||
let checked = if resolution.kind.is_single() || inputs.count_only {
|
||||
all.clone()
|
||||
} else {
|
||||
let when_all_unhealthy = if request.health_check_probe {
|
||||
WhenAllUnhealthy::Nothing
|
||||
} else {
|
||||
WhenAllUnhealthy::KeepAll
|
||||
};
|
||||
self.health
|
||||
.health_check_filter(settings, all.clone(), when_all_unhealthy)
|
||||
.await
|
||||
};
|
||||
let mut healthy: Vec<Arc<Deployment>> = checked
|
||||
.iter()
|
||||
.filter(|d| !cooling.iter().any(|(id, _)| *id == d.id))
|
||||
.cloned()
|
||||
.collect();
|
||||
// The safety net when health-check routing drives cooldowns (an allowed_fails_policy):
|
||||
// never on a listed pass-through pick or the healthy count.
|
||||
if healthy.is_empty()
|
||||
&& !listed
|
||||
&& !inputs.count_only
|
||||
&& settings.enable_health_check_routing
|
||||
&& settings.allowed_fails_policy.is_some()
|
||||
{
|
||||
healthy = checked;
|
||||
}
|
||||
if resolution.kind.is_single()
|
||||
&& let Some(deployment) = all.first()
|
||||
{
|
||||
// `_core_async_get_healthy_deployments`: a deployment named by id is checked alone.
|
||||
if deployment.model_info.blocked {
|
||||
return Err(Error::Paused {
|
||||
model: deployment.params.model.clone(),
|
||||
});
|
||||
}
|
||||
if request.pass_through.is_some() && !uses_pass_through(deployment) {
|
||||
return Err(Error::NotPassThrough {
|
||||
id: deployment.id.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
healthy.retain(|d| !d.model_info.blocked);
|
||||
// `continuation_deployments`: a mid-stream continuation never lands on a deployment the host
|
||||
// named as unable to take its prefill.
|
||||
if let Some(excluded) = request
|
||||
.body
|
||||
.get("continuation_excluded_ids")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
{
|
||||
healthy.retain(|d| !excluded.iter().any(|id| id.as_str() == Some(d.id.as_str())));
|
||||
}
|
||||
let request_order = request.target_order.filter(|_| !inputs.fallback_hop);
|
||||
// `async_callback_filter_deployments` runs after the cooldown and blocked filters on every
|
||||
// async candidate list but a deployment named by id and a listed pass-through pick.
|
||||
if settings.callback_filters
|
||||
&& !inputs.count_only
|
||||
&& !resolution.kind.is_single()
|
||||
&& request.pass_through != Some(PassThrough::Listed)
|
||||
{
|
||||
let mut ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
ctx.excluded_ids = &inputs.excluded_ids;
|
||||
ctx.skipped_ids = &inputs.skipped_ids;
|
||||
ctx.target_order = inputs.target_order.or(request_order);
|
||||
healthy = self.host.callback_filters(&ctx, healthy).await?;
|
||||
}
|
||||
if inputs.count_only {
|
||||
return Ok(Candidates {
|
||||
name: name.to_owned(),
|
||||
model_group: resolution.group,
|
||||
kind: resolution.kind,
|
||||
all,
|
||||
healthy,
|
||||
wildcard_request,
|
||||
});
|
||||
}
|
||||
|
||||
if listed {
|
||||
if settings.enable_pre_call_checks && !resolution.kind.is_single() {
|
||||
healthy = self
|
||||
.pre_call_checks(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
&all,
|
||||
healthy,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
// Only a priority group narrows a listed pass-through pick by order.
|
||||
let priority = resolution.kind == ResolveKind::RoutingGroup
|
||||
&& registry.routing_groups().iter().any(|g| {
|
||||
g.group_name == resolution.group && crate::routing_groups::is_priority(g)
|
||||
});
|
||||
if priority {
|
||||
let order = request
|
||||
.target_order
|
||||
.or_else(|| healthy.iter().filter_map(|d| d.order()).min());
|
||||
if let Some(order) = order {
|
||||
healthy.retain(|d| d.order() == Some(order));
|
||||
}
|
||||
}
|
||||
return Ok(Candidates {
|
||||
name: name.to_owned(),
|
||||
model_group: resolution.group,
|
||||
kind: resolution.kind,
|
||||
all,
|
||||
healthy,
|
||||
wildcard_request,
|
||||
});
|
||||
}
|
||||
|
||||
if !settings.plugins.is_empty() && !healthy.is_empty() {
|
||||
let ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
healthy = self
|
||||
.host
|
||||
.select(&SelectStage::Plugins, &ctx, healthy)
|
||||
.await?;
|
||||
}
|
||||
|
||||
if settings.enable_pre_call_checks && !resolution.kind.is_single() {
|
||||
healthy = self
|
||||
.pre_call_checks(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
&all,
|
||||
healthy,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
if request.filter_tags && !healthy.is_empty() {
|
||||
let mut ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
ctx.pool = &all;
|
||||
healthy = self.host.select(&SelectStage::Tags, &ctx, healthy).await?;
|
||||
}
|
||||
|
||||
if !resolution.kind.is_single() {
|
||||
healthy = crate::routing_plugins::narrow(healthy, plugin_models.as_deref())?;
|
||||
}
|
||||
|
||||
// `_without_declared_text_only_deployments`: an image request skips the deployments that
|
||||
// declare no vision, unless that leaves nothing.
|
||||
if request.has_image && !resolution.kind.is_single() && !healthy.is_empty() {
|
||||
let ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
let text_only = self
|
||||
.host
|
||||
.select(&SelectStage::TextOnly, &ctx, healthy.clone())
|
||||
.await?;
|
||||
let vision: Vec<Arc<Deployment>> = healthy
|
||||
.iter()
|
||||
.filter(|d| !text_only.iter().any(|t| t.id == d.id))
|
||||
.cloned()
|
||||
.collect();
|
||||
if !vision.is_empty() {
|
||||
healthy = vision;
|
||||
}
|
||||
}
|
||||
|
||||
// `_without_decisions_deployments_refusing_media`: a Decisions call carrying audio or
|
||||
// videos skips the deployments that refuse them, unless that leaves nothing.
|
||||
if request.decisions_media && !resolution.kind.is_single() && !healthy.is_empty() {
|
||||
let ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&resolution.group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
let refusing = self
|
||||
.host
|
||||
.select(&SelectStage::MediaRefusing, &ctx, healthy.clone())
|
||||
.await?;
|
||||
let taking: Vec<Arc<Deployment>> = healthy
|
||||
.iter()
|
||||
.filter(|d| !refusing.iter().any(|r| r.id == d.id))
|
||||
.cloned()
|
||||
.collect();
|
||||
if !taking.is_empty() {
|
||||
healthy = taking;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(pinned) = request.pinned_deployment_id.as_deref()
|
||||
&& healthy.iter().any(|d| d.id == pinned)
|
||||
{
|
||||
healthy.retain(|d| d.id == pinned);
|
||||
}
|
||||
|
||||
let order = inputs
|
||||
.target_order
|
||||
.or(request_order)
|
||||
.or_else(|| healthy.iter().filter_map(|d| d.order()).min());
|
||||
if let Some(order) = order {
|
||||
healthy.retain(|d| d.order() == Some(order));
|
||||
}
|
||||
|
||||
healthy.retain(|d| {
|
||||
!request.excluded_ids.contains(&d.id) && !inputs.excluded_ids.contains(&d.id)
|
||||
});
|
||||
|
||||
if !inputs.skipped_ids.is_empty() || !request.retry_skipped_ids.is_empty() {
|
||||
let kept: Vec<Arc<Deployment>> = healthy
|
||||
.iter()
|
||||
.filter(|d| {
|
||||
!inputs.skipped_ids.contains(&d.id)
|
||||
&& !request.retry_skipped_ids.contains(&d.id)
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
if !kept.is_empty() {
|
||||
healthy = kept;
|
||||
}
|
||||
}
|
||||
|
||||
if request.pass_through == Some(PassThrough::Healthy)
|
||||
&& !resolution.kind.is_single()
|
||||
&& !healthy.is_empty()
|
||||
{
|
||||
healthy.retain(|d| uses_pass_through(d));
|
||||
if healthy.is_empty() {
|
||||
return Err(Error::NoPassThroughDeployments {
|
||||
model: name.to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Candidates {
|
||||
name: name.to_owned(),
|
||||
model_group: resolution.group,
|
||||
kind: resolution.kind,
|
||||
all,
|
||||
healthy,
|
||||
wildcard_request,
|
||||
})
|
||||
}
|
||||
|
||||
/// `_pre_call_checks` over `healthy`, with the host's facts.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn pre_call_checks(
|
||||
&self,
|
||||
snapshot: &RouterStateSnapshot,
|
||||
op: &Op,
|
||||
request: &RouterRequest,
|
||||
model_group: &str,
|
||||
name: &str,
|
||||
pool: &[Arc<Deployment>],
|
||||
healthy: Vec<Arc<Deployment>>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
let mut ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
model_group,
|
||||
name,
|
||||
StrategyScope::Default,
|
||||
);
|
||||
ctx.pool = pool;
|
||||
let facts = self.host.pre_call_facts(&ctx, &healthy).await?;
|
||||
crate::precall::filter(model_group, healthy, pool, &facts)
|
||||
}
|
||||
|
||||
/// Pick one of `candidates.healthy` with the group's strategy.
|
||||
pub(crate) async fn pick(
|
||||
&self,
|
||||
snapshot: &RouterStateSnapshot,
|
||||
op: &Op,
|
||||
request: &RouterRequest,
|
||||
candidates: &Candidates,
|
||||
) -> Result<Arc<Deployment>, Error> {
|
||||
let no_deployment = || Error::NoHealthyDeployments {
|
||||
model: candidates.failed_model(&snapshot.registry),
|
||||
};
|
||||
// A deployment named by its id is returned as is, before any strategy runs.
|
||||
if candidates.kind.is_single() {
|
||||
return candidates
|
||||
.healthy
|
||||
.first()
|
||||
.cloned()
|
||||
.ok_or_else(no_deployment);
|
||||
}
|
||||
let (strategy, scope) = routing_context(
|
||||
&snapshot.registry,
|
||||
&snapshot.settings.routing_strategy,
|
||||
request.routing_strategy.as_deref(),
|
||||
&candidates.kind,
|
||||
&candidates.model_group,
|
||||
);
|
||||
if strategy == RoutingStrategy::Unset {
|
||||
return Err(no_deployment());
|
||||
}
|
||||
if strategy == RoutingStrategy::SimpleShuffle {
|
||||
let weights = crate::weights::validate(request.router_weights.as_ref())
|
||||
.map_err(|_| Error::InvalidRouterWeights)?;
|
||||
return simple_shuffle(
|
||||
&candidates.healthy,
|
||||
&candidates.model_group,
|
||||
weights.as_ref(),
|
||||
self.rng.as_ref(),
|
||||
)?
|
||||
.ok_or_else(no_deployment);
|
||||
}
|
||||
let ctx = self.context(
|
||||
snapshot,
|
||||
op,
|
||||
request,
|
||||
&candidates.model_group,
|
||||
&candidates.name,
|
||||
scope,
|
||||
);
|
||||
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()
|
||||
.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()
|
||||
.unwrap();
|
||||
assert_eq!(picked.id, "b");
|
||||
}
|
||||
}
|
||||
133
litellm-rust/crates/router-core/src/settings.rs
Normal file
133
litellm-rust/crates/router-core/src/settings.rs
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
//! [`Settings`]: one typed, frozen view of the configuration a caller may read, built from the
|
||||
//! static config and the current dynamic snapshot.
|
||||
|
||||
use std::{collections::BTreeMap, sync::Arc, time::Duration};
|
||||
|
||||
use crate::{
|
||||
config::{
|
||||
AllowedFailsPolicy, DEFAULT_COOLDOWN_TIME, PluginSpec, RetryPolicy, RouterConfig,
|
||||
RouterDynamicConfig, RoutingStrategy,
|
||||
},
|
||||
error::AttemptError,
|
||||
};
|
||||
|
||||
#[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 model_group_retry_policy: Arc<BTreeMap<String, RetryPolicy>>,
|
||||
pub num_retries: u32,
|
||||
pub max_fallbacks: u32,
|
||||
pub retry_after: Duration,
|
||||
pub timeout: Option<Duration>,
|
||||
pub drop_params: Option<bool>,
|
||||
pub enable_pre_call_checks: bool,
|
||||
pub routing_plugins: bool,
|
||||
pub callback_filters: bool,
|
||||
pub enable_health_check_routing: bool,
|
||||
pub health_check_staleness: Duration,
|
||||
pub background_health_check_model_groups: Option<Arc<[String]>>,
|
||||
pub enable_weighted_failover: bool,
|
||||
pub fallback_on_unsupported_params: bool,
|
||||
pub provider_default_deployment_ids: Arc<[String]>,
|
||||
pub dynamic: Arc<RouterDynamicConfig>,
|
||||
}
|
||||
|
||||
/// `DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER`.
|
||||
pub const DEFAULT_HEALTH_CHECK_STALENESS: Duration = Duration::from_secs(300 * 2);
|
||||
|
||||
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(),
|
||||
model_group_retry_policy: Arc::new(config.model_group_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_pre_call_checks: config.enable_pre_call_checks,
|
||||
routing_plugins: config.routing_plugins,
|
||||
callback_filters: config.callback_filters,
|
||||
enable_health_check_routing: config.enable_health_check_routing,
|
||||
health_check_staleness: config
|
||||
.health_check_staleness_threshold
|
||||
.map_or(DEFAULT_HEALTH_CHECK_STALENESS, seconds),
|
||||
background_health_check_model_groups: config
|
||||
.background_health_check_model_groups
|
||||
.clone()
|
||||
.map(Into::into),
|
||||
enable_weighted_failover: config.enable_weighted_failover,
|
||||
fallback_on_unsupported_params: config.fallback_on_unsupported_params.unwrap_or(true),
|
||||
provider_default_deployment_ids: config.provider_default_deployment_ids.clone().into(),
|
||||
dynamic,
|
||||
}
|
||||
}
|
||||
|
||||
/// `allowed_fails_for`: the policy's count for this failure.
|
||||
pub fn allowed_fails_for(&self, failure: &AttemptError) -> Option<u32> {
|
||||
self.allowed_fails_policy
|
||||
.as_ref()
|
||||
.and_then(|policy| policy.allowed_fails_for(failure))
|
||||
}
|
||||
|
||||
/// How long an allowed-fails counter lives: `settings().cooldown_time or
|
||||
/// DEFAULT_COOLDOWN_TIME_SECONDS`, so a router cooldown of 0 still counts over the default.
|
||||
pub fn fails_window(&self) -> Duration {
|
||||
if self.cooldown_time.is_zero() {
|
||||
DEFAULT_COOLDOWN_TIME
|
||||
} else {
|
||||
self.cooldown_time
|
||||
}
|
||||
}
|
||||
|
||||
/// `_is_allowed_fails_set_on_router`.
|
||||
pub fn allowed_fails_set(&self) -> bool {
|
||||
self.allowed_fails.is_some() && !self.dynamic.allowed_fails_is_default
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::passed_to_the_router(Some(3), false, true)]
|
||||
#[case::the_process_default(Some(3), true, false)]
|
||||
#[case::none(None, false, false)]
|
||||
fn allowed_fails_counts_as_set_unless_it_is_the_process_default(
|
||||
#[case] allowed_fails: Option<u32>,
|
||||
#[case] allowed_fails_is_default: bool,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
let dynamic = RouterDynamicConfig {
|
||||
allowed_fails,
|
||||
allowed_fails_is_default,
|
||||
..RouterDynamicConfig::default()
|
||||
};
|
||||
let settings = Settings::new(&RouterConfig::default(), Arc::new(dynamic));
|
||||
assert_eq!(settings.allowed_fails_set(), expected);
|
||||
}
|
||||
}
|
||||
110
litellm-rust/crates/router-core/src/strategy.rs
Normal file
110
litellm-rust/crates/router-core/src/strategy.rs
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
//! Which routing strategy picks for a request, and whose state it reads: a port of
|
||||
//! `Router._get_routing_context`. A per-request `routing_strategy` override wins, then the
|
||||
//! routing group the resolved model belongs to (called as the group, or as a member of a
|
||||
//! non-priority group), then the router's own strategy.
|
||||
|
||||
use crate::{
|
||||
config::RoutingStrategy,
|
||||
registry::{Registry, ResolveKind},
|
||||
};
|
||||
|
||||
/// The per-request overrides `_OVERRIDABLE_ROUTING_STRATEGIES` honours: simple-shuffle and every
|
||||
/// strategy with a default selector. Anything else is ignored.
|
||||
pub const OVERRIDABLE: [&str; 6] = [
|
||||
"simple-shuffle",
|
||||
"least-busy",
|
||||
"usage-based-routing",
|
||||
"usage-based-routing-v2",
|
||||
"latency-based-routing",
|
||||
"cost-based-routing",
|
||||
];
|
||||
|
||||
/// Whose selector state a strategy pick reads.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum StrategyScope {
|
||||
/// The router's own strategy (the implicit `"default"` group).
|
||||
Default,
|
||||
/// A routing group's own selector.
|
||||
Group(String),
|
||||
/// A per-request override: the default selector when it names the router's strategy, else
|
||||
/// one selector per strategy kept for the router's life.
|
||||
Override,
|
||||
}
|
||||
|
||||
impl StrategyScope {
|
||||
/// The scope as the host step spells it.
|
||||
pub fn wire(&self) -> (&'static str, Option<&str>) {
|
||||
match self {
|
||||
Self::Default => ("default", None),
|
||||
Self::Group(name) => ("group", Some(name)),
|
||||
Self::Override => ("override", None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A request's `routing_strategy` override, when it is one Python honours.
|
||||
pub fn request_override(raw: Option<&str>) -> Option<RoutingStrategy> {
|
||||
let raw = raw?;
|
||||
OVERRIDABLE
|
||||
.contains(&raw)
|
||||
.then(|| raw.parse().unwrap_or_default())
|
||||
}
|
||||
|
||||
/// The strategy and selector scope for a pick of `model_group` (the resolved group, after
|
||||
/// aliases), reached through a resolution of `kind`.
|
||||
pub fn routing_context(
|
||||
registry: &Registry,
|
||||
default: &RoutingStrategy,
|
||||
override_strategy: Option<&str>,
|
||||
kind: &ResolveKind,
|
||||
model_group: &str,
|
||||
) -> (RoutingStrategy, StrategyScope) {
|
||||
if let Some(strategy) = request_override(override_strategy) {
|
||||
return (strategy, StrategyScope::Override);
|
||||
}
|
||||
let group = match kind {
|
||||
ResolveKind::RoutingGroup => registry
|
||||
.routing_groups()
|
||||
.iter()
|
||||
.find(|group| group.group_name == model_group),
|
||||
ResolveKind::Alias | ResolveKind::Group => registry.member_routing_group(model_group),
|
||||
ResolveKind::SpecificDeployment
|
||||
| ResolveKind::Id
|
||||
| ResolveKind::Default
|
||||
| ResolveKind::Wildcard => None,
|
||||
};
|
||||
match group {
|
||||
None => (default.clone(), StrategyScope::Default),
|
||||
Some(group) if is_priority(&group.routing_strategy) => {
|
||||
(RoutingStrategy::SimpleShuffle, StrategyScope::Default)
|
||||
}
|
||||
Some(group) => (
|
||||
group.routing_strategy.clone(),
|
||||
StrategyScope::Group(group.group_name.clone()),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_priority(strategy: &RoutingStrategy) -> bool {
|
||||
matches!(strategy, RoutingStrategy::Other(name) if name == "priority")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::honoured(Some("least-busy"), Some(RoutingStrategy::LeastBusy))]
|
||||
#[case::shuffle(Some("simple-shuffle"), Some(RoutingStrategy::SimpleShuffle))]
|
||||
#[case::lar1_ignored(Some("lar1"), None)]
|
||||
#[case::unknown_ignored(Some("bogus"), None)]
|
||||
#[case::unset(None, None)]
|
||||
fn only_overridable_strategies_override(
|
||||
#[case] raw: Option<&str>,
|
||||
#[case] expected: Option<RoutingStrategy>,
|
||||
) {
|
||||
assert_eq!(request_override(raw), expected);
|
||||
}
|
||||
}
|
||||
81
litellm-rust/crates/router-core/src/weights.rs
Normal file
81
litellm-rust/crates/router-core/src/weights.rs
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
//! `_router_weights` as `validate_router_weights` reads it (`litellm/types/router_weights.py`):
|
||||
//! model group to deployment id to a finite weight of at least 0 (an int or a float, never a
|
||||
//! bool or a string), with non-blank ids and at least one positive weight in every non-empty group.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
/// Validated request weights.
|
||||
pub type RouterWeights = BTreeMap<String, BTreeMap<String, f64>>;
|
||||
|
||||
/// The request's weights didn't validate; the host raises the validator's own error.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct InvalidRouterWeights;
|
||||
|
||||
fn identifier(text: &str) -> bool {
|
||||
text.chars().any(|c| !c.is_whitespace())
|
||||
}
|
||||
|
||||
fn weight(value: &Value) -> Option<f64> {
|
||||
match value {
|
||||
Value::Number(number) => number.as_f64().filter(|w| w.is_finite() && *w >= 0.0),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `validate_router_weights(raw)`: `None` when the request passed none.
|
||||
pub fn validate(raw: Option<&Value>) -> Result<Option<RouterWeights>, InvalidRouterWeights> {
|
||||
let groups = match raw {
|
||||
None | Some(Value::Null) => return Ok(None),
|
||||
Some(Value::Object(groups)) => groups,
|
||||
Some(_) => return Err(InvalidRouterWeights),
|
||||
};
|
||||
let mut validated = RouterWeights::new();
|
||||
for (group, weights) in groups {
|
||||
let Value::Object(weights) = weights else {
|
||||
return Err(InvalidRouterWeights);
|
||||
};
|
||||
if !identifier(group) {
|
||||
return Err(InvalidRouterWeights);
|
||||
}
|
||||
let mut entries = BTreeMap::new();
|
||||
for (id, value) in weights {
|
||||
if !identifier(id) {
|
||||
return Err(InvalidRouterWeights);
|
||||
}
|
||||
entries.insert(id.clone(), weight(value).ok_or(InvalidRouterWeights)?);
|
||||
}
|
||||
if !entries.is_empty() && !entries.values().any(|w| *w > 0.0) {
|
||||
return Err(InvalidRouterWeights);
|
||||
}
|
||||
validated.insert(group.clone(), entries);
|
||||
}
|
||||
Ok(Some(validated))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::all_zero(json!({"g": {"a": 0.0, "b": 0.0}}))]
|
||||
#[case::negative(json!({"g": {"a": -1.0}}))]
|
||||
#[case::text(json!({"g": {"a": "1"}}))]
|
||||
#[case::boolean(json!({"g": {"a": true}}))]
|
||||
#[case::blank_id(json!({"g": {" ": 1.0}}))]
|
||||
#[case::not_a_mapping(json!(["g"]))]
|
||||
fn invalid_weights_are_refused(#[case] raw: Value) {
|
||||
assert_eq!(validate(Some(&raw)), Err(InvalidRouterWeights));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty_group(json!({"g": {}}))]
|
||||
#[case::int_weight(json!({"g": {"a": 1, "b": 0}}))]
|
||||
fn valid_weights_are_kept(#[case] raw: Value) {
|
||||
assert!(validate(Some(&raw)).unwrap().is_some());
|
||||
}
|
||||
}
|
||||
1757
litellm-rust/crates/router-core/tests/call_loop.rs
Normal file
1757
litellm-rust/crates/router-core/tests/call_loop.rs
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -620,7 +620,7 @@ class DualCache(BaseCache):
|
|||
increment_list: list["RedisPipelineIncrementOperation"],
|
||||
local_only: bool = False,
|
||||
parent_otel_span: Span | None = None,
|
||||
**kwargs,
|
||||
**kwargs: object, # kwargs-ok: accepted and unused, as the cache interface's other calls take them
|
||||
) -> list[float] | None:
|
||||
result: list[float] | None = None
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -202,7 +202,7 @@ def response_total_token_count(response: object) -> int:
|
|||
|
||||
def apply_remaining_usage_headers(
|
||||
additional_headers: dict[str, object],
|
||||
remaining_usage: dict[str, int],
|
||||
remaining_usage: Mapping[str, object],
|
||||
) -> None:
|
||||
for header, value in remaining_usage.items():
|
||||
if value is not None and header not in additional_headers:
|
||||
|
|
|
|||
|
|
@ -108,6 +108,17 @@ def as_object_list(value: object) -> Sequence[object] | None:
|
|||
return value if isinstance(value, list) else None
|
||||
|
||||
|
||||
_Key = TypeVar("_Key") # rebind-ok: a TypeVar cannot be Final
|
||||
_Value = TypeVar("_Value") # rebind-ok: a TypeVar cannot be Final
|
||||
|
||||
|
||||
def without(
|
||||
mapping: Mapping[_Key, _Value], *keys: object
|
||||
) -> dict[_Key, _Value]: # mutable-ok: a fresh copy the caller owns
|
||||
"""``mapping`` without ``keys``, as a new dict in the same order; ``mapping`` is unchanged."""
|
||||
return {key: value for key, value in mapping.items() if key not in keys}
|
||||
|
||||
|
||||
def decoded_json_object(value: object) -> Mapping[str, object]:
|
||||
mapping: Final = as_str_mapping(value)
|
||||
if mapping is not None:
|
||||
|
|
|
|||
|
|
@ -1081,4 +1081,8 @@ def convert_to_model_response_object(
|
|||
end_time=end_time,
|
||||
convert_tool_call_to_json_mode=convert_tool_call_to_json_mode,
|
||||
)
|
||||
raise Exception(f"Invalid response object {traceback.format_exc()}\n\nreceived_args={received_args}")
|
||||
# The conversion's own traceback, without the errors being handled when it ran: a router
|
||||
# retry runs inside the handler of the attempt before it, and that attempt's error text
|
||||
# (a rate limit) would otherwise decide how this one is classified.
|
||||
own_traceback: Final = "".join(traceback.format_exception(type(e), e, e.__traceback__, chain=False))
|
||||
raise Exception(f"Invalid response object {own_traceback}\n\nreceived_args={received_args}")
|
||||
|
|
|
|||
|
|
@ -378,7 +378,8 @@ def filter_team_based_models(
|
|||
|
||||
ids_to_remove: Final = set()
|
||||
if isinstance(healthy_deployments, dict):
|
||||
return healthy_deployments
|
||||
# A deployment asked for by id is filtered like any other: another team's is not served.
|
||||
return [] if _owner_team(healthy_deployments) not in (None, request_team_id) else healthy_deployments # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the legacy untyped deployment dict
|
||||
for deployment in healthy_deployments:
|
||||
_model_info = deployment.get("model_info") or {}
|
||||
model_team_id = _model_info.get("team_id")
|
||||
|
|
@ -431,6 +432,16 @@ def _deployment_supports_web_search(deployment: dict) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
def is_web_search_request(request_kwargs: Mapping[str, object]) -> bool:
|
||||
tools: Final[object] = request_kwargs.get("tools")
|
||||
return isinstance(tools, (list, tuple)) and any(
|
||||
isinstance(tool, dict)
|
||||
and isinstance(tool_type := tool.get("type"), str)
|
||||
and tool_type in _WEB_SEARCH_TOOL_TYPES
|
||||
for tool in tools
|
||||
)
|
||||
|
||||
|
||||
def filter_web_search_deployments(
|
||||
healthy_deployments: list[dict] | dict,
|
||||
request_kwargs: dict | None = None,
|
||||
|
|
@ -445,14 +456,7 @@ def filter_web_search_deployments(
|
|||
if isinstance(healthy_deployments, dict):
|
||||
return healthy_deployments
|
||||
|
||||
tools: Final[object] = request_kwargs.get("tools")
|
||||
is_web_search_request: Final = isinstance(tools, (list, tuple)) and any(
|
||||
isinstance(tool, dict)
|
||||
and isinstance(tool_type := tool.get("type"), str)
|
||||
and tool_type in _WEB_SEARCH_TOOL_TYPES
|
||||
for tool in tools
|
||||
)
|
||||
if not is_web_search_request:
|
||||
if not is_web_search_request(request_kwargs):
|
||||
return healthy_deployments
|
||||
|
||||
# Filter out deployments that don't support web search
|
||||
|
|
@ -607,3 +611,9 @@ def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mappin
|
|||
)
|
||||
verbose_router_logger.warning(warning)
|
||||
return warning
|
||||
|
||||
|
||||
def _owner_team(deployment: Mapping[str, object]) -> object:
|
||||
"""The team a deployment belongs to (``model_info.team_id``), or None."""
|
||||
model_info: Final = deployment.get("model_info")
|
||||
return model_info.get("team_id") if isinstance(model_info, Mapping) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a deployment dict's model_info
|
||||
|
|
|
|||
|
|
@ -670,6 +670,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
# Synthesized compaction block from compact_20260112 polyfill (streaming).
|
||||
self.compaction_block = compaction_block
|
||||
self.iterations_usage = iterations_usage
|
||||
# The provider's error behind this stream's error frame, for a router that accounts for it.
|
||||
self.stream_failure: Exception | None = None
|
||||
self._refusal_text: str = ""
|
||||
self._paused_turn: bool = False
|
||||
self._claude_refusal_stop_details: AnthropicStopDetails | None = None
|
||||
|
|
@ -1396,6 +1398,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
verbose_logger.exception("Anthropic Adapter - mid-stream error: %s", e)
|
||||
logging_obj: Final = self.litellm_logging_obj
|
||||
provider_error: Final = _provider_error(e)
|
||||
self.stream_failure = provider_error
|
||||
if logging_obj is not None and logging_obj.on_detached_stream_failure is not None:
|
||||
if provider_error is e:
|
||||
raise
|
||||
|
|
@ -1619,6 +1622,10 @@ class AnthropicSSEStream(AsyncIterator[bytes]):
|
|||
self._byte_stream: Final[AsyncIterator[bytes]] = anthropic_wrapper.async_anthropic_sse_wrapper()
|
||||
self._hidden_params: dict[str, object] = {}
|
||||
|
||||
@property
|
||||
def stream_failure(self) -> Exception | None:
|
||||
return self._anthropic_wrapper.stream_failure
|
||||
|
||||
@property
|
||||
def chunks(self) -> "list[ModelResponseStream] | None":
|
||||
return self._anthropic_wrapper.chunks
|
||||
|
|
|
|||
|
|
@ -141,6 +141,8 @@ class AnthropicResponsesStreamWrapper:
|
|||
if litellm_logging_obj is not None:
|
||||
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
|
||||
self._unreported_failure: Exception | None = None
|
||||
# The failure behind this stream's error frame, for a router that accounts for it.
|
||||
self.stream_failure: Exception | None = None
|
||||
self._current_block_index: int = -1
|
||||
# Map item_id -> content_block_index so we can stop the right block later
|
||||
self._item_id_to_block_index: dict[str, int] = {}
|
||||
|
|
@ -467,6 +469,7 @@ class AnthropicResponsesStreamWrapper:
|
|||
llm_provider=getattr(self.responses_stream, "custom_llm_provider", None),
|
||||
model=self.model,
|
||||
)
|
||||
self.stream_failure = self._unreported_failure
|
||||
self._chunk_queue.append(_anthropic_error_chunk(status_code, message))
|
||||
|
||||
async def _report_unreported_failure(self) -> None:
|
||||
|
|
@ -578,5 +581,9 @@ class AnthropicResponsesSSEByteStream:
|
|||
async def aclose(self) -> None:
|
||||
await self._frames.aclose()
|
||||
|
||||
@property
|
||||
def stream_failure(self) -> Exception | None:
|
||||
return self._wrapper.stream_failure
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return self._wrapper.partial_response_on_disconnect()
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ outside code used before the entry points live in ``DeprecatedRouterMembers`` un
|
|||
import inspect
|
||||
import weakref
|
||||
from collections.abc import Mapping
|
||||
from types import CellType, FunctionType, MappingProxyType
|
||||
from types import FunctionType, MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm import router_python_core
|
||||
|
|
@ -35,7 +35,7 @@ from litellm.types.router import (
|
|||
RetryPolicy,
|
||||
)
|
||||
|
||||
_CORE_SIGNATURE: Final = inspect.signature(PythonCore.__init__)
|
||||
_CORE_SIGNATURE: Final = inspect.signature(PythonCore.__init__) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # reflection reads only the parameters' names
|
||||
|
||||
|
||||
def router_arguments(
|
||||
|
|
@ -65,142 +65,40 @@ else:
|
|||
_replaced_core_members: Final[ # mutable-ok: originals of patched core members, restored on delete
|
||||
dict[tuple[type, str], list[object]]
|
||||
] = {}
|
||||
_FACADE_ONLY: Final = frozenset(
|
||||
{
|
||||
"__module__",
|
||||
"__qualname__",
|
||||
"__doc__",
|
||||
"__dict__",
|
||||
"__weakref__",
|
||||
"__init__",
|
||||
"__init_subclass__",
|
||||
"__annotations__",
|
||||
"__firstlineno__",
|
||||
"__static_attributes__",
|
||||
"__classcell__",
|
||||
"_core_class",
|
||||
}
|
||||
)
|
||||
|
||||
def _with_class_cell(value: object, cell: "CellType") -> object:
|
||||
"""``value`` with its ``__class__`` cell (what zero-argument ``super()`` reads) replaced by ``cell``."""
|
||||
if isinstance(value, staticmethod | classmethod):
|
||||
return type(value)(_function_with_class_cell(value.__func__, cell))
|
||||
if isinstance(value, property):
|
||||
return property(
|
||||
*(
|
||||
None if f is None else _function_with_class_cell(f, cell)
|
||||
for f in (value.fget, value.fset, value.fdel)
|
||||
),
|
||||
value.__doc__,
|
||||
)
|
||||
return _function_with_class_cell(value, cell)
|
||||
|
||||
def _function_with_class_cell(value: object, cell: "CellType") -> object:
|
||||
"""A plain function with its ``__class__`` cell replaced by ``cell``; anything else unchanged."""
|
||||
if not isinstance(value, FunctionType) or "__class__" not in value.__code__.co_freevars:
|
||||
return value
|
||||
closure: Final = tuple(
|
||||
cell if name == "__class__" else original
|
||||
for name, original in zip(value.__code__.co_freevars, value.__closure__ or (), strict=True)
|
||||
)
|
||||
rebound: Final = FunctionType(value.__code__, value.__globals__, value.__name__, value.__defaults__, closure)
|
||||
rebound.__kwdefaults__ = value.__kwdefaults__
|
||||
rebound.__qualname__ = value.__qualname__
|
||||
rebound.__dict__.update(value.__dict__)
|
||||
return rebound
|
||||
|
||||
class _Absent:
|
||||
"""A facade-only member, hidden from the core class: reading it raises the usual AttributeError."""
|
||||
|
||||
def __set_name__(self, owner: type, name: str) -> None:
|
||||
self._name = name
|
||||
|
||||
def __get__(self, instance: object, owner: type | None = None) -> object:
|
||||
if instance is None:
|
||||
raise AttributeError(
|
||||
f"type object {getattr(owner, '__name__', owner)!r} has no attribute {self._name!r}"
|
||||
)
|
||||
raise AttributeError(f"{type(instance).__name__!r} object has no attribute {self._name!r}")
|
||||
|
||||
def _core_class_for(facade: type, parent: type) -> type:
|
||||
"""
|
||||
A core class carrying a Router subclass's own members, so the core runs its overrides. It
|
||||
subclasses the facade too, so an override's ``super(Sub, self)`` finds the core it runs on.
|
||||
"""
|
||||
cell: Final = CellType()
|
||||
members: Final = {
|
||||
name: _with_class_cell(value, cell) for name, value in vars(facade).items() if name not in _FACADE_ONLY
|
||||
}
|
||||
namespace: Final = {
|
||||
# The facade's own plumbing sits in the core class's MRO ahead of the parent core: put
|
||||
# the parent core's versions back, and hide the facade-only members.
|
||||
"__init__": parent.__init__,
|
||||
"__getattribute__": object.__getattribute__,
|
||||
"__setattr__": parent.__setattr__,
|
||||
"__delattr__": parent.__delattr__,
|
||||
"__dir__": parent.__dir__,
|
||||
**{name: _Absent() for name in ("_attach", "_core", "from_config", "_core_class")},
|
||||
**members,
|
||||
"__module__": facade.__module__,
|
||||
}
|
||||
core_class: Final = _CoreType(f"{facade.__name__}Core", (facade, parent), namespace)
|
||||
cell.cell_contents = core_class
|
||||
for name in members:
|
||||
type.__delattr__(facade, name)
|
||||
return core_class
|
||||
|
||||
class _RouterType(type):
|
||||
"""
|
||||
Class-level reads of core members (``Router._helper(...)``) go to the router's core class, and
|
||||
so do class-level patches, which put the original back when the patch is removed.
|
||||
``isinstance`` and ``dir`` see the core class too.
|
||||
Class-level reads of core members (``Router._helper(...)``) go to ``PythonCore``, and so do
|
||||
class-level patches, which put the original back when the patch is removed.
|
||||
``isinstance`` and ``dir`` see ``PythonCore`` too.
|
||||
"""
|
||||
|
||||
def __getattr__(cls, name: str) -> object:
|
||||
if name == "_core_class":
|
||||
raise AttributeError(name)
|
||||
return getattr(cls._core_class, name)
|
||||
return getattr(PythonCore, name)
|
||||
|
||||
def __setattr__(cls, name: str, value: object) -> None:
|
||||
core_class: Final = vars(cls).get("_core_class") or getattr(cls, "_core_class", None)
|
||||
if core_class is None or any(name in vars(base) for base in cls.__mro__) or not hasattr(core_class, name):
|
||||
if any(name in vars(base) for base in cls.__mro__) or not hasattr(PythonCore, name):
|
||||
super().__setattr__(name, value)
|
||||
return
|
||||
_replaced_core_members.setdefault((core_class, name), []).append(vars(core_class).get(name, _MISSING))
|
||||
setattr(core_class, name, value)
|
||||
_replaced_core_members.setdefault((PythonCore, name), []).append(vars(PythonCore).get(name, _MISSING))
|
||||
setattr(PythonCore, name, value)
|
||||
|
||||
def __delattr__(cls, name: str) -> None:
|
||||
core_class: Final = getattr(cls, "_core_class", None)
|
||||
replaced: Final = _replaced_core_members.get((core_class, name)) if core_class is not None else None
|
||||
replaced: Final = _replaced_core_members.get((PythonCore, name))
|
||||
if not replaced or any(name in vars(base) for base in cls.__mro__):
|
||||
super().__delattr__(name)
|
||||
return
|
||||
original: Final = replaced.pop()
|
||||
if original is _MISSING:
|
||||
delattr(core_class, name)
|
||||
delattr(PythonCore, name)
|
||||
return
|
||||
setattr(core_class, name, original)
|
||||
setattr(PythonCore, name, original)
|
||||
|
||||
def __instancecheck__(cls, instance: object) -> bool:
|
||||
return type.__instancecheck__(cls, instance) or isinstance(instance, cls._core_class)
|
||||
return type.__instancecheck__(cls, instance) or isinstance(instance, PythonCore)
|
||||
|
||||
def __dir__(cls) -> list[str]: # mutable-ok: dir() returns a list
|
||||
return sorted({*type.__dir__(cls), *dir(cls._core_class)})
|
||||
|
||||
class _CoreType(_RouterType):
|
||||
"""The metaclass of a Router subclass's core class: a plain class, none of the facade's redirects."""
|
||||
|
||||
def __getattr__(cls, name: str) -> object:
|
||||
raise AttributeError(f"type object {cls.__name__!r} has no attribute {name!r}")
|
||||
|
||||
__setattr__ = type.__setattr__
|
||||
__delattr__ = type.__delattr__
|
||||
__instancecheck__ = type.__instancecheck__
|
||||
__dir__ = type.__dir__
|
||||
|
||||
def _is_core(instance: object) -> bool:
|
||||
return isinstance(type(instance), _CoreType)
|
||||
return sorted({*type.__dir__(cls), *dir(PythonCore)})
|
||||
|
||||
def _patchable_method(core: object, name: str) -> bool:
|
||||
"""A method of the core: an instance-level replacement of it is also kept on the facade, so
|
||||
|
|
@ -212,27 +110,20 @@ else:
|
|||
return hasattr(type(core), "__getattr__") and inspect.ismethod(getattr(core, name, None))
|
||||
|
||||
class Router(DeprecatedRouterMembers, metaclass=_RouterType):
|
||||
_core_class: type[PythonCore] = PythonCore
|
||||
|
||||
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: class keywords, passed on to super()
|
||||
if isinstance(cls, _CoreType):
|
||||
# A core class built below: the parent core's hooks wire it, as for any core class.
|
||||
super().__init_subclass__(**kwargs)
|
||||
return
|
||||
# DeprecatedRouterMembers wires a subclass's overrides of old members; here that runs
|
||||
# on the core class built from it, which is where the router's own calls land.
|
||||
super(DeprecatedRouterMembers, cls).__init_subclass__(**kwargs)
|
||||
type.__setattr__(cls, "_core_class", _core_class_for(cls, cls._core_class))
|
||||
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: class keywords, refused with the class
|
||||
raise TypeError(
|
||||
"litellm.Router cannot be subclassed. Use the router's entry points instead: "
|
||||
"register_plugin() for custom routing, apply() for settings, observe() for outcomes. "
|
||||
"See https://docs.litellm.ai/docs/routing_plugins."
|
||||
)
|
||||
|
||||
def __init__(self, *args: object, **kwargs: object) -> None: # kwargs-ok: Router's constructor arguments
|
||||
self._attach(
|
||||
make_core(build_router_config(router_arguments(*args, **kwargs)), python_host, type(self)._core_class)
|
||||
)
|
||||
self._attach(make_core(build_router_config(router_arguments(*args, **kwargs)), python_host))
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: RouterConfig, host: PythonHost = python_host) -> "Router":
|
||||
router: Final = cls.__new__(cls)
|
||||
router._attach(make_core(config, host, cls._core_class))
|
||||
router._attach(make_core(config, host))
|
||||
return router
|
||||
|
||||
def _attach(self, core: RouterCore) -> None:
|
||||
|
|
@ -253,8 +144,7 @@ else:
|
|||
try:
|
||||
return object.__getattribute__(self, name)
|
||||
except AttributeError:
|
||||
if _is_core(self):
|
||||
raise
|
||||
pass
|
||||
facade: Final = object.__getattribute__(self, "__dict__")
|
||||
core: Final = facade.get("_core")
|
||||
if core is None:
|
||||
|
|
@ -273,9 +163,6 @@ else:
|
|||
) from None
|
||||
|
||||
def __setattr__(self, name: str, value: object) -> None:
|
||||
if _is_core(self):
|
||||
super().__setattr__(name, value)
|
||||
return
|
||||
if _class_descriptor(type(self), name):
|
||||
object.__setattr__(self, name, value)
|
||||
return
|
||||
|
|
@ -289,9 +176,6 @@ else:
|
|||
self.__dict__[name] = value
|
||||
|
||||
def __delattr__(self, name: str) -> None:
|
||||
if _is_core(self):
|
||||
super().__delattr__(name)
|
||||
return
|
||||
if _class_descriptor(type(self), name):
|
||||
object.__delattr__(self, name)
|
||||
return
|
||||
|
|
@ -299,9 +183,7 @@ else:
|
|||
delattr(self._core, name)
|
||||
|
||||
def __dir__(self) -> list[str]: # mutable-ok: dir() returns a list
|
||||
if _is_core(self):
|
||||
return super().__dir__()
|
||||
return sorted({*super().__dir__(), *dir(self.__dict__.get("_core", type(self)._core_class))})
|
||||
return sorted({*super().__dir__(), *dir(self.__dict__.get("_core", PythonCore))})
|
||||
|
||||
Router.__init__.__signature__ = _CORE_SIGNATURE.replace( # pyright: ignore[reportFunctionMemberAccess] # inspect reads it
|
||||
parameters=[parameter for parameter in _CORE_SIGNATURE.parameters.values() if parameter.name != "_host"]
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ as they were.
|
|||
# members are still untyped (bare ``dict`` / ``list``) or mutable (the live objects); the overloads
|
||||
# repeat them unchanged, and typing the members themselves is the core's job.
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
|
|
@ -45,6 +46,7 @@ from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, overload
|
|||
from litellm.router_config import detached
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from functools import _lru_cache_wrapper
|
||||
from typing import Any # noqa: TID251 # only to spell the cores' untyped (Unknown) containers and one Any type
|
||||
|
||||
from opentelemetry.trace import Span
|
||||
|
|
@ -179,9 +181,9 @@ def deployment_usable_by_team(deployment: "Mapping[str, object] | Deployment", t
|
|||
Whether ``team_id`` may use ``deployment``: a team-scoped deployment (``model_info.team_id``
|
||||
set) only by that team, any other deployment by everyone. Needs no router instance.
|
||||
"""
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_rust.surface.registry import HostRegistry
|
||||
|
||||
return PythonCore._core_deployment_usable_by_team(deployment, team_id)
|
||||
return HostRegistry._core_deployment_usable_by_team(deployment, team_id)
|
||||
|
||||
|
||||
def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str:
|
||||
|
|
@ -201,11 +203,77 @@ def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str
|
|||
return hashlib.sha256("".join(parts).encode()).hexdigest()
|
||||
|
||||
|
||||
# The arguments ``Router.__init__`` accepts, in its signature's order. Every core reads this list
|
||||
# (the Python core's ``get_valid_args`` reads its constructor's signature, and the core suite checks
|
||||
# that the two agree).
|
||||
ROUTER_INIT_ARGS: Final[tuple[str, ...]] = (
|
||||
"model_list",
|
||||
"assistants_config",
|
||||
"search_tools",
|
||||
"guardrail_list",
|
||||
"redis_url",
|
||||
"redis_host",
|
||||
"redis_port",
|
||||
"redis_password",
|
||||
"redis_db",
|
||||
"cache_responses",
|
||||
"cache_kwargs",
|
||||
"caching_groups",
|
||||
"client_ttl",
|
||||
"polling_interval",
|
||||
"default_priority",
|
||||
"num_retries",
|
||||
"max_fallbacks",
|
||||
"timeout",
|
||||
"stream_timeout",
|
||||
"default_litellm_params",
|
||||
"default_max_parallel_requests",
|
||||
"set_verbose",
|
||||
"debug_level",
|
||||
"default_fallbacks",
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
"model_group_alias",
|
||||
"enable_pre_call_checks",
|
||||
"enable_tag_filtering",
|
||||
"tag_filtering_match_any",
|
||||
"tag_routing_prefix",
|
||||
"plugins",
|
||||
"retry_after",
|
||||
"retry_policy",
|
||||
"model_group_retry_policy",
|
||||
"allowed_fails",
|
||||
"allowed_fails_policy",
|
||||
"cooldown_time",
|
||||
"disable_cooldowns",
|
||||
"fallback_on_unsupported_params",
|
||||
"routing_strategy",
|
||||
"optional_pre_call_checks",
|
||||
"routing_strategy_args",
|
||||
"routing_groups",
|
||||
"provider_budget_config",
|
||||
"alerting_config",
|
||||
"router_general_settings",
|
||||
"deployment_affinity_ttl_seconds",
|
||||
"model_group_affinity_config",
|
||||
"ignore_invalid_deployments",
|
||||
"enable_health_check_routing",
|
||||
"health_check_staleness_threshold",
|
||||
"health_check_ignore_transient_errors",
|
||||
"background_health_check_model_groups",
|
||||
"enable_weighted_failover",
|
||||
"enable_mid_stream_fallback_continuation",
|
||||
"fallback_access_check",
|
||||
"fallback_budget_check",
|
||||
"auto_router_capability_limit",
|
||||
"session_affinity_repin_on_failover",
|
||||
)
|
||||
|
||||
|
||||
def router_init_args() -> tuple[str, ...]:
|
||||
"""The arguments ``Router.__init__`` accepts (was ``Router.get_valid_args``). Needs no router instance."""
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
return tuple(PythonCore.get_valid_args())
|
||||
return ROUTER_INIT_ARGS
|
||||
|
||||
|
||||
def mark_advisor_orchestration_failure(exception: BaseException) -> None:
|
||||
|
|
@ -269,9 +337,9 @@ def deployment_metadata(request_kwargs: Mapping[str, object]) -> Mapping[str, ob
|
|||
|
||||
|
||||
def _json_default_stable_id(obj: object) -> str:
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_rust.steps.prerouting import _json_default_stable_id as stable_id
|
||||
|
||||
return PythonCore._json_default_stable_id(obj)
|
||||
return stable_id(obj)
|
||||
|
||||
|
||||
class RouterRegistry:
|
||||
|
|
@ -316,6 +384,13 @@ class RouterRegistry:
|
|||
return self._router._core_cached_deployment_model_info(id, model_name)
|
||||
return self._router._core_get_deployment_model_info(model_id=id, model_name=model_name)
|
||||
|
||||
def deployment_model_info_cache(self) -> "_lru_cache_wrapper[ModelMapInfo | None]":
|
||||
"""
|
||||
The lru cache ``deployment_model_info(..., cached=True)`` reads, called as ``(id, model_name)``
|
||||
and cleared with ``cache_clear()`` (was the ``cached_deployment_model_info`` attribute).
|
||||
"""
|
||||
return self._router._core_cached_deployment_model_info
|
||||
|
||||
def router_model_info(
|
||||
self,
|
||||
deployment: "dict[str, object] | Deployment | None", # mutable-ok: the core takes a dict
|
||||
|
|
@ -532,10 +607,8 @@ class RouterRegistry:
|
|||
``should_include_deployment``), or, without ``model_name``, whether the team may use it at
|
||||
all (was ``deployment_usable_by_team``).
|
||||
"""
|
||||
from litellm.router_python_core import PythonCore
|
||||
|
||||
if model_name is None:
|
||||
return PythonCore._core_deployment_usable_by_team(deployment, team_id)
|
||||
return deployment_usable_by_team(deployment, team_id)
|
||||
return self._router._core_should_include_deployment(model_name=model_name, model=deployment, team_id=team_id)
|
||||
|
||||
@overload
|
||||
|
|
@ -1074,13 +1147,22 @@ PLAIN_SETTING_ATTRIBUTES: Final[frozenset[str]] = frozenset({"num_retries", "max
|
|||
ASSIGNABLE_ALIASES: Final[Mapping[str, str]] = MappingProxyType({"model_group_alias": "_core_model_group_alias"})
|
||||
|
||||
_MISSING: Final = object()
|
||||
# The settings fields ``RouterSettings`` copies, so an in-place edit of what a read returns never
|
||||
# reaches the router (only ``apply()`` changes settings).
|
||||
_COPIED_SETTING_FIELDS: Final = (
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
"default_litellm_params",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RouterSettings:
|
||||
"""
|
||||
The router's configuration as one frozen snapshot (``settings()``). Fields hold the router's
|
||||
own objects, not copies: mutating a list in here mutates the router's.
|
||||
own objects, except the fallback lists and ``default_litellm_params``, which are copies: settings
|
||||
change only through ``apply()``, so editing one of those in place changes nothing on any core.
|
||||
A field the router never set reads as ``None``.
|
||||
|
||||
Each field has the type of the Router attribute it reads (``_core_<field>``), the type the old
|
||||
|
|
@ -1119,6 +1201,12 @@ class RouterSettings:
|
|||
_router: "Router | None" = field(default=None, repr=False, compare=False)
|
||||
_unset: frozenset[str] = field(default=frozenset(), repr=False, compare=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
for name in _COPIED_SETTING_FIELDS:
|
||||
value: object = getattr(self, name) # pyright: ignore[reportAny] # a settings field
|
||||
if isinstance(value, list | dict):
|
||||
object.__setattr__(self, name, copy.deepcopy(value)) # pyright: ignore[reportUnknownArgumentType] # a settings list or dict
|
||||
|
||||
@classmethod
|
||||
def of(cls, router: "Router") -> "RouterSettings":
|
||||
read: Final = {name: getattr(router, "_core_" + name, _MISSING) for name in SETTING_FIELDS}
|
||||
|
|
@ -1188,16 +1276,16 @@ class RouterHealth:
|
|||
@staticmethod
|
||||
def is_blocked(deployment: "Deployment") -> bool:
|
||||
"""Whether an admin paused this deployment (was ``is_deployment_blocked``)."""
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_rust.surface.health_routing import is_blocked as deployment_is_blocked
|
||||
|
||||
return PythonCore._core_is_deployment_blocked(deployment)
|
||||
return deployment_is_blocked(deployment)
|
||||
|
||||
@staticmethod
|
||||
def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool:
|
||||
"""Whether every deployment in the list is paused (was ``are_all_deployments_blocked``)."""
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_rust.surface.health_routing import all_blocked as deployments_all_blocked
|
||||
|
||||
return PythonCore._core_are_all_deployments_blocked(deployments)
|
||||
return deployments_all_blocked(deployments)
|
||||
|
||||
def fully_blocked_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set
|
||||
"""Groups whose every deployment is paused (was ``get_fully_blocked_model_names``)."""
|
||||
|
|
|
|||
215
litellm/router_attempt.py
Normal file
215
litellm/router_attempt.py
Normal file
|
|
@ -0,0 +1,215 @@
|
|||
"""
|
||||
One router attempt, as the host runs it: the request for the deployment the router core picked.
|
||||
|
||||
The Rust router core picks a deployment and hands the attempt to its host; ``PythonHost`` builds
|
||||
the keyword arguments for that one deployment here and calls ``litellm.<op>`` with them. Nothing
|
||||
here selects, retries, falls back or records cooldowns: the core does all of that. Nothing here
|
||||
reads a Python router either, so this module outlives ``router_python_core.py``.
|
||||
"""
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
|
||||
from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter
|
||||
from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs
|
||||
|
||||
# Request keys the router reads to route the call; the deployment call never receives them.
|
||||
ROUTING_KEYS: Final = frozenset(
|
||||
{
|
||||
"model",
|
||||
"num_retries",
|
||||
"disable_fallbacks",
|
||||
"specific_deployment",
|
||||
"_router_weights",
|
||||
"_retry_skipped_deployment_ids",
|
||||
"_target_order",
|
||||
"model_group_retry_policy",
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
"mock_testing_fallbacks",
|
||||
"mock_testing_context_fallbacks",
|
||||
"mock_testing_content_policy_fallbacks",
|
||||
"mock_testing_rate_limit_error",
|
||||
}
|
||||
)
|
||||
|
||||
# The positional parameters of each endpoint, in litellm's order: a call that passes them
|
||||
# positionally is read back into keywords before it is routed.
|
||||
POSITIONAL: Final = MappingProxyType(
|
||||
{
|
||||
"acompletion": ("model", "messages"),
|
||||
"completion": ("model", "messages"),
|
||||
"atext_completion": ("model", "prompt"),
|
||||
"text_completion": ("model", "prompt"),
|
||||
"aembedding": ("model", "input"),
|
||||
"embedding": ("model", "input"),
|
||||
"aimage_generation": ("prompt", "model"),
|
||||
"image_generation": ("prompt", "model"),
|
||||
"atranscription": ("model", "file"),
|
||||
"aspeech": ("model", "input", "voice"),
|
||||
"arerank": ("model", "query", "documents"),
|
||||
"amoderation": ("input", "model"),
|
||||
"aresponses": ("input", "model"),
|
||||
"aanthropic_messages": (),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The call's arguments as one keyword request. Raises ``TypeError`` like the endpoint would."""
|
||||
names: Final = POSITIONAL.get(op, ())
|
||||
if len(args) > len(names):
|
||||
raise TypeError(f"{op}() takes {len(names)} positional arguments but {len(args)} were given")
|
||||
request: Final = dict(kwargs)
|
||||
for name, value in zip(names, args, strict=False):
|
||||
if name in request:
|
||||
raise TypeError(f"{op}() got multiple values for argument {name!r}")
|
||||
request[name] = value
|
||||
return request
|
||||
|
||||
|
||||
# Endpoints the Python router serves through ``_ageneric_api_call_with_fallbacks``, which writes
|
||||
# the routing metadata under ``litellm_metadata``.
|
||||
_GENERIC_API_OPS: Final = frozenset({"aanthropic_messages", "aresponses"})
|
||||
|
||||
|
||||
def metadata_field(op: str) -> str:
|
||||
"""``metadata`` or ``litellm_metadata``: where the router writes what it routed the call as."""
|
||||
return "litellm_metadata" if op in _GENERIC_API_OPS else get_router_metadata_variable_name(function_name=op)
|
||||
|
||||
|
||||
def prepare_call(op: str, request: Mapping[str, object], model_group_alias: str | None) -> Mapping[str, object]:
|
||||
"""What the router adds to a request once, before its first attempt
|
||||
(``_update_kwargs_before_fallbacks``): a trace id and the model group it was asked for."""
|
||||
field: Final = metadata_field(op)
|
||||
metadata: Final = dict(_as_mapping(request.get(field)))
|
||||
metadata.update({"model_group": request.get("model"), "model_group_alias": model_group_alias})
|
||||
return {**request, "litellm_trace_id": request.get("litellm_trace_id") or str(uuid.uuid4()), field: metadata}
|
||||
|
||||
|
||||
def _wildcard_match(pattern: object, requested: str) -> "re.Match[str] | None":
|
||||
if not isinstance(pattern, str) or "*" not in pattern:
|
||||
return None
|
||||
regex: Final = PatternMatchRouter().pattern_to_regex(pattern)
|
||||
return re.match(regex, requested) or re.match(regex, f"{inferred_provider(requested)}/{requested}")
|
||||
|
||||
|
||||
def for_requested_model(deployment: Mapping[str, object], requested: str | None) -> Mapping[str, object]:
|
||||
"""A wildcard deployment as the router serves ``requested`` through it: its ``model`` with
|
||||
the wildcards filled in from the requested name (``PatternMatchRouter.route``), by the team's
|
||||
public pattern (``_core_team_pattern_routers``) or its own. Any other deployment comes back
|
||||
unchanged."""
|
||||
params: Final = _as_mapping(deployment.get("litellm_params"))
|
||||
model: Final = params.get("model")
|
||||
if requested is None or not isinstance(model, str):
|
||||
return deployment
|
||||
public_name: Final = _as_mapping(deployment.get("model_info")).get("team_public_model_name")
|
||||
match: Final = _wildcard_match(public_name, requested) or _wildcard_match(deployment.get("model_name"), requested)
|
||||
if match is None:
|
||||
return deployment
|
||||
return {
|
||||
**deployment,
|
||||
"litellm_params": {**params, "model": PatternMatchRouter.set_deployment_model_name(match, model)},
|
||||
}
|
||||
|
||||
|
||||
def attempt_kwargs(
|
||||
op: str,
|
||||
request: Mapping[str, object],
|
||||
deployment: Mapping[str, object],
|
||||
*,
|
||||
model_group: str,
|
||||
default_litellm_params: Mapping[str, object],
|
||||
timeout: float | None,
|
||||
) -> Mapping[str, object]:
|
||||
"""The keyword arguments for ``litellm.<op>`` against one deployment
|
||||
(``_update_kwargs_with_deployment`` and the call sites that merge ``litellm_params``)."""
|
||||
params: Final = _as_mapping(deployment.get("litellm_params"))
|
||||
model_info: Final = _model_info_for(request, deployment, model_group)
|
||||
field: Final = metadata_field(op)
|
||||
metadata: Final = dict(_as_mapping(request.get(field)))
|
||||
metadata.update(
|
||||
{
|
||||
"model_group": model_group,
|
||||
"deployment": params.get("model"),
|
||||
"model_info": model_info,
|
||||
"api_base": params.get("api_base"),
|
||||
"deployment_model_name": deployment.get("model_name"),
|
||||
}
|
||||
)
|
||||
tags: Final = _as_list(params.get("tags"))
|
||||
if tags:
|
||||
merged: Final = list(_as_list(metadata.get("tags")))
|
||||
merged.extend(tag for tag in tags if tag not in merged)
|
||||
metadata["tags"] = merged
|
||||
call: Final = {key: value for key, value in request.items() if key not in ROUTING_KEYS}
|
||||
for key, value in default_litellm_params.items():
|
||||
call.setdefault(key, value)
|
||||
call[field] = metadata
|
||||
call["model_info"] = model_info
|
||||
if "timeout" not in call:
|
||||
deployment_timeout: Final = params.get("timeout")
|
||||
chosen: Final = deployment_timeout if deployment_timeout is not None else timeout
|
||||
if chosen is not None:
|
||||
call["timeout"] = chosen
|
||||
return without_router_only_kwargs({**params, **call, "model": params.get("model")})
|
||||
|
||||
|
||||
def _model_info_for(
|
||||
request: Mapping[str, object], deployment: Mapping[str, object], model_group: str
|
||||
) -> dict[str, object]: # mutable-ok: a fresh dict for the call
|
||||
"""The deployment's ``model_info``, or, for a request carrying its own credentials, the
|
||||
request-scoped deployment's (``_handle_clientside_credential``): its own stable id, and the
|
||||
configured deployment's id as ``original_model_id``."""
|
||||
from litellm.litellm_core_utils import clientside_credential_handler as clientside
|
||||
from litellm.router_api import deployment_id
|
||||
|
||||
model_info: Final = dict(_as_mapping(deployment.get("model_info")))
|
||||
if not clientside.is_clientside_credential(request_kwargs=dict(request)): # pyright: ignore[reportUnknownMemberType] # an untyped helper
|
||||
return model_info
|
||||
dynamic: Final = _as_mapping(
|
||||
clientside.get_dynamic_litellm_params( # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # an untyped helper
|
||||
litellm_params=dict(_as_mapping(deployment.get("litellm_params"))), request_kwargs=dict(request)
|
||||
)
|
||||
)
|
||||
original: Final = model_info.get("id")
|
||||
model_info["id"] = deployment_id(model_group, dynamic)
|
||||
model_info["original_model_id"] = original
|
||||
return model_info
|
||||
|
||||
|
||||
def stamp_attempt_error(error: BaseException, deployment: Mapping[str, object]) -> None:
|
||||
"""What the router stamps on a failed attempt's error: the deployment's own ``num_retries``
|
||||
(``_set_deployment_num_retries_on_exception``) and its id (``failed_deployment_id``)."""
|
||||
params: Final = _as_mapping(deployment.get("litellm_params"))
|
||||
if getattr(error, "num_retries", None) is None and params.get("num_retries") is not None:
|
||||
try:
|
||||
error.num_retries = int(str(params["num_retries"])) # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the router stamps the error in place
|
||||
except ValueError:
|
||||
pass
|
||||
deployment_id: Final = _as_mapping(deployment.get("model_info")).get("id")
|
||||
if deployment_id and not getattr(error, "failed_deployment_id", None):
|
||||
try:
|
||||
error.failed_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: see above
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
|
||||
def _as_mapping(value: object) -> Mapping[str, object]:
|
||||
if isinstance(value, Mapping):
|
||||
entries: Final[Mapping[object, object]] = value # pyright: ignore[reportUnknownVariableType] # a request's untyped mapping
|
||||
return {str(key): item for key, item in entries.items()}
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def _as_list(value: object) -> list[object]: # mutable-ok: a fresh list
|
||||
if isinstance(value, list):
|
||||
items: Final[list[object]] = value # pyright: ignore[reportUnknownVariableType] # a request's untyped list # mutable-ok: read, then copied
|
||||
return list(items)
|
||||
return []
|
||||
|
|
@ -79,7 +79,12 @@ def detached(value: object) -> object:
|
|||
|
||||
|
||||
def build_router_config(arguments: Mapping[str, object]) -> RouterConfig:
|
||||
"""The constructor arguments, detached from the caller's objects: settings are static."""
|
||||
"""The constructor arguments, detached from the caller's objects: settings are static.
|
||||
|
||||
A bool `cooldown_time` is refused: `True` is not a number of seconds, and
|
||||
reading it as 1 s or as the default would both be a guess."""
|
||||
if isinstance(arguments.get("cooldown_time"), bool):
|
||||
raise TypeError("Router cooldown_time must be a number of seconds, not a bool")
|
||||
return RouterConfig(arguments=MappingProxyType({name: detached(value) for name, value in arguments.items()}))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
"""
|
||||
``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one.
|
||||
|
||||
``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. The Rust core
|
||||
(litellm-rust/crates/router-core) implements the same operations. The core is picked once, when
|
||||
the router is built, from its ``RouterConfig``.
|
||||
``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. ``RustCore``
|
||||
(``litellm/router_rust/core.py`` over litellm-rust/crates/router-core) is a separate implementation
|
||||
that reaches outside itself only through its host; it does not use ``PythonCore``. The core is
|
||||
picked once, when the router is built, from its ``RouterConfig``.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Protocol, runtime_checkable
|
||||
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.router_config import RouterConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -20,7 +22,6 @@ if TYPE_CHECKING:
|
|||
RouterSettings,
|
||||
RoutingObservers,
|
||||
)
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_python_host import PythonHost
|
||||
from litellm.types.router import CustomRoutingStrategyBase
|
||||
|
||||
|
|
@ -50,7 +51,26 @@ class RouterCore(Protocol):
|
|||
def retire(self) -> None: ...
|
||||
|
||||
|
||||
def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore:
|
||||
def make_core(config: RouterConfig, host: "PythonHost") -> RouterCore:
|
||||
"""
|
||||
The Rust core when the route catalog's router rule allows it and the Rust core serves this
|
||||
config, else the Python core. Decided once, for the router's life, from the config alone: the
|
||||
Rust core never builds, wraps or falls back to a Python core.
|
||||
"""
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.rust_bridge.catalog import Route, RouteContext, decision
|
||||
from litellm.rust_bridge.configuration import Decision
|
||||
|
||||
return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature
|
||||
chosen: Final = decision(RouteContext(Route.ROUTER))
|
||||
if chosen is not Decision.PYTHON:
|
||||
from litellm.router_rust.core import RustCore
|
||||
from litellm.rust_bridge._native import RustBridgeDeclined
|
||||
|
||||
try:
|
||||
return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported
|
||||
except RustBridgeDeclined as declined:
|
||||
# A config the Rust core cannot take keeps the router on Python while both cores exist.
|
||||
if chosen is Decision.RUST_REQUIRED:
|
||||
raise
|
||||
verbose_router_logger.info("Router served by the Python core: %s", declined)
|
||||
return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature
|
||||
|
|
|
|||
|
|
@ -5,8 +5,10 @@ DeprecatedRouterMembers: the Router members that outside code used before the en
|
|||
entry points (``registry()``, ``settings()``, ``health()``, ``select()``, ``apply()``,
|
||||
``observe()``, ``register_plugin()``, ``retire()``; see litellm/router_api.py). Every member here
|
||||
returns what it returned before, with the type it had before, and names its replacement in a
|
||||
``DeprecationWarning``. Reads go to the same Router implementation (``Router._core_<name>``) the
|
||||
entry point uses; writes go through ``apply`` / ``aapply`` so the registry version moves.
|
||||
``DeprecationWarning``. Reads go through that entry point, so they work on either router core
|
||||
(Python or Rust) and, on the Python core, return the same ``_core_<name>`` object as before.
|
||||
Settings writes go through ``apply`` / ``aapply`` so the registry version moves; any other old
|
||||
attribute write sets the core's ``_core_<name>``, which each core keeps where its reads find it.
|
||||
|
||||
Access from the router itself and from its own modules (``litellm.router*``) is not counted:
|
||||
those are the router's internals and move behind the core interface separately. Every other access
|
||||
|
|
@ -17,11 +19,10 @@ deleted (planned about a month after the core swap).
|
|||
# pyright: reportPrivateUsage=false, reportMissingTypeArgument=false, reportUnknownParameterType=false
|
||||
# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
|
||||
# pyright: reportUnknownLambdaType=false
|
||||
# This shim is part of the router (it reads Router's ``_core_*`` members, as Router's own methods do)
|
||||
# This shim is part of the router (it writes Router's ``_core_*`` members, as Router's own methods do)
|
||||
# and repeats each old member's signature exactly, untyped parts included (``dict``, ``list``), so
|
||||
# callers keep the types they had. Typing those parts is the core's job, not the shim's.
|
||||
|
||||
import contextvars
|
||||
import functools
|
||||
import inspect
|
||||
import sys
|
||||
|
|
@ -31,7 +32,7 @@ from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, S
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload
|
||||
|
||||
from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet
|
||||
from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet, RoutingContext
|
||||
from litellm.router_strategy.least_busy import picks_reserve
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -40,6 +41,7 @@ if TYPE_CHECKING:
|
|||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm.caching.caching import RedisCache
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
|
@ -131,6 +133,9 @@ class _DeprecatedAttribute(Generic[_T]):
|
|||
_note(self.name, self.replacement)
|
||||
router: Final = _as_router(obj)
|
||||
try:
|
||||
if obj.__dict__.get("_core") is None:
|
||||
# No core yet (a Router made without __init__): the value written through this name.
|
||||
return getattr(obj, "_core_" + self.name) # pyright: ignore[reportAny] # the pending write, as given
|
||||
return self.read(router)
|
||||
except AttributeError as e:
|
||||
raise _public_attribute_error(router, "_core_" + self.name, self.name, e) from None
|
||||
|
|
@ -178,243 +183,174 @@ def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coro
|
|||
return looked_up
|
||||
|
||||
|
||||
# (id(router), member name) pairs whose subclass override is running in this context.
|
||||
_OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar(
|
||||
"router_overrides_running", default=frozenset()
|
||||
)
|
||||
|
||||
|
||||
def _override_key(router: object, name: str) -> tuple[int, str]:
|
||||
"""The running-override key for ``router``: a Router and its core share one, so an old member
|
||||
entered on the Router is seen by the override routing that runs on the core."""
|
||||
core: Final = getattr(router, "__dict__", {}).get("_core")
|
||||
return (id(router if core is None else core), name)
|
||||
|
||||
|
||||
def _enter_override(name: str, override: Callable[..., object]) -> Callable[..., object]:
|
||||
def _async_picker(router: "Router") -> Callable[..., Awaitable[object]]:
|
||||
"""
|
||||
A direct call of a subclass override (``router.<name>()``) marks it running, so its
|
||||
``super().<name>()`` reaches the parent once instead of re-entering the override through
|
||||
``_core_<name>``.
|
||||
The router's async pick (the old ``async_get_available_deployment`` attribute). A core that
|
||||
keeps its pick as a member (the Python core's, or one assigned through this attribute) returns
|
||||
that member, so a patched pick reads back as patched; any other core picks through ``select``.
|
||||
"""
|
||||
if inspect.iscoroutinefunction(override):
|
||||
async_override: Final[Callable[..., Awaitable[object]]] = override
|
||||
member: Final[object] = getattr(router, "_core_async_get_available_deployment", None)
|
||||
if callable(member):
|
||||
return member # pyright: ignore[reportReturnType] # the core's own async pick
|
||||
|
||||
@functools.wraps(override)
|
||||
async def entered_async(
|
||||
router: object,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)})
|
||||
try:
|
||||
return await async_override(router, *args, **kwargs)
|
||||
finally:
|
||||
_OVERRIDES_RUNNING.reset(token)
|
||||
|
||||
return entered_async
|
||||
|
||||
@functools.wraps(override)
|
||||
def entered(
|
||||
router: object,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
async def async_get_available_deployment(
|
||||
model: str,
|
||||
request_kwargs: dict, # mutable-ok: the old parameter type, unchanged
|
||||
messages: list[dict[str, str]] | None = None, # mutable-ok: the old parameter type, unchanged
|
||||
input: str | list | None = None, # mutable-ok: the old parameter type, unchanged
|
||||
specific_deployment: bool | None = False,
|
||||
) -> object:
|
||||
token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)})
|
||||
try:
|
||||
return override(router, *args, **kwargs)
|
||||
finally:
|
||||
_OVERRIDES_RUNNING.reset(token)
|
||||
return await router.select(
|
||||
model, request_kwargs, messages=messages, input=input, specific_deployment=specific_deployment
|
||||
)
|
||||
|
||||
return entered
|
||||
return async_get_available_deployment
|
||||
|
||||
|
||||
def _private_alias(name: str, member: Callable[..., object]) -> Callable[..., object]:
|
||||
def _unreserved_picker(router: "Router") -> Callable[..., Awaitable[object]]:
|
||||
"""``_async_picker``, picking without holding a least-busy reservation (a lookup makes no
|
||||
call). A pick assigned through the attribute reads back as the object assigned."""
|
||||
picker: Final = _async_picker(router)
|
||||
if inspect.ismethod(picker) or picker.__name__ == "async_get_available_deployment":
|
||||
return _without_reserving(picker)
|
||||
return picker
|
||||
|
||||
|
||||
def _slack_alerting_logger(router: "Router") -> "SlackAlerting":
|
||||
"""``observe().slack_alerting_logger``; a router without one has no such attribute, as before."""
|
||||
logger: Final = router.observe().slack_alerting_logger
|
||||
if logger is None:
|
||||
raise AttributeError(
|
||||
f"'{type(router).__name__}' object has no attribute 'slack_alerting_logger'",
|
||||
name="slack_alerting_logger",
|
||||
obj=router,
|
||||
)
|
||||
return logger
|
||||
|
||||
|
||||
def _run_ready(step: Coroutine[object, object, _T]) -> _T:
|
||||
"""
|
||||
A private alias of an old member (``_get_all_deployments``). It was bound to Router's own
|
||||
method, so it never reached a subclass override of ``name``, and still doesn't.
|
||||
The result of a coroutine that finishes without suspending: ``select(kind='pass_through')``
|
||||
picks synchronously on every core, so the old synchronous member can return its result.
|
||||
"""
|
||||
return _enter_override(name, member)
|
||||
|
||||
|
||||
def _route_override(cls: type, name: str, override: Callable[..., object]) -> Callable[..., object]:
|
||||
"""
|
||||
Make the router's own calls to ``_core_<name>`` reach a subclass override of ``<name>``.
|
||||
|
||||
When the override calls ``super().<name>()``, the call comes back to ``_core_<name>``; it is
|
||||
then sent on to the parent's implementation instead of back into the override.
|
||||
"""
|
||||
core: Final = "_core_" + name
|
||||
|
||||
def parent(router: object) -> Callable[..., object]:
|
||||
return getattr(super(cls, router), core)
|
||||
|
||||
if inspect.iscoroutinefunction(override):
|
||||
async_override: Final[Callable[..., Awaitable[object]]] = override
|
||||
|
||||
@functools.wraps(override)
|
||||
async def routed_async(
|
||||
router: object,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
key: Final = _override_key(router, name)
|
||||
running: Final = _OVERRIDES_RUNNING.get()
|
||||
if key in running:
|
||||
parent_async: Final[Callable[..., Awaitable[object]]] = getattr(super(cls, router), core)
|
||||
return await parent_async(*args, **kwargs)
|
||||
token: Final = _OVERRIDES_RUNNING.set(running | {key})
|
||||
try:
|
||||
return await async_override(router, *args, **kwargs)
|
||||
finally:
|
||||
_OVERRIDES_RUNNING.reset(token)
|
||||
|
||||
return routed_async
|
||||
|
||||
@functools.wraps(override)
|
||||
def routed(
|
||||
router: object,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: passes the override's arguments through unchanged
|
||||
) -> object:
|
||||
key: Final = _override_key(router, name)
|
||||
running: Final = _OVERRIDES_RUNNING.get()
|
||||
if key in running:
|
||||
return parent(router)(*args, **kwargs)
|
||||
token: Final = _OVERRIDES_RUNNING.set(running | {key})
|
||||
try:
|
||||
return override(router, *args, **kwargs)
|
||||
finally:
|
||||
_OVERRIDES_RUNNING.reset(token)
|
||||
|
||||
return routed
|
||||
try:
|
||||
step.send(None)
|
||||
except StopIteration as done:
|
||||
return done.value # pyright: ignore[reportAny] # the coroutine's own result
|
||||
step.close()
|
||||
raise RuntimeError("select(kind='pass_through') suspended; it must pick without awaiting")
|
||||
|
||||
|
||||
class DeprecatedRouterMembers:
|
||||
"""Old Router members, kept for about a month. Each names its replacement."""
|
||||
|
||||
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: class keywords, passed on to super()
|
||||
# A Router subclass that overrides an old member (say async_pre_routing_hook) keeps
|
||||
# having its override called by the router, which now calls _core_<name> internally.
|
||||
super().__init_subclass__(**kwargs)
|
||||
for name, value in list(vars(cls).items()):
|
||||
if (
|
||||
name.startswith("_core_")
|
||||
or name not in vars(DeprecatedRouterMembers)
|
||||
or not hasattr(cls, "_core_" + name)
|
||||
):
|
||||
continue
|
||||
if isinstance(value, staticmethod):
|
||||
# Its super() call reaches the old member here, which calls Router's own implementation.
|
||||
setattr(cls, "_core_" + name, value)
|
||||
elif inspect.isfunction(value):
|
||||
setattr(cls, "_core_" + name, _route_override(cls, name, value))
|
||||
setattr(cls, name, _enter_override(name, value))
|
||||
|
||||
# -- registry: attributes --------------------------------------------------------------------
|
||||
model_list = _attr("registry().all()", lambda r: r._core_model_list)
|
||||
model_names = _attr("registry().names(kind='group_set')", lambda r: r._core_model_names)
|
||||
deployment_names = _attr("registry().names(kind='deployment_names')", lambda r: r._core_deployment_names)
|
||||
team_public_model_names = _attr("registry().names(kind='team_public')", lambda r: r._core_team_public_model_names)
|
||||
model_list = _attr("registry().all()", lambda r: r.registry().all())
|
||||
model_names = _attr("registry().names(kind='group_set')", lambda r: r.registry().names(kind="group_set"))
|
||||
deployment_names = _attr(
|
||||
"registry().names(kind='deployment_names')", lambda r: r.registry().names(kind="deployment_names")
|
||||
)
|
||||
team_public_model_names = _attr(
|
||||
"registry().names(kind='team_public')", lambda r: r.registry().names(kind="team_public")
|
||||
)
|
||||
model_name_to_deployment_indices = _attr(
|
||||
"registry().index('model_name')", lambda r: r._core_model_name_to_deployment_indices
|
||||
"registry().index('model_name')", lambda r: r.registry().index("model_name")
|
||||
)
|
||||
team_model_to_deployment_indices = _attr(
|
||||
"registry().index('team_model')", lambda r: r._core_team_model_to_deployment_indices
|
||||
"registry().index('team_model')", lambda r: r.registry().index("team_model")
|
||||
)
|
||||
cached_deployment_model_info = _attr(
|
||||
"registry().deployment_model_info(id, model_name, cached=True)",
|
||||
lambda r: r._core_cached_deployment_model_info,
|
||||
lambda r: r.registry().deployment_model_info_cache(),
|
||||
)
|
||||
model_group_alias = _attr("registry().aliases()", lambda r: r._core_model_group_alias)
|
||||
pattern_router = _attr("registry().patterns()", lambda r: r._core_pattern_router)
|
||||
team_pattern_routers = _attr("registry().patterns(team=True)", lambda r: r._core_team_pattern_routers)
|
||||
routing_groups = _attr("registry().routing_groups()", lambda r: r._core_routing_groups)
|
||||
model_group_alias = _attr("registry().aliases()", lambda r: r.registry().aliases())
|
||||
pattern_router = _attr("registry().patterns()", lambda r: r.registry().patterns())
|
||||
team_pattern_routers = _attr("registry().patterns(team=True)", lambda r: r.registry().patterns(team=True))
|
||||
routing_groups = _attr("registry().routing_groups()", lambda r: r.registry().routing_groups())
|
||||
_routing_groups = routing_groups
|
||||
auto_routers = _attr("registry().plugins('auto')", lambda r: r._core_auto_routers)
|
||||
complexity_routers = _attr("registry().plugins('complexity')", lambda r: r._core_complexity_routers)
|
||||
quality_routers = _attr("registry().plugins('quality')", lambda r: r._core_quality_routers)
|
||||
adaptive_routers = _attr("registry().plugins('adaptive')", lambda r: r._core_adaptive_routers)
|
||||
routing_plugins = _attr("registry().plugins('routing')", lambda r: r._core_routing_plugins)
|
||||
auto_routers = _attr("registry().plugins('auto')", lambda r: r.registry().plugins("auto"))
|
||||
complexity_routers = _attr("registry().plugins('complexity')", lambda r: r.registry().plugins("complexity"))
|
||||
quality_routers = _attr("registry().plugins('quality')", lambda r: r.registry().plugins("quality"))
|
||||
adaptive_routers = _attr("registry().plugins('adaptive')", lambda r: r.registry().plugins("adaptive"))
|
||||
routing_plugins = _attr("registry().plugins('routing')", lambda r: r.registry().plugins("routing"))
|
||||
prompt_caching_pin_cache = _attr(
|
||||
"registry().plugins('prompt_caching_pins')", lambda r: r._core_prompt_caching_pin_cache
|
||||
"registry().plugins('prompt_caching_pins')", lambda r: r.registry().plugins("prompt_caching_pins")
|
||||
)
|
||||
|
||||
# -- health and selection: attributes --------------------------------------------------------
|
||||
cache = _attr("health().store", lambda r: r._core_cache)
|
||||
cooldown_cache = _attr("health().cooldowns", lambda r: r._core_cooldown_cache)
|
||||
health_state_cache = _attr("health().health_checks", lambda r: r._core_health_state_cache)
|
||||
deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map)
|
||||
async_get_available_deployment = _attr(
|
||||
"select(model, request_kwargs)",
|
||||
lambda r: _without_reserving(r._core_async_get_available_deployment),
|
||||
)
|
||||
cache = _attr("health().store", lambda r: r.health().store)
|
||||
cooldown_cache = _attr("health().cooldowns", lambda r: r.health().cooldowns)
|
||||
health_state_cache = _attr("health().health_checks", lambda r: r.health().health_checks)
|
||||
deployment_latency_map = _attr("health().latency", lambda r: r.health().latency)
|
||||
async_get_available_deployment = _attr("select(model, request_kwargs)", lambda r: _unreserved_picker(r))
|
||||
|
||||
# -- observers ------------------------------------------------------------------------------
|
||||
slack_alerting_logger = _attr("observe().slack_alerting_logger", lambda r: r._core_slack_alerting_logger)
|
||||
slack_alerting_logger = _attr("observe().slack_alerting_logger", _slack_alerting_logger)
|
||||
|
||||
# -- settings -------------------------------------------------------------------------------
|
||||
routing_strategy = _setting("routing_strategy", lambda r: r._core_routing_strategy)
|
||||
context_window_fallbacks = _setting("context_window_fallbacks", lambda r: r._core_context_window_fallbacks)
|
||||
content_policy_fallbacks = _setting("content_policy_fallbacks", lambda r: r._core_content_policy_fallbacks)
|
||||
default_deployment = _setting("default_deployment", lambda r: r._core_default_deployment)
|
||||
default_litellm_params = _setting("default_litellm_params", lambda r: r._core_default_litellm_params)
|
||||
routing_strategy = _setting("routing_strategy", lambda r: r.settings().routing_strategy)
|
||||
context_window_fallbacks = _setting("context_window_fallbacks", lambda r: r.settings().context_window_fallbacks)
|
||||
content_policy_fallbacks = _setting("content_policy_fallbacks", lambda r: r.settings().content_policy_fallbacks)
|
||||
default_deployment = _setting("default_deployment", lambda r: r.settings().default_deployment)
|
||||
default_litellm_params = _setting("default_litellm_params", lambda r: r.settings().default_litellm_params)
|
||||
default_max_parallel_requests = _setting(
|
||||
"default_max_parallel_requests", lambda r: r._core_default_max_parallel_requests
|
||||
"default_max_parallel_requests", lambda r: r.settings().default_max_parallel_requests
|
||||
)
|
||||
cooldown_time = _setting("cooldown_time", lambda r: r.settings().cooldown_time)
|
||||
allowed_fails = _setting("allowed_fails", lambda r: r.settings().allowed_fails)
|
||||
allowed_fails_policy = _setting("allowed_fails_policy", lambda r: r.settings().allowed_fails_policy)
|
||||
disable_cooldowns = _setting("disable_cooldowns", lambda r: r.settings().disable_cooldowns)
|
||||
enable_pre_call_checks = _setting("enable_pre_call_checks", lambda r: r.settings().enable_pre_call_checks)
|
||||
enable_health_check_routing = _setting(
|
||||
"enable_health_check_routing", lambda r: r.settings().enable_health_check_routing
|
||||
)
|
||||
cooldown_time = _setting("cooldown_time", lambda r: r._core_cooldown_time)
|
||||
allowed_fails = _setting("allowed_fails", lambda r: r._core_allowed_fails)
|
||||
allowed_fails_policy = _setting("allowed_fails_policy", lambda r: r._core_allowed_fails_policy)
|
||||
disable_cooldowns = _setting("disable_cooldowns", lambda r: r._core_disable_cooldowns)
|
||||
enable_pre_call_checks = _setting("enable_pre_call_checks", lambda r: r._core_enable_pre_call_checks)
|
||||
enable_health_check_routing = _setting("enable_health_check_routing", lambda r: r._core_enable_health_check_routing)
|
||||
health_check_ignore_transient_errors = _setting(
|
||||
"health_check_ignore_transient_errors", lambda r: r._core_health_check_ignore_transient_errors
|
||||
"health_check_ignore_transient_errors", lambda r: r.settings().health_check_ignore_transient_errors
|
||||
)
|
||||
background_health_check_model_groups = _setting(
|
||||
"background_health_check_model_groups", lambda r: r._core_background_health_check_model_groups
|
||||
"background_health_check_model_groups", lambda r: r.settings().background_health_check_model_groups
|
||||
)
|
||||
enable_tag_filtering = _setting("enable_tag_filtering", lambda r: r._core_enable_tag_filtering)
|
||||
tag_filtering_match_any = _setting("tag_filtering_match_any", lambda r: r._core_tag_filtering_match_any)
|
||||
tag_routing_prefix = _setting("tag_routing_prefix", lambda r: r._core_tag_routing_prefix)
|
||||
router_general_settings = _setting("router_general_settings", lambda r: r._core_router_general_settings)
|
||||
provider_budget_config = _setting("provider_budget_config", lambda r: r._core_provider_budget_config)
|
||||
enable_tag_filtering = _setting("enable_tag_filtering", lambda r: r.settings().enable_tag_filtering)
|
||||
tag_filtering_match_any = _setting("tag_filtering_match_any", lambda r: r.settings().tag_filtering_match_any)
|
||||
tag_routing_prefix = _setting("tag_routing_prefix", lambda r: r.settings().tag_routing_prefix)
|
||||
router_general_settings = _setting("router_general_settings", lambda r: r.settings().router_general_settings)
|
||||
provider_budget_config = _setting("provider_budget_config", lambda r: r.settings().provider_budget_config)
|
||||
provider_default_deployment_ids = _setting(
|
||||
"provider_default_deployment_ids", lambda r: r._core_provider_default_deployment_ids
|
||||
"provider_default_deployment_ids", lambda r: r.settings().provider_default_deployment_ids
|
||||
)
|
||||
fallback_access_check = _setting("fallback_access_check", lambda r: r._core_fallback_access_check)
|
||||
fallback_budget_check = _setting("fallback_budget_check", lambda r: r._core_fallback_budget_check)
|
||||
fallback_access_check = _setting("fallback_access_check", lambda r: r.settings().fallback_access_check)
|
||||
fallback_budget_check = _setting("fallback_budget_check", lambda r: r.settings().fallback_budget_check)
|
||||
# These three are declared with the attribute type Router gave them rather than the descriptor's:
|
||||
# protocols in router_utils and the proxy match Router on them, and protocol matching does not
|
||||
# see through a descriptor. Reads and writes still go through _setting at runtime.
|
||||
fallbacks: list | list[dict[str, list[str] | list | None]] | None = _setting( # pyright: ignore[reportAssignmentType] # see above # mutable-ok: the old type
|
||||
"fallbacks", lambda r: r._core_fallbacks
|
||||
"fallbacks", lambda r: r.settings().fallbacks
|
||||
)
|
||||
kubernetes_pod_discovery: "KubernetesPodDiscovery" = _setting( # pyright: ignore[reportAssignmentType] # see above
|
||||
"kubernetes_pod_discovery", lambda r: r._core_kubernetes_pod_discovery
|
||||
"kubernetes_pod_discovery", lambda r: r.settings().kubernetes_pod_discovery
|
||||
)
|
||||
search_tools: "list[SearchToolTypedDict]" = _setting( # pyright: ignore[reportAssignmentType] # see above # mutable-ok: the old type
|
||||
"search_tools", lambda r: r._core_search_tools
|
||||
"search_tools", lambda r: r.settings().search_tools
|
||||
)
|
||||
guardrail_list = _setting("guardrail_list", lambda r: r._core_guardrail_list)
|
||||
cache_responses = _setting("cache_responses", lambda r: r._core_cache_responses)
|
||||
guardrail_list = _setting("guardrail_list", lambda r: r.settings().guardrail_list)
|
||||
cache_responses = _setting("cache_responses", lambda r: r.settings().cache_responses)
|
||||
|
||||
# -- registry: one deployment ---------------------------------------------------------------
|
||||
def get_deployment(self, model_id: str) -> "Deployment | None":
|
||||
_note("get_deployment", "registry().deployment(id)")
|
||||
return _as_router(self)._core_get_deployment(model_id=model_id)
|
||||
return _as_router(self).registry().deployment(model_id)
|
||||
|
||||
def has_model_id(self, candidate_id: str) -> bool:
|
||||
_note("has_model_id", "registry().has_deployment(id)")
|
||||
return _as_router(self)._core_has_model_id(candidate_id)
|
||||
return _as_router(self).registry().has_deployment(candidate_id)
|
||||
|
||||
def get_model_info(self, id: str) -> dict | None: # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_info", "registry().deployment_dict(id)")
|
||||
return _as_router(self)._core_get_model_info(id=id)
|
||||
return _as_router(self).registry().deployment_dict(id)
|
||||
|
||||
def get_deployment_model_info(self, model_id: str, model_name: str) -> "ModelMapInfo | None":
|
||||
_note("get_deployment_model_info", "registry().deployment_model_info(id, model_name)")
|
||||
return _as_router(self)._core_get_deployment_model_info(model_id=model_id, model_name=model_name)
|
||||
return _as_router(self).registry().deployment_model_info(model_id, model_name)
|
||||
|
||||
@overload
|
||||
def get_router_model_info(
|
||||
|
|
@ -434,62 +370,56 @@ class DeprecatedRouterMembers:
|
|||
id: str | None = None,
|
||||
) -> "ModelMapInfo":
|
||||
_note("get_router_model_info", "registry().router_model_info(...)")
|
||||
router: Final = _as_router(self)
|
||||
# The core's overloads split (deployment, id) by which one is None; this forwards either form unchanged.
|
||||
return router._core_get_router_model_info( # pyright: ignore[reportCallIssue] # see above
|
||||
deployment=deployment, # pyright: ignore[reportArgumentType] # see above
|
||||
received_model_name=received_model_name,
|
||||
id=id, # pyright: ignore[reportArgumentType] # see above
|
||||
)
|
||||
return _as_router(self).registry().router_model_info(deployment, received_model_name, id)
|
||||
|
||||
def get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]:
|
||||
_note("get_discovered_model_info", "registry().discovered_model_info(id)")
|
||||
return _as_router(self)._core_get_discovered_model_info(model_id)
|
||||
return _as_router(self).registry().discovered_model_info(model_id)
|
||||
|
||||
# -- registry: many deployments -------------------------------------------------------------
|
||||
def get_model_list(
|
||||
self, model_name: str | None = None, team_id: str | None = None
|
||||
) -> "list[DeploymentTypedDict] | None": # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_list", "registry().deployments(model_name, team_id=...)")
|
||||
return _as_router(self)._core_get_model_list(model_name=model_name, team_id=team_id)
|
||||
return _as_router(self).registry().deployments(model_name, team_id=team_id)
|
||||
|
||||
def get_all_deployments(
|
||||
self, model_name: str, model_alias: str | None = None, team_id: str | None = None
|
||||
) -> "list[DeploymentTypedDict]": # mutable-ok: the old return type, unchanged
|
||||
_note("get_all_deployments", "registry().deployments(model_name, source='all')")
|
||||
return _as_router(self)._core_get_all_deployments(
|
||||
model_name=model_name, model_alias=model_alias, team_id=team_id
|
||||
return (
|
||||
_as_router(self).registry().deployments(model_name, source="all", model_alias=model_alias, team_id=team_id)
|
||||
)
|
||||
|
||||
_get_all_deployments = _private_alias("get_all_deployments", get_all_deployments)
|
||||
_get_all_deployments = get_all_deployments
|
||||
|
||||
def config_deployments(self) -> Iterator[Mapping[str, object]]:
|
||||
_note("config_deployments", "registry().deployments(source='config')")
|
||||
return _as_router(self)._core_config_deployments()
|
||||
return _as_router(self).registry().deployments(source="config")
|
||||
|
||||
def deployments_named(self, model_name: str) -> "tuple[DeploymentTypedDict, ...] | None":
|
||||
_note("deployments_named", "registry().deployments(model_name, source='named')")
|
||||
return _as_router(self)._core_deployments_named(model_name)
|
||||
return _as_router(self).registry().deployments(model_name, source="named")
|
||||
|
||||
def get_model_list_from_model_alias(
|
||||
self, model_name: str | None = None
|
||||
) -> "list[DeploymentTypedDict]": # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_list_from_model_alias", "registry().deployments(model_name, source='alias')")
|
||||
return _as_router(self)._core_get_model_list_from_model_alias(model_name=model_name)
|
||||
return _as_router(self).registry().deployments(model_name, source="alias")
|
||||
|
||||
def get_model_list_from_routing_groups(self, model_name: str | None = None) -> "Sequence[DeploymentTypedDict]":
|
||||
_note("get_model_list_from_routing_groups", "registry().deployments(model_name, source='routing_groups')")
|
||||
return _as_router(self)._core_get_model_list_from_routing_groups(model_name=model_name)
|
||||
return _as_router(self).registry().deployments(model_name, source="routing_groups")
|
||||
|
||||
def get_model_list_of_routed_group(
|
||||
self, model_group: str
|
||||
) -> "list[DeploymentTypedDict]": # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_list_of_routed_group", "registry().deployments(model_group, source='routed_group')")
|
||||
return _as_router(self)._core_get_model_list_of_routed_group(model_group)
|
||||
return _as_router(self).registry().deployments(model_group, source="routed_group")
|
||||
|
||||
def wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
|
||||
_note("wildcard_route_deployments", "registry().deployments(model, source='wildcard')")
|
||||
return _as_router(self)._core_wildcard_route_deployments(model, team_id)
|
||||
return _as_router(self).registry().deployments(model, source="wildcard", team_id=team_id)
|
||||
|
||||
def get_routable_deployments(
|
||||
self,
|
||||
|
|
@ -500,17 +430,21 @@ class DeprecatedRouterMembers:
|
|||
_note(
|
||||
"get_routable_deployments", "registry().deployments(model, source='routable', request=..., call_type=...)"
|
||||
)
|
||||
return _as_router(self)._core_get_routable_deployments(model, request_kwargs, call_type)
|
||||
return (
|
||||
_as_router(self)
|
||||
.registry()
|
||||
.deployments(model, source="routable", request=request_kwargs, call_type=call_type)
|
||||
)
|
||||
|
||||
def get_deployment_by_model_group_name(self, model_group_name: str) -> "Deployment | None":
|
||||
_note("get_deployment_by_model_group_name", "registry().deployments(name, source='first_in_group')")
|
||||
return _as_router(self)._core_get_deployment_by_model_group_name(model_group_name=model_group_name)
|
||||
return _as_router(self).registry().deployments(model_group_name, source="first_in_group")
|
||||
|
||||
def deployments_for_request(
|
||||
self, model: str, request_kwargs: Mapping[str, object]
|
||||
) -> "Sequence[DeploymentTypedDict]":
|
||||
_note("deployments_for_request", "registry().deployments(model, source='for_request', request=...)")
|
||||
return _as_router(self)._core_deployments_for_request(model, request_kwargs)
|
||||
return _as_router(self).registry().deployments(model, source="for_request", request=request_kwargs)
|
||||
|
||||
def should_include_deployment(
|
||||
self,
|
||||
|
|
@ -519,135 +453,137 @@ class DeprecatedRouterMembers:
|
|||
team_id: str | None = None,
|
||||
) -> bool:
|
||||
_note("should_include_deployment", "registry().includes(deployment, model_name=..., team_id=...)")
|
||||
return _as_router(self)._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id)
|
||||
return _as_router(self).registry().includes(model, model_name=model_name, team_id=team_id)
|
||||
|
||||
@staticmethod
|
||||
def deployment_usable_by_team(model: "Mapping[str, object] | Deployment", team_id: str | None) -> bool:
|
||||
_note("deployment_usable_by_team", "registry().includes(deployment, team_id=...)")
|
||||
from litellm.router_python_core import PythonCore
|
||||
from litellm.router_api import deployment_usable_by_team as usable_by_team
|
||||
|
||||
return PythonCore._core_deployment_usable_by_team(model, team_id)
|
||||
return usable_by_team(model, team_id)
|
||||
|
||||
_deployment_usable_by_team = deployment_usable_by_team
|
||||
|
||||
# -- registry: names -------------------------------------------------------------------------
|
||||
def get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_names", "registry().names(team_id=...)")
|
||||
return _as_router(self)._core_get_model_names(team_id=team_id)
|
||||
return _as_router(self).registry().names(team_id=team_id)
|
||||
|
||||
def get_model_ids(
|
||||
self, model_name: str | None = None, exclude_team_models: bool = False
|
||||
) -> list[str]: # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_ids", "registry().names(model_name, kind='ids')")
|
||||
return _as_router(self)._core_get_model_ids(model_name=model_name, exclude_team_models=exclude_team_models)
|
||||
return _as_router(self).registry().names(model_name, kind="ids", exclude_team_models=exclude_team_models)
|
||||
|
||||
def resolved_litellm_models(
|
||||
self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False
|
||||
) -> tuple[str, ...]:
|
||||
_note("resolved_litellm_models", "registry().names(model_name, kind='litellm_models' or 'litellm_base_models')")
|
||||
return _as_router(self)._core_resolved_litellm_models(
|
||||
model_name, team_id=team_id, use_base_model=use_base_model
|
||||
return (
|
||||
_as_router(self)
|
||||
.registry()
|
||||
.names(model_name, kind="litellm_base_models" if use_base_model else "litellm_models", team_id=team_id)
|
||||
)
|
||||
|
||||
def wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]:
|
||||
_note("wildcard_routed_models", "registry().names(model, kind='wildcard')")
|
||||
return _as_router(self)._core_wildcard_routed_models(model, team_id)
|
||||
return _as_router(self).registry().names(model, kind="wildcard", team_id=team_id)
|
||||
|
||||
def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]:
|
||||
_note("get_candidate_model_ids_for_route", "registry().names(model, kind='route_candidates')")
|
||||
return _as_router(self)._core_get_candidate_model_ids_for_route(model, team_id=team_id)
|
||||
return _as_router(self).registry().names(model, kind="route_candidates", team_id=team_id)
|
||||
|
||||
# -- registry: resolving names ---------------------------------------------------------------
|
||||
def routable_model_group(self, model_name: str) -> str:
|
||||
_note("routable_model_group", "registry().resolve(name)")
|
||||
return _as_router(self)._core_routable_model_group(model_name)
|
||||
return _as_router(self).registry().resolve(model_name)
|
||||
|
||||
def get_model_from_alias(self, model: str) -> str | None:
|
||||
_note("get_model_from_alias", "registry().resolve(name, kind='alias')")
|
||||
return _as_router(self)._core_get_model_from_alias(model=model)
|
||||
return _as_router(self).registry().resolve(model, kind="alias")
|
||||
|
||||
_get_model_from_alias = _private_alias("get_model_from_alias", get_model_from_alias)
|
||||
_get_model_from_alias = get_model_from_alias
|
||||
|
||||
def resolve_served_model_name(self, model_name: str, served_model: str) -> str:
|
||||
_note("resolve_served_model_name", "registry().resolve(name, kind='served', served_model=...)")
|
||||
return _as_router(self)._core_resolve_served_model_name(model_name, served_model)
|
||||
return _as_router(self).registry().resolve(model_name, kind="served", served_model=served_model)
|
||||
|
||||
def resolve_model_name_from_model_id(
|
||||
self, model_id: str | None, custom_llm_provider: str | None = None
|
||||
) -> str | None:
|
||||
_note("resolve_model_name_from_model_id", "registry().resolve(id, kind='from_id')")
|
||||
return _as_router(self)._core_resolve_model_name_from_model_id(
|
||||
model_id, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
return _as_router(self).registry().resolve(model_id, kind="from_id", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
def map_team_model(
|
||||
self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None
|
||||
) -> str | None:
|
||||
_note("map_team_model", "registry().resolve(name, kind='team_model', team_id=...)")
|
||||
return _as_router(self)._core_map_team_model(team_model_name, team_id, custom_llm_provider=custom_llm_provider)
|
||||
return (
|
||||
_as_router(self)
|
||||
.registry()
|
||||
.resolve(team_model_name, kind="team_model", team_id=team_id, custom_llm_provider=custom_llm_provider)
|
||||
)
|
||||
|
||||
def is_recognized_model(self, model: str) -> bool:
|
||||
_note("is_recognized_model", "registry().resolve(name, kind='recognized')")
|
||||
return _as_router(self)._core_is_recognized_model(model)
|
||||
return _as_router(self).registry().resolve(model, kind="recognized")
|
||||
|
||||
def get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None:
|
||||
_note("get_routable_upstream_model", "registry().resolve(name, kind='upstream')")
|
||||
return _as_router(self)._core_get_routable_upstream_model(model_name, team_id=team_id)
|
||||
return _as_router(self).registry().resolve(model_name, kind="upstream", team_id=team_id)
|
||||
|
||||
def get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]:
|
||||
_note("get_search_tool_support", "registry().resolve(name, kind='search_tool_support')")
|
||||
return _as_router(self)._core_get_search_tool_support(model_name, team_id=team_id)
|
||||
return _as_router(self).registry().resolve(model_name, kind="search_tool_support", team_id=team_id)
|
||||
|
||||
def get_model_group(self, id: str) -> list | None: # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_group", "registry().resolve(id, kind='group_of_id')")
|
||||
return _as_router(self)._core_get_model_group(id=id)
|
||||
return _as_router(self).registry().resolve(id, kind="group_of_id")
|
||||
|
||||
def is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool:
|
||||
_note("is_model_access_group_for_wildcard_route", "registry().resolve(group, kind='wildcard_access_group')")
|
||||
return _as_router(self)._core_is_model_access_group_for_wildcard_route(model_access_group=model_access_group)
|
||||
return _as_router(self).registry().resolve(model_access_group, kind="wildcard_access_group")
|
||||
|
||||
_is_model_access_group_for_wildcard_route = _private_alias(
|
||||
"is_model_access_group_for_wildcard_route", is_model_access_group_for_wildcard_route
|
||||
)
|
||||
_is_model_access_group_for_wildcard_route = is_model_access_group_for_wildcard_route
|
||||
|
||||
def team_model_has_alternatives(self, deployment_id: str) -> bool:
|
||||
_note("team_model_has_alternatives", "registry().resolve(id, kind='team_model_alternatives')")
|
||||
return _as_router(self)._core_team_model_has_alternatives(deployment_id)
|
||||
return _as_router(self).registry().resolve(deployment_id, kind="team_model_alternatives")
|
||||
|
||||
def routing_group_has_alternatives(self, model_group: str | None) -> bool:
|
||||
_note("routing_group_has_alternatives", "registry().resolve(group, kind='routing_group_alternatives')")
|
||||
return _as_router(self)._core_routing_group_has_alternatives(model_group)
|
||||
return _as_router(self).registry().resolve(model_group, kind="routing_group_alternatives")
|
||||
|
||||
# -- registry: groups, access groups, credentials, plugins -----------------------------------
|
||||
def get_model_group_info(
|
||||
self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None
|
||||
) -> "ModelGroupInfo | None":
|
||||
_note("get_model_group_info", "registry().group(name)")
|
||||
return _as_router(self)._core_get_model_group_info(model_group, include_hidden=include_hidden, team_id=team_id)
|
||||
return _as_router(self).registry().group(model_group, include_hidden=include_hidden, team_id=team_id)
|
||||
|
||||
def cached_model_group_info(self, model_group: str) -> "ModelGroupInfo | None":
|
||||
_note("cached_model_group_info", "registry().group(name, field='cached_info')")
|
||||
return _as_router(self)._core_cached_model_group_info(model_group)
|
||||
return _as_router(self).registry().group(model_group, field="cached_info")
|
||||
|
||||
def get_configured_display_name(self, model_name: str) -> str | None:
|
||||
_note("get_configured_display_name", "registry().group(name, field='display_name')")
|
||||
return _as_router(self)._core_get_configured_display_name(model_name)
|
||||
return _as_router(self).registry().group(model_name, field="display_name")
|
||||
|
||||
def get_configured_mode(self, model_name: str) -> str | None:
|
||||
_note("get_configured_mode", "registry().group(name, field='mode')")
|
||||
return _as_router(self)._core_get_configured_mode(model_name)
|
||||
return _as_router(self).registry().group(model_name, field="mode")
|
||||
|
||||
def get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]:
|
||||
_note("get_configured_service_tiers", "registry().group(name, field='service_tiers')")
|
||||
return _as_router(self)._core_get_configured_service_tiers(model_name, team_id=team_id)
|
||||
return _as_router(self).registry().group(model_name, field="service_tiers", team_id=team_id)
|
||||
|
||||
def get_configured_token_limits(self, model_name: str) -> tuple[int | None, int | None]:
|
||||
_note("get_configured_token_limits", "registry().group(name, field='token_limits')")
|
||||
return _as_router(self)._core_get_configured_token_limits(model_name)
|
||||
return _as_router(self).registry().group(model_name, field="token_limits")
|
||||
|
||||
def get_model_listing_info(self, model_name: str) -> "DeploymentModelListingInfo | None":
|
||||
_note("get_model_listing_info", "registry().group(name, field='listing_info')")
|
||||
return _as_router(self)._core_get_model_listing_info(model_name)
|
||||
return _as_router(self).registry().group(model_name, field="listing_info")
|
||||
|
||||
def get_model_access_groups(
|
||||
self,
|
||||
|
|
@ -656,31 +592,31 @@ class DeprecatedRouterMembers:
|
|||
team_id: str | None = None,
|
||||
) -> dict[str, list[str]]: # mutable-ok: the old return type, unchanged
|
||||
_note("get_model_access_groups", "registry().access_groups(...)")
|
||||
return _as_router(self)._core_get_model_access_groups(
|
||||
model_name=model_name, model_access_group=model_access_group, team_id=team_id
|
||||
return (
|
||||
_as_router(self)
|
||||
.registry()
|
||||
.access_groups(model_name=model_name, model_access_group=model_access_group, team_id=team_id)
|
||||
)
|
||||
|
||||
def get_deployment_credentials(self, model_id: str) -> dict | None: # mutable-ok: the old return type, unchanged
|
||||
_note("get_deployment_credentials", "registry().credentials(id)")
|
||||
return _as_router(self)._core_get_deployment_credentials(model_id=model_id)
|
||||
return _as_router(self).registry().credentials(model_id)
|
||||
|
||||
def get_deployment_credentials_with_provider(
|
||||
self, model_id: str, team_id: str | None = None
|
||||
) -> "dict[str, Any] | None": # pyright: ignore[reportExplicitAny] # the old return type, unchanged # mutable-ok: same
|
||||
_note("get_deployment_credentials_with_provider", "registry().credentials(id, form='with_provider')")
|
||||
return _as_router(self)._core_get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id)
|
||||
return _as_router(self).registry().credentials(model_id, form="with_provider", team_id=team_id)
|
||||
|
||||
def get_credential_deployment(self, model_id: str, team_id: str | None = None) -> "Deployment | None":
|
||||
_note("get_credential_deployment", "registry().credentials(id, form='deployment')")
|
||||
return _as_router(self)._core_get_credential_deployment(model_id=model_id, team_id=team_id)
|
||||
return _as_router(self).registry().credentials(model_id, form="deployment", team_id=team_id)
|
||||
|
||||
def get_router_deployment_budget_limiter(self) -> "RouterBudgetLimiting | None":
|
||||
_note("get_router_deployment_budget_limiter", "registry().plugins('budget_limiter')")
|
||||
return _as_router(self)._core_get_router_deployment_budget_limiter()
|
||||
return _as_router(self).registry().plugins("budget_limiter")
|
||||
|
||||
_get_router_deployment_budget_limiter = _private_alias(
|
||||
"get_router_deployment_budget_limiter", get_router_deployment_budget_limiter
|
||||
)
|
||||
_get_router_deployment_budget_limiter = get_router_deployment_budget_limiter
|
||||
|
||||
@staticmethod
|
||||
def generate_model_id(model_group: str, litellm_params: dict) -> str: # mutable-ok: the old parameter type
|
||||
|
|
@ -692,15 +628,15 @@ class DeprecatedRouterMembers:
|
|||
# -- settings: methods -----------------------------------------------------------------------
|
||||
def get_settings(self) -> dict[str, object]: # mutable-ok: Router._core_get_settings's type, a fresh dict
|
||||
_note("get_settings", "settings().as_dict()")
|
||||
return _as_router(self)._core_get_settings()
|
||||
return dict(_as_router(self).settings().as_dict())
|
||||
|
||||
def get_allowed_fails_from_policy(self, exception: Exception) -> int | None:
|
||||
_note("get_allowed_fails_from_policy", "settings().allowed_fails_for(exception)")
|
||||
return _as_router(self)._core_get_allowed_fails_from_policy(exception=exception)
|
||||
return _as_router(self).settings().allowed_fails_for(exception)
|
||||
|
||||
def get_available_guardrail(self, guardrail_name: str) -> "GuardrailTypedDict":
|
||||
_note("get_available_guardrail", "settings().guardrail(name)")
|
||||
return _as_router(self)._core_get_available_guardrail(guardrail_name)
|
||||
return _as_router(self).settings().guardrail(guardrail_name)
|
||||
|
||||
# -- health ----------------------------------------------------------------------------------
|
||||
@staticmethod
|
||||
|
|
@ -723,31 +659,31 @@ class DeprecatedRouterMembers:
|
|||
|
||||
def get_fully_blocked_model_names(self) -> set[str]: # mutable-ok: the old return type, unchanged
|
||||
_note("get_fully_blocked_model_names", "health().fully_blocked_model_names()")
|
||||
return _as_router(self)._core_get_fully_blocked_model_names()
|
||||
return _as_router(self).health().fully_blocked_model_names()
|
||||
|
||||
async def async_get_fully_unhealthy_model_names(self) -> set[str]: # mutable-ok: the old return type, unchanged
|
||||
_note("async_get_fully_unhealthy_model_names", "health().fully_unhealthy_model_names()")
|
||||
return await _as_router(self)._core_async_get_fully_unhealthy_model_names()
|
||||
return await _as_router(self).health().fully_unhealthy_model_names()
|
||||
|
||||
async def get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]:
|
||||
_note("get_model_group_usage", "health().usage(group)")
|
||||
return await _as_router(self)._core_get_model_group_usage(model_group=model_group)
|
||||
return await _as_router(self).health().usage(model_group)
|
||||
|
||||
async def get_remaining_model_group_usage(
|
||||
self, model_group: str
|
||||
) -> dict[str, int]: # mutable-ok: the old return type, unchanged
|
||||
_note("get_remaining_model_group_usage", "health().usage(group, kind='remaining')")
|
||||
return await _as_router(self)._core_get_remaining_model_group_usage(model_group=model_group)
|
||||
return await _as_router(self).health().usage(model_group, kind="remaining")
|
||||
|
||||
def team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None":
|
||||
_note("team_ptu_ceiling", "health().team_ptu_ceiling(team_id, group)")
|
||||
return _as_router(self)._core_team_ptu_ceiling(team_id, model_group)
|
||||
return _as_router(self).health().team_ptu_ceiling(team_id, model_group)
|
||||
|
||||
def split_request_priority(
|
||||
self, model: str, kwargs: Mapping[str, object]
|
||||
) -> tuple[int | None, Mapping[str, object]]:
|
||||
_note("split_request_priority", "health().split_priority(model, kwargs)")
|
||||
return _as_router(self)._core_split_request_priority(model, kwargs)
|
||||
return _as_router(self).health().split_priority(model, kwargs)
|
||||
|
||||
# -- selection -------------------------------------------------------------------------------
|
||||
async def async_get_healthy_deployments(
|
||||
|
|
@ -761,9 +697,10 @@ class DeprecatedRouterMembers:
|
|||
health_check_probe: bool = False,
|
||||
) -> list[dict] | dict: # mutable-ok: the old return type, unchanged
|
||||
_note("async_get_healthy_deployments", "select(model, request_kwargs, kind='healthy')")
|
||||
return await _as_router(self)._core_async_get_healthy_deployments(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
return await _as_router(self).select(
|
||||
model,
|
||||
request_kwargs,
|
||||
kind="healthy",
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
|
|
@ -771,7 +708,7 @@ class DeprecatedRouterMembers:
|
|||
health_check_probe=health_check_probe,
|
||||
)
|
||||
|
||||
def get_available_deployment_for_pass_through( # noqa: ANN201 # the old member had no return annotation; its inferred type is kept by inferring it from the core
|
||||
def get_available_deployment_for_pass_through( # noqa: ANN201 # the old member had no return annotation; its inferred type is kept by inferring it from select
|
||||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, str]] | None = None, # mutable-ok: the old parameter type, unchanged
|
||||
|
|
@ -780,12 +717,15 @@ class DeprecatedRouterMembers:
|
|||
request_kwargs: dict | None = None, # mutable-ok: the old parameter type, unchanged
|
||||
):
|
||||
_note("get_available_deployment_for_pass_through", "select(model, kind='pass_through')")
|
||||
return _as_router(self)._core_get_available_deployment_for_pass_through(
|
||||
model=model,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
request_kwargs=request_kwargs,
|
||||
return _run_ready(
|
||||
_as_router(self).select(
|
||||
model,
|
||||
request_kwargs,
|
||||
kind="pass_through",
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
)
|
||||
|
||||
# -- writes ----------------------------------------------------------------------------------
|
||||
|
|
@ -805,7 +745,7 @@ class DeprecatedRouterMembers:
|
|||
_note("update_redis_cache", "apply(ChangeSet(redis_cache=cache))")
|
||||
_as_router(self).apply(ChangeSet(redis_cache=cache))
|
||||
|
||||
_update_redis_cache = _private_alias("update_redis_cache", update_redis_cache)
|
||||
_update_redis_cache = update_redis_cache
|
||||
|
||||
async def arefresh_model_info(self, *, client: "AsyncHTTPHandler | None" = None) -> None:
|
||||
_note("arefresh_model_info", "aapply(ChangeSet(refresh_model_info=True))")
|
||||
|
|
@ -822,8 +762,15 @@ class DeprecatedRouterMembers:
|
|||
|
||||
def bind_to_request_logging(self, logging_obj: "LiteLLMLogging") -> None:
|
||||
_note("bind_to_request_logging", "observe(logging_obj)")
|
||||
# Straight to the core, as before: observe() skips binding when logging_obj is None, the old member did not.
|
||||
_as_router(self)._core_bind_to_request_logging(logging_obj)
|
||||
router: Final = _as_router(self)
|
||||
if logging_obj is None: # pyright: ignore[reportUnnecessaryComparison] # callers pass None despite the type
|
||||
# observe() skips binding when logging_obj is None; the old member bound it, and failed on None.
|
||||
raise AttributeError(
|
||||
"'NoneType' object has no attribute 'add_dynamic_function_callbacks'",
|
||||
name="add_dynamic_function_callbacks",
|
||||
obj=None,
|
||||
)
|
||||
router.observe(logging_obj)
|
||||
|
||||
def discard(self) -> None:
|
||||
_note("discard", "retire()")
|
||||
|
|
@ -836,7 +783,7 @@ class DeprecatedRouterMembers:
|
|||
request_kwargs: dict[str, object] | None = None, # mutable-ok: the old parameter type, unchanged
|
||||
) -> None:
|
||||
_note("arm_routing_read_prefetch", "RoutingContext.arm_read_prefetch")
|
||||
_as_router(self)._core_arm_routing_read_prefetch(model, request_kwargs)
|
||||
RoutingContext(_as_router(self)).arm_read_prefetch(model, request_kwargs)
|
||||
|
||||
async def async_pre_routing_hook(
|
||||
self,
|
||||
|
|
@ -847,7 +794,7 @@ class DeprecatedRouterMembers:
|
|||
specific_deployment: bool | None = False,
|
||||
) -> "PreRoutingHookResponse | None":
|
||||
_note("async_pre_routing_hook", "RoutingContext.pre_routing_hook")
|
||||
return await _as_router(self)._core_async_pre_routing_hook(
|
||||
return await RoutingContext(_as_router(self)).pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
|
|
@ -857,7 +804,7 @@ class DeprecatedRouterMembers:
|
|||
|
||||
def log_retry(self, kwargs: dict, e: Exception) -> dict: # mutable-ok: the old signature, unchanged
|
||||
_note("log_retry", "RoutingContext.log_retry")
|
||||
return _as_router(self)._core_log_retry(kwargs=kwargs, e=e)
|
||||
return RoutingContext(_as_router(self)).log_retry(kwargs=kwargs, e=e)
|
||||
|
||||
|
||||
# Router-family functions and classes that took the Router (``litellm_router_instance=``,
|
||||
|
|
|
|||
|
|
@ -127,6 +127,7 @@ from litellm.litellm_core_utils.clientside_credential_handler import (
|
|||
)
|
||||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_str_mapping,
|
||||
coerce_token_limit,
|
||||
get_litellm_metadata_from_kwargs,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
|
|
@ -322,6 +323,7 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
mark_deployment_selection_failure,
|
||||
set_cooldown_deployments,
|
||||
)
|
||||
from litellm.router_utils.deployment_numbers import refuse_unservable_numbers
|
||||
from litellm.router_utils.embedding_batch_size_check import (
|
||||
async_log_rejection_before_provider_call,
|
||||
clear_rejected_before_provider_call,
|
||||
|
|
@ -366,6 +368,11 @@ from litellm.router_utils.fallback_event_handlers import (
|
|||
routed_deployment_id,
|
||||
run_async_fallback,
|
||||
)
|
||||
from litellm.router_utils.global_callback_holders import (
|
||||
live_check_owners,
|
||||
routers_holding_global_callbacks,
|
||||
routers_holding_global_callbacks_lock,
|
||||
)
|
||||
from litellm.router_utils.handle_error import (
|
||||
async_raise_no_deployment_exception,
|
||||
send_llm_exception_alert,
|
||||
|
|
@ -1233,8 +1240,9 @@ class RoutingArgs(enum.Enum):
|
|||
# as the per-request one built from a caller-supplied user_config, drops out on its
|
||||
# own rather than leaving entries behind that nothing can withdraw.
|
||||
_live_routers: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet()
|
||||
_routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet()
|
||||
_routers_holding_global_callbacks_lock: Final = threading.Lock()
|
||||
# Shared with the other router cores of the process (litellm/router_utils/global_callback_holders.py).
|
||||
_routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = routers_holding_global_callbacks # pyright: ignore[reportAssignmentType] # every member serves the holder members
|
||||
_routers_holding_global_callbacks_lock: Final = routers_holding_global_callbacks_lock
|
||||
|
||||
|
||||
def _replay_live_router_model_cost() -> None:
|
||||
|
|
@ -1878,6 +1886,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
# cost-map entries from the list it is serving at that moment, whether
|
||||
# those deployments came from model_list or were added later from the DB.
|
||||
_live_routers.add(self)
|
||||
live_check_owners.add(self)
|
||||
|
||||
if allowed_fails is not None:
|
||||
self._core_allowed_fails = allowed_fails
|
||||
|
|
@ -2103,6 +2112,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
# Stop contributing to cost-map rebuilds straight away rather than waiting
|
||||
# for this router to be collected.
|
||||
_live_routers.discard(self)
|
||||
live_check_owners.discard(self)
|
||||
self.retire_routing_state()
|
||||
self._release_global_callbacks()
|
||||
|
||||
|
|
@ -3414,7 +3424,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls]
|
||||
if any(
|
||||
router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or []))
|
||||
for router in tuple(_live_routers)
|
||||
for router in (*_live_routers, *(owner for owner in live_check_owners if not isinstance(owner, PythonCore)))
|
||||
):
|
||||
return
|
||||
for cb in tuple(self._host.callback_list("callbacks")):
|
||||
|
|
@ -7150,6 +7160,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
or self._stream_can_fall_back(model, initial_kwargs)
|
||||
)
|
||||
buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline
|
||||
# An error this loop raises off a frame, not the provider's.
|
||||
raised_here: BaseException | None = None # rebind-ok: set when the loop raises off a frame
|
||||
try:
|
||||
async for chunk in source_iterator:
|
||||
if _anthropic_stream_forwards_ping_live(chunk, has_generated_content):
|
||||
|
|
@ -7180,25 +7192,27 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
and self._refusal_fallback_available(model, initial_kwargs)
|
||||
):
|
||||
refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details)
|
||||
raise MidStreamFallbackError(
|
||||
raised_here = MidStreamFallbackError(
|
||||
message=refusal_error.message,
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
original_exception=refusal_error,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
raise raised_here
|
||||
if not has_generated_content and error_event is None:
|
||||
buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk)
|
||||
continue
|
||||
if recoverable_frame_error is not None:
|
||||
assert error_event is not None
|
||||
raise MidStreamFallbackError(
|
||||
raised_here = MidStreamFallbackError(
|
||||
message=error_event[1],
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
original_exception=recoverable_frame_error,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
raise raised_here
|
||||
for buffered_chunk in buffered_lifecycle_chunks:
|
||||
yield buffered_chunk
|
||||
buffered_lifecycle_chunks = ()
|
||||
|
|
@ -7206,6 +7220,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
for buffered_chunk in buffered_lifecycle_chunks:
|
||||
yield buffered_chunk
|
||||
except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate
|
||||
if stream_error is not raised_here:
|
||||
await self._areport_raised_messages_stream_failure(stream_error, initial_kwargs, source_iterator)
|
||||
async for item in self._aanthropic_messages_recover_stream_error(
|
||||
stream_error,
|
||||
has_generated_content,
|
||||
|
|
@ -7225,6 +7241,52 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator)
|
||||
return wrapper
|
||||
|
||||
async def _areport_raised_messages_stream_failure(
|
||||
self, stream_error: Exception, attempt_kwargs: Mapping[str, object], source: object
|
||||
) -> None:
|
||||
"""A /v1/messages stream whose provider iterator raised, counted against the deployment
|
||||
that served it and cooled down by the rules, the same as a non-stream call's failure, so a
|
||||
retry picks another deployment. A stream that logs itself (litellm's own stream logging
|
||||
reaches the router's failure callback), or a failure the request's logging object already
|
||||
logged, is not counted twice."""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
if isinstance(getattr(source, "logging_obj", None), Logging):
|
||||
return
|
||||
logging_obj: Final = attempt_kwargs.get("litellm_logging_obj")
|
||||
if isinstance(logging_obj, Logging) and logging_obj.model_call_details.get("has_logged_async_failure") is True:
|
||||
return
|
||||
bucket: Final = as_str_mapping(attempt_kwargs.get(get_metadata_variable_name_from_kwargs(attempt_kwargs))) or {}
|
||||
model_info: Final = as_str_mapping(bucket.get("model_info")) or {}
|
||||
deployment_id: Final = model_info.get("id")
|
||||
if deployment_id is None:
|
||||
return
|
||||
deployment: Final = self.get_deployment(model_id=str(deployment_id))
|
||||
deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {}
|
||||
provider_error: Final = (
|
||||
stream_error.original_exception
|
||||
if isinstance(stream_error, MidStreamFallbackError) and stream_error.original_exception is not None
|
||||
else stream_error
|
||||
)
|
||||
now: Final = datetime.now() # noqa: DTZ005 # naive, as the logging pipeline's times
|
||||
logged: Final = {
|
||||
"call_type": "anthropic_messages",
|
||||
"exception": provider_error,
|
||||
"litellm_params": {
|
||||
**{
|
||||
name: deployment_params[name] for name in ("api_base", "cooldown_time") if name in deployment_params
|
||||
},
|
||||
"metadata": dict(bucket),
|
||||
"model_info": dict(model_info),
|
||||
},
|
||||
}
|
||||
try:
|
||||
_ = self.deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # an untyped callback
|
||||
await self.async_deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # see above
|
||||
except Exception as error: # noqa: BLE001 # a recording never fails the stream
|
||||
verbose_router_logger.warning("a stream's failure was not recorded: %s", error)
|
||||
|
||||
async def _aanthropic_messages_recover_stream_error(
|
||||
self,
|
||||
stream_error: Exception,
|
||||
|
|
@ -11235,6 +11297,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
model_info=_model_info,
|
||||
)
|
||||
)
|
||||
refuse_unservable_numbers(deployment.to_json(exclude_none=True, warnings=False)) # pyright: ignore[reportUnknownMemberType] # pydantic's to_json
|
||||
PythonCore._register_deployment_pricing(deployment)
|
||||
|
||||
## Check if LLM Deployment is allowed for this deployment
|
||||
|
|
@ -11738,6 +11801,16 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
return model_info_is_active_for_environment(model_info=deployment.model_info)
|
||||
|
||||
def set_model_list(self, model_list: list):
|
||||
"""Replace every deployment. All or nothing: when an entry raises, the deployments served
|
||||
before the call are put back and the error is raised."""
|
||||
previous: Final = copy.deepcopy(list(self.__dict__.get("_core_model_list") or ()))
|
||||
try:
|
||||
self._replace_model_list(model_list) # pyright: ignore[reportUnknownMemberType] # the legacy untyped model list
|
||||
except Exception:
|
||||
self._replace_model_list(previous) # pyright: ignore[reportUnknownMemberType] # see above
|
||||
raise
|
||||
|
||||
def _replace_model_list(self, model_list: list) -> None: # pyright: ignore[reportUnknownParameterType, reportMissingTypeArgument] # the legacy untyped model list
|
||||
original_model_list: Final = copy.deepcopy(model_list)
|
||||
self._discovered_model_info_cache.flush_cache()
|
||||
with self._deployment_index_lock:
|
||||
|
|
@ -12031,6 +12104,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
# add to model list
|
||||
_deployment: Final = served.to_json(exclude_none=True)
|
||||
refuse_unservable_numbers(_deployment)
|
||||
# initialize client
|
||||
self._add_deployment(deployment=served)
|
||||
|
||||
|
|
@ -12174,6 +12248,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
_deployment_model_id: Final = deployment.model_info.id or ""
|
||||
_deployment_on_router: Final[Deployment | None] = self._core_get_deployment(model_id=_deployment_model_id)
|
||||
served: Final = without_model_capability_flags(deployment)
|
||||
# `_create_deployment` lists the model name before it can still fail; a dropped deployment
|
||||
# takes its name back off the list.
|
||||
names_before: Final = len(self._core_deployment_names)
|
||||
try:
|
||||
if _deployment_on_router is not None:
|
||||
# deployment with this model_id exists on the router
|
||||
|
|
@ -12223,6 +12300,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
deployment.model_info.id,
|
||||
e,
|
||||
)
|
||||
del self._core_deployment_names[names_before:]
|
||||
self._restore_deployment_after_failed_upsert(
|
||||
previous_deployment=_deployment_on_router, model_id=_deployment_model_id
|
||||
)
|
||||
|
|
@ -14713,9 +14791,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
"""
|
||||
if team_id is not None and self._team_deployments_named(team_id, model_group):
|
||||
return self._core_get_all_deployments(model_name=model_group, team_id=team_id)
|
||||
own: Final = cast( # cast-ok: the private alias returns what get_all_deployments returns
|
||||
list[DeploymentTypedDict], self._get_all_deployments(model_name=model_group)
|
||||
)
|
||||
own: Final = self._get_all_deployments(model_name=model_group)
|
||||
named: Final = [*own, *self._core_get_model_list_from_routing_groups(model_name=model_group)]
|
||||
return named or self._get_wildcard_deployments(model_name=model_group)
|
||||
|
||||
|
|
@ -15923,10 +15999,15 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
# IF TEAM ID SPECIFIED ON MODEL, AND REQUEST CONTAINS USER_API_KEY_TEAM_ID, FILTER OUT MODELS THAT ARE NOT IN THE TEAM
|
||||
## THIS PREVENTS WRITING FILES OF OTHER TEAMS TO MODELS THAT ARE TEAM-ONLY MODELS
|
||||
by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict
|
||||
healthy_deployments = filter_team_based_models(
|
||||
healthy_deployments=healthy_deployments,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
if by_id is not None and not healthy_deployments:
|
||||
# Another team's deployment asked for by id: refused under its group's name.
|
||||
group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above
|
||||
model = group_name # rebind-ok: the error names the group
|
||||
|
||||
if verbose_router_logger.isEnabledFor(logging.DEBUG):
|
||||
verbose_router_logger.debug("healthy_deployments after team filter: %s", healthy_deployments)
|
||||
|
|
@ -17236,6 +17317,16 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
model,
|
||||
)
|
||||
)
|
||||
# As the async pick: a team-owned deployment serves only its own team's callers, also when
|
||||
# asked for by id (then refused under its group's name).
|
||||
by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict
|
||||
healthy_deployments = filter_team_based_models( # rebind-ok: the team filter narrows # pyright: ignore[reportUnknownVariableType] # legacy selector type
|
||||
healthy_deployments=healthy_deployments, # pyright: ignore[reportArgumentType] # legacy selector type
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
if by_id is not None and not healthy_deployments:
|
||||
group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above
|
||||
model = group_name # rebind-ok: the error names the group
|
||||
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
|
||||
routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs)
|
||||
pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,10 @@ come from ``dynamic()`` (``litellm/router_config.py``). A Rust core gets the sam
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
import importlib
|
||||
import time
|
||||
import types
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
|
|
@ -19,12 +22,13 @@ from typing import (
|
|||
from typing_extensions import overload
|
||||
|
||||
import litellm
|
||||
from litellm.router_config import RouterDynamicConfig, live_dynamic_config
|
||||
from litellm.router_config import RouterConfig, RouterDynamicConfig, live_dynamic_config
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.caching.caching import Cache
|
||||
from litellm.caching.caching import Cache, RedisCache
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
|
||||
|
|
@ -97,7 +101,7 @@ class PythonHost:
|
|||
def invoke(self, op: str) -> Endpoint: ...
|
||||
def invoke(self, op: str) -> Endpoint:
|
||||
"""The litellm function named ``op`` (``acompletion``, ``get_model_info``, ...), looked up at call time."""
|
||||
endpoint: Final[object] = getattr(litellm, op, None)
|
||||
endpoint: Final[object] = getattr(litellm, op, None) if "." not in op else _dotted(op)
|
||||
if endpoint is None:
|
||||
raise AttributeError(f"litellm has no endpoint {op!r}")
|
||||
if not callable(endpoint):
|
||||
|
|
@ -107,6 +111,10 @@ class PythonHost:
|
|||
def sleep(self, seconds: float) -> Awaitable[None]:
|
||||
return asyncio.sleep(seconds)
|
||||
|
||||
def sleep_sync(self, seconds: float) -> None:
|
||||
"""The backoff of a sync endpoint's retry."""
|
||||
time.sleep(seconds)
|
||||
|
||||
def dynamic(self) -> RouterDynamicConfig:
|
||||
"""The process settings, read live on each access (``current_dynamic_config()`` takes a snapshot)."""
|
||||
return cast(RouterDynamicConfig, live_dynamic_config) # cast-ok: the live view reads the same fields
|
||||
|
|
@ -116,7 +124,7 @@ class PythonHost:
|
|||
return litellm.logging_callback_manager
|
||||
|
||||
def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself
|
||||
value: Final = getattr(litellm, name)
|
||||
value: Final[object] = getattr(litellm, name) # pyright: ignore[reportAny] # an untyped litellm global
|
||||
return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global
|
||||
|
||||
def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given
|
||||
|
|
@ -131,5 +139,78 @@ class PythonHost:
|
|||
def suppress_debug_info(self) -> None:
|
||||
litellm.suppress_debug_info = True
|
||||
|
||||
def routing_cache(self, config: RouterConfig) -> "DualCache":
|
||||
"""The store a router keeps its routing state in (cooldowns, fail counts, usage): in memory,
|
||||
plus Redis when the router's arguments name one. The host builds and owns it; a Rust core
|
||||
reads and writes it through cache steps."""
|
||||
from litellm.caching.caching import DualCache, InMemoryCache
|
||||
|
||||
return DualCache(redis_cache=redis_cache_for(config), in_memory_cache=InMemoryCache())
|
||||
|
||||
|
||||
def redis_cache_for(config: RouterConfig) -> "RedisCache | None":
|
||||
"""The Redis cache the router's ``redis_*`` and ``cache_kwargs`` arguments describe, or None."""
|
||||
import json
|
||||
|
||||
from litellm.caching.caching import RedisCache, RedisClusterCache
|
||||
from litellm.utils import get_secret
|
||||
|
||||
url: Final = config.get("redis_url")
|
||||
host: Final = config.get("redis_host")
|
||||
port: Final = config.get("redis_port")
|
||||
if url is None and (host is None or port is None):
|
||||
return None
|
||||
settings: Final[dict[str, object]] = {} # mutable-ok: built here, then passed as keywords
|
||||
if url is not None:
|
||||
settings["url"] = url
|
||||
if host is not None:
|
||||
settings["host"] = host
|
||||
if port is not None:
|
||||
settings["port"] = str(port)
|
||||
if config.get("redis_password") is not None:
|
||||
settings["password"] = config.get("redis_password")
|
||||
if config.get("redis_db") is not None:
|
||||
settings["db"] = str(config.get("redis_db"))
|
||||
extra: Final = config.get("cache_kwargs")
|
||||
if isinstance(extra, dict):
|
||||
settings.update({str(key): value for key, value in extra.items()}) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the cache_kwargs argument is untyped
|
||||
nodes: object = settings.get("startup_nodes") # rebind-ok: the env cluster nodes when the arguments name none
|
||||
if not nodes:
|
||||
env_nodes: Final[object] = get_secret("REDIS_CLUSTER_NODES") # pyright: ignore[reportUnknownVariableType] # get_secret is untyped
|
||||
if isinstance(env_nodes, str):
|
||||
parsed: Final[object] = json.loads(env_nodes) # pyright: ignore[reportAny] # plain JSON
|
||||
nodes = parsed # rebind-ok: see above
|
||||
if nodes:
|
||||
return RedisClusterCache(**{**settings, "startup_nodes": nodes}) # pyright: ignore[reportArgumentType] # RedisCache's keyword arguments
|
||||
return RedisCache(**settings) # pyright: ignore[reportArgumentType] # RedisCache's keyword arguments
|
||||
|
||||
|
||||
python_host: Final = PythonHost()
|
||||
|
||||
|
||||
def _dotted(op: str) -> object:
|
||||
"""``op`` named by its path under ``litellm`` (``vector_stores.main.asearch``), read at call time.
|
||||
|
||||
Each segment is a module attribute, a submodule, or a key of a mapping such as a module's
|
||||
table of generated endpoints.
|
||||
"""
|
||||
node: object = litellm # rebind-ok: walks the path one segment at a time
|
||||
path: str = "litellm" # rebind-ok: the module path walked so far
|
||||
for segment in op.split("."):
|
||||
path = f"{path}.{segment}"
|
||||
if isinstance(node, Mapping):
|
||||
node = cast(Mapping[str, object], node).get(segment) # cast-ok: endpoint tables are keyed by name
|
||||
elif isinstance(node, types.ModuleType) and (submodule := _submodule(path)) is not None:
|
||||
node = submodule
|
||||
else:
|
||||
node = cast(object, getattr(node, segment, None)) # cast-ok: any attribute
|
||||
if node is None:
|
||||
return None
|
||||
return node
|
||||
|
||||
|
||||
def _submodule(path: str) -> types.ModuleType | None:
|
||||
try:
|
||||
return importlib.import_module(path)
|
||||
except ModuleNotFoundError:
|
||||
return None
|
||||
|
|
|
|||
13
litellm/router_rust/__init__.py
Normal file
13
litellm/router_rust/__init__.py
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
"""
|
||||
The Rust router core's Python side: ``RustCore`` (``core``) and the host-side modules router-core
|
||||
(litellm-rust/crates/router-core) reaches through its steps.
|
||||
|
||||
- ``steps``: answers to host steps that need live Python objects (strategies, pre-routing,
|
||||
pre-call facts, plugins, callbacks, caller access)
|
||||
- ``invoke``: the host side of ``Host::invoke``, one attempt against one deployment (request
|
||||
shaping, endpoints, streams, slots, retry and fallback metadata, prompt management, websockets)
|
||||
- ``surface``: the Python ``Router`` API kept over router-core's state (registry views, model
|
||||
info, settings, health, retries, fallbacks, alerting); a Rust host needs none of it
|
||||
|
||||
Nothing is imported here, so importing one module does not import the core.
|
||||
"""
|
||||
4943
litellm/router_rust/core.py
Normal file
4943
litellm/router_rust/core.py
Normal file
File diff suppressed because it is too large
Load diff
13
litellm/router_rust/invoke/README.md
Normal file
13
litellm/router_rust/invoke/README.md
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
# invoke
|
||||
|
||||
The host side of `Host::invoke`: running one attempt against the one deployment router-core picked. Nothing here selects, retries, falls back or records cooldowns.
|
||||
|
||||
- `shaping`: the kwargs an attempt sends to its deployment
|
||||
- `endpoints`: what each public router method does before, around and after the routed call
|
||||
- `streaming`: stream wrappers, so a stream that fails before its first token fails the attempt
|
||||
- `slots`: a deployment's `max_parallel_requests`
|
||||
- `walk`: the retry and fallback metadata stamped on each attempt
|
||||
- `prompt_management`: prompt-management groups
|
||||
- `websocket`: admission for each turn of a Responses websocket
|
||||
|
||||
A Rust host replaces this with native endpoint calls.
|
||||
1
litellm/router_rust/invoke/__init__.py
Normal file
1
litellm/router_rust/invoke/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""The host side of ``Host::invoke``: one router attempt against one deployment."""
|
||||
2454
litellm/router_rust/invoke/endpoints.py
Normal file
2454
litellm/router_rust/invoke/endpoints.py
Normal file
File diff suppressed because it is too large
Load diff
129
litellm/router_rust/invoke/prompt_management.py
Normal file
129
litellm/router_rust/invoke/prompt_management.py
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
"""
|
||||
Prompt-management groups on the Rust router core (``_is_prompt_management_model`` and
|
||||
``_prompt_management_factory``): a group of exactly one deployment whose model names a prompt
|
||||
management integration (``langfuse/...``) or ``litellm_agent/...``. The request is not routed with
|
||||
retries and fallbacks: the deployment is picked, its prompt rendered by the request's logging
|
||||
object, and the rendered model called directly, or routed when it names a group.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import without
|
||||
from litellm.router_rust.invoke.endpoints import PLANS, EndpointCore
|
||||
from litellm.router_rust.invoke.shaping import as_dict, update_kwargs_with_deployment
|
||||
|
||||
_PROMPT_MANAGEMENT_PARAMS: Final = frozenset(
|
||||
{"bitbucket_config", "dotprompt_config", "prompt_id", "prompt_variables", "prompt_label", "prompt_version"}
|
||||
)
|
||||
|
||||
|
||||
def is_prompt_management_model(core: EndpointCore, model: str) -> bool:
|
||||
"""`_is_prompt_management_model`: one deployment, whose model's provider prefix is a known
|
||||
custom-logger-compatible callback."""
|
||||
model_list: Final = core.group_deployments(model_name=model)
|
||||
if model_list is None or len(model_list) != 1:
|
||||
return False
|
||||
litellm_model: Final = as_dict(model_list[0].get("litellm_params")).get("model")
|
||||
if not isinstance(litellm_model, str) or "/" not in litellm_model:
|
||||
return False
|
||||
known: Final = core.host.dynamic().known_custom_logger_compatible_callbacks
|
||||
return litellm_model.split("/")[0] in known
|
||||
|
||||
|
||||
def _logging_and_kwargs(
|
||||
kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
) -> tuple[object, dict[str, object]]: # mutable-ok: same
|
||||
"""The request's logging object, else the one `function_setup` builds (with the kwargs it returns)."""
|
||||
from litellm.utils import Rules, function_setup, get_utc_datetime
|
||||
|
||||
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
||||
if logging_obj is not None:
|
||||
return logging_obj, kwargs
|
||||
return function_setup( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers
|
||||
**{ # pyright: ignore[reportArgumentType] # the request's keyword payload
|
||||
"original_function": "acompletion",
|
||||
"rules_obj": Rules(),
|
||||
"start_time": get_utc_datetime(),
|
||||
**kwargs,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def prompt_management_acompletion(
|
||||
core: EndpointCore,
|
||||
model: str,
|
||||
messages: object,
|
||||
request: Mapping[str, object],
|
||||
) -> object:
|
||||
"""`_prompt_management_factory`, on the request as `acompletion` prepared it."""
|
||||
from litellm.utils import get_non_default_completion_params
|
||||
|
||||
# `acompletion` hands the routing function along; function_setup consumes it.
|
||||
logging_obj, kwargs = _logging_and_kwargs({**request, "original_function": core.router.acompletion})
|
||||
caller_timeout: Final = kwargs.get("timeout") or kwargs.get("request_timeout")
|
||||
caller_stream_timeout: Final = kwargs.get("stream_timeout")
|
||||
specific_deployment: Final = kwargs.pop("specific_deployment", None)
|
||||
picked: Final = await core.select(
|
||||
model,
|
||||
kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs
|
||||
kind="pick",
|
||||
specific_deployment=bool(specific_deployment),
|
||||
messages=messages,
|
||||
)
|
||||
deployment: Final = as_dict(picked)
|
||||
update_kwargs_with_deployment(
|
||||
core.shaping,
|
||||
deployment,
|
||||
kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs
|
||||
cache=core.routing_cache,
|
||||
splits_stream_timeout=True,
|
||||
)
|
||||
params: Final = as_dict(deployment.get("litellm_params"))
|
||||
data: Final = dict(params)
|
||||
litellm_model: Final = data.get("model")
|
||||
is_litellm_agent_model: Final = isinstance(litellm_model, str) and litellm_model.startswith("litellm_agent/")
|
||||
prompt_id: Final = kwargs.get("prompt_id") or params.get("prompt_id")
|
||||
prompt_variables: Final = kwargs.get("prompt_variables") or params.get("prompt_variables")
|
||||
prompt_label: Final = kwargs.get("prompt_label") or params.get("prompt_label")
|
||||
if not is_litellm_agent_model and (prompt_id is None or not isinstance(prompt_id, str)):
|
||||
raise ValueError(f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}")
|
||||
if prompt_variables is not None and not isinstance(prompt_variables, dict):
|
||||
raise ValueError(
|
||||
f"Prompt variables is set but not a dictionary. Got={prompt_variables}, type={type(prompt_variables)}"
|
||||
)
|
||||
rendered_model, rendered_messages, optional_params = logging_obj.get_chat_completion_prompt( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportAttributeAccessIssue, reportOptionalMemberAccess] # the request's Logging object
|
||||
model=litellm_model, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it
|
||||
non_default_params=get_non_default_completion_params(kwargs=kwargs), # pyright: ignore[reportArgumentType] # the request kwargs
|
||||
prompt_id=prompt_id, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it
|
||||
prompt_variables=prompt_variables,
|
||||
prompt_label=prompt_label, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it
|
||||
request_kwargs=kwargs,
|
||||
injected_for_every_deployment=True,
|
||||
)
|
||||
filtered_data: Final = without(data, *_PROMPT_MANAGEMENT_PARAMS)
|
||||
final: Final[dict[str, object]] = {**filtered_data, **kwargs, **optional_params} # pyright: ignore[reportUnknownVariableType] # mutable-ok: the call's kwargs
|
||||
final["model"] = rendered_model
|
||||
final["messages"] = rendered_messages
|
||||
final["litellm_logging_obj"] = logging_obj
|
||||
final["prompt_id"] = prompt_id
|
||||
final["prompt_variables"] = prompt_variables
|
||||
final["prompt_label"] = prompt_label
|
||||
target_list: Final = core.group_deployments(model_name=str(rendered_model)) # pyright: ignore[reportUnknownArgumentType] # the core's host registry
|
||||
if target_list is None or len(target_list) == 0: # a direct call to the rendered model
|
||||
_ = final.pop("original_function")
|
||||
return await core.host.invoke("acompletion")(**final) # pyright: ignore[reportUnknownMemberType] # the core's host
|
||||
timeouts: Final = {
|
||||
key: value
|
||||
for key, value in (
|
||||
("timeout", caller_timeout if caller_timeout is not None else filtered_data.get("timeout")),
|
||||
(
|
||||
"stream_timeout",
|
||||
caller_stream_timeout if caller_stream_timeout is not None else filtered_data.get("stream_timeout"),
|
||||
),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
routed: Final = without(final, "timeout", "stream_timeout")
|
||||
return await core.drive(PLANS["acompletion"], {**routed, **timeouts}, prepared_already=True)
|
||||
635
litellm/router_rust/invoke/shaping.py
Normal file
635
litellm/router_rust/invoke/shaping.py
Normal file
|
|
@ -0,0 +1,635 @@
|
|||
"""
|
||||
The keyword arguments one router attempt sends to one deployment, as the Python router builds them.
|
||||
|
||||
The Rust router core picks the deployment; the host shapes the request for it here
|
||||
(``_update_kwargs_with_deployment``, ``_update_kwargs_with_default_litellm_params``,
|
||||
``_merge_tools_from_deployment``, the timeout resolvers and the reasoning-effort override), with
|
||||
the same precedence and the same oddities. Nothing here reads a Python router.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
from collections.abc import Mapping, MutableMapping, Sequence
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm import exceptions as litellm_exceptions
|
||||
from litellm.constants import (
|
||||
BAGGAGE_MAPPED_TAGS_METADATA_KEY,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
ROUTING_REQUEST_TAGS_METADATA_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name
|
||||
from litellm.litellm_core_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
is_clientside_credential, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
)
|
||||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.litellm_core_utils.core_helpers import without
|
||||
from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, inferred_provider
|
||||
from litellm.litellm_core_utils.request_timeout_resolver import get_service_tier_timeout
|
||||
from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin
|
||||
from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy
|
||||
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults
|
||||
from litellm.router_api import deployment_id as generate_deployment_id
|
||||
from litellm.router_strategy.tag_based_routing import (
|
||||
baggage_stamp_for_serving_deployment,
|
||||
get_tags_from_request_kwargs,
|
||||
tags_with_serving_deployment_tags,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
deployment_has_io_token_limits, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
refund_stale_reservation_before_retry,
|
||||
set_io_token_rate_limit_request_kwargs,
|
||||
)
|
||||
from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage
|
||||
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
|
||||
from litellm.types.router import Deployment, GenericLiteLLMParams, LiteLLM_Params, validate_litellm_params_strictly
|
||||
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.router_python_host import PythonHost
|
||||
|
||||
# The SDK-client and max_parallel_requests reads of the routing cache: process-local, so the family
|
||||
# only names a Redis span when a read misses through to Redis.
|
||||
ROUTER_CLIENTS_TARGET: Final = "router_clients"
|
||||
_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt", "anthropic"})
|
||||
_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params"
|
||||
# Generic handlers whose stream timeout is split from the request timeout.
|
||||
STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"})
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class Shaping:
|
||||
"""What the request shaping reads from the router, resolved once when the core is built.
|
||||
|
||||
``timeout`` is the router's ``timeout`` (its own setting, else the process ``request_timeout``),
|
||||
``request_timeout`` the explicitly configured process ``request_timeout``, kept only when the
|
||||
router has a timeout of its own, as the Python router keeps them."""
|
||||
|
||||
default_litellm_params: Mapping[str, object]
|
||||
timeout: object
|
||||
explicit_timeout: object
|
||||
request_timeout: float | None
|
||||
stream_timeout: object
|
||||
cache_responses: bool
|
||||
|
||||
|
||||
def as_dict(value: object) -> dict[str, object]: # mutable-ok: a fresh dict of a request mapping
|
||||
try:
|
||||
return dict(_MAPPING.validate_python(value))
|
||||
except ValidationError:
|
||||
return {}
|
||||
|
||||
|
||||
def _mapping(value: object) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _MAPPING.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def metadata_bucket(
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
name: str,
|
||||
) -> dict[str, object]: # mutable-ok: the request's own bucket
|
||||
"""``kwargs.setdefault(name, {})``: the request's own metadata dict, created when absent."""
|
||||
bucket: Final = kwargs.get(name)
|
||||
if isinstance(bucket, dict):
|
||||
return bucket # pyright: ignore[reportUnknownVariableType] # the request's untyped metadata dict
|
||||
if name not in kwargs:
|
||||
fresh: Final[dict[str, object]] = {} # mutable-ok: the request kwargs the Python router mutates in place
|
||||
kwargs[name] = fresh # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
return fresh
|
||||
# A non-dict value under the key: `setdefault` returns it as is, and `.update` would raise.
|
||||
raise AttributeError(f"'{type(bucket).__name__}' object has no attribute 'update'")
|
||||
|
||||
|
||||
# --- request params the deployment sets --------------------------------------------------------
|
||||
|
||||
|
||||
def _drops_malformed_request_tools(
|
||||
deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object]
|
||||
) -> bool:
|
||||
if isinstance(request.get("tools") or [], (list, tuple)) or not deployment_params.get("tools"):
|
||||
return False
|
||||
set_defaults: Final = {name: value for name, value in defaults.items() if value is not None}
|
||||
return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools")
|
||||
|
||||
|
||||
def _malformed_request_params_the_deployment_sets(
|
||||
deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object]
|
||||
) -> frozenset[str]:
|
||||
overrides: Final = {
|
||||
key: value
|
||||
for key, value in request.items()
|
||||
if key in SUPPLEMENTAL_KWARGS_KEYS and deployment_params.get(key) is not None
|
||||
}
|
||||
if not overrides:
|
||||
return frozenset()
|
||||
try:
|
||||
validate_litellm_params_strictly(GenericLiteLLMParams, overrides)
|
||||
except ValidationError as error:
|
||||
set_defaults: Final = {name: value for name, value in defaults.items() if value is not None}
|
||||
policy: Final = drop_policy({**deployment_params, **set_defaults, **request})
|
||||
malformed: Final = (str(detail["loc"][0]) for detail in error.errors() if detail["loc"])
|
||||
return frozenset(key for key in malformed if policy.drops(key))
|
||||
return frozenset()
|
||||
|
||||
|
||||
def _merged_tools(
|
||||
deployment_tools: Sequence[object], request_tools: object
|
||||
) -> list[object] | None: # mutable-ok: sent as the request's tools
|
||||
if not isinstance(request_tools, (list, tuple)):
|
||||
return None
|
||||
items: Final[Sequence[object]] = request_tools # pyright: ignore[reportUnknownVariableType] # a request's untyped tools
|
||||
return [*deployment_tools, *items] if deployment_tools or items else None
|
||||
|
||||
|
||||
def merge_tools_from_deployment(
|
||||
deployment_params: Mapping[str, object],
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
defaults: Mapping[str, object],
|
||||
) -> None:
|
||||
"""``_merge_tools_from_deployment``: deployment tools first, then the request's; the request's
|
||||
``tool_choice`` else the deployment's."""
|
||||
if _drops_malformed_request_tools(deployment_params, defaults, kwargs):
|
||||
for name in TOOL_FIELDS:
|
||||
kwargs.pop(name, None)
|
||||
deployment_tools: Final = deployment_params.get("tools") or [] # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
merged: Final = _merged_tools(
|
||||
deployment_tools if isinstance(deployment_tools, list) else [], # pyright: ignore[reportUnknownArgumentType] # the deployment's untyped tools
|
||||
kwargs.get("tools") or [],
|
||||
)
|
||||
if merged is not None:
|
||||
kwargs["tools"] = merged # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
if "tool_choice" not in kwargs and deployment_params.get("tool_choice") is not None:
|
||||
kwargs["tool_choice"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
deployment_params["tool_choice"]
|
||||
)
|
||||
|
||||
|
||||
def _deployment_sets_litellm_param(deployment_params: Mapping[str, object], key: str) -> bool:
|
||||
value: Final = deployment_params.get(key)
|
||||
if value is None:
|
||||
return False
|
||||
field: Final = LiteLLM_Params.model_fields.get(key)
|
||||
return field is None or value != field.default # pyright: ignore[reportAny] # litellm's untyped helpers and request values
|
||||
|
||||
|
||||
def _forwarded_alias_marker_keys_the_deployment_sets(
|
||||
deployment_params: Mapping[str, object], forwarded_keys: object
|
||||
) -> tuple[str, ...]:
|
||||
if not isinstance(forwarded_keys, tuple):
|
||||
return ()
|
||||
keys: Final[tuple[object, ...]] = forwarded_keys # pyright: ignore[reportUnknownVariableType] # the router's own marker
|
||||
return tuple(key for key in keys if isinstance(key, str) and _deployment_sets_litellm_param(deployment_params, key))
|
||||
|
||||
|
||||
# --- reasoning effort -------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _pop_effort_from_nested_carrier(params: dict[str, object], carrier: str) -> None: # mutable-ok: copy-on-write
|
||||
nested: Final = _mapping(params.get(carrier)) if isinstance(params.get(carrier), Mapping) else None
|
||||
if nested is None:
|
||||
return
|
||||
sanitized: Final = without(nested, "effort")
|
||||
if sanitized:
|
||||
params[carrier] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
sanitized
|
||||
)
|
||||
else:
|
||||
params.pop(carrier, None)
|
||||
|
||||
|
||||
def deployment_params_with_request_reasoning_override(
|
||||
deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object]
|
||||
) -> dict[str, object]: # mutable-ok: the call's own copy of the deployment params
|
||||
"""``_deployment_params_with_request_reasoning_override``: a request ``reasoning_effort`` drops
|
||||
the deployment's ``thinking`` and nested ``effort`` carriers."""
|
||||
sanitized: Final = dict(deployment_params)
|
||||
if request_kwargs.get("reasoning_effort") is None:
|
||||
return sanitized
|
||||
sanitized.pop("thinking", None)
|
||||
_pop_effort_from_nested_carrier(sanitized, "output_config")
|
||||
_pop_effort_from_nested_carrier(sanitized, "reasoning")
|
||||
extra_body: Final = sanitized.get("extra_body")
|
||||
if isinstance(extra_body, Mapping):
|
||||
sanitized_extra_body: Final = as_dict(extra_body) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
|
||||
sanitized_extra_body.pop("reasoning_effort", None)
|
||||
sanitized_extra_body.pop("thinking", None)
|
||||
_pop_effort_from_nested_carrier(sanitized_extra_body, "output_config")
|
||||
_pop_effort_from_nested_carrier(sanitized_extra_body, "reasoning")
|
||||
if sanitized_extra_body:
|
||||
sanitized["extra_body"] = sanitized_extra_body
|
||||
else:
|
||||
sanitized.pop("extra_body", None)
|
||||
return sanitized
|
||||
|
||||
|
||||
def _is_classifier_internal_call(kwargs: Mapping[str, object]) -> bool:
|
||||
return any(
|
||||
isinstance(candidate, Mapping)
|
||||
and (_mapping(candidate) or {}).get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
|
||||
for candidate in (kwargs.get("metadata"), kwargs.get("litellm_metadata"))
|
||||
)
|
||||
|
||||
|
||||
def _declared_param_allowlist(params: Mapping[str, object]) -> frozenset[str]:
|
||||
declared: Final = params.get("allowed_openai_params")
|
||||
if not isinstance(declared, (list, tuple, set, frozenset)):
|
||||
return frozenset()
|
||||
entries: Final[Sequence[object]] = list(declared) # pyright: ignore[reportUnknownArgumentType] # the deployment's untyped allowlist
|
||||
return frozenset(entry for entry in entries if isinstance(entry, str))
|
||||
|
||||
|
||||
def _deployment_supports_param(deployment: Mapping[str, object], group: str, param: str, host: "PythonHost") -> bool:
|
||||
params: Final = as_dict(deployment.get("litellm_params"))
|
||||
if not params:
|
||||
return True
|
||||
model_info: Final = as_dict(deployment.get("model_info"))
|
||||
base_model: Final = model_info.get("base_model") or params.get("base_model")
|
||||
try:
|
||||
model_label: Final = str(params.get("model") or group)
|
||||
provider_param: Final = params.get("custom_llm_provider")
|
||||
authenticating: Final = declared_authenticating_provider(
|
||||
model_label, provider_param if isinstance(provider_param, str) else None
|
||||
)
|
||||
if authenticating is not None:
|
||||
model, custom_llm_provider = model_label.removeprefix(f"{authenticating}/"), authenticating
|
||||
else:
|
||||
model, custom_llm_provider, _, _ = host.invoke( # rebind-ok: rebound as the Python router rebinds it
|
||||
"get_llm_provider"
|
||||
)(
|
||||
model=model_label,
|
||||
custom_llm_provider=provider_param,
|
||||
litellm_params=LiteLLM_Params.model_validate(params),
|
||||
)
|
||||
supported: Final = host.invoke("get_supported_openai_params")(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
base_model=base_model if isinstance(base_model, str) else None,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request
|
||||
return True
|
||||
return supported is None or param in supported
|
||||
|
||||
|
||||
def deployment_accepts_param(deployment: Mapping[str, object], group: str, param: str, host: "PythonHost") -> bool:
|
||||
params: Final = as_dict(deployment.get("litellm_params"))
|
||||
if not params:
|
||||
return True
|
||||
if param in _declared_param_allowlist(params):
|
||||
return True
|
||||
provider_param: Final = params.get("custom_llm_provider")
|
||||
if declared_authenticating_provider(
|
||||
str(params.get("model") or ""), provider_param if isinstance(provider_param, str) else None
|
||||
):
|
||||
return True
|
||||
return _deployment_supports_param(deployment, group, param, host)
|
||||
|
||||
|
||||
def drop_unsupported_classifier_reasoning_effort(
|
||||
deployment: Mapping[str, object],
|
||||
model: str,
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
host: "PythonHost",
|
||||
) -> None:
|
||||
"""``_drop_unsupported_classifier_reasoning_effort``."""
|
||||
if kwargs.get("reasoning_effort") is None or not _is_classifier_internal_call(kwargs):
|
||||
return
|
||||
if deployment_accepts_param(deployment, model, "reasoning_effort", host):
|
||||
return
|
||||
kwargs.pop("reasoning_effort", None)
|
||||
proxy_server_request: Final = kwargs.get("proxy_server_request")
|
||||
if not isinstance(proxy_server_request, dict):
|
||||
return
|
||||
body: Final[object] = proxy_server_request.get("body") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the proxy's untyped request
|
||||
if isinstance(body, dict):
|
||||
body.pop("reasoning_effort", None) # pyright: ignore[reportUnknownMemberType] # same
|
||||
|
||||
|
||||
# --- timeouts ---------------------------------------------------------------------------------
|
||||
|
||||
|
||||
def stream_timeout(shaping: Shaping, kwargs: Mapping[str, object], data: Mapping[str, object]) -> object:
|
||||
"""``_get_stream_timeout``."""
|
||||
return (
|
||||
kwargs.get("stream_timeout", None)
|
||||
or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier")
|
||||
or data.get("stream_timeout", None)
|
||||
or shaping.stream_timeout
|
||||
or shaping.request_timeout
|
||||
or shaping.default_litellm_params.get("stream_timeout", None)
|
||||
)
|
||||
|
||||
|
||||
def non_stream_timeout(shaping: Shaping, kwargs: Mapping[str, object], data: Mapping[str, object]) -> object:
|
||||
"""``_get_non_stream_timeout``."""
|
||||
return (
|
||||
kwargs.get("timeout", None)
|
||||
or kwargs.get("request_timeout", None)
|
||||
or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier")
|
||||
or data.get("timeout", None)
|
||||
or data.get("request_timeout", None)
|
||||
or shaping.request_timeout
|
||||
or shaping.timeout
|
||||
or shaping.default_litellm_params.get("timeout", None)
|
||||
)
|
||||
|
||||
|
||||
def call_timeout(shaping: Shaping, kwargs: Mapping[str, object], data: Mapping[str, object]) -> object:
|
||||
"""``_get_timeout``."""
|
||||
stream: Final = stream_timeout(shaping, kwargs, data) if kwargs.get("stream") else None
|
||||
return stream if stream is not None else non_stream_timeout(shaping, kwargs, data)
|
||||
|
||||
|
||||
# --- default litellm params -------------------------------------------------------------------
|
||||
|
||||
|
||||
def update_kwargs_with_default_litellm_params(
|
||||
shaping: Shaping,
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
metadata_variable_name: str,
|
||||
) -> None:
|
||||
"""``_update_kwargs_with_default_litellm_params``: unset non-None defaults fill the request;
|
||||
the default metadata is merged over the request's."""
|
||||
defaults: Final = dict(without_retrieve_stream_defaults(shaping.default_litellm_params, kwargs))
|
||||
metadata_defaults: Final = defaults.pop("metadata", {}) or {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
for key, value in defaults.items():
|
||||
if value is None:
|
||||
continue
|
||||
kwargs.setdefault(key, value)
|
||||
metadata_bucket(kwargs, metadata_variable_name).update(as_dict(metadata_defaults)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
|
||||
|
||||
|
||||
# --- clientside credentials -------------------------------------------------------------------
|
||||
|
||||
|
||||
def _clientside_deployment(
|
||||
deployment: Mapping[str, object], kwargs: Mapping[str, object], function_name: str | None
|
||||
) -> tuple[Mapping[str, object], str, object]:
|
||||
"""``_handle_clientside_credential``: the request-scoped model_info, model and api_base."""
|
||||
model_info: Final = as_dict(deployment.get("model_info"))
|
||||
litellm_params: Final = as_dict(deployment.get("litellm_params"))
|
||||
dynamic_litellm_params: Final = get_dynamic_litellm_params( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
litellm_params=litellm_params, request_kwargs=dict(kwargs)
|
||||
)
|
||||
metadata_variable_name: Final = get_router_metadata_variable_name(function_name=function_name)
|
||||
model_group: Final = as_dict(kwargs.get(metadata_variable_name)).get("model_group")
|
||||
model_id: Final = generate_deployment_id(model_group, dynamic_litellm_params) # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # the group as the request carries it
|
||||
original_model_id: Final = model_info.get("id")
|
||||
model_info["id"] = model_id
|
||||
model_info["original_model_id"] = original_model_id
|
||||
built: Final = Deployment(
|
||||
model_name=model_group, # pyright: ignore[reportArgumentType] # the group as the request carries it
|
||||
litellm_params=LiteLLM_Params.model_validate(dynamic_litellm_params),
|
||||
model_info=model_info, # pyright: ignore[reportArgumentType] # model_info as configured
|
||||
)
|
||||
return built.model_info.model_dump(), built.litellm_params.model, built.litellm_params.api_base
|
||||
|
||||
|
||||
# --- the deployment --------------------------------------------------------------------------
|
||||
|
||||
|
||||
def update_kwargs_with_deployment(
|
||||
shaping: Shaping,
|
||||
deployment: Mapping[str, object],
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
*,
|
||||
cache: "DualCache | None",
|
||||
function_name: str | None = None,
|
||||
splits_stream_timeout: bool = False,
|
||||
opens_responses_socket: bool = False,
|
||||
) -> None:
|
||||
"""``_update_kwargs_with_deployment``, in place on the attempt's kwargs."""
|
||||
deployment_params: Final = as_dict(deployment.get("litellm_params"))
|
||||
for key in _forwarded_alias_marker_keys_the_deployment_sets(
|
||||
deployment_params, kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ())
|
||||
):
|
||||
kwargs.pop(key, None)
|
||||
kwargs.pop("daily_reset_time", None)
|
||||
defaults: Final = shaping.default_litellm_params
|
||||
merge_tools_from_deployment(deployment_params, kwargs, defaults)
|
||||
|
||||
model_info: Mapping[str, object] = as_dict( # rebind-ok: rebound as the Python router rebinds it
|
||||
deployment.get("model_info")
|
||||
)
|
||||
deployment_model: object = deployment_params.get("model") # rebind-ok: same
|
||||
deployment_api_base: object = deployment_params.get("api_base") # rebind-ok: same
|
||||
deployment_model_name: Final = deployment.get("model_name")
|
||||
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
|
||||
|
||||
for key in _malformed_request_params_the_deployment_sets(deployment_params, defaults, kwargs):
|
||||
kwargs.pop(key)
|
||||
declared_provider: Final = deployment_params.get("custom_llm_provider")
|
||||
declared_model: Final = deployment_params.get("model")
|
||||
deployment_provider: Final = (
|
||||
declared_provider
|
||||
if isinstance(declared_provider, str)
|
||||
else inferred_provider(declared_model if isinstance(declared_model, str) else None)
|
||||
)
|
||||
forwarded_authorization: Final = (
|
||||
ProviderSpecificHeaderUtils.forwarded_authorization(kwargs.get("provider_specific_header"), deployment_provider)
|
||||
if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS
|
||||
else None
|
||||
)
|
||||
if is_clientside_credential(request_kwargs=dict(kwargs)) or forwarded_authorization is not None:
|
||||
model_info, deployment_model, deployment_api_base = ( # rebind-ok: rebound as the Python router rebinds it
|
||||
_clientside_deployment(
|
||||
deployment,
|
||||
{**kwargs, **({"api_key": forwarded_authorization} if forwarded_authorization is not None else {})},
|
||||
function_name,
|
||||
)
|
||||
)
|
||||
|
||||
metadata_variable_name: Final = get_router_metadata_variable_name(function_name=function_name)
|
||||
routing: Final = deployment.get(KUBERNETES_POD_ROUTING_KEY)
|
||||
metadata: Final = metadata_bucket(kwargs, metadata_variable_name)
|
||||
metadata.update(
|
||||
{
|
||||
"deployment": deployment_model,
|
||||
"model_info": model_info,
|
||||
"api_base": deployment_api_base,
|
||||
"deployment_model_name": deployment_model_name,
|
||||
KUBERNETES_POD_ROUTING_KEY: routing,
|
||||
}
|
||||
)
|
||||
other_bucket: Final = "metadata" if metadata_variable_name == "litellm_metadata" else "litellm_metadata"
|
||||
other_bucket_metadata: Final = kwargs.get(other_bucket)
|
||||
if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata:
|
||||
kwargs[other_bucket] = {**as_dict(other_bucket_metadata), KUBERNETES_POD_ROUTING_KEY: routing} # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
|
||||
if cache is not None:
|
||||
refund_stale_reservation_before_retry(cache, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict
|
||||
set_io_token_rate_limit_request_kwargs(
|
||||
kwargs, # pyright: ignore[reportArgumentType] # the request's kwargs dict
|
||||
store_in_context=deployment_has_io_token_limits(deployment), # pyright: ignore[reportArgumentType] # a deployment record
|
||||
opens_responses_socket=opens_responses_socket,
|
||||
metadata_channel=metadata_variable_name,
|
||||
)
|
||||
metadata.setdefault(
|
||||
ROUTING_REQUEST_TAGS_METADATA_KEY,
|
||||
tuple(get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)), # pyright: ignore[reportArgumentType] # the request's kwargs dict
|
||||
)
|
||||
deployment_tags: Final = deployment_params.get("tags")
|
||||
merged_tags: Final = tags_with_serving_deployment_tags(metadata, deployment_tags)
|
||||
if merged_tags is not None:
|
||||
metadata["tags"] = list(merged_tags)
|
||||
baggage_stamp: Final = baggage_stamp_for_serving_deployment(metadata, deployment_tags)
|
||||
if baggage_stamp is not None:
|
||||
metadata[BAGGAGE_MAPPED_TAGS_METADATA_KEY] = baggage_stamp
|
||||
|
||||
credential_name: Final = deployment_params.get("litellm_credential_name")
|
||||
if credential_name:
|
||||
credential_tag: Final = f"Credential: {credential_name}"
|
||||
existing_tags: Final = metadata.get("tags") or [] # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values
|
||||
if isinstance(existing_tags, list) and credential_tag not in existing_tags:
|
||||
existing_tags.append(credential_tag) # pyright: ignore[reportUnknownMemberType] # the request's own tags list
|
||||
metadata["tags"] = existing_tags
|
||||
|
||||
from litellm.llms.github_copilot.per_user_auth import github_copilot_per_user_credential_name
|
||||
|
||||
configured_per_user_name: Final = github_copilot_per_user_credential_name(deployment_params)
|
||||
caller_credential_name: Final = kwargs.get("litellm_credential_name")
|
||||
if (
|
||||
configured_per_user_name is not None
|
||||
and isinstance(caller_credential_name, str)
|
||||
and caller_credential_name != configured_per_user_name
|
||||
):
|
||||
raise litellm_exceptions.BadRequestError(
|
||||
message="litellm_credential_name cannot be overridden on a deployment that uses per-user GitHub OAuth",
|
||||
model=str(deployment_model_name),
|
||||
llm_provider="",
|
||||
)
|
||||
|
||||
kwargs["model_info"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
model_info
|
||||
)
|
||||
_set_timeouts(shaping, deployment_params, kwargs, function_name, splits_stream_timeout)
|
||||
|
||||
if router_config_requests_stream_usage(kwargs, (deployment_params, defaults)):
|
||||
kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
STREAM_USAGE_INJECTED_FIELD
|
||||
] = True
|
||||
update_kwargs_with_default_litellm_params(shaping, kwargs, metadata_variable_name)
|
||||
pinned_model_id: Final = deployment_model_id_pin(deployment_params)
|
||||
if pinned_model_id is not None:
|
||||
kwargs["model_id"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
pinned_model_id
|
||||
)
|
||||
|
||||
|
||||
def _set_timeouts(
|
||||
shaping: Shaping,
|
||||
deployment_params: Mapping[str, object],
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place
|
||||
function_name: str | None,
|
||||
splits_stream_timeout: bool,
|
||||
) -> None:
|
||||
if function_name == "_ageneric_api_call_with_fallbacks":
|
||||
from litellm.passthrough.timeout_utils import resolve_llm_passthrough_timeout
|
||||
|
||||
router_timeout: Final = (
|
||||
shaping.request_timeout
|
||||
if shaping.request_timeout is not None
|
||||
else float(shaping.explicit_timeout)
|
||||
if isinstance(shaping.explicit_timeout, (int, float))
|
||||
else None
|
||||
)
|
||||
router_stream_timeout: Final = (
|
||||
shaping.stream_timeout
|
||||
if shaping.stream_timeout is not None
|
||||
else shaping.request_timeout
|
||||
if shaping.request_timeout is not None
|
||||
else shaping.default_litellm_params.get("stream_timeout")
|
||||
)
|
||||
passthrough_timeout: Final = partial(
|
||||
resolve_llm_passthrough_timeout,
|
||||
litellm_params=deployment_params,
|
||||
router_timeout=router_timeout,
|
||||
router_stream_timeout=router_stream_timeout, # pyright: ignore[reportArgumentType] # a configured default, passed as the Python router passes it
|
||||
)
|
||||
timeout_request: Final = {**kwargs, "stream": True} if requests_streamed_retrieval(kwargs) else dict(kwargs)
|
||||
kwargs.update(
|
||||
CompletionTimeout.stream_request_kwargs(
|
||||
passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=timeout_request)
|
||||
)
|
||||
if splits_stream_timeout and timeout_request.get("stream")
|
||||
else {"timeout": passthrough_timeout(kwargs=timeout_request)}
|
||||
)
|
||||
elif splits_stream_timeout:
|
||||
kwargs.update(
|
||||
CompletionTimeout.stream_request_kwargs(
|
||||
non_stream_timeout(shaping, kwargs, deployment_params),
|
||||
stream_timeout(shaping, kwargs, deployment_params) if kwargs.get("stream") else None,
|
||||
)
|
||||
)
|
||||
else:
|
||||
kwargs["timeout"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
call_timeout(shaping, kwargs, deployment_params)
|
||||
)
|
||||
|
||||
|
||||
def stamp_deployment_num_retries(error: BaseException, deployment: Mapping[str, object]) -> None:
|
||||
"""``_set_deployment_num_retries_on_exception``."""
|
||||
if getattr(error, "num_retries", None) is not None:
|
||||
return
|
||||
value: Final = as_dict(deployment.get("litellm_params")).get("num_retries")
|
||||
if value is None:
|
||||
return
|
||||
try:
|
||||
error.num_retries = int(str(value)) if not isinstance(value, int) else value # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
|
||||
def stamp_failed_deployment_id(error: BaseException, model_info: object) -> None:
|
||||
"""``_set_failed_deployment_id_on_exception`` with the effective model_info."""
|
||||
if getattr(error, "failed_deployment_id", None):
|
||||
return
|
||||
deployment_id: Final = as_dict(model_info).get("id")
|
||||
if deployment_id:
|
||||
try:
|
||||
error.failed_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
except Exception: # noqa: BLE001, S110 # as the Python router: a read-only error keeps its attributes
|
||||
pass
|
||||
|
||||
|
||||
def stamp_retry_skip_deployment_id(error: BaseException, kwargs: Mapping[str, object]) -> None:
|
||||
"""``_stamp_retry_skip_deployment_id``."""
|
||||
model_info: Final = kwargs.get("model_info")
|
||||
if not isinstance(model_info, Mapping):
|
||||
return
|
||||
info: Final = as_dict(model_info) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values
|
||||
configured: Final = (
|
||||
info.get("original_model_id") if isinstance(error, litellm_exceptions.UnsupportedEndpointError) else None
|
||||
)
|
||||
deployment_id: Final = configured or info.get("id")
|
||||
if isinstance(deployment_id, str) and deployment_id:
|
||||
error.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read on retry # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them
|
||||
|
||||
|
||||
NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
__all__ = [
|
||||
"STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS",
|
||||
"Shaping",
|
||||
"as_dict",
|
||||
"call_timeout",
|
||||
"deployment_params_with_request_reasoning_override",
|
||||
"drop_unsupported_classifier_reasoning_effort",
|
||||
"merge_tools_from_deployment",
|
||||
"metadata_bucket",
|
||||
"stamp_deployment_num_retries",
|
||||
"stamp_failed_deployment_id",
|
||||
"stamp_retry_skip_deployment_id",
|
||||
"update_kwargs_with_default_litellm_params",
|
||||
"update_kwargs_with_deployment",
|
||||
]
|
||||
107
litellm/router_rust/invoke/slots.py
Normal file
107
litellm/router_rust/invoke/slots.py
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
"""
|
||||
A deployment's ``max_parallel_requests`` slot on the Rust router core (``_deployment_slot`` and
|
||||
``_get_client(client_type="max_parallel_requests")``): the attempt holds one of the deployment's
|
||||
slots around its provider call, and an open chat stream keeps it until the stream is exhausted,
|
||||
closed, or falls back. A caller arriving while every slot is taken gets a 429, which the core
|
||||
records as that attempt's failure.
|
||||
|
||||
The limit object is the one the Python router caches (``MaxParallelRequestsLimit``), kept in the
|
||||
core's routing cache under the same key, so its in-flight count is shared by every attempt.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET, as_dict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_rust.invoke.endpoints import EndpointCore, Plan
|
||||
|
||||
# The attempt helpers that run inside `_deployment_slot` (`_acompletion` holds the same slot
|
||||
# itself). Sync helpers, rerank, moderation and the search helper hold none.
|
||||
_SLOTTED_KINDS: Final = frozenset(
|
||||
{
|
||||
"chat",
|
||||
"text",
|
||||
"embedding",
|
||||
"image",
|
||||
"transcription",
|
||||
"speech",
|
||||
"adapter",
|
||||
"generic",
|
||||
"batch",
|
||||
"cancel_batch",
|
||||
"vector_store_create",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@with_service_target(ROUTER_CLIENTS_TARGET)
|
||||
def max_parallel_requests_limit(
|
||||
core: "EndpointCore", deployment: Mapping[str, object], kwargs: Mapping[str, object]
|
||||
) -> object:
|
||||
"""`_get_client(client_type="max_parallel_requests")`: the cached limit, built on first use
|
||||
(`InitalizeCachedClient.set_max_parallel_requests_client`)."""
|
||||
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs
|
||||
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
|
||||
from litellm.utils import calculate_max_parallel_requests
|
||||
|
||||
model_id: Final = as_dict(deployment.get("model_info")).get("id")
|
||||
key: Final = f"{model_id}_max_parallel_requests_client"
|
||||
span: Final = get_parent_otel_span_from_kwargs(dict(kwargs))
|
||||
cache: Final = core.routing_cache
|
||||
cached: Final[object] = cache.get_cache(key=key, local_only=True, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # the cache holds untyped values
|
||||
if cached is not None:
|
||||
return cached # pyright: ignore[reportUnknownVariableType] # the cache holds untyped values
|
||||
params: Final = as_dict(deployment.get("litellm_params"))
|
||||
calculated: Final = calculate_max_parallel_requests(
|
||||
rpm=params.get("rpm"), # pyright: ignore[reportArgumentType] # the configured value, as the Python router reads it
|
||||
max_parallel_requests=params.get("max_parallel_requests"), # pyright: ignore[reportArgumentType] # same
|
||||
tpm=params.get("tpm"), # pyright: ignore[reportArgumentType] # same
|
||||
default_max_parallel_requests=core.endpoints.default_max_parallel_requests, # pyright: ignore[reportArgumentType] # the setting as given, as the Python router passes it
|
||||
)
|
||||
if calculated:
|
||||
cache.set_cache( # pyright: ignore[reportUnknownMemberType] # the routing cache
|
||||
key=key,
|
||||
value=MaxParallelRequestsLimit(
|
||||
max_parallel_requests=calculated,
|
||||
model_id=str(model_id),
|
||||
model_group=str(deployment.get("model_name", "")),
|
||||
),
|
||||
local_only=True,
|
||||
)
|
||||
return cache.get_cache(key=key, local_only=True, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # see above
|
||||
|
||||
|
||||
def enter_deployment_slot(
|
||||
core: "EndpointCore", plan: "Plan", deployment: Mapping[str, object], kwargs: Mapping[str, object]
|
||||
) -> contextlib.ExitStack:
|
||||
"""The attempt's slot, entered (a full deployment raises its 429 here); an empty stack when
|
||||
the plan's helper holds none or the deployment has no limit."""
|
||||
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
|
||||
|
||||
slot: Final = contextlib.ExitStack()
|
||||
if plan.sync or plan.kind not in _SLOTTED_KINDS:
|
||||
return slot
|
||||
limit: Final = max_parallel_requests_limit(core, deployment, kwargs)
|
||||
if isinstance(limit, MaxParallelRequestsLimit):
|
||||
slot.enter_context(limit)
|
||||
return slot
|
||||
|
||||
|
||||
def deployment_slot(
|
||||
core: "EndpointCore", deployment: Mapping[str, object], kwargs: Mapping[str, object]
|
||||
) -> contextlib.ExitStack:
|
||||
"""`_deployment_slot` outside a plan's attempt (a file created on every deployment)."""
|
||||
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
|
||||
|
||||
slot: Final = contextlib.ExitStack()
|
||||
limit: Final = max_parallel_requests_limit(core, deployment, kwargs)
|
||||
if isinstance(limit, MaxParallelRequestsLimit):
|
||||
slot.enter_context(limit)
|
||||
return slot
|
||||
|
||||
|
||||
__all__ = ["deployment_slot", "enter_deployment_slot", "max_parallel_requests_limit"]
|
||||
1839
litellm/router_rust/invoke/streaming.py
Normal file
1839
litellm/router_rust/invoke/streaming.py
Normal file
File diff suppressed because it is too large
Load diff
398
litellm/router_rust/invoke/walk.py
Normal file
398
litellm/router_rust/invoke/walk.py
Normal file
|
|
@ -0,0 +1,398 @@
|
|||
"""
|
||||
The request metadata the Python router stamps as it walks retries and fallbacks, kept for a
|
||||
``RustCore`` call. router-core decides the walk and reports it (each attempt's ``scope``,
|
||||
``max_retries`` and ``model_group_size``, and the ``fallback`` and ``retry_logged`` events); this
|
||||
module only holds what each level's metadata dict holds in the Python router and lays it over the
|
||||
attempt's kwargs.
|
||||
|
||||
In the Python router one metadata dict serves a level's retries, and every fallback hop copies
|
||||
its parent's (``run_async_fallback``'s ``_hop_metadata``), so a hop's breadcrumbs never reach
|
||||
its parent. Each scope here is one such dict. The root level's dict is the caller's own (a proxy
|
||||
reads a failed request's walk off it), so a fresh call's root writes reach that dict as well.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Coroutine, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets
|
||||
|
||||
# `RETRY_BREADCRUMB_LIMIT`: how many failures `previous_models` keeps.
|
||||
RETRY_BREADCRUMB_LIMIT: Final = 4
|
||||
|
||||
# The keys the walk stamps on each attempt's metadata.
|
||||
_WALK_KEYS: Final = (
|
||||
"original_model_group",
|
||||
"attempted_fallbacks",
|
||||
"attempted_retries",
|
||||
"max_retries",
|
||||
"model_group_size",
|
||||
"previous_models",
|
||||
"request_retry_count",
|
||||
)
|
||||
# What a breadcrumb reads from the last attempt a scope made.
|
||||
_SERVED_KEYS: Final = ("model_info", "deployment")
|
||||
# The walk state `run_async_fallback` hands every hop as kwargs.
|
||||
_HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors")
|
||||
|
||||
|
||||
class _Scope:
|
||||
__slots__ = ("kwargs", "metadata")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
metadata: dict[str, object], # mutable-ok: the level's metadata, stamped as the walk runs
|
||||
kwargs: Mapping[str, object],
|
||||
) -> None:
|
||||
self.metadata: Final = metadata
|
||||
self.kwargs: Final = kwargs
|
||||
|
||||
|
||||
def _breadcrumb(
|
||||
metadata: Mapping[str, object], model: str, error: BaseException, params_base_model: Callable[[str], str | None]
|
||||
) -> dict[str, object]: # mutable-ok: a fresh RetryAttemptRecord
|
||||
"""``_core_log_retry``'s record of the failure."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_base_model_from_litellm_call_metadata
|
||||
|
||||
model_info: Final = as_str_mapping(metadata.get("model_info")) or {}
|
||||
deployment_id: Final = model_info.get("id")
|
||||
known_id: Final = deployment_id if isinstance(deployment_id, str) else None
|
||||
base_model: Final = get_base_model_from_litellm_call_metadata(metadata=metadata) or (
|
||||
params_base_model(known_id) if known_id is not None else None
|
||||
)
|
||||
deployment: Final = metadata.get("deployment")
|
||||
attempted: Final = metadata.get("attempted_retries")
|
||||
return {
|
||||
"model_group": model,
|
||||
"deployment_id": known_id,
|
||||
"model": deployment if isinstance(deployment, str) else None,
|
||||
"base_model": base_model,
|
||||
"exception_type": type(error).__name__,
|
||||
"exception_string": str(error),
|
||||
"attempted_retries": attempted if type(attempted) is int else None,
|
||||
}
|
||||
|
||||
|
||||
def _kept_breadcrumbs(earlier: object) -> tuple[object, ...]:
|
||||
"""The breadcrumbs a new one joins: the latest ``RETRY_BREADCRUMB_LIMIT - 1``."""
|
||||
if not isinstance(earlier, list | tuple):
|
||||
return ()
|
||||
stamped: Final[tuple[object, ...]] = tuple(earlier) # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the breadcrumbs as stamped
|
||||
return stamped[-(RETRY_BREADCRUMB_LIMIT - 1) :]
|
||||
|
||||
|
||||
def log_retry_into(
|
||||
kwargs: dict[str, object], # mutable-ok: the request kwargs whose metadata the record joins
|
||||
error: Exception,
|
||||
params_base_model: Callable[[str], str | None],
|
||||
) -> dict[str, object]: # mutable-ok: the same kwargs, as ``_core_log_retry`` returns them
|
||||
"""
|
||||
``_core_log_retry`` (``RoutingContext.log_retry``): the failure joins the request metadata's
|
||||
``previous_models`` and counts toward ``request_retry_count``, with a ``litellm.routing.retry``
|
||||
phase event when a pre-routing strategy picked the group.
|
||||
"""
|
||||
from litellm.integrations.otel.runtime import phase_event
|
||||
|
||||
field: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
|
||||
metadata: Final = kwargs[field]
|
||||
view: Final = as_str_mapping(metadata)
|
||||
if view is None:
|
||||
# Python reads it with ``.get``.
|
||||
raise AttributeError(f"{type(metadata).__name__!r} object has no attribute 'get'")
|
||||
model_group: Final = kwargs.get("model")
|
||||
record: Final = _breadcrumb(view, "", error, params_base_model)
|
||||
record["model_group"] = model_group if isinstance(model_group, str) else None
|
||||
deployment_id: Final = record["deployment_id"]
|
||||
breadcrumbs: Final = (*_kept_breadcrumbs(view.get("previous_models")), record)
|
||||
earlier: Final = view.get("request_retry_count")
|
||||
request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1
|
||||
decision: Final = view.get("routing_decision")
|
||||
router_name: Final = (as_str_mapping(decision) or {}).get("router_model_name")
|
||||
if isinstance(router_name, str):
|
||||
phase_event(
|
||||
"litellm.routing.retry",
|
||||
{
|
||||
"litellm.routing.router_model_name": router_name,
|
||||
"litellm.retry.count": request_retry_count,
|
||||
"error.type": type(error).__name__,
|
||||
**({"litellm.deployment.model_group": model_group} if isinstance(model_group, str) else {}),
|
||||
**({"litellm.deployment.id": deployment_id} if isinstance(deployment_id, str) else {}),
|
||||
},
|
||||
)
|
||||
metadata["previous_models"] = breadcrumbs # pyright: ignore[reportIndexIssue] # the logging object already holds this dict
|
||||
metadata["request_retry_count"] = request_retry_count # pyright: ignore[reportIndexIssue] # same dict, read by the cap
|
||||
return kwargs
|
||||
|
||||
|
||||
class WalkMetadata:
|
||||
"""The metadata each level of one call's walk holds, by scope."""
|
||||
|
||||
__slots__ = ("_attempted", "_field", "_live", "_params_base_model", "_request", "_scopes")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request: Mapping[str, object],
|
||||
field: str,
|
||||
params_base_model: Callable[[str], str | None],
|
||||
*,
|
||||
live: bool = False,
|
||||
) -> None:
|
||||
model: Final = request.get("model")
|
||||
root: Final[dict[str, object]] = {"attempted_fallbacks": 0} # mutable-ok: the root level's metadata
|
||||
if model is not None:
|
||||
root["original_model_group"] = model
|
||||
self._request: Final = request
|
||||
self._field: Final = field
|
||||
self._params_base_model: Final = params_base_model
|
||||
self._scopes: Final[dict[int, _Scope]] = {0: _Scope(root, {})} # mutable-ok: filled as the walk hops
|
||||
self._attempted: AttemptedFallbackTargets | None = None
|
||||
# A call whose retry and fallback wrappers stamp the request itself; a stream's recovery
|
||||
# re-enters below them.
|
||||
self._live: Final = live
|
||||
if live:
|
||||
self._stamp_fallback_entry()
|
||||
|
||||
def _stamp_fallback_entry(self) -> None:
|
||||
"""The head of ``async_function_with_fallbacks``, in the request's own buckets. ``run_async_fallback``'s
|
||||
hops carry their ``attempted_targets`` and never reach it."""
|
||||
from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets
|
||||
|
||||
request: Final = self._request
|
||||
if isinstance(request.get("attempted_targets"), AttemptedFallbackTargets):
|
||||
return
|
||||
sibling: Final = request.get("metadata" if self._field == "litellm_metadata" else "litellm_metadata")
|
||||
if isinstance(sibling, dict):
|
||||
sibling.pop("attempted_fallbacks", None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict
|
||||
sibling.pop("original_model_group", None) # pyright: ignore[reportUnknownMemberType] # see above
|
||||
bucket: Final = request.get(self._field)
|
||||
if not isinstance(bucket, dict):
|
||||
return
|
||||
bucket["attempted_fallbacks"] = 0
|
||||
model: Final = request.get("model")
|
||||
if model is not None:
|
||||
bucket["original_model_group"] = model
|
||||
|
||||
def _retry_bucket(self, scope: _Scope) -> dict[str, object] | None: # mutable-ok: the request's own dict
|
||||
"""``_metadata`` of ``async_function_with_retries`` when it is a dict of the request's: the
|
||||
root level's, or a hop's when the bucket is not the one ``run_async_fallback`` copies."""
|
||||
if not self._live:
|
||||
return None
|
||||
name: Final = "litellm_metadata" if "litellm_metadata" in self._request else "metadata"
|
||||
if scope is not self._scopes[0] and name == self._field:
|
||||
return None
|
||||
bucket: Final = self._request.get(name)
|
||||
# `... or {}`: an empty bucket leaves the stamps on a dict of their own.
|
||||
return bucket if isinstance(bucket, dict) and bucket else None # pyright: ignore[reportUnknownVariableType] # the request's untyped metadata dict
|
||||
|
||||
def _write_through(self, scope: _Scope, values: Mapping[str, object]) -> None:
|
||||
bucket: Final = self._retry_bucket(scope)
|
||||
if bucket is not None:
|
||||
bucket.update(values)
|
||||
|
||||
def retries_opened(self, event: Mapping[str, object]) -> None:
|
||||
"""The head of ``async_function_with_retries``: the group's size and the retry tracking, before any pick."""
|
||||
scope: Final = self._scope(event.get("scope"))
|
||||
stamped: Final = {"attempted_retries": 0, "max_retries": event.get("max_retries")}
|
||||
scope.metadata.update(stamped)
|
||||
self._write_through(scope, stamped)
|
||||
bucket: Final = self._retry_bucket(scope)
|
||||
if bucket is not None and isinstance(bucket.get("model_group"), str):
|
||||
bucket["model_group_size"] = event.get("model_group_size")
|
||||
|
||||
def retry_budget(self, event: Mapping[str, object]) -> None:
|
||||
"""``_metadata["max_retries"] = num_retries``, once a level's first failure settled its overrides."""
|
||||
scope: Final = self._scope(event.get("scope"))
|
||||
stamped: Final = {"max_retries": event.get("max_retries")}
|
||||
scope.metadata.update(stamped)
|
||||
self._write_through(scope, stamped)
|
||||
|
||||
def _scope(self, value: object) -> _Scope:
|
||||
found: Final = self._scopes.get(value) if isinstance(value, int) else None
|
||||
return found if found is not None else self._scopes[0]
|
||||
|
||||
def targets(self) -> "AttemptedFallbackTargets":
|
||||
"""The walk's ``attempted_targets``: the targets tried and the ones skipped, and why."""
|
||||
from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets
|
||||
|
||||
if self._attempted is None:
|
||||
self._attempted = AttemptedFallbackTargets()
|
||||
return self._attempted
|
||||
|
||||
def hop(self, event: Mapping[str, object]) -> None:
|
||||
"""A ``fallback`` event: the new level copies its parent's metadata and kwargs."""
|
||||
parent: Final = self._scope(event.get("parent"))
|
||||
attempted: Final = self.targets()
|
||||
for key in (event.get("from"), event.get("to")):
|
||||
if isinstance(key, str):
|
||||
attempted.record(key)
|
||||
metadata: Final = dict(parent.metadata)
|
||||
original: Final = metadata.pop("original_model_group", self._request.get("model"))
|
||||
metadata.pop("model_group", None)
|
||||
metadata.pop("attempted_fallbacks", None)
|
||||
metadata["original_model_group"] = original
|
||||
# A dict target's non-string `model` is what `kwargs["model"]` holds on the hop.
|
||||
metadata["model_group"] = event["unroutable"] if "unroutable" in event else event.get("to")
|
||||
metadata["attempted_fallbacks"] = event.get("depth")
|
||||
kwargs: Final = {
|
||||
**parent.kwargs,
|
||||
"fallback_depth": event.get("depth"),
|
||||
"max_fallbacks": event.get("max_fallbacks"),
|
||||
"attempted_targets": attempted,
|
||||
}
|
||||
if self._request.get("include_fallback_errors") is True:
|
||||
kwargs["include_fallback_errors"] = True
|
||||
scope: Final = event.get("scope")
|
||||
if isinstance(scope, int):
|
||||
self._scopes[scope] = _Scope(metadata, kwargs)
|
||||
|
||||
def hop_kwargs(self, event: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a fresh dict for the loggers
|
||||
"""The kwargs the level held after a hop (``kwargs["model"] = target``), as the
|
||||
fallback loggers see them."""
|
||||
scope: Final = self._scope(event.get("scope"))
|
||||
model: Final = scope.metadata.get("model_group", event.get("to"))
|
||||
return {**self._request, **scope.kwargs, "model": model, self._field: dict(scope.metadata)}
|
||||
|
||||
def log_retry(self, event: Mapping[str, object], error: BaseException | None) -> None:
|
||||
"""A ``retry_logged`` event: the failure joins the scope's ``previous_models``."""
|
||||
scope: Final = self._scope(event.get("scope"))
|
||||
metadata: Final = scope.metadata
|
||||
attempted: Final = event.get("attempted_retries")
|
||||
if isinstance(attempted, int):
|
||||
metadata["attempted_retries"] = attempted
|
||||
self._write_through(scope, {"attempted_retries": attempted})
|
||||
if error is None:
|
||||
return
|
||||
record: Final = _breadcrumb(metadata, str(event.get("model")), error, self._params_base_model)
|
||||
earlier: Final = metadata.get("previous_models")
|
||||
kept: Final = _kept_breadcrumbs(earlier)
|
||||
count: Final = metadata.get("request_retry_count")
|
||||
metadata["previous_models"] = (*kept, record)
|
||||
metadata["request_retry_count"] = (count if type(count) is int and count >= 0 else 0) + 1
|
||||
self._write_through(
|
||||
scope,
|
||||
{"previous_models": metadata["previous_models"], "request_retry_count": metadata["request_retry_count"]},
|
||||
)
|
||||
|
||||
def deployment_that_just_failed(self, step: Mapping[str, object]) -> str | None:
|
||||
"""``_deployment_that_just_failed``: the deployment the metadata an attempt inherits names,
|
||||
before the attempt stamps its own; ``None`` when the walk's level has run none yet."""
|
||||
model_info: Final = as_str_mapping(self._scope(step.get("scope")).metadata.get("model_info"))
|
||||
model_id: Final = model_info.get("id") if model_info is not None else None
|
||||
return model_id if isinstance(model_id, str) else None
|
||||
|
||||
def request_retry_count(self, step: Mapping[str, object]) -> object:
|
||||
"""The level's ``request_retry_count``, as the retry loop has counted it so far."""
|
||||
return self._scope(step.get("scope")).metadata.get("request_retry_count")
|
||||
|
||||
def stamp(self, step: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The attempt's kwargs with its level's walk metadata and hop kwargs laid over them."""
|
||||
scope: Final = self._scope(step.get("scope"))
|
||||
metadata: Final = scope.metadata
|
||||
metadata["attempted_retries"] = step.get("retry")
|
||||
metadata["max_retries"] = step.get("max_retries")
|
||||
metadata["model_group_size"] = step.get("model_group_size")
|
||||
self._write_through(
|
||||
scope, {"attempted_retries": metadata["attempted_retries"], "max_retries": metadata["max_retries"]}
|
||||
)
|
||||
served: Final = dict(as_str_mapping(kwargs.get(self._field)) or {})
|
||||
served.update({key: metadata[key] for key in _WALK_KEYS if key in metadata})
|
||||
metadata.update({key: served.get(key) for key in _SERVED_KEYS})
|
||||
return {**kwargs, **{key: scope.kwargs[key] for key in _HOP_KWARGS if key in scope.kwargs}, self._field: served}
|
||||
|
||||
|
||||
async def _log_failed_fallbacks(ctx: "RoutingContext", detail: str, parent_otel_span: "Span | None") -> None:
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments_with_debug_info
|
||||
|
||||
cooldown_info: Final = await async_get_cooldown_deployments_with_debug_info(
|
||||
ctx=ctx, parent_otel_span=parent_otel_span
|
||||
)
|
||||
verbose_router_logger.error(
|
||||
"litellm.router.py::async_function_with_fallbacks() - "
|
||||
"Error occurred while trying to do fallbacks - %s\n"
|
||||
"Debug Information:\nCooldown Deployments=%s",
|
||||
detail,
|
||||
cooldown_info,
|
||||
)
|
||||
|
||||
|
||||
class FallbackOutcomes:
|
||||
"""What ``run_async_fallback`` does when a hop returns: it logs the hop to the process's
|
||||
custom loggers and, under ``include_fallback_errors``, lists the level's errors (its own
|
||||
first, then each hop that failed) in the response's ``x-litellm-fallback-errors``."""
|
||||
|
||||
__slots__ = ("_failed", "_header_errors", "_include", "_pending")
|
||||
|
||||
def __init__(self, request: Mapping[str, object]) -> None:
|
||||
self._include: Final = request.get("include_fallback_errors", False) is True
|
||||
self._failed: Final[dict[object, list[object]]] = {} # mutable-ok: each level's failed hops, in order
|
||||
self._header_errors: Final[list[list[object]]] = [] # mutable-ok: each successful hop's errors, innermost first
|
||||
self._pending: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: logger calls not yet awaited
|
||||
|
||||
def done(
|
||||
self,
|
||||
event: Mapping[str, object],
|
||||
level_error: BaseException | None,
|
||||
hop_error: BaseException | None,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> None:
|
||||
from litellm.litellm_core_utils.add_retry_fallback_headers import get_fallback_error_info
|
||||
from litellm.router_utils import fallback_event_handlers as handlers
|
||||
|
||||
if level_error is None or not isinstance(level_error, Exception):
|
||||
return
|
||||
parent: Final = event.get("parent")
|
||||
failed: Final = self._failed.setdefault(parent, [])
|
||||
original: Final = str(event.get("from"))
|
||||
if event.get("succeeded") is True:
|
||||
if self._include:
|
||||
self._header_errors.append([get_fallback_error_info(level_error), *failed])
|
||||
self._pending.append(
|
||||
handlers.log_success_fallback_event( # pyright: ignore[reportUnknownMemberType] # takes the kwargs as an untyped dict
|
||||
original_model_group=original, kwargs=dict(kwargs), original_exception=level_error
|
||||
)
|
||||
)
|
||||
return
|
||||
if isinstance(hop_error, Exception):
|
||||
failed.append(get_fallback_error_info(hop_error))
|
||||
self._pending.append(
|
||||
handlers.log_failure_fallback_event( # pyright: ignore[reportUnknownMemberType] # see above
|
||||
original_model_group=original, kwargs=dict(kwargs), original_exception=level_error
|
||||
)
|
||||
)
|
||||
|
||||
def failed_fallbacks(self, ctx: "RoutingContext", detail: str, parent_otel_span: "Span | None") -> None:
|
||||
"""The ``except`` around the regular fallbacks of ``async_function_with_fallbacks_common_utils``:
|
||||
``detail`` is the failure's text as it read then, the cooldowns are read when the pending calls run."""
|
||||
self._pending.append(_log_failed_fallbacks(ctx, detail, parent_otel_span))
|
||||
|
||||
def take_pending(self) -> tuple[Coroutine[object, object, None], ...]:
|
||||
pending: Final = tuple(self._pending)
|
||||
self._pending.clear()
|
||||
return pending
|
||||
|
||||
def add_headers(self, response: object, attempted_fallbacks: int) -> object:
|
||||
"""``add_fallback_headers_to_response`` for each successful hop, innermost first."""
|
||||
from litellm.litellm_core_utils.add_retry_fallback_headers import add_fallback_headers_to_response
|
||||
|
||||
headed: object = response # rebind-ok: gains each level's errors
|
||||
for errors in self._header_errors:
|
||||
headed = add_fallback_headers_to_response(
|
||||
response=headed,
|
||||
attempted_fallbacks=attempted_fallbacks,
|
||||
fallback_errors=errors, # pyright: ignore[reportArgumentType] # FallbackErrorInfo dicts
|
||||
)
|
||||
return headed
|
||||
|
||||
|
||||
async def await_all(pending: Sequence[Coroutine[object, object, None]]) -> None:
|
||||
for call in pending:
|
||||
await call
|
||||
|
||||
|
||||
__all__ = ["RETRY_BREADCRUMB_LIMIT", "FallbackOutcomes", "WalkMetadata", "await_all"]
|
||||
127
litellm/router_rust/invoke/websocket.py
Normal file
127
litellm/router_rust/invoke/websocket.py
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
"""
|
||||
Responses websocket turns on the Rust router core (``async_admit_deployment_turn``): the socket is
|
||||
opened on one deployment, and the core hands the handler a ``turn_admission`` callable that re-runs
|
||||
that deployment's admission before every later turn: the io-token turn budget, the deployment
|
||||
budget limiter, the routing strategy (on that deployment alone) and the pre-call checks.
|
||||
|
||||
Which strategy and selector apply is the Rust core's routing context (``routing_context``); the
|
||||
selectors, limiters and checks are the live Python objects the host keeps.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import without
|
||||
from litellm.router_rust.invoke.shaping import as_dict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_rust.invoke.endpoints import EndpointCore
|
||||
|
||||
|
||||
def turn_admission(
|
||||
core: "EndpointCore", model_group: str, deployment: Mapping[str, object], kwargs: Mapping[str, object]
|
||||
) -> Callable[[], object]:
|
||||
"""The ``turn_admission`` the handler calls before each later turn."""
|
||||
return partial(admit_deployment_turn, core, model_group, deployment, kwargs)
|
||||
|
||||
|
||||
def _deployment_as_configured_now(core: "EndpointCore", deployment: Mapping[str, object]) -> Mapping[str, object]:
|
||||
model_info: Final = deployment.get("model_info")
|
||||
model_id: Final = as_dict(model_info).get("id") if isinstance(model_info, dict) else None # pyright: ignore[reportUnknownArgumentType] # the configured model_info dict
|
||||
current: Final = core.configured_deployment(model_id) if isinstance(model_id, str) else None
|
||||
return deployment if current is None else current.model_dump(exclude_none=True)
|
||||
|
||||
|
||||
def _budget_limiter(core: "EndpointCore") -> object:
|
||||
"""The router's deployment-budget callback; the Rust core declines every config that has one."""
|
||||
from litellm.router_rust.core import RustCoreUnsupported
|
||||
|
||||
try:
|
||||
return core.registry().plugins("budget_limiter")
|
||||
except RustCoreUnsupported:
|
||||
return None
|
||||
|
||||
|
||||
async def admit_deployment_turn(
|
||||
core: "EndpointCore", model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
"""`async_admit_deployment_turn`."""
|
||||
from litellm import exceptions as litellm_exceptions
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
async_admit_io_token_turn,
|
||||
without_io_token_limits,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
deployment: Final = _deployment_as_configured_now(core, opened_on)
|
||||
rate_limit_check: Final = next(
|
||||
(callback for callback in core.routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)),
|
||||
None,
|
||||
)
|
||||
io_token_admission: Final[Mapping[str, object]] = (
|
||||
MappingProxyType({})
|
||||
if rate_limit_check is None
|
||||
else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict
|
||||
)
|
||||
budget_limiter: Final[object] = _budget_limiter(core)
|
||||
if budget_limiter is not None:
|
||||
try:
|
||||
await budget_limiter.async_filter_deployments( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # RouterBudgetLimiting
|
||||
model=model_group,
|
||||
healthy_deployments=[dict(deployment)],
|
||||
messages=None,
|
||||
request_kwargs=dict(request_kwargs),
|
||||
)
|
||||
except ValueError as e:
|
||||
raise litellm_exceptions.RateLimitError(
|
||||
message=str(e),
|
||||
llm_provider="",
|
||||
model=model_group,
|
||||
category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT,
|
||||
) from e
|
||||
await _admit_turn_through_routing_strategy(core, model_group, deployment, request_kwargs)
|
||||
await core.router.async_routing_strategy_pre_call_checks(
|
||||
deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None
|
||||
)
|
||||
return io_token_admission
|
||||
|
||||
|
||||
async def _admit_turn_through_routing_strategy(
|
||||
core: "EndpointCore", model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
"""`_admit_turn_through_routing_strategy`: the group's strategy picks among this deployment alone."""
|
||||
from litellm import exceptions as litellm_exceptions
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.router_rust.steps.selection import pick_async
|
||||
from litellm.router_utils.handle_error import async_raise_no_deployment_exception
|
||||
|
||||
routing_kwargs: Final = without(request_kwargs, "litellm_logging_obj")
|
||||
override: Final = routing_kwargs.get("routing_strategy")
|
||||
strategy, scope, group = core.handle.routing_context(model_group, override if isinstance(override, str) else None)
|
||||
selectors: Final = core.selectors
|
||||
selector: Final = selectors.selector_for(strategy, scope, group, routing_kwargs)
|
||||
if strategy == "simple-shuffle" or selector is None:
|
||||
return
|
||||
selected: Final = await pick_async(
|
||||
strategy, selector, model_group, [dict(deployment)], None, None, dict(routing_kwargs)
|
||||
)
|
||||
if selected is None:
|
||||
no_deployment: Final = await async_raise_no_deployment_exception(
|
||||
ctx=RoutingContext(core), # pyright: ignore[reportArgumentType] # the error reads registry() and health()
|
||||
model=model_group,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
raise litellm_exceptions.RateLimitError(
|
||||
message=str(no_deployment),
|
||||
llm_provider="",
|
||||
model=model_group,
|
||||
category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT,
|
||||
) from no_deployment
|
||||
if selectors.is_override(strategy, selector):
|
||||
await selector.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger
|
||||
|
||||
|
||||
__all__ = ["admit_deployment_turn", "turn_admission"]
|
||||
12
litellm/router_rust/steps/README.md
Normal file
12
litellm/router_rust/steps/README.md
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
# steps
|
||||
|
||||
Answers to router-core's host steps that need live Python objects. router-core makes the routing decision; these modules supply the Python-only facts or run the Python-only objects it asks about.
|
||||
|
||||
- `selection`: stateful strategies (latency, tpm/rpm, cost, usage) on `select` steps
|
||||
- `prerouting`: auto, complexity, adaptive and quality routers, and session binding, on `pre_route` steps
|
||||
- `precall`: token counts, context windows and param support for pre-call checks
|
||||
- `plugins`: `Router(plugins=[...])` objects
|
||||
- `callbacks`: the optional pre-call checks (budgets, affinity, rate limits) as `CustomLogger`s
|
||||
- `access`: the caller facts read from auth metadata
|
||||
|
||||
A Rust host answers these steps natively, or declines the config at build.
|
||||
1
litellm/router_rust/steps/__init__.py
Normal file
1
litellm/router_rust/steps/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""Answers to router-core host steps that need live Python objects."""
|
||||
85
litellm/router_rust/steps/access.py
Normal file
85
litellm/router_rust/steps/access.py
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
"""The caller facts the Rust core routes on, read from the request's auth metadata: whether the
|
||||
caller is a proxy admin, the models its key and team allow (which narrow an access-group-only
|
||||
caller to its groups' deployments) and, for a proxy admin, the model the request names."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request, request_user_api_key_auth
|
||||
|
||||
_NAMES: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def _names(value: object) -> tuple[str, ...]:
|
||||
names: Final = _NAMES.validate_python(value) if isinstance(value, list | tuple | set | frozenset) else ()
|
||||
return tuple(str(name) for name in names)
|
||||
|
||||
|
||||
def _admin_model(request: Mapping[str, object]) -> object:
|
||||
"""The model ``filter_team_based_models`` compares team deployments against."""
|
||||
metadata: Final = as_str_mapping(request.get("metadata")) or {}
|
||||
litellm_metadata: Final = as_str_mapping(request.get("litellm_metadata")) or {}
|
||||
return request.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group")
|
||||
|
||||
|
||||
def caller_fields(request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""``proxy_admin``, ``allowed_models`` and ``admin_model`` for the Rust request document."""
|
||||
auth: Final[object] = request_user_api_key_auth(request)
|
||||
if auth is None:
|
||||
return {}
|
||||
allowed: Final = sorted({*_names(getattr(auth, "models", None)), *_names(getattr(auth, "team_models", None))})
|
||||
admin: Final = is_proxy_admin_request(request)
|
||||
model: Final = _admin_model(request) if admin else None
|
||||
return {
|
||||
**({"allowed_models": allowed} if allowed else {}),
|
||||
**({"proxy_admin": True} if admin else {}),
|
||||
**({"admin_model": model} if isinstance(model, str) else {}),
|
||||
}
|
||||
|
||||
|
||||
def _router_weights_error(raw: object) -> BaseException:
|
||||
"""The error ``validate_router_weights`` raises for weights the Rust core refused."""
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.types.router_weights import validate_router_weights
|
||||
|
||||
try:
|
||||
_ = validate_router_weights(raw)
|
||||
except ValidationError as error:
|
||||
return error
|
||||
return RuntimeError(f"router-core refused router weights the validator accepts: {raw!r}")
|
||||
|
||||
|
||||
def selection_refusal(step: Mapping[str, object], model: str, request: Mapping[str, object]) -> BaseException | None:
|
||||
"""The error for a pick the Rust core refused over who the caller is or what it asked for:
|
||||
a public name several teams share, deployments reserved for other teams, invalid request
|
||||
router weights or a non-numeric deployment weight. None for any other failure."""
|
||||
import litellm.exceptions as litellm_exceptions
|
||||
|
||||
kind: Final = step.get("error_kind")
|
||||
named: Final = str(step.get("error_model") or model)
|
||||
if kind == "multiple_teams":
|
||||
return litellm_exceptions.BadRequestError(
|
||||
message=(
|
||||
f"Model name '{named}' matches deployments from multiple teams. "
|
||||
"Specify the deployment ID directly to disambiguate."
|
||||
),
|
||||
model=named,
|
||||
llm_provider="",
|
||||
)
|
||||
if kind == "reserved":
|
||||
return litellm_exceptions.BadRequestError(
|
||||
message=f"Deployment {named} is reserved for another team until {step.get('error_detail')}",
|
||||
model=named,
|
||||
llm_provider="",
|
||||
)
|
||||
if kind == "invalid_router_weights":
|
||||
return _router_weights_error(request.get("_router_weights"))
|
||||
if kind == "non_numeric_weight":
|
||||
return TypeError(f"Deployment {step.get('error_detail')} must be numeric")
|
||||
return None
|
||||
650
litellm/router_rust/steps/callbacks.py
Normal file
650
litellm/router_rust/steps/callbacks.py
Normal file
|
|
@ -0,0 +1,650 @@
|
|||
"""
|
||||
Host-side optional pre-call checks for ``RustCore``: the ``CustomLogger`` checks a router builds
|
||||
from ``optional_pre_call_checks``, ``provider_budget_config``, per-deployment budgets and
|
||||
``model_group_affinity_config``, registered in the host's ``callbacks`` list exactly as the Python
|
||||
router's ``add_optional_pre_call_checks`` registers them, since they are features users configure.
|
||||
|
||||
The Rust core asks the host to run every routing callback's ``async_filter_deployments`` through a
|
||||
``callback_filters`` select stage (``async_callback_filter_deployments``) and decides itself what
|
||||
the surviving deployments or the raised error mean: a pick with nothing left, a retry, a fallback.
|
||||
|
||||
This module also keeps the process-wide holder rules (``litellm/router_utils/global_callback_holders``):
|
||||
a retiring router hands the checks it vacates to a survivor, a runtime removal keeps a check
|
||||
another live router still has, and ``observe(request_scoped=True)`` takes every routing callback
|
||||
off the host lists while the router keeps running them (``_routing_callbacks``). It also answers
|
||||
which outcome callbacks the Python router would still have running for a request
|
||||
(``OptionalChecks.outcome_callbacks``), which router-core records by. It never reaches a Python
|
||||
router core.
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from weakref import WeakSet
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import PROMPT_CACHE_PIN_CACHE_MAX_SIZE
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.routing_common_utils import get_request_team_id
|
||||
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
|
||||
from litellm.router_utils.global_callback_holders import (
|
||||
live_check_owners,
|
||||
routers_holding_global_callbacks,
|
||||
routers_holding_global_callbacks_lock,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
||||
DeploymentAffinityCheck,
|
||||
completed_affinity_response,
|
||||
warn_on_unknown_model_group_affinity_flags,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck
|
||||
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
|
||||
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.router_python_host import PythonHost
|
||||
from litellm.router_rust.steps.selection import StrategySelectors
|
||||
from litellm.router_rust.surface.registry import HostRegistry
|
||||
from litellm.router_utils.cooldown_cache import CooldownCache
|
||||
from litellm.router_utils.global_callback_holders import GlobalCallbackHolder
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
_AFFINITY_CHECKS: Final = (
|
||||
"deployment_affinity",
|
||||
"responses_api_deployment_check",
|
||||
"session_affinity",
|
||||
"encrypted_content_affinity",
|
||||
)
|
||||
# `_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS`: the checks a runtime list drops when it leaves them out.
|
||||
_RUNTIME_TOGGLEABLE: Final[Mapping[str, type[CustomLogger]]] = {
|
||||
"prompt_caching": PromptCachingDeploymentCheck,
|
||||
"enforce_model_rate_limits": ModelRateLimitingCheck,
|
||||
"encrypted_content_affinity": EncryptedContentAffinityCheck,
|
||||
}
|
||||
|
||||
|
||||
def _callbacks(host: "PythonHost") -> list[object]: # mutable-ok: the process list itself
|
||||
return host.callback_list("callbacks") # pyright: ignore[reportArgumentType] # not in the CallbackList literal
|
||||
|
||||
|
||||
def register_router_selector(host: "PythonHost", selector: CustomLogger) -> None:
|
||||
"""``_register_router_selector``: least-busy also counts calls as they start."""
|
||||
if isinstance(selector, LeastBusyLoggingHandler):
|
||||
if isinstance(host.callback_list("input_callback"), list): # pyright: ignore[reportUnnecessaryIsInstance] # a process list a user may have reassigned
|
||||
host.callbacks.add_litellm_input_callback(selector)
|
||||
else:
|
||||
host.set_callback_list("input_callback", [selector])
|
||||
if isinstance(_callbacks(host), list): # pyright: ignore[reportUnnecessaryIsInstance] # see above
|
||||
host.callbacks.add_litellm_callback(selector)
|
||||
|
||||
|
||||
def _move_before_deployment_affinity(
|
||||
callback_list: list[object], # mutable-ok: reordered in place
|
||||
callback: object,
|
||||
) -> None:
|
||||
if callback not in callback_list:
|
||||
return
|
||||
callback_list.remove(callback)
|
||||
index: Final = next(
|
||||
(i for i, entry in enumerate(callback_list) if isinstance(entry, DeploymentAffinityCheck)),
|
||||
len(callback_list),
|
||||
)
|
||||
callback_list.insert(index, callback)
|
||||
|
||||
|
||||
class OptionalChecks:
|
||||
"""One router's optional pre-call checks (``optional_callbacks``) and the stores they read."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
owner: "GlobalCallbackHolder",
|
||||
host: "PythonHost",
|
||||
cache: DualCache,
|
||||
hosted: "HostRegistry",
|
||||
ctx: object,
|
||||
arguments: Mapping[str, object],
|
||||
) -> None:
|
||||
self._owner: Final = owner
|
||||
self._host: Final = host
|
||||
self._cache: Final = cache
|
||||
self._hosted: Final = hosted
|
||||
self._ctx: Final = ctx
|
||||
self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it
|
||||
self.router_budget_logger: RouterBudgetLimiting | None = None
|
||||
self.provider_budget_config: object = arguments.get("provider_budget_config")
|
||||
ttl: Final = arguments.get("deployment_affinity_ttl_seconds")
|
||||
self.deployment_affinity_ttl_seconds: object = 3600 if ttl is None else ttl
|
||||
self.session_affinity_repin_on_failover: object = bool(arguments.get("session_affinity_repin_on_failover"))
|
||||
affinity_config: Final = arguments.get("model_group_affinity_config")
|
||||
self.model_group_affinity_config: dict[str, list[str]] | None = ( # mutable-ok: the argument as passed
|
||||
affinity_config if isinstance(affinity_config, dict) else None # pyright: ignore[reportAttributeAccessIssue] # the argument as passed
|
||||
)
|
||||
warn_on_unknown_model_group_affinity_flags(self.model_group_affinity_config)
|
||||
self.prompt_caching_pin_cache: DualCache = DualCache(
|
||||
redis_cache=cache.redis_cache,
|
||||
in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE),
|
||||
)
|
||||
self._request_scoped_selectors: tuple[CustomLogger, ...] = ()
|
||||
self._request_scoped_policies: tuple[CustomLogger, ...] = ()
|
||||
# ``_slack_alerting_loggers``: the Slack alerting logger ``alerting_config`` builds, if any.
|
||||
self.slack_alerting_loggers: tuple[CustomLogger, ...] = ()
|
||||
# Whether the outcome callbacks would still be on the host lists (``_release_global_callbacks``).
|
||||
self._outcome_callbacks_listed: bool = True
|
||||
# The logging objects ``_core_bind_to_request_logging`` would have bound them to.
|
||||
self._bound_requests: Final[WeakSet[Logging]] = WeakSet()
|
||||
|
||||
# -- construction ---------------------------------------------------------------------------
|
||||
|
||||
def build(self, checks: object, model_list: object) -> None:
|
||||
"""The constructor's checks: a budget appends ``router_budget_limiting`` to the list given (in
|
||||
place, as Python does), the list is added, a per-group affinity config adds the affinity
|
||||
check, and ``apply_default_settings`` adds the empty default list."""
|
||||
listed: list[object] | None = checks if isinstance(checks, list) else None # pyright: ignore[reportUnknownVariableType] # mutable-ok: appended in place, as Python does # rebind-ok: the argument as passed
|
||||
if RouterBudgetLimiting.should_init_router_budget_limiter( # pyright: ignore[reportUnknownMemberType] # untyped model_list
|
||||
model_list=model_list if isinstance(model_list, list) else None, # pyright: ignore[reportUnknownArgumentType] # the argument as passed
|
||||
provider_budget_config=self.provider_budget_config, # pyright: ignore[reportArgumentType] # the argument as passed
|
||||
):
|
||||
if listed is not None:
|
||||
listed.append("router_budget_limiting")
|
||||
else:
|
||||
listed = ["router_budget_limiting"] # rebind-ok: the list Python builds
|
||||
if listed is not None:
|
||||
self.add(listed) # pyright: ignore[reportUnknownArgumentType] # names as given
|
||||
if self.model_group_affinity_config:
|
||||
self.ensure_deployment_affinity_callback()
|
||||
self.add([])
|
||||
|
||||
def _is_priority_group(self, model: str) -> bool:
|
||||
"""``_is_priority_routing_group``."""
|
||||
hosted: Final = self._hosted
|
||||
resolved: Final = hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # the host registry's alias lookup
|
||||
group: Final = hosted.get_routing_group(resolved)
|
||||
return group is not None and group.routing_strategy == "priority"
|
||||
|
||||
async def repin_session_after_success(
|
||||
self, litellm_params: Mapping[str, object], response: object, cooldowns: "CooldownCache"
|
||||
) -> None:
|
||||
"""``_repin_session_after_success``: with ``session_affinity_repin_on_failover`` on, a completed
|
||||
success moves the session's pin to the deployment that served it, among the group's
|
||||
deployments not cooling down."""
|
||||
if not self.session_affinity_repin_on_failover:
|
||||
return
|
||||
try:
|
||||
if not completed_affinity_response(response):
|
||||
return
|
||||
callback: Final = next(
|
||||
(logger for logger in self.optional_callbacks or () if isinstance(logger, DeploymentAffinityCheck)),
|
||||
None,
|
||||
)
|
||||
if callback is None:
|
||||
return
|
||||
model_group: Final = callback.get_repin_model_group(litellm_params)
|
||||
if model_group is None or self._is_priority_group(model_group):
|
||||
return
|
||||
registered_ids: Final = self._hosted._core_get_candidate_model_ids_for_route( # pyright: ignore[reportPrivateUsage] # the host registry's candidate lookup
|
||||
model=model_group, team_id=get_request_team_id(litellm_params)
|
||||
)
|
||||
cooldown_ids: Final = frozenset(
|
||||
model_id
|
||||
for model_id, _ in await cooldowns.async_get_active_cooldowns(
|
||||
model_ids=list(registered_ids), parent_otel_span=None
|
||||
)
|
||||
)
|
||||
await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids)
|
||||
except Exception as error: # noqa: BLE001 # optional pin updates must not interrupt success accounting
|
||||
verbose_router_logger.debug("Session re-pinning failed: %s", error)
|
||||
|
||||
def _affinity_check(self, user_key: bool, responses_api: bool, session_id: bool) -> DeploymentAffinityCheck:
|
||||
return DeploymentAffinityCheck(
|
||||
cache=self._cache,
|
||||
ttl_seconds=self.deployment_affinity_ttl_seconds, # pyright: ignore[reportArgumentType] # the argument as passed
|
||||
enable_user_key_affinity=user_key,
|
||||
enable_responses_api_affinity=responses_api,
|
||||
enable_session_id_affinity=session_id,
|
||||
model_group_affinity_config=self.model_group_affinity_config,
|
||||
is_priority_group=self._is_priority_group,
|
||||
session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, # pyright: ignore[reportArgumentType] # same
|
||||
)
|
||||
|
||||
def _own(self, callback: CustomLogger) -> None:
|
||||
if self.optional_callbacks is None:
|
||||
self.optional_callbacks = []
|
||||
self.optional_callbacks.append(callback)
|
||||
self._host.callbacks.add_litellm_callback(callback)
|
||||
|
||||
def ensure_deployment_affinity_callback(self) -> None:
|
||||
"""``_ensure_deployment_affinity_callback``: the affinity check with every global flag off."""
|
||||
if any(isinstance(cb, DeploymentAffinityCheck) for cb in self.optional_callbacks or ()):
|
||||
return
|
||||
self._own(self._affinity_check(False, False, False))
|
||||
|
||||
def _add_affinity(self, checks: Sequence[object]) -> None:
|
||||
user_key: Final = "deployment_affinity" in checks
|
||||
responses_api: Final = "responses_api_deployment_check" in checks
|
||||
session_id: Final = "session_affinity" in checks
|
||||
if not (user_key or responses_api or session_id):
|
||||
return
|
||||
if self.optional_callbacks is None:
|
||||
self.optional_callbacks = []
|
||||
existing: Final = next(
|
||||
(cb for cb in self.optional_callbacks if isinstance(cb, DeploymentAffinityCheck)),
|
||||
None,
|
||||
)
|
||||
if existing is None:
|
||||
self._own(self._affinity_check(user_key, responses_api, session_id))
|
||||
return
|
||||
existing.enable_user_key_affinity = existing.enable_user_key_affinity or user_key
|
||||
existing.enable_responses_api_affinity = existing.enable_responses_api_affinity or responses_api
|
||||
existing.enable_session_id_affinity = existing.enable_session_id_affinity or session_id
|
||||
existing.ttl_seconds = self.deployment_affinity_ttl_seconds # pyright: ignore[reportAttributeAccessIssue] # the argument as passed
|
||||
if self.model_group_affinity_config:
|
||||
existing.model_group_affinity_config = self.model_group_affinity_config
|
||||
|
||||
def _add_encrypted_content_affinity(self, enable_global_affinity: bool) -> None:
|
||||
"""``_add_encrypted_content_affinity_check``."""
|
||||
if not (
|
||||
enable_global_affinity
|
||||
or EncryptedContentAffinityCheck.has_model_group_affinity_enabled(self.model_group_affinity_config)
|
||||
):
|
||||
return
|
||||
if self.optional_callbacks is None:
|
||||
self.optional_callbacks = []
|
||||
existing: Final = next(
|
||||
(cb for cb in self.optional_callbacks if isinstance(cb, EncryptedContentAffinityCheck)), None
|
||||
)
|
||||
if existing is not None:
|
||||
existing.ctx = self._ctx # pyright: ignore[reportAttributeAccessIssue] # the core's RoutingContext serves registry()
|
||||
existing.enable_global_affinity = existing.enable_global_affinity or enable_global_affinity
|
||||
existing.model_group_affinity_config = self.model_group_affinity_config or {}
|
||||
check: EncryptedContentAffinityCheck = existing # rebind-ok: the check to move
|
||||
else:
|
||||
check = EncryptedContentAffinityCheck( # rebind-ok: the check to move
|
||||
ctx=self._ctx, # pyright: ignore[reportArgumentType] # see above
|
||||
enable_global_affinity=enable_global_affinity,
|
||||
model_group_affinity_config=self.model_group_affinity_config,
|
||||
)
|
||||
self._own(check)
|
||||
_move_before_deployment_affinity(self.optional_callbacks, check)
|
||||
_move_before_deployment_affinity(_callbacks(self._host), check)
|
||||
|
||||
def add(self, checks: Sequence[object] | None) -> None:
|
||||
"""``add_optional_pre_call_checks``."""
|
||||
if checks is None:
|
||||
return
|
||||
self._add_affinity(checks)
|
||||
self._add_encrypted_content_affinity("encrypted_content_affinity" in checks)
|
||||
for name in checks:
|
||||
if name in _AFFINITY_CHECKS:
|
||||
continue
|
||||
callback: CustomLogger | None = None
|
||||
if name == "prompt_caching":
|
||||
callback = PromptCachingDeploymentCheck(
|
||||
cache=self.prompt_caching_pin_cache, is_priority_group=self._is_priority_group
|
||||
)
|
||||
elif name == "router_budget_limiting":
|
||||
if self.budget_limiter() is not None:
|
||||
continue
|
||||
callback = RouterBudgetLimiting(
|
||||
dual_cache=self._cache,
|
||||
provider_budget_config=self.provider_budget_config, # pyright: ignore[reportArgumentType] # the argument as passed
|
||||
model_list=self._hosted._core_model_list, # pyright: ignore[reportPrivateUsage, reportArgumentType] # the live model_list
|
||||
)
|
||||
self.router_budget_logger = callback
|
||||
elif name == "enforce_model_rate_limits":
|
||||
callback = ModelRateLimitingCheck(dual_cache=self._cache)
|
||||
if callback is None:
|
||||
continue
|
||||
if self.optional_callbacks is not None and any(
|
||||
isinstance(existing, type(callback)) for existing in self.optional_callbacks
|
||||
):
|
||||
continue
|
||||
self._own(callback)
|
||||
|
||||
def set(self, checks: object) -> None:
|
||||
"""``set_optional_pre_call_checks``: drop the toggleable checks the list leaves out, then add it."""
|
||||
if checks is None:
|
||||
return
|
||||
listed: Final[tuple[object, ...]] = tuple(checks) if isinstance(checks, list | tuple) else () # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # names as given
|
||||
requested: Final = frozenset(listed)
|
||||
for name, callback_cls in _RUNTIME_TOGGLEABLE.items():
|
||||
if name not in requested:
|
||||
self._remove_of_type(callback_cls)
|
||||
self.add(listed)
|
||||
|
||||
def _remove_of_type(self, callback_cls: type[CustomLogger]) -> None:
|
||||
"""``_remove_optional_callbacks_of_type``: kept on the host while any other live router of
|
||||
the process has a check of that type."""
|
||||
if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks):
|
||||
return
|
||||
self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls]
|
||||
others: Final = tuple(owner for owner in live_check_owners if id(owner) != id(self._owner))
|
||||
if any(any(type(cb) is callback_cls for cb in (owner.optional_callbacks or ())) for owner in others):
|
||||
return
|
||||
callbacks: Final = _callbacks(self._host)
|
||||
for cb in tuple(callbacks):
|
||||
if type(cb) is callback_cls:
|
||||
self._host.callbacks.remove_callback_from_list_by_object(callbacks, cb, require_self=False) # pyright: ignore[reportUnknownMemberType] # the manager's untyped helper
|
||||
|
||||
# -- budgets --------------------------------------------------------------------------------
|
||||
|
||||
def budget_limiter(self) -> RouterBudgetLimiting | None:
|
||||
"""``_core_get_router_deployment_budget_limiter``: exact type, so a proxy subclass is not taken."""
|
||||
if self.router_budget_logger is not None:
|
||||
return self.router_budget_logger
|
||||
for cb in self.optional_callbacks or ():
|
||||
if isinstance(cb, RouterBudgetLimiting) and type(cb) is RouterBudgetLimiting:
|
||||
self.router_budget_logger = cb
|
||||
return cb
|
||||
return None
|
||||
|
||||
def sync_deployment_budget(self, deployment: "Deployment") -> None:
|
||||
"""``_sync_deployment_budget_config``, after a deployment is added."""
|
||||
model_id: Final = deployment.model_info.id
|
||||
if model_id is None:
|
||||
return
|
||||
limiter: RouterBudgetLimiting | None = self.budget_limiter() # rebind-ok: built when first needed
|
||||
params: Final = deployment.litellm_params
|
||||
if params.max_budget is None or params.budget_duration is None:
|
||||
if limiter is not None:
|
||||
limiter.unregister_deployment_budget(model_id=model_id)
|
||||
return
|
||||
if limiter is None:
|
||||
self.add(["router_budget_limiting"])
|
||||
limiter = self.budget_limiter() # rebind-ok: the one just built
|
||||
if limiter is not None:
|
||||
limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) # pyright: ignore[reportUnknownMemberType] # pydantic's untyped dump
|
||||
|
||||
def unregister_deployment_budget(self, model_id: str) -> None:
|
||||
limiter: Final = self.budget_limiter()
|
||||
if limiter is not None:
|
||||
limiter.unregister_deployment_budget(model_id=model_id)
|
||||
|
||||
# -- the routing callbacks ------------------------------------------------------------------
|
||||
|
||||
def loggers(self) -> tuple[CustomLogger, ...]:
|
||||
"""``_optional_loggers``."""
|
||||
return tuple(cb for cb in self.optional_callbacks or () if isinstance(cb, CustomLogger))
|
||||
|
||||
def _policies_not_set_globally(self) -> tuple[CustomLogger, ...]:
|
||||
keys: Final = self._host.callbacks.get_custom_logger_key
|
||||
global_keys: Final = frozenset(keys(cb) for cb in _callbacks(self._host) if isinstance(cb, CustomLogger))
|
||||
return tuple(policy for policy in self._request_scoped_policies if keys(policy) not in global_keys)
|
||||
|
||||
def bind_to_request_logging(self, logging_obj: "Logging") -> None:
|
||||
"""``_core_bind_to_request_logging``'s routing half: the request-scoped selectors replace the
|
||||
host's same-keyed callbacks for the request, and the policies the host does not run are
|
||||
added to it. The outcome callbacks are not bound: the Rust core records every attempt's
|
||||
outcome itself, so the request is noted for ``outcome_callbacks`` instead."""
|
||||
self._bound_requests.add(logging_obj)
|
||||
for selector in self._request_scoped_selectors:
|
||||
logging_obj.replace_global_callback(selector)
|
||||
for policy in self._policies_not_set_globally():
|
||||
logging_obj.add_dynamic_callback(policy)
|
||||
|
||||
def outcome_callbacks(
|
||||
self, request: Mapping[str, object], *, sync: bool, make_call: bool = True
|
||||
) -> Mapping[str, bool]:
|
||||
"""Router-core's ``OutcomeCallbacks`` for ``request``, empty (its default) while the outcome
|
||||
callbacks are listed. Past that only the request's bound logging object runs them, and
|
||||
``_core_bind_to_request_logging`` binds neither ``sync_deployment_callback_on_success`` nor any
|
||||
success callback a sync call runs. ``make_call``: the attempt runs through ``make_call``."""
|
||||
if self._outcome_callbacks_listed:
|
||||
return {}
|
||||
bound: Final = request.get("litellm_logging_obj") in self._bound_requests
|
||||
return {"success": bound and not sync, "failure": bound, "make_call": make_call}
|
||||
|
||||
def routing_callbacks(self) -> Iterable[object]:
|
||||
"""``_routing_callbacks``: the host's callbacks, or after ``observe(request_scoped=True)`` the
|
||||
host's callbacks with the router's own selectors and policies in place of theirs."""
|
||||
if not (self._request_scoped_selectors or self._request_scoped_policies):
|
||||
return tuple(_callbacks(self._host))
|
||||
keys: Final = self._host.callbacks.get_custom_logger_key
|
||||
replaced: Final = frozenset(map(keys, self._request_scoped_selectors))
|
||||
return (
|
||||
*self._host.callbacks.without_custom_logger_keys(_callbacks(self._host), replaced), # pyright: ignore[reportArgumentType] # the process list holds any callback
|
||||
*self._request_scoped_selectors,
|
||||
*self._policies_not_set_globally(),
|
||||
)
|
||||
|
||||
async def filter_deployments(
|
||||
self,
|
||||
model: str,
|
||||
deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract
|
||||
messages: object,
|
||||
request: Mapping[str, object],
|
||||
) -> list[object]: # mutable-ok: as the filters return it
|
||||
"""``async_callback_filter_deployments``: each routing callback narrows the deployments in
|
||||
turn; a filter's error is raised for the Rust core to decide on."""
|
||||
remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed
|
||||
base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts
|
||||
request_kwargs: Final = dict(request)
|
||||
for callback in tuple(self.routing_callbacks()):
|
||||
# The base filter returns its input unchanged: only overriding filters are run.
|
||||
if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above
|
||||
continue
|
||||
remaining = await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportAssignmentType, reportUnknownVariableType] # see above
|
||||
model=model,
|
||||
healthy_deployments=remaining,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given
|
||||
request_kwargs=request_kwargs,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
return remaining
|
||||
|
||||
# -- release --------------------------------------------------------------------------------
|
||||
|
||||
def serve_through_request_logging(self, selectors: Sequence[CustomLogger], dropped: Sequence[object]) -> None:
|
||||
"""``_core_serve_through_request_logging``: keep running ``selectors`` (and request-scoped
|
||||
adaptive hooks) and the optional checks, but take everything off the host lists."""
|
||||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||||
|
||||
adaptive: Final = tuple(
|
||||
AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) # pyright: ignore[reportArgumentType] # an adaptive router
|
||||
for tagged in chain.from_iterable(self._hosted._core_adaptive_routers.values()) # pyright: ignore[reportPrivateUsage] # the host registry's live registry
|
||||
)
|
||||
self._request_scoped_selectors = (*selectors, *adaptive)
|
||||
self._request_scoped_policies = (*self.loggers(), *self.slack_alerting_loggers)
|
||||
self.release(dropped)
|
||||
|
||||
def own_adaptive_hooks(self) -> tuple[CustomLogger, ...]:
|
||||
"""``_adaptive_router_hooks``: the host's adaptive post-call hooks of this router's routers."""
|
||||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||||
|
||||
own: Final = frozenset(
|
||||
id(tagged.strategy)
|
||||
for tagged in chain.from_iterable(self._hosted._core_adaptive_routers.values()) # pyright: ignore[reportPrivateUsage] # see above
|
||||
)
|
||||
return tuple(
|
||||
hook
|
||||
for hook in self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
|
||||
if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own
|
||||
)
|
||||
|
||||
def retire_state(self, selectors: Sequence[object]) -> None:
|
||||
"""``retire_routing_state``: stop the background work of the selectors and the budget limiter."""
|
||||
for strategy in (*selectors, *self.loggers()):
|
||||
if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting):
|
||||
strategy.retire()
|
||||
|
||||
def release(self, dropped: Sequence[object]) -> None:
|
||||
"""``_release_global_callbacks``: drop ``dropped``, the optional checks and the adaptive hooks
|
||||
from the host lists, then hand each vacated check to a surviving router's same-keyed one."""
|
||||
self._outcome_callbacks_listed = False
|
||||
with routers_holding_global_callbacks_lock:
|
||||
routers_holding_global_callbacks.discard(self._owner)
|
||||
adaptive: Final = self.own_adaptive_hooks()
|
||||
slack: Final = self.slack_alerting_loggers
|
||||
vacated: Final = self._host.callbacks.drop_callbacks(
|
||||
(*dropped, *(self.optional_callbacks or ()), *slack, *adaptive), method_owners=(self._owner, *slack)
|
||||
)
|
||||
with routers_holding_global_callbacks_lock:
|
||||
survivors: Final = tuple(routers_holding_global_callbacks)
|
||||
candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) # pyright: ignore[reportPrivateUsage] # the holder protocol
|
||||
for candidate in self._host.callbacks.with_custom_logger_keys(candidates, vacated):
|
||||
register_router_selector(self._host, candidate)
|
||||
adaptive_keys: Final = frozenset(map(self._host.callbacks.get_custom_logger_key, adaptive))
|
||||
survivor: Final = next((router for router in survivors if router._core_adaptive_routers), None) # pyright: ignore[reportPrivateUsage] # see above
|
||||
if survivor is not None and vacated & adaptive_keys:
|
||||
survivor._sync_adaptive_router_hooks() # pyright: ignore[reportPrivateUsage] # see above
|
||||
|
||||
|
||||
_ENCRYPTED_PIN: Final = "_encrypted_content_affinity_pinned"
|
||||
|
||||
|
||||
def _request_with_failed_ids(request: Mapping[str, object], step: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The request as the filters read it in the Python router: the Rust walk holds what the Python
|
||||
one keeps in the attempt's kwargs, ``_excluded_deployment_ids`` from a weighted failover hop,
|
||||
``_retry_skipped_deployment_ids`` from the retry loop (``deployment_ids_failed_in_request``) and
|
||||
``_target_order`` from an order fallback."""
|
||||
excluded: Final = tuple(str(id) for id in _ids_of(step.get("excluded_ids")))
|
||||
skipped: Final = tuple(str(id) for id in _ids_of(step.get("skipped_ids")))
|
||||
target_order: Final = step.get("target_order")
|
||||
if not excluded and not skipped and target_order is None:
|
||||
return request
|
||||
return {
|
||||
**request,
|
||||
**({"_excluded_deployment_ids": list(excluded)} if excluded else {}),
|
||||
**({"_retry_skipped_deployment_ids": skipped} if skipped else {}),
|
||||
**({"_target_order": target_order} if target_order is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
def _ids_of(value: object) -> Sequence[object]:
|
||||
return value if isinstance(value, list | tuple) else () # pyright: ignore[reportUnknownVariableType] # the step's id list
|
||||
|
||||
|
||||
def filtered_ids(remaining: Iterable[object], candidates: Sequence[str]) -> list[str]: # mutable-ok: a JSON list
|
||||
"""The ids of the deployments the filters left, in their order, among the candidates."""
|
||||
ids: Final[list[str]] = [] # mutable-ok: built here
|
||||
for deployment in remaining:
|
||||
found = (as_str_mapping((as_str_mapping(deployment) or {}).get("model_info")) or {}).get("id")
|
||||
if isinstance(found, str) and found in candidates and found not in ids:
|
||||
ids.append(found)
|
||||
return ids
|
||||
|
||||
|
||||
class RustOptionalChecks:
|
||||
"""``RustCore`` members for its optional pre-call checks and routing callbacks."""
|
||||
|
||||
# RustCore sets these as Final attributes in its constructor.
|
||||
@property
|
||||
def _opt(self) -> OptionalChecks:
|
||||
checks: Final[OptionalChecks] = self.__dict__["_checks"] # pyright: ignore[reportAny] # RustCore's own attribute
|
||||
return checks
|
||||
|
||||
@property
|
||||
def _sel(self) -> "StrategySelectors":
|
||||
selectors: Final[StrategySelectors] = self.__dict__["_selectors"] # pyright: ignore[reportAny] # same
|
||||
return selectors
|
||||
|
||||
@property
|
||||
def _reg(self) -> "HostRegistry":
|
||||
hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # same
|
||||
return hosted
|
||||
|
||||
@property
|
||||
def optional_callbacks(self) -> "list[object] | None": # mutable-ok: the router's live list
|
||||
return self._opt.optional_callbacks
|
||||
|
||||
@property
|
||||
def router_budget_logger(self) -> RouterBudgetLimiting | None:
|
||||
return self._opt.router_budget_logger
|
||||
|
||||
@property
|
||||
def deployment_affinity_ttl_seconds(self) -> object:
|
||||
return self._opt.deployment_affinity_ttl_seconds
|
||||
|
||||
@property
|
||||
def session_affinity_repin_on_failover(self) -> object:
|
||||
return self._opt.session_affinity_repin_on_failover
|
||||
|
||||
@property
|
||||
def model_group_affinity_config(self) -> "dict[str, list[str]] | None": # mutable-ok: the argument as passed
|
||||
return self._opt.model_group_affinity_config
|
||||
|
||||
def add_optional_pre_call_checks(self, optional_pre_call_checks: "Sequence[object] | None") -> None:
|
||||
self._opt.add(optional_pre_call_checks)
|
||||
|
||||
def apply_default_settings(self) -> None:
|
||||
"""Add the (empty) default checks; an enabled continuation registers the prefill check."""
|
||||
from litellm.router_rust.invoke.streaming import register_continuation_prefill_check
|
||||
|
||||
self._opt.add([])
|
||||
if getattr(self, "enable_mid_stream_fallback_continuation", False):
|
||||
register_continuation_prefill_check(self._opt._host) # pyright: ignore[reportPrivateUsage] # the checks' host
|
||||
|
||||
# The members another router's retire reads (`global_callback_holders`).
|
||||
|
||||
def _global_callback_candidates(self) -> tuple[CustomLogger, ...]:
|
||||
return (*self._sel.callback_selectors(), *self._opt.loggers(), *self._opt.slack_alerting_loggers)
|
||||
|
||||
@property
|
||||
def _core_adaptive_routers(self) -> Mapping[str, object]:
|
||||
return self._reg._core_adaptive_routers # pyright: ignore[reportPrivateUsage] # the host registry's live registry
|
||||
|
||||
def _sync_adaptive_router_hooks(self) -> None:
|
||||
strategies: Final = self._reg.strategies
|
||||
if strategies is not None:
|
||||
strategies._sync_adaptive_router_hooks() # pyright: ignore[reportPrivateUsage] # the strategies' own hook sync
|
||||
|
||||
def _join_callback_holders(self) -> None:
|
||||
with routers_holding_global_callbacks_lock:
|
||||
routers_holding_global_callbacks.add(self)
|
||||
live_check_owners.add(self)
|
||||
|
||||
def outcome_callbacks(
|
||||
self, request: Mapping[str, object], *, sync: bool, make_call: bool = True
|
||||
) -> Mapping[str, bool]:
|
||||
return self._opt.outcome_callbacks(request, sync=sync, make_call=make_call)
|
||||
|
||||
def _serve_through_request_logging(self) -> None:
|
||||
selectors: Final = self._sel
|
||||
self._opt.serve_through_request_logging(
|
||||
selectors.callback_selectors(), (*selectors.callback_selectors(), *selectors.override_selectors())
|
||||
)
|
||||
|
||||
def _retire_callbacks(self) -> None:
|
||||
"""``_core_discard``'s callback half: retire the routing state, release the callbacks. The
|
||||
selectors stay: a retired router still picks with them."""
|
||||
live_check_owners.discard(self)
|
||||
selectors: Final = (*self._sel.callback_selectors(), *self._sel.override_selectors())
|
||||
self._opt.retire_state(selectors)
|
||||
self._opt.release(selectors)
|
||||
|
||||
async def _answer_callback_filters(
|
||||
self, step: Mapping[str, object], request: Mapping[str, object], messages: object
|
||||
) -> tuple[object, BaseException | None]:
|
||||
"""A ``callback_filters`` step: the remaining ids, or a filter's error as an attempt
|
||||
failure with the error itself, which the caller keeps to raise."""
|
||||
from litellm.router_rust.core import attempt_failure
|
||||
|
||||
candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list
|
||||
records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts
|
||||
narrowed: Final = _request_with_failed_ids(request, step)
|
||||
try:
|
||||
remaining: Final = await self._opt.filter_deployments(
|
||||
str(step.get("model_group")),
|
||||
records, # pyright: ignore[reportUnknownArgumentType] # see above
|
||||
messages,
|
||||
narrowed,
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # handed to the Rust core as the try's failure
|
||||
return attempt_failure(error), error
|
||||
# A filter's pin lands on the request kwargs the pick goes on to read.
|
||||
if narrowed is not request and narrowed.get(_ENCRYPTED_PIN) and isinstance(request, dict):
|
||||
request[_ENCRYPTED_PIN] = True # rebind-ok: the request dict the pick goes on to read
|
||||
return filtered_ids(remaining, candidates), None # pyright: ignore[reportUnknownArgumentType] # see above
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OptionalChecks",
|
||||
"RustOptionalChecks",
|
||||
"filtered_ids",
|
||||
"register_router_selector",
|
||||
]
|
||||
87
litellm/router_rust/steps/plugins.py
Normal file
87
litellm/router_rust/steps/plugins.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""
|
||||
Host side of the Rust core's routing-plugin pipeline (``Router(plugins=[...])``).
|
||||
|
||||
The Rust core decides where the pipeline applies (only the async pick runs it; the sync pick
|
||||
refuses a router with plugins) and which deployments the plugins' candidate models keep
|
||||
(``router-core/src/routing_plugins.rs``). The plugins themselves are live Python objects, so this
|
||||
module runs them on a ``routing_plugins`` select step, the way the Python router core's
|
||||
``_run_routing_plugins`` does, and reports the candidate models left on the request.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Protocol
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs
|
||||
from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias
|
||||
from litellm.types.router import RoutingContext, RoutingPlugin
|
||||
|
||||
SYNC_PATH_MESSAGE: Final = (
|
||||
"Router(plugins=[...]) is configured, but this call resolved to the synchronous "
|
||||
"deployment-selection path, which never runs the routing-plugin pipeline. This "
|
||||
"happens for sync Router methods (e.g. Router.completion()) and for async calls "
|
||||
"with a routing_strategy that has no async-native selector (e.g. legacy "
|
||||
"'usage-based-routing', v1). Silently skipping "
|
||||
"configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. "
|
||||
"Use an async Router method with a supported routing_strategy (simple-shuffle, "
|
||||
"usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), "
|
||||
"or remove `plugins` from the Router config."
|
||||
)
|
||||
_CANDIDATES_KEY: Final = "_routing_plugin_candidate_models"
|
||||
|
||||
|
||||
class PluginRegistry(Protocol):
|
||||
def _core_get_model_list(
|
||||
self, model_name: str | None = None, team_id: str | None = None
|
||||
) -> Sequence[Mapping[str, object]] | None: ...
|
||||
|
||||
|
||||
def _resolved_litellm_models(
|
||||
registry: PluginRegistry, model: str
|
||||
) -> list[str]: # mutable-ok: the plugins' candidate list
|
||||
"""``_core_resolved_litellm_models``: the provider models a call to ``model`` is served by."""
|
||||
return [
|
||||
litellm_model
|
||||
for deployment in registry._core_get_model_list(model_name=model) or () # pyright: ignore[reportPrivateUsage] # the host registry's listing
|
||||
if isinstance(litellm_model := (as_str_mapping(deployment.get("litellm_params")) or {}).get("model"), str)
|
||||
and litellm_model
|
||||
]
|
||||
|
||||
|
||||
def candidate_models(request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The answer to a ``routing_plugins`` step: the candidate models the plugins left on the
|
||||
request, ``None`` when none narrowed it."""
|
||||
metadata: Final = as_str_mapping(request.get(get_metadata_variable_name_from_kwargs(dict(request))))
|
||||
return {"models": metadata.get(_CANDIDATES_KEY) if metadata is not None else None}
|
||||
|
||||
|
||||
async def run_plugins(
|
||||
plugins: Sequence[RoutingPlugin],
|
||||
registry: PluginRegistry,
|
||||
model_group_alias: Mapping[str, object],
|
||||
model: str,
|
||||
request: dict[str, object], # mutable-ok: the plugins stamp their outcome on the caller's request
|
||||
messages: object,
|
||||
) -> None:
|
||||
"""``_run_routing_plugins``: run ``plugins`` in order over ``model``'s candidate models, then
|
||||
stamp their signals (and a narrowed candidate list) on the request's metadata."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
|
||||
|
||||
registered: Final = resolve_model_group_alias(model_group_alias, model) or model
|
||||
models: Final = _resolved_litellm_models(registry, registered)
|
||||
metadata_key: Final = get_metadata_variable_name_from_kwargs(request)
|
||||
metadata: Final = request.setdefault(metadata_key, {})
|
||||
raw_messages: Final[object] = messages or list[object]()
|
||||
context = RoutingContext( # rebind-ok: each plugin returns the context the next one reads
|
||||
raw_messages=raw_messages, # pyright: ignore[reportArgumentType] # the request's messages
|
||||
structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request) or [], # pyright: ignore[reportArgumentType] # the request's messages
|
||||
candidate_models=list(models),
|
||||
metadata=metadata, # pyright: ignore[reportArgumentType] # the request's metadata dict
|
||||
)
|
||||
for plugin in plugins:
|
||||
context = await plugin.run(context)
|
||||
metadata["routing_plugin_signals"] = context.signals # pyright: ignore[reportIndexIssue] # the request's metadata dict
|
||||
if len(context.candidate_models) < len(models):
|
||||
metadata[_CANDIDATES_KEY] = context.candidate_models # pyright: ignore[reportIndexIssue] # see above
|
||||
|
||||
|
||||
__all__ = ["SYNC_PATH_MESSAGE", "PluginRegistry", "candidate_models", "run_plugins"]
|
||||
494
litellm/router_rust/steps/precall.py
Normal file
494
litellm/router_rust/steps/precall.py
Normal file
|
|
@ -0,0 +1,494 @@
|
|||
"""
|
||||
Host-side facts for the Rust core's ``enable_pre_call_checks`` filter.
|
||||
|
||||
The Rust core decides which deployments the pre-call checks drop and which error a group with
|
||||
nothing left raises (``router-core/src/precall.rs``). What needs live Python objects is answered
|
||||
here, on a ``pre_call`` select step: the prompt's token count (messages, Responses input and
|
||||
instructions, tools and an Anthropic ``system`` block), each deployment's context window from
|
||||
its router model info, and which requested params each deployment takes. Ported from the Python
|
||||
router core's ``_pre_call_checks`` helpers; it never reaches a Python router core.
|
||||
|
||||
``pre_call_checks`` is the one decision made here: ``RoutingContext.pre_call_checks`` hands over
|
||||
deployment records rather than a step, and those need not be in the Rust registry.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
|
||||
from litellm.litellm_core_utils.token_counter import (
|
||||
offload_token_count,
|
||||
served_openai_models,
|
||||
single_encoding_openai_model,
|
||||
split_embedding_input_tokens,
|
||||
)
|
||||
from litellm.router_python_host import python_host
|
||||
from litellm.router_strategy.complexity_router.context_compaction import compaction_pending
|
||||
from litellm.types.router import LiteLLM_Params
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_python_host import PythonHost
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
# The chat params a deployment must take when the request sends them (`_PRE_CALL_CHECK_FILTERED_PARAMS`).
|
||||
_FILTERED_PARAMS: Final = frozenset({"response_format"})
|
||||
|
||||
|
||||
class PreCallRegistry(Protocol):
|
||||
"""The host registry reads the facts need."""
|
||||
|
||||
def _core_get_model_info(self, id: str) -> Mapping[str, object] | None: ... # the registry's name
|
||||
|
||||
def _core_deployments_named(self, model_name: str) -> Sequence[Mapping[str, object]] | None: ...
|
||||
|
||||
def _core_get_router_model_info(
|
||||
self, deployment: Mapping[str, object], received_model_name: str
|
||||
) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
def count_tokens(
|
||||
host: "PythonHost",
|
||||
messages: object,
|
||||
input: object, # the request field's name
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
tokenizer_model: str,
|
||||
deployments: Sequence[Mapping[str, object]] = (),
|
||||
) -> int:
|
||||
"""``_count_pre_call_check_tokens``: messages, or Responses input normalized to messages, plus
|
||||
the tools and Anthropic ``system`` block the message list never carries. Embedding media an
|
||||
item of ``deployments`` reads as media is left out of the input: it never counts as text."""
|
||||
from litellm.llms.anthropic.messages.native.utils import anthropic_system_to_openai_message
|
||||
|
||||
extras: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else MappingProxyType({})
|
||||
raw_instructions: Final = extras.get("instructions")
|
||||
instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
|
||||
raw_tools: Final = extras.get("tools")
|
||||
tools: Final = raw_tools if _is_list(raw_tools) and raw_tools else None
|
||||
system_message: Final = anthropic_system_to_openai_message(extras.get("system"))
|
||||
counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # the host's overloads name litellm's untyped returns
|
||||
if messages is not None:
|
||||
counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages
|
||||
return counter(model=tokenizer_model, messages=counted_messages, tools=tools)
|
||||
if input is not None:
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments))
|
||||
input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped
|
||||
input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it
|
||||
responses_api_request={"instructions": instructions} if instructions is not None else {},
|
||||
only_since_compaction=True,
|
||||
)
|
||||
return counter(model=tokenizer_model, messages=input_messages, tools=tools)
|
||||
raise ValueError("Either messages or input must be provided to count tokens")
|
||||
|
||||
|
||||
def _max_input_tokens(registry: PreCallRegistry, model: str, deployment: Mapping[str, object]) -> int | None:
|
||||
"""``_deployment_max_input_tokens``: the declared context window, or None when there is none or
|
||||
the deployment cannot be resolved."""
|
||||
try:
|
||||
model_info: Final = registry._core_get_router_model_info( # pyright: ignore[reportPrivateUsage] # the host registry's lookup
|
||||
deployment=deployment, received_model_name=model
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others
|
||||
verbose_router_logger.debug("litellm.router.py::_deployment_max_input_tokens: skipping deployment. Got - %s", e)
|
||||
return None
|
||||
max_input_tokens: Final = model_info.get("max_input_tokens")
|
||||
return max_input_tokens if isinstance(max_input_tokens, int) else None
|
||||
|
||||
|
||||
def _items(value: object) -> tuple[object, ...]:
|
||||
"""The items of a list, tuple or set; nothing for any other value."""
|
||||
return tuple(value) if isinstance(value, list | tuple | set | frozenset) else () # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the items are read as objects
|
||||
|
||||
|
||||
def _is_list(value: object) -> bool:
|
||||
return isinstance(value, list)
|
||||
|
||||
|
||||
def _declared_param_allowlist(params: Mapping[str, object]) -> frozenset[str]:
|
||||
declared: Final = params.get("allowed_openai_params")
|
||||
return frozenset(entry for entry in _items(declared) if isinstance(entry, str))
|
||||
|
||||
|
||||
def _capability_model_and_provider(deployment_params: Mapping[str, object], group: str) -> tuple[str, str]:
|
||||
model_label: Final = str(deployment_params.get("model") or group)
|
||||
custom_llm_provider: Final = deployment_params.get("custom_llm_provider")
|
||||
authenticating_provider: Final = declared_authenticating_provider(
|
||||
model_label,
|
||||
custom_llm_provider, # pyright: ignore[reportArgumentType] # passed as configured, as the Python core does
|
||||
)
|
||||
if authenticating_provider is not None:
|
||||
return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider
|
||||
model, provider, _, _ = python_host.invoke("get_llm_provider")( # pyright: ignore[reportUnknownMemberType] # see count_tokens
|
||||
model=model_label,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=LiteLLM_Params.model_validate(dict(deployment_params)),
|
||||
)
|
||||
return model, provider
|
||||
|
||||
|
||||
def _supports_param(deployment: Mapping[str, object], group: str, param: str) -> bool:
|
||||
"""``_deployment_supports_param``."""
|
||||
deployment_params: Final = as_str_mapping(deployment.get("litellm_params"))
|
||||
if not deployment_params:
|
||||
return True
|
||||
deployment_model_info: Final = as_str_mapping(deployment.get("model_info"))
|
||||
base_model: Final = (
|
||||
deployment_model_info.get("base_model") if deployment_model_info else None
|
||||
) or deployment_params.get("base_model")
|
||||
try:
|
||||
model, custom_llm_provider = _capability_model_and_provider(deployment_params, group)
|
||||
supported: Final = python_host.invoke("get_supported_openai_params")( # pyright: ignore[reportUnknownMemberType] # see count_tokens
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
base_model=base_model if isinstance(base_model, str) else None,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request
|
||||
verbose_router_logger.debug(
|
||||
"litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e
|
||||
)
|
||||
return True
|
||||
return supported is None or param in supported
|
||||
|
||||
|
||||
def _lacks_params(deployment: Mapping[str, object], group: str, params: frozenset[str]) -> bool:
|
||||
return any(not _supports_param(deployment, group, param) for param in params)
|
||||
|
||||
|
||||
def _rejects_params(
|
||||
deployment: Mapping[str, object], group: str, params: frozenset[str], request_kwargs: Mapping[str, object]
|
||||
) -> bool:
|
||||
"""``_deployment_rejects_params``: lacks a param neither the deployment nor the request drops or allows."""
|
||||
call_settings: Final[Mapping[str, object]] = {
|
||||
**(as_str_mapping(deployment.get("litellm_params")) or {}),
|
||||
**request_kwargs,
|
||||
}
|
||||
if call_settings.get("drop_params") is True:
|
||||
return False
|
||||
dropped: Final = call_settings.get("additional_drop_params")
|
||||
dropped_names: Final = frozenset(
|
||||
entry for entry in (_items(dropped) if _is_list(dropped) else ()) if isinstance(entry, str)
|
||||
)
|
||||
return _lacks_params(deployment, group, params - _declared_param_allowlist(call_settings) - dropped_names)
|
||||
|
||||
|
||||
def _requested_params(request_kwargs: Mapping[str, object]) -> frozenset[str]:
|
||||
"""``_pre_call_check_requested_chat_params``."""
|
||||
from litellm import utils as litellm_utils
|
||||
|
||||
return _filtered_params(
|
||||
litellm_utils.get_non_default_params(passed_params=dict(request_kwargs)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # an untyped helper
|
||||
)
|
||||
|
||||
|
||||
def _filtered_params(requested: Mapping[object, object]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
param
|
||||
for param, value in requested.items()
|
||||
if isinstance(param, str) and param in _FILTERED_PARAMS and not isinstance(value, str)
|
||||
)
|
||||
|
||||
|
||||
class PreCallFacts:
|
||||
"""The facts of one ``pre_call`` step."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: "PythonHost",
|
||||
registry: PreCallRegistry,
|
||||
step: Mapping[str, object],
|
||||
records: Callable[[str], Mapping[str, object]],
|
||||
request: Mapping[str, object] | None,
|
||||
messages: object,
|
||||
input: object, # the request field's name
|
||||
) -> None:
|
||||
self._host: Final = host
|
||||
self._registry: Final = registry
|
||||
self._model: Final = str(step.get("model_group"))
|
||||
self._name: Final = str(step.get("model"))
|
||||
self._candidates: Final = [str(id) for id in _items(step.get("candidates"))]
|
||||
self._pool: Final = [str(id) for id in _items(step.get("pool"))]
|
||||
self._records: Final = records
|
||||
self._request: Final = request
|
||||
self._messages: Final = messages
|
||||
self._input: Final = input
|
||||
self._has_input: Final = messages is not None or input is not None
|
||||
self._windows: Final = {id: _max_input_tokens(registry, self._model, records(id)) for id in self._candidates}
|
||||
|
||||
def _record(self, id: str) -> Mapping[str, object]: # a deployment id
|
||||
return self._records(id)
|
||||
|
||||
def _needs_count(self) -> bool:
|
||||
return self._has_input and any(limit is not None for limit in self._windows.values())
|
||||
|
||||
def _tokenizer_model(self) -> str:
|
||||
named: Final = self._registry._core_deployments_named(self._model) or () # pyright: ignore[reportPrivateUsage] # the host registry's index
|
||||
return (
|
||||
single_encoding_openai_model(
|
||||
served_openai_models(named) or served_openai_models(self._record(id) for id in self._candidates)
|
||||
)
|
||||
or ""
|
||||
)
|
||||
|
||||
def _candidate_records(self) -> Sequence[Mapping[str, object]]:
|
||||
return tuple(self._record(id) for id in self._candidates)
|
||||
|
||||
def _count(self) -> int | None:
|
||||
try:
|
||||
return count_tokens(
|
||||
self._host,
|
||||
self._messages,
|
||||
self._input,
|
||||
self._request,
|
||||
self._tokenizer_model(),
|
||||
self._candidate_records(),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request
|
||||
verbose_router_logger.error(
|
||||
"litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e
|
||||
)
|
||||
return None
|
||||
|
||||
async def _count_async(self) -> int | None:
|
||||
try:
|
||||
return await offload_token_count(count_tokens)(
|
||||
self._host,
|
||||
self._messages,
|
||||
self._input,
|
||||
self._request,
|
||||
self._tokenizer_model(),
|
||||
self._candidate_records(),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request
|
||||
verbose_router_logger.error(
|
||||
"litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e
|
||||
)
|
||||
return None
|
||||
|
||||
def _answer(self, input_tokens: int | None) -> Mapping[str, object]:
|
||||
if not self._has_input:
|
||||
return {"has_input": False}
|
||||
request: Final[Mapping[str, object]] = self._request if self._request is not None else {}
|
||||
params: Final = (
|
||||
_requested_params(request)
|
||||
if self._request is not None and self._host.dynamic().drop_params is False
|
||||
else frozenset[str]()
|
||||
)
|
||||
deployments: Final[dict[str, dict[str, object]]] = {} # mutable-ok: built here, then returned
|
||||
for id in dict.fromkeys([*self._candidates, *self._pool]):
|
||||
record = self._record(id)
|
||||
deployments[id] = {
|
||||
"max_input_tokens": self._windows.get(id),
|
||||
"lacks_params": bool(params) and id in self._candidates and _lacks_params(record, self._model, params),
|
||||
"rejects_params": bool(params) and _rejects_params(record, self._model, params, request),
|
||||
}
|
||||
allowed_region: Final = request.get("allowed_model_region")
|
||||
return {
|
||||
"has_input": True,
|
||||
"countable": not compaction_pending(self._request),
|
||||
"input_tokens": input_tokens,
|
||||
"allowed_model_region": allowed_region,
|
||||
"requested_params": sorted(params),
|
||||
"deployments": deployments,
|
||||
}
|
||||
|
||||
def answer(self) -> Mapping[str, object]:
|
||||
"""The facts on the sync path: the prompt is counted inline."""
|
||||
return self._answer(self._count() if self._needs_count() else None)
|
||||
|
||||
async def answer_async(self) -> Mapping[str, object]:
|
||||
"""The facts on the async path: the prompt is counted off the event loop."""
|
||||
return self._answer(await self._count_async() if self._needs_count() else None)
|
||||
|
||||
|
||||
def pre_call_checks(
|
||||
host: "PythonHost",
|
||||
registry: PreCallRegistry,
|
||||
model: str,
|
||||
deployments: Sequence[Mapping[str, object]],
|
||||
messages: object,
|
||||
input: object, # the request field's name
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
input_token_count: int | None = None,
|
||||
skip_inline_token_count: bool = False,
|
||||
request_deployments: Sequence[Mapping[str, object]] | None = None,
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
"""``_pre_call_checks`` over the given deployments, as ``router-core/src/precall.rs`` decides it."""
|
||||
import litellm.exceptions as litellm_exceptions
|
||||
|
||||
pool: Final = (
|
||||
tuple(request_deployments)
|
||||
if request_deployments is not None
|
||||
else (tuple(registry._core_deployments_named(model) or ()) or tuple(deployments)) # pyright: ignore[reportPrivateUsage] # the host registry's index
|
||||
)
|
||||
candidates: Final = {f"candidate-{index}": record for index, record in enumerate(deployments)}
|
||||
group: Final = {f"pool-{index}": record for index, record in enumerate(pool)}
|
||||
records: Final = {**candidates, **group}
|
||||
step: Final = {"model_group": model, "model": model, "candidates": list(candidates), "pool": list(group)}
|
||||
facts: Final = PreCallFacts(host, registry, step, records.__getitem__, request_kwargs, messages, input)
|
||||
counted: Final = (
|
||||
input_token_count
|
||||
if input_token_count is not None or skip_inline_token_count or not facts._needs_count() # pyright: ignore[reportPrivateUsage] # this module's own class
|
||||
else facts._count() # pyright: ignore[reportPrivateUsage] # see above
|
||||
)
|
||||
answer: Final = facts._answer(counted) # pyright: ignore[reportPrivateUsage] # see above
|
||||
rows: Final = as_str_mapping(answer.get("deployments")) or {}
|
||||
allowed_region: Final = answer.get("allowed_model_region")
|
||||
|
||||
def row(id: str) -> Mapping[str, object]: # a candidate or pool key
|
||||
return as_str_mapping(rows.get(id)) or {}
|
||||
|
||||
def over_window(id: str) -> bool:
|
||||
limit: Final = row(id).get("max_input_tokens")
|
||||
return isinstance(limit, int) and answer.get("countable") is True and counted is not None and counted > limit
|
||||
|
||||
def out_of_region(id: str) -> bool:
|
||||
params: Final = as_str_mapping(records[id].get("litellm_params")) or {}
|
||||
return allowed_region is not None and params.get("region_name") != allowed_region
|
||||
|
||||
too_small: Final = tuple(id for id in candidates if over_window(id))
|
||||
invalid: Final = frozenset(too_small) | frozenset(id for id in candidates if out_of_region(id))
|
||||
param_candidates: Final = tuple(id for id in candidates if id not in invalid)
|
||||
lacking: Final = tuple(row(id).get("lacks_params") is True for id in param_candidates)
|
||||
rejecting: Final = tuple(row(id).get("rejects_params") is True for id in param_candidates)
|
||||
failing: Final = rejecting if lacking and all(lacking) else lacking
|
||||
excluded: Final = invalid | frozenset(id for id, fails in zip(param_candidates, failing, strict=True) if fails)
|
||||
if len(excluded) == len(candidates):
|
||||
requested: Final = [str(param) for param in _items(answer.get("requested_params"))]
|
||||
if not too_small and group and all(row(id).get("rejects_params") is True for id in group):
|
||||
raise litellm_exceptions.UnsupportedParamsError(
|
||||
message=(
|
||||
f"no deployment of model={model} supports {requested}. "
|
||||
"To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, "
|
||||
f"or send allowed_openai_params={requested} to pass them through"
|
||||
),
|
||||
model=model,
|
||||
)
|
||||
if too_small:
|
||||
detail: Final = "".join(
|
||||
f"Model={_error_model(records[id])}, Max Input Tokens={row(id).get('max_input_tokens')}, Got={counted}"
|
||||
for id in too_small
|
||||
)
|
||||
raise litellm_exceptions.ContextWindowExceededError(
|
||||
message="litellm._pre_call_checks: Context Window exceeded for given call. No models have context "
|
||||
f"window large enough for this call.\n{detail}",
|
||||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
return tuple(record for id, record in candidates.items() if id not in excluded)
|
||||
|
||||
|
||||
def _error_model(deployment: Mapping[str, object]) -> object:
|
||||
"""The name a context-window error gives a deployment: its base model, else its model."""
|
||||
params: Final = as_str_mapping(deployment.get("litellm_params")) or {}
|
||||
model_info: Final = as_str_mapping(deployment.get("model_info")) or {}
|
||||
return model_info.get("base_model") or params.get("base_model") or params.get("model")
|
||||
|
||||
|
||||
def request_carries_image(messages: object, request: Mapping[str, object]) -> bool:
|
||||
"""``_request_carries_image``: the messages (or the request's own messages or input) hold an image."""
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content
|
||||
|
||||
return request_contains_image_content(messages or request.get("messages") or request.get("input"))
|
||||
|
||||
|
||||
def declares_no_vision(deployment: Mapping[str, object], group: str) -> bool:
|
||||
"""``_deployment_declares_no_vision``: ``model_info.supports_vision`` is ``False``, or, when it
|
||||
is unset, the cost map says the model has vision explicitly disabled."""
|
||||
from litellm import utils as litellm_utils
|
||||
|
||||
model_info: Final = as_str_mapping(deployment.get("model_info")) or {}
|
||||
declared: Final = model_info.get("supports_vision")
|
||||
if declared is not None:
|
||||
return declared is False
|
||||
params: Final = as_str_mapping(deployment.get("litellm_params")) or {}
|
||||
if not params:
|
||||
return False
|
||||
base_model: Final = model_info.get("base_model") or params.get("base_model")
|
||||
try:
|
||||
model, custom_llm_provider = _capability_model_and_provider(params, group)
|
||||
return litellm_utils.is_vision_explicitly_disabled(
|
||||
base_model if isinstance(base_model, str) else model, custom_llm_provider
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request
|
||||
verbose_router_logger.debug(
|
||||
"litellm.router.py::_deployment_declares_no_vision: keeping deployment for model=%s. Got - %s", group, e
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PreCallFacts",
|
||||
"PreCallRegistry",
|
||||
"count_tokens",
|
||||
"declares_no_vision",
|
||||
"pre_call_checks",
|
||||
"request_carries_image",
|
||||
]
|
||||
|
||||
|
||||
def _route_capability_resolvers() -> Mapping[str, Callable[[str, "LlmProviders"], object | None]]:
|
||||
"""``_ROUTE_CAPABILITY_RESOLVERS``: the calls only some providers serve, by handler name."""
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
return {
|
||||
"adecisions": ProviderConfigManager.get_provider_decisions_config, # pyright: ignore[reportUnknownMemberType] # litellm's config lookup
|
||||
"decisions": ProviderConfigManager.get_provider_decisions_config, # pyright: ignore[reportUnknownMemberType] # see above
|
||||
}
|
||||
|
||||
|
||||
# ``_DECISIONS_MEDIA_CALL_TYPES``: the calls whose audio or videos some deployments refuse.
|
||||
DECISIONS_MEDIA_CALL_TYPES: Final = frozenset({"adecisions", "decisions", "asystemone", "systemone"})
|
||||
|
||||
|
||||
def serves_route_capability(handler: str | None) -> bool:
|
||||
"""Whether only some providers serve the call ``handler`` names."""
|
||||
return handler is not None and handler in _route_capability_resolvers()
|
||||
|
||||
|
||||
def cannot_serve_route(deployment: Mapping[str, object], group: str, handler: str) -> bool:
|
||||
"""Not ``_deployment_serves_route``: the deployment's provider has no config for the call."""
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
resolver: Final = _route_capability_resolvers().get(handler)
|
||||
params: Final = as_str_mapping(deployment.get("litellm_params"))
|
||||
if resolver is None or params is None:
|
||||
return False
|
||||
try:
|
||||
model, provider_name = _capability_model_and_provider(params, group)
|
||||
provider: Final = next((member for member in LlmProviders if member.value == provider_name), None)
|
||||
return provider is None or resolver(model, provider) is None
|
||||
except Exception as e: # noqa: BLE001 # best-effort capability filter, as the Python core keeps it
|
||||
verbose_router_logger.debug(
|
||||
"litellm.router.py::_deployment_serves_route: keeping deployment for model=%s. Got - %s", group, e
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def refuses_decisions_media(
|
||||
deployment: Mapping[str, object], group: str, sent_kwargs: Mapping[str, object], *, audio: bool, videos: bool
|
||||
) -> bool:
|
||||
"""``_decisions_deployment_refuses_media``."""
|
||||
from litellm.decisions.call import deployment_refuses_media
|
||||
|
||||
model_info: Final = as_str_mapping(deployment.get("model_info")) or {}
|
||||
params: Final = as_str_mapping(deployment.get("litellm_params")) or {}
|
||||
try:
|
||||
model, custom_llm_provider = _capability_model_and_provider(params, group)
|
||||
return deployment_refuses_media(
|
||||
model, custom_llm_provider, model_info, {**params, **sent_kwargs}, audio=audio, videos=videos
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort filter: a deployment it cannot resolve answers for itself
|
||||
verbose_router_logger.debug(
|
||||
"litellm.router.py::_decisions_deployment_refuses_media: keeping a deployment of model=%s. Got - %s",
|
||||
group,
|
||||
e,
|
||||
)
|
||||
return False
|
||||
1107
litellm/router_rust/steps/prerouting.py
Normal file
1107
litellm/router_rust/steps/prerouting.py
Normal file
File diff suppressed because it is too large
Load diff
445
litellm/router_rust/steps/selection.py
Normal file
445
litellm/router_rust/steps/selection.py
Normal file
|
|
@ -0,0 +1,445 @@
|
|||
"""
|
||||
Host-side strategy selectors for ``RustCore``.
|
||||
|
||||
The Rust core decides which strategy picks and whose state it reads (``strategy::routing_context``:
|
||||
a request override, then the routing group, then the router's own strategy) and asks the host to
|
||||
run the stateful strategies it does not implement natively through a ``select`` step. The
|
||||
strategies themselves are the standalone ``litellm/router_strategy`` loggers, which double as
|
||||
``CustomLogger`` callbacks that record traffic in the host's routing cache. This module builds,
|
||||
registers and retires them the way the Python router's ``routing_strategy_init``,
|
||||
``_init_routing_groups`` and ``_get_override_strategy_selector`` do. It never reaches a Python
|
||||
router core.
|
||||
"""
|
||||
|
||||
import threading
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, TypeAlias
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
|
||||
from litellm.router_strategy.least_busy import (
|
||||
CALL_IDENTITY,
|
||||
SELECTION_RESERVATION_MAX_KEYS,
|
||||
LeastBusyLoggingHandler,
|
||||
call_id_of,
|
||||
)
|
||||
from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
|
||||
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
|
||||
from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler
|
||||
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
|
||||
from litellm.types.router import RoutingGroup, RoutingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.router_python_host import PythonHost
|
||||
|
||||
Selector: TypeAlias = (
|
||||
LeastBusyLoggingHandler
|
||||
| LowestCostLoggingHandler
|
||||
| LowestLatencyLoggingHandler
|
||||
| LowestTPMLoggingHandler
|
||||
| LowestTPMLoggingHandler_v2
|
||||
)
|
||||
|
||||
# The legacy router attribute that holds the default group's selector, per strategy.
|
||||
SELECTOR_ATTRIBUTES: Final = MappingProxyType(
|
||||
{
|
||||
"least-busy": "leastbusy_logger",
|
||||
"usage-based-routing": "lowesttpm_logger",
|
||||
"usage-based-routing-v2": "lowesttpm_logger_v2",
|
||||
"latency-based-routing": "lowestlatency_logger",
|
||||
"cost-based-routing": "lowestcost_logger",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def normalize_strategy(strategy: object) -> str | None:
|
||||
"""A strategy as the string Python compares against (``_normalize_strategy``)."""
|
||||
if isinstance(strategy, RoutingStrategy):
|
||||
return strategy.value
|
||||
return strategy if isinstance(strategy, str) else None
|
||||
|
||||
|
||||
def build_selector(
|
||||
strategy: str | None,
|
||||
args: Mapping[str, object],
|
||||
cache: "DualCache",
|
||||
reservations: "InMemoryCache | None" = None,
|
||||
) -> Selector | None:
|
||||
"""``_build_strategy_selector`` without registering: ``None`` for simple-shuffle and every
|
||||
strategy without a selector. Every least-busy selector of a router keeps its pick
|
||||
reservations in the router's one ``reservations`` store."""
|
||||
routing_args: Final = dict(args)
|
||||
match strategy:
|
||||
case "least-busy":
|
||||
return LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations)
|
||||
case "usage-based-routing":
|
||||
return LowestTPMLoggingHandler(router_cache=cache, routing_args=routing_args)
|
||||
case "usage-based-routing-v2":
|
||||
return LowestTPMLoggingHandler_v2(router_cache=cache, routing_args=routing_args)
|
||||
case "latency-based-routing":
|
||||
return LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args)
|
||||
case "cost-based-routing":
|
||||
return LowestCostLoggingHandler(router_cache=cache, routing_args={})
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _callbacks(host: "PythonHost") -> list[object]: # mutable-ok: the process list itself
|
||||
"""``litellm.callbacks``, which the host serves like its other callback lists."""
|
||||
return host.callback_list("callbacks") # pyright: ignore[reportArgumentType] # not in the CallbackList literal
|
||||
|
||||
|
||||
class StrategySelectors:
|
||||
"""The selectors of one router: the default group's, each routing group's, and the lazily
|
||||
built per-request override ones."""
|
||||
|
||||
def __init__(self, host: "PythonHost", cache: "DualCache") -> None:
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
|
||||
self._host: Final = host
|
||||
self._cache: Final = cache
|
||||
# `_core_least_busy_reservations`: the picks least-busy reserved, shared by its selectors.
|
||||
self.reservations: Final = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS)
|
||||
self._strategy: str | None = None
|
||||
# The router's `routing_strategy`, which a direct `routing_strategy_init` leaves as it was.
|
||||
self._configured: str | None = None
|
||||
self._default: Selector | None = None
|
||||
self._groups: dict[str, Selector] = {} # mutable-ok: replaced as routing groups change
|
||||
self._group_strategies: frozenset[str] = frozenset()
|
||||
self._overrides: dict[str, Selector | None] = {} # mutable-ok: filled lazily per strategy
|
||||
self._lock: Final = threading.Lock()
|
||||
|
||||
# -- building ------------------------------------------------------------------------------
|
||||
|
||||
def _register(self, selector: CustomLogger) -> None:
|
||||
"""``_register_router_selector``: least-busy also counts calls as they start."""
|
||||
host: Final = self._host
|
||||
if isinstance(selector, LeastBusyLoggingHandler):
|
||||
if isinstance(host.callback_list("input_callback"), list): # pyright: ignore[reportUnnecessaryIsInstance] # a process list a user may have reassigned
|
||||
host.callbacks.add_litellm_input_callback(selector)
|
||||
else:
|
||||
host.set_callback_list("input_callback", [selector])
|
||||
if isinstance(_callbacks(host), list): # pyright: ignore[reportUnnecessaryIsInstance] # see above
|
||||
host.callbacks.add_litellm_callback(selector)
|
||||
|
||||
def _unregister(self, selectors: Sequence[object]) -> None:
|
||||
live: Final = [selector for selector in selectors if selector is not None]
|
||||
for selector in live:
|
||||
if isinstance(selector, BaseRoutingStrategy):
|
||||
selector.retire()
|
||||
if live:
|
||||
self._host.callbacks.drop_callbacks(live)
|
||||
|
||||
def init_default(self, strategy: object, args: Mapping[str, object], *, configured: bool = True) -> None:
|
||||
"""``routing_strategy_init``: drop the old default and override selectors, build the new one.
|
||||
``configured``: ``strategy`` is the router's ``routing_strategy`` from now on."""
|
||||
self._unregister([self._default, *self._overrides.values()])
|
||||
self._overrides = {}
|
||||
self._default = None
|
||||
self._strategy = normalize_strategy(strategy)
|
||||
if configured:
|
||||
self._configured = self._strategy
|
||||
self._default = build_selector(self._strategy, args, self._cache, self.reservations)
|
||||
if self._default is not None:
|
||||
self._register(self._default)
|
||||
|
||||
def rebuild_default(self, args: Mapping[str, object]) -> None:
|
||||
"""``_apply_updated_routing_strategy_args``: relink the default selector to new args, keeping
|
||||
the old one when they are invalid."""
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
||||
if self._default is None or self._strategy != self._configured:
|
||||
return
|
||||
try:
|
||||
rebuilt: Final = build_selector(self._strategy, args, self._cache, self.reservations)
|
||||
except (TypeError, ValidationError):
|
||||
verbose_router_logger.exception(
|
||||
"Invalid routing_strategy_args %s for '%s'; keeping the previous ones", dict(args), self._strategy
|
||||
)
|
||||
return
|
||||
self._unregister([self._default])
|
||||
self._default = rebuilt
|
||||
if rebuilt is not None:
|
||||
self._register(rebuilt)
|
||||
|
||||
def init_groups(self, groups: Sequence[tuple[str, str, Mapping[str, object]]]) -> None:
|
||||
"""``_replace_routing_groups``: one selector per group, from (name, strategy, args)."""
|
||||
self._unregister(list(self._groups.values()))
|
||||
built: Final = {
|
||||
name: selector
|
||||
for name, strategy, args in groups
|
||||
if (
|
||||
selector := build_selector(
|
||||
None if strategy == "priority" else strategy, args, self._cache, self.reservations
|
||||
)
|
||||
)
|
||||
is not None
|
||||
}
|
||||
for selector in built.values():
|
||||
self._register(selector)
|
||||
self._groups = built
|
||||
self._group_strategies = frozenset(strategy for _, strategy, _ in groups)
|
||||
|
||||
def retire(self) -> None:
|
||||
self._unregister([self._default, *self._groups.values(), *self._overrides.values()])
|
||||
self.clear()
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Forget every selector, after the caller released them from the host lists."""
|
||||
self._default = None
|
||||
self._groups = {}
|
||||
self._overrides = {}
|
||||
|
||||
def callback_selectors(self) -> tuple[Selector, ...]:
|
||||
"""``_callback_selectors``: the default group's selector and each routing group's."""
|
||||
return tuple(selector for selector in (self._default, *self._groups.values()) if selector is not None)
|
||||
|
||||
def override_selectors(self) -> tuple[Selector, ...]:
|
||||
"""The per-request override selectors built so far."""
|
||||
return tuple(selector for selector in self._overrides.values() if selector is not None)
|
||||
|
||||
@property
|
||||
def default(self) -> Selector | None:
|
||||
"""The default selector of the router's ``routing_strategy``, if it has one."""
|
||||
return self._default if self._strategy == self._configured else None
|
||||
|
||||
def attribute(self, name: str) -> Selector | None:
|
||||
"""A legacy ``*_logger`` attribute: the default selector when it is that strategy's."""
|
||||
return self._default if SELECTOR_ATTRIBUTES.get(self._strategy or "") == name else None
|
||||
|
||||
def latency_args(self) -> Mapping[str, object] | None:
|
||||
"""The default latency selector's full routing args, as ``get_settings`` reports them."""
|
||||
if not isinstance(self._default, LowestLatencyLoggingHandler):
|
||||
return None
|
||||
reported: Final[object] = self._default.routing_args.json() # pyright: ignore[reportUnknownMemberType] # RoutingArgs.json is untyped
|
||||
return as_str_mapping(reported)
|
||||
|
||||
# -- picking -------------------------------------------------------------------------------
|
||||
|
||||
def selector_for(
|
||||
self, strategy: str, scope: str, group: str | None, request: Mapping[str, object] | None
|
||||
) -> Selector | None:
|
||||
"""The selector a Rust ``select`` step names (``_get_routing_context``)."""
|
||||
if scope == "group":
|
||||
return self._groups.get(group or "")
|
||||
if scope == "override":
|
||||
return self._override(strategy, request)
|
||||
return self._default if strategy == self._strategy else None
|
||||
|
||||
def _override(self, strategy: str, request: Mapping[str, object] | None) -> Selector | None:
|
||||
"""``_get_override_strategy_selector`` and ``_bind_override_selector_to_request``."""
|
||||
if strategy == self._configured:
|
||||
return self.default
|
||||
with self._lock:
|
||||
if strategy not in self._overrides:
|
||||
self._overrides[strategy] = build_selector(strategy, {}, self._cache, self.reservations)
|
||||
selector: Final = self._overrides[strategy]
|
||||
if selector is not None and request is not None and strategy not in self._registered_strategies():
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj: Final = request.get("litellm_logging_obj")
|
||||
if isinstance(logging_obj, Logging):
|
||||
logging_obj.add_dynamic_callback(selector)
|
||||
return selector
|
||||
|
||||
def _registered_strategies(self) -> frozenset[str]:
|
||||
return frozenset({*((self._configured,) if self._configured else ()), *self._group_strategies})
|
||||
|
||||
def is_override(self, strategy: str, selector: object) -> bool:
|
||||
"""Whether ``selector`` is a per-request override one, whose pre-call check the router runs
|
||||
itself because it is not in the callback lists."""
|
||||
return selector is not None and self._overrides.get(strategy) is selector
|
||||
|
||||
|
||||
async def pick_async(
|
||||
strategy: str,
|
||||
selector: Selector,
|
||||
model: str,
|
||||
deployments: list[dict[str, object]], # mutable-ok: the selectors take a list of deployment dicts
|
||||
messages: object,
|
||||
input: object,
|
||||
request_kwargs: dict[str, object] | None, # mutable-ok: the latency selector reads it as given
|
||||
reserve: bool = False,
|
||||
) -> object:
|
||||
"""``_select_deployment_async``; ``reserve``: least-busy holds the pick for the request's call."""
|
||||
match selector:
|
||||
case LeastBusyLoggingHandler() if strategy == "least-busy":
|
||||
return await selector.async_get_available_deployments(
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
reserve=reserve,
|
||||
call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})),
|
||||
) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
case LowestTPMLoggingHandler() if strategy == "usage-based-routing":
|
||||
return selector.get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given
|
||||
input=input, # pyright: ignore[reportArgumentType] # same
|
||||
)
|
||||
case LowestTPMLoggingHandler_v2() | LowestCostLoggingHandler() if strategy in (
|
||||
"usage-based-routing-v2",
|
||||
"cost-based-routing",
|
||||
):
|
||||
return await selector.async_get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # same
|
||||
input=input, # pyright: ignore[reportArgumentType] # same
|
||||
)
|
||||
case LowestLatencyLoggingHandler() if strategy == "latency-based-routing":
|
||||
return await selector.async_get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # same
|
||||
input=input, # pyright: ignore[reportArgumentType] # same
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def pick_sync(
|
||||
strategy: str,
|
||||
selector: Selector,
|
||||
model: str,
|
||||
deployments: list[dict[str, object]], # mutable-ok: see pick_async
|
||||
messages: object,
|
||||
input: object,
|
||||
request_kwargs: dict[str, object] | None, # mutable-ok: see pick_async
|
||||
reserve: bool = False,
|
||||
) -> object:
|
||||
"""``_select_deployment_sync``: cost-based routing has no sync pick."""
|
||||
match selector:
|
||||
case LeastBusyLoggingHandler() if strategy == "least-busy":
|
||||
return selector.get_available_deployments(
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
reserve=reserve,
|
||||
call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})),
|
||||
) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
case LowestTPMLoggingHandler() | LowestTPMLoggingHandler_v2() if strategy in (
|
||||
"usage-based-routing",
|
||||
"usage-based-routing-v2",
|
||||
):
|
||||
return selector.get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # same
|
||||
input=input, # pyright: ignore[reportArgumentType] # same
|
||||
)
|
||||
case LowestLatencyLoggingHandler() if strategy == "latency-based-routing":
|
||||
return selector.get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts
|
||||
model_group=model,
|
||||
healthy_deployments=deployments,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # same
|
||||
input=input, # pyright: ignore[reportArgumentType] # same
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
async def routing_pre_call_checks(callbacks: Iterable[object], deployment: Mapping[str, object]) -> None:
|
||||
"""``async_routing_strategy_pre_call_checks``: every ``CustomLogger`` among the router's routing
|
||||
callbacks (``_routing_callbacks``) checks the picked deployment before the call
|
||||
(usage-based-routing-v2 counts its rpm here). A check's error fails the attempt; the core
|
||||
records that failure like any other."""
|
||||
for callback in tuple(callbacks):
|
||||
if isinstance(callback, CustomLogger):
|
||||
await callback.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict
|
||||
|
||||
|
||||
def routing_pre_call_checks_sync(callbacks: Iterable[object], deployment: Mapping[str, object]) -> None:
|
||||
"""``routing_strategy_pre_call_checks``: the sync check of each ``CustomLogger`` callback."""
|
||||
for callback in tuple(callbacks):
|
||||
if isinstance(callback, CustomLogger):
|
||||
callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # see above
|
||||
|
||||
|
||||
def picked_id(picked: object) -> str | None:
|
||||
"""The id of the deployment a selector returned."""
|
||||
info: Final = as_str_mapping((as_str_mapping(picked) or {}).get("model_info")) or {}
|
||||
found: Final = info.get("id")
|
||||
return str(found) if found is not None else None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SELECTOR_ATTRIBUTES",
|
||||
"Selector",
|
||||
"StrategySelectors",
|
||||
"build_selector",
|
||||
"group_documents",
|
||||
"group_settings",
|
||||
"member_map",
|
||||
"normalize_strategy",
|
||||
"parse_groups",
|
||||
"pick_async",
|
||||
"pick_sync",
|
||||
"picked_id",
|
||||
"routing_pre_call_checks",
|
||||
"routing_pre_call_checks_sync",
|
||||
]
|
||||
|
||||
|
||||
def parse_groups(
|
||||
groups_input: object, model_names: frozenset[str], aliases: frozenset[str]
|
||||
) -> tuple[RoutingGroup, ...]:
|
||||
"""``_init_routing_groups``' validation: ``parse_routing_groups``, then a priority group may not
|
||||
share its name with a model group or alias (a plain group only warns: its name is shadowed)."""
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.routing_groups import (
|
||||
parse_routing_groups, # pyright: ignore[reportUnknownVariableType] # takes untyped dicts
|
||||
)
|
||||
|
||||
if not groups_input or not isinstance(groups_input, list | tuple):
|
||||
return ()
|
||||
groups: Final = parse_routing_groups( # pyright: ignore[reportUnknownVariableType] # returns RoutingGroup models
|
||||
list(groups_input), # pyright: ignore[reportUnknownArgumentType] # validated by parse_routing_groups
|
||||
known_model_names=model_names,
|
||||
)
|
||||
for group in groups:
|
||||
if group.group_name in model_names or group.group_name in aliases:
|
||||
if group.routing_strategy == "priority":
|
||||
raise ValueError("Priority routing group names must not shadow a model or alias")
|
||||
verbose_router_logger.warning(
|
||||
"routing_groups: group_name '%s' is shadowed by an existing model_name or model_group_alias; "
|
||||
"the group's strategy still applies to its members, but the name is not callable until renamed.",
|
||||
group.group_name,
|
||||
)
|
||||
return groups
|
||||
|
||||
|
||||
def group_documents(groups: Sequence[RoutingGroup]) -> list[Mapping[str, object]]: # mutable-ok: a JSON list
|
||||
"""The routing groups as the Rust core reads them."""
|
||||
return [
|
||||
{
|
||||
"group_name": group.group_name,
|
||||
"models": list(group.models),
|
||||
"routing_strategy": group.routing_strategy,
|
||||
**({"model_priorities": dict(group.model_priorities)} if group.model_priorities is not None else {}),
|
||||
}
|
||||
for group in groups
|
||||
]
|
||||
|
||||
|
||||
def group_settings(groups: Sequence[RoutingGroup]) -> list[object]: # mutable-ok: the settings snapshot's list
|
||||
"""``get_settings()['routing_groups']``."""
|
||||
return [
|
||||
group.model_dump(exclude={"model_priorities"}) if group.model_priorities is None else group.model_dump()
|
||||
for group in groups
|
||||
]
|
||||
|
||||
|
||||
def member_map(groups: Sequence[RoutingGroup]) -> Mapping[str, str]:
|
||||
"""``_model_to_group``: each member model group of a non-priority routing group, to the group
|
||||
(a model listed by two groups maps to the last)."""
|
||||
kept: Final = [group for group in groups if group.routing_strategy != "priority"]
|
||||
return dict(chain.from_iterable(dict.fromkeys(group.models, group.group_name).items() for group in kept))
|
||||
14
litellm/router_rust/surface/README.md
Normal file
14
litellm/router_rust/surface/README.md
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
# surface
|
||||
|
||||
The Python `Router` API, kept over router-core's state so Python callers see today's behavior. router-core holds the state and the rules; these modules are views and adapters.
|
||||
|
||||
- `registry`, `model_list`, `indexes`, `members`: the deployment list, indexes and lookups; writes become changesets
|
||||
- `legacy_members`: the old public members outside the entry points
|
||||
- `model_info`: model and group info, answering the `model_info` step's lookups
|
||||
- `replay`: re-registers pricing after a price data reload
|
||||
- `settings`: `update_settings` rules and the settings view
|
||||
- `health`, `health_routing`, `retries`: outcome callbacks, health-check routing and retry answers
|
||||
- `fallbacks`: fallback config validation
|
||||
- `alerting`: Slack alerting
|
||||
|
||||
A Rust host needs none of this.
|
||||
1
litellm/router_rust/surface/__init__.py
Normal file
1
litellm/router_rust/surface/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
"""The Python ``Router`` API kept over router-core's state, for Python callers."""
|
||||
33
litellm/router_rust/surface/alerting.py
Normal file
33
litellm/router_rust/surface/alerting.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
"""The Rust core's Slack alerting: ``alerting_config`` built into a ``SlackAlerting`` logger on the
|
||||
host, as the Python core's ``_initialize_alerting`` builds it."""
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.router_python_host import PythonHost
|
||||
from litellm.types.router import AlertingConfig
|
||||
|
||||
|
||||
def initialize_alerting(host: "PythonHost", alerting_config: "AlertingConfig | None") -> "SlackAlerting | None":
|
||||
"""``_initialize_alerting``: a ``SlackAlerting`` logger with the config's threshold and webhook,
|
||||
added to the host's callbacks, and its slow-response check to the host's success callbacks.
|
||||
No config builds nothing."""
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
|
||||
if alerting_config is None:
|
||||
return None
|
||||
logger: Final = SlackAlerting(
|
||||
alerting_threshold=alerting_config.alerting_threshold,
|
||||
alerting=["slack"],
|
||||
default_webhook_url=alerting_config.webhook_url,
|
||||
)
|
||||
host.callbacks.add_litellm_callback(logger)
|
||||
host.callbacks.add_litellm_success_callback(logger.response_taking_too_long_callback) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # SlackAlerting's untyped callback
|
||||
verbose_router_logger.info("\033[94m\nInitialized Alerting for litellm.Router\033[0m\n")
|
||||
return logger
|
||||
|
||||
|
||||
__all__ = ["initialize_alerting"]
|
||||
129
litellm/router_rust/surface/fallbacks.py
Normal file
129
litellm/router_rust/surface/fallbacks.py
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
"""
|
||||
Fallback configuration checks ``RustCore`` runs on the host: ``validate_fallbacks``, with the
|
||||
Python router's error texts. router-core walks the fallbacks; this module only refuses a list it
|
||||
could not walk the way the Python router would, before the core is built.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.router_utils.fallback_event_handlers import check_non_standard_fallback_format
|
||||
|
||||
|
||||
def refuse_malformed_fallbacks(fallback_param: Sequence[object] | None) -> None:
|
||||
"""``Router.validate_fallbacks``: every entry is a one-key dict whose targets are a list (or a
|
||||
client-side dict target)."""
|
||||
if fallback_param is None:
|
||||
return
|
||||
for fallback_dict in fallback_param:
|
||||
if not isinstance(fallback_dict, dict):
|
||||
raise ValueError(f"Item '{fallback_dict}' is not a dictionary.") # noqa: TRY004 # the Python router's error class
|
||||
if len(fallback_dict) != 1: # pyright: ignore[reportUnknownArgumentType] # only the size is read
|
||||
raise ValueError(
|
||||
f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." # pyright: ignore[reportUnknownArgumentType] # see above
|
||||
)
|
||||
((model_group, targets),) = fallback_dict.items() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # an entry as configured
|
||||
entry: list[object] = [fallback_dict] # mutable-ok: the one-entry list the check reads
|
||||
if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=entry):
|
||||
continue
|
||||
raise TypeError(
|
||||
f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. "
|
||||
f"Use {{'{model_group}': ['{targets}']}}"
|
||||
)
|
||||
|
||||
|
||||
def rejected_fallback_target(targets: Sequence[object]) -> Exception | None:
|
||||
"""The error ``run_async_fallback`` raises before its walk when a chain, from any list and of
|
||||
any level, has a dict target a request may not set (a server-owned federation field, an
|
||||
``oidc/`` api key): every target of that walk is refused."""
|
||||
from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference
|
||||
from litellm.types.router import reject_server_owned_wif_params
|
||||
|
||||
for target in targets:
|
||||
if isinstance(target, dict):
|
||||
try:
|
||||
reject_server_owned_wif_params(target) # pyright: ignore[reportUnknownArgumentType] # a chain as configured
|
||||
reject_api_key_secret_reference(target) # pyright: ignore[reportUnknownArgumentType] # see above
|
||||
except Exception as error: # noqa: BLE001 # the walk's own refusal, reported as its failure
|
||||
return error
|
||||
return None
|
||||
|
||||
|
||||
# The `mock_testing_*` flags the router reads, and the mock failure each asks router-core for.
|
||||
MOCK_FLAGS: Final = MappingProxyType(
|
||||
{
|
||||
"mock_testing_fallbacks": "fallbacks",
|
||||
"mock_testing_context_fallbacks": "context_window",
|
||||
"mock_testing_content_policy_fallbacks": "content_policy",
|
||||
"mock_testing_rate_limit_error": "rate_limit",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def mock_failure(request: Mapping[str, object]) -> str | None:
|
||||
"""The mock failure a request asks for (``MockRouterTestingParams.from_kwargs``, then
|
||||
``_handle_mock_testing_rate_limit_error``): the first fallback flag that is ``True`` (a string
|
||||
is read as a bool), else the rate limit flag when it is ``True`` itself."""
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
for flag in ("mock_testing_fallbacks", "mock_testing_context_fallbacks", "mock_testing_content_policy_fallbacks"):
|
||||
value = request.get(flag)
|
||||
if (str_to_bool(value) if isinstance(value, str) else value) is True:
|
||||
return MOCK_FLAGS[flag]
|
||||
return "rate_limit" if request.get("mock_testing_rate_limit_error") is True else None
|
||||
|
||||
|
||||
def mock_error(
|
||||
mock: str, model: str, num_retries: object, request: Mapping[str, object], configured: Mapping[str, object]
|
||||
) -> Exception:
|
||||
"""The exception the Python router raises for a mock failure; ``configured`` holds the
|
||||
router's fallback lists, which the request's own replace."""
|
||||
import litellm.exceptions as litellm_exceptions
|
||||
|
||||
def listed(name: str) -> object:
|
||||
return request.get(name, configured.get(name))
|
||||
|
||||
if mock == "rate_limit":
|
||||
return litellm_exceptions.RateLimitError(
|
||||
model=model,
|
||||
llm_provider="",
|
||||
message=f"This is a mock exception for model={model}, to trigger a rate limit error.",
|
||||
num_retries=num_retries if isinstance(num_retries, int) else None,
|
||||
)
|
||||
if mock == "context_window":
|
||||
return litellm_exceptions.ContextWindowExceededError(
|
||||
model=model,
|
||||
llm_provider="",
|
||||
message=f"This is a mock exception for model={model}, to trigger a fallback. \
|
||||
Context_Window_Fallbacks={listed('context_window_fallbacks')}",
|
||||
)
|
||||
if mock == "content_policy":
|
||||
return litellm_exceptions.ContentPolicyViolationError(
|
||||
model=model,
|
||||
llm_provider="",
|
||||
message=f"This is a mock exception for model={model}, to trigger a fallback. \
|
||||
Context_Policy_Fallbacks={listed('content_policy_fallbacks')}",
|
||||
)
|
||||
return litellm_exceptions.InternalServerError(
|
||||
model=model,
|
||||
llm_provider="",
|
||||
message=f"This is a mock exception for model={model}, to trigger a fallback. Fallbacks={listed('fallbacks')}",
|
||||
)
|
||||
|
||||
|
||||
class RustFallbacks:
|
||||
"""Fallback members served on the host."""
|
||||
|
||||
def validate_fallbacks(self, fallback_param: Sequence[object] | None) -> None:
|
||||
refuse_malformed_fallbacks(fallback_param)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MOCK_FLAGS",
|
||||
"RustFallbacks",
|
||||
"mock_error",
|
||||
"mock_failure",
|
||||
"refuse_malformed_fallbacks",
|
||||
"rejected_fallback_target",
|
||||
]
|
||||
558
litellm/router_rust/surface/health.py
Normal file
558
litellm/router_rust/surface/health.py
Normal file
|
|
@ -0,0 +1,558 @@
|
|||
"""
|
||||
The deployment outcome callbacks ``RustCore`` serves as members: ``deployment_callback_on_failure``,
|
||||
``deployment_callback_on_success`` and ``sync_deployment_callback_on_success``. router-core counts
|
||||
the outcome and decides the cooldown; this module reads litellm's logging kwargs, drops the
|
||||
failures the Python router does not hold against the deployment, and drives the Rust call. The
|
||||
stream wrappers report a stream's outcome through the same members when the stream ends.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_rust.steps.callbacks import OptionalChecks
|
||||
from litellm.router_rust.surface.registry import HostRegistry
|
||||
from litellm.router_utils.cooldown_cache import CooldownCache
|
||||
from litellm.rust_bridge._native import (
|
||||
RouterCall,
|
||||
RouterCoreHandle,
|
||||
_FailureReportSteps, # pyright: ignore[reportPrivateUsage] # the stub's step shapes
|
||||
_SuccessReportSteps, # pyright: ignore[reportPrivateUsage] # see above
|
||||
)
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
|
||||
def _cast_status(exception_status: object) -> int | None:
|
||||
if isinstance(exception_status, int) and not isinstance(exception_status, bool):
|
||||
return exception_status
|
||||
if isinstance(exception_status, str):
|
||||
try:
|
||||
return int(exception_status)
|
||||
except ValueError:
|
||||
return 500
|
||||
return None
|
||||
|
||||
|
||||
def failure_status(exception_status: object) -> int | None:
|
||||
"""``cast_exception_status_to_int``: a number as given, a numeric string read as one, any
|
||||
other string 500. A number no status can hold reads as no status."""
|
||||
status: Final = _cast_status(exception_status)
|
||||
return status if status is not None and 0 <= status <= 65535 else None
|
||||
|
||||
|
||||
def _logged_deployment_id(kwargs: Mapping[str, object]) -> object:
|
||||
"""``litellm_params.model_info.id`` of a logged call."""
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
return (as_str_mapping(litellm_params.get("model_info")) or {}).get("id")
|
||||
|
||||
|
||||
def first_cooldown(*sources: Mapping[str, object] | None) -> int | float | None:
|
||||
"""``_first_present(*sources, key="cooldown_time")``, the number as it was written."""
|
||||
for source in sources:
|
||||
value = source.get("cooldown_time") if source is not None else None
|
||||
if isinstance(value, int | float) and not isinstance(value, bool):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _configured_cooldown(model_info: object, litellm_params: Mapping[str, object]) -> float | None:
|
||||
"""``first_cooldown`` of a logged call, as the float router-core takes."""
|
||||
seconds: Final = first_cooldown(as_str_mapping(model_info), litellm_params)
|
||||
return None if seconds is None else float(seconds)
|
||||
|
||||
|
||||
def python_cooldown_time(seconds: float, *sources: float | None) -> float:
|
||||
"""``seconds`` as ``deployment_callback_on_failure`` hands them on. Router-core's duration is
|
||||
always a float; Python's is the number that won its precedence, an int or a float as it was
|
||||
written. ``sources`` run in that precedence, and only the type is read from them."""
|
||||
return next((source for source in sources if source is not None and source == seconds), seconds)
|
||||
|
||||
|
||||
def _fallback_cooldown(deployment: "Deployment | None") -> float | None:
|
||||
"""``first_cooldown`` of the deployment a failed hop tried, as ``_core_get_model_info`` lists it."""
|
||||
if deployment is None:
|
||||
return None
|
||||
listed: Final = deployment.model_dump()
|
||||
return first_cooldown(as_str_mapping(listed.get("model_info")), as_str_mapping(listed.get("litellm_params")))
|
||||
|
||||
|
||||
def _deployment_cooldown(deployment: Mapping[str, object]) -> float | None:
|
||||
"""``_configured_cooldown`` of a deployment as listed."""
|
||||
return _configured_cooldown(deployment.get("model_info"), as_str_mapping(deployment.get("litellm_params")) or {})
|
||||
|
||||
|
||||
def logged_call(op: str, kwargs: Mapping[str, object], started: datetime | None = None) -> Mapping[str, object]:
|
||||
"""The logging kwargs ``deployment_callback_on_failure`` would read for a call made with
|
||||
``kwargs``: the call's logging object's (a proxy request carries one), else what the request says.
|
||||
Read before the failure logger stamps the call's ``end_time``, so it is stamped here."""
|
||||
ended: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline's api_call_start_time
|
||||
details: Final = as_str_mapping(getattr(kwargs.get("litellm_logging_obj"), "model_call_details", None))
|
||||
call_type: Final = op.rsplit(".", 1)[-1]
|
||||
if details is not None:
|
||||
return {"call_type": call_type, **details, "end_time": ended}
|
||||
fields: Final = ("metadata", "litellm_metadata", "client_side_timeout", "timeout")
|
||||
return {
|
||||
"call_type": call_type,
|
||||
"litellm_params": {name: kwargs.get(name) for name in fields},
|
||||
"start_time": started,
|
||||
"end_time": ended,
|
||||
}
|
||||
|
||||
|
||||
def _authenticates_per_caller(litellm_params: object) -> bool:
|
||||
"""``github_copilot_per_user_credential_name`` in ``_should_run_cooldown_logic``. Per-user params
|
||||
with no credential raise there, which ends Python's failure callback short of any cooldown."""
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.llms.github_copilot.per_user_auth import github_copilot_per_user_credential_name
|
||||
|
||||
try:
|
||||
return github_copilot_per_user_credential_name(litellm_params) is not None
|
||||
except BadRequestError:
|
||||
return True
|
||||
|
||||
|
||||
def failure_already_logged(kwargs: Mapping[str, object], event: Literal["async_failure", "sync_failure"]) -> bool:
|
||||
"""``should_run_logging``: a logging object logs one failure a request. The attempts of a request
|
||||
that shares one (``litellm_logging_obj``) after the first logged failure run no failure callback;
|
||||
without one, every litellm call logs on its own."""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
||||
return isinstance(logging_obj, Logging) and logging_obj.model_call_details.get(f"has_logged_{event}") is True
|
||||
|
||||
|
||||
def _success_target(kwargs: Mapping[str, object]) -> tuple[str, Mapping[str, object]] | None:
|
||||
"""The deployment id and metadata a success callback counts under, or None."""
|
||||
params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
metadata: Final = as_str_mapping(params.get("metadata"))
|
||||
if metadata is None:
|
||||
return None
|
||||
model_info: Final = as_str_mapping(params.get("model_info")) or {}
|
||||
deployment_id: Final = model_info.get("id")
|
||||
if metadata.get("model_group") is None or deployment_id is None:
|
||||
return None
|
||||
return str(deployment_id), metadata
|
||||
|
||||
|
||||
class RustHealthCallbacks:
|
||||
"""Outcome callbacks read into the Rust core's health state."""
|
||||
|
||||
_handle: "RouterCoreHandle"
|
||||
_cooldown_given: object
|
||||
|
||||
def _drive_sync_steps(self, call: "RouterCall[_FailureReportSteps | _SuccessReportSteps]") -> Mapping[str, object]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def _drive_cache_steps(
|
||||
self, call: "RouterCall[_FailureReportSteps | _SuccessReportSteps]"
|
||||
) -> Mapping[str, object]:
|
||||
raise NotImplementedError
|
||||
|
||||
def _sync_allowed_fails(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def _core_get_allowed_fails_from_policy(self, exception: Exception) -> int | None:
|
||||
"""The router ``allowed_fails_policy``'s allowance for this exception, decided by router-core."""
|
||||
from litellm.router_rust.core import attempt_failure
|
||||
|
||||
return self._handle.allowed_fails_for(dict(attempt_failure(exception)))
|
||||
|
||||
def _deployment_of(self, deployment_id: object) -> "Deployment | None":
|
||||
# RustCore sets `_hosted` as an attribute in its constructor.
|
||||
hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # RustCore's own attribute
|
||||
return hosted._core_get_deployment(model_id=deployment_id) if isinstance(deployment_id, str) else None # pyright: ignore[reportPrivateUsage] # the host registry's id lookup
|
||||
|
||||
def deployment_tracks_usage(self, deployment_id: object, model_group: object) -> bool:
|
||||
"""``_cached_deployment_tracks_usage``. A deployment whose model is not mapped raises there,
|
||||
which ends Python's usage write, so it reads as not tracked."""
|
||||
# RustCore sets `_hosted` as an attribute in its constructor.
|
||||
hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # RustCore's own attribute
|
||||
if not isinstance(deployment_id, str) or not isinstance(model_group, str):
|
||||
return False
|
||||
try:
|
||||
return hosted.cached_deployment_tracks_usage(deployment_id, model_group)
|
||||
except Exception: # noqa: BLE001 # as `deployment_callback_on_success` swallows it
|
||||
return False
|
||||
|
||||
def cooldown_event(self, deployment_id: str, exception_status: object, cooldown_time: float | None) -> None:
|
||||
"""The end of ``set_cooldown_deployments``, once the cooldown is written: the cooldown event
|
||||
Prometheus reads, as a task of the running loop. Outside one this raises ``RuntimeError``."""
|
||||
import asyncio
|
||||
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.router_utils import cooldown_handlers
|
||||
|
||||
_ = asyncio.get_running_loop()
|
||||
_ = asyncio.create_task( # fire and forget, as set_cooldown_deployments schedules it
|
||||
cooldown_handlers.router_cooldown_event_callback(
|
||||
ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext
|
||||
deployment_id=deployment_id,
|
||||
exception_status=exception_status, # pyright: ignore[reportArgumentType] # the status as raised
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
)
|
||||
|
||||
def cooldown_event_on_own_loop(
|
||||
self, deployment_id: str, exception_status: object, cooldown_time: float | None
|
||||
) -> None:
|
||||
"""``cooldown_event`` for a sync walk ``function_with_fallbacks`` would run on its private
|
||||
loop: ``set_cooldown_deployments`` finds that loop however the caller runs, and the event's
|
||||
task keeps what it raises to itself."""
|
||||
import asyncio
|
||||
|
||||
from litellm.litellm_core_utils import asyncify
|
||||
|
||||
async def fire() -> None:
|
||||
self.cooldown_event(deployment_id, exception_status, cooldown_time)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
asyncify.run_async_function(fire) # pyright: ignore[reportUnknownMemberType] # the untyped sync bridge
|
||||
|
||||
def python_cooldown(self, seconds: float, configured: float | None = None, error: object = None) -> float:
|
||||
"""The ``time_to_cooldown`` Python hands ``set_cooldown_deployments``, for router-core's
|
||||
``seconds``. A failure callback's precedence is the ``configured`` number, then the ``retry-after``
|
||||
header of ``error``, then the router's own; a path passing neither hands on the router's own."""
|
||||
from litellm.router_rust.core import retry_after_seconds
|
||||
|
||||
retry_after: Final = retry_after_seconds(error) if isinstance(error, BaseException) else None
|
||||
default: Final = self._cooldown_given
|
||||
return python_cooldown_time(
|
||||
seconds,
|
||||
configured,
|
||||
None if retry_after is None else int(retry_after),
|
||||
default if isinstance(default, int | float) else None,
|
||||
)
|
||||
|
||||
def cooled_down(
|
||||
self,
|
||||
done: Mapping[str, object],
|
||||
deployment_id: object,
|
||||
exception_status: object,
|
||||
*,
|
||||
own_loop: bool = False,
|
||||
configured: float | None = None,
|
||||
error: object = None,
|
||||
) -> bool:
|
||||
"""Whether a failure report cooled the deployment down, with its cooldown event, which
|
||||
carries ``python_cooldown`` of ``configured`` and ``error``."""
|
||||
cooled: Final = bool(done.get("cooled"))
|
||||
if cooled and isinstance(deployment_id, str):
|
||||
seconds: Final = done.get("cooldown_time")
|
||||
event: Final = self.cooldown_event_on_own_loop if own_loop else self.cooldown_event
|
||||
event(
|
||||
deployment_id,
|
||||
exception_status,
|
||||
self.python_cooldown(seconds, configured, error) if isinstance(seconds, float) else None,
|
||||
)
|
||||
return cooled
|
||||
|
||||
def per_user_deployment(self, deployment_id: object) -> bool:
|
||||
"""The failed deployment authenticates per caller (``_should_run_cooldown_logic``), which is
|
||||
all ``set_cooldown_deployments`` reads when it is called without the request's kwargs."""
|
||||
deployment: Final = self._deployment_of(deployment_id)
|
||||
return deployment is not None and _authenticates_per_caller(deployment.litellm_params)
|
||||
|
||||
def failure_exits(
|
||||
self, exception: object, logged: Mapping[str, object], deployment_id: object
|
||||
) -> Mapping[str, bool]:
|
||||
"""What the Python router does not hold against the deployment, for every failure it logs.
|
||||
``outside_deployment``: ``deployment_callback_on_failure``'s early exits, counting nothing.
|
||||
``caller_scoped``: ``_should_run_cooldown_logic``'s per-caller exits, counted but not cooled;
|
||||
its exits on the exception's class are router-core's (``FailureKind``)."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
_is_per_user_provider_request, # pyright: ignore[reportPrivateUsage] # the cooldown check's own rule
|
||||
is_background_response_poll_not_found,
|
||||
is_caller_scoped_auth_failure,
|
||||
is_caller_timeout_408,
|
||||
is_failure_outside_selected_deployment,
|
||||
is_video_lookup_not_found,
|
||||
)
|
||||
|
||||
status: Final[object] = getattr(exception, "status_code", "")
|
||||
litellm_params: Final = as_str_mapping(logged.get("litellm_params")) or {}
|
||||
outside: Final = (
|
||||
is_failure_outside_selected_deployment(exception if isinstance(exception, BaseException) else None)
|
||||
or (isinstance(exception, Exception) and is_background_response_poll_not_found(exception, litellm_params))
|
||||
or is_video_lookup_not_found(exception, logged)
|
||||
or is_caller_timeout_408(logged, status) # pyright: ignore[reportArgumentType] # the status as raised
|
||||
or is_caller_scoped_auth_failure(self._deployment_of(deployment_id), status) # pyright: ignore[reportArgumentType] # see above
|
||||
)
|
||||
return {
|
||||
"outside_deployment": outside,
|
||||
"caller_scoped": _is_per_user_provider_request(logged) or self.per_user_deployment(deployment_id),
|
||||
}
|
||||
|
||||
def trigger_fallback_cooldown(
|
||||
self, kwargs: Mapping[str, object], exception: object, generic: bool, *, own_loop: bool = False
|
||||
) -> None:
|
||||
"""The failed fallback hop's ``_trigger_cooldown_for_failed_deployment``, which runs only when
|
||||
the request's logging object blocked the hop's failure callbacks. ``generic``: the call's kwargs
|
||||
carry ``original_generic_function``. ``own_loop``: see ``cooldown_event_on_own_loop``."""
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
_failure_callbacks_blocked_logging_object, # pyright: ignore[reportPrivateUsage] # the fallback path's own rule
|
||||
)
|
||||
|
||||
blocked: Final = _failure_callbacks_blocked_logging_object(kwargs)
|
||||
if blocked is None or not isinstance(exception, Exception):
|
||||
return
|
||||
try:
|
||||
call: Final = self._fallback_failure_call(kwargs, exception, blocked.model_call_details, generic)
|
||||
if call is not None:
|
||||
_ = self.cooled_down(
|
||||
self._drive_sync_steps(call),
|
||||
getattr(exception, "failed_deployment_id", None),
|
||||
getattr(exception, "status_code", ""),
|
||||
own_loop=own_loop,
|
||||
configured=_fallback_cooldown(
|
||||
self._deployment_of(getattr(exception, "failed_deployment_id", None))
|
||||
),
|
||||
error=exception,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best effort, as the Python helper: it never breaks the fallback response
|
||||
verbose_router_logger.debug("Error triggering cooldown for fallback deployment: %s", e)
|
||||
|
||||
def _fallback_failure_call(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
exception: Exception,
|
||||
model_call_details: Mapping[str, object],
|
||||
generic: bool,
|
||||
) -> "RouterCall[_FailureReportSteps] | None":
|
||||
"""The Rust call recording a failed hop, or None past the exits of
|
||||
``_trigger_cooldown_for_failed_deployment``: another list than ``failure_exits``, checked per
|
||||
failed hop. The selection tag is not on it; a stream stall and a generic call's 404 are."""
|
||||
from litellm.router_rust.core import attempt_failure
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
_is_per_user_provider_request, # pyright: ignore[reportPrivateUsage] # the cooldown check's own rule
|
||||
cast_exception_status_to_int,
|
||||
is_advisor_orchestration_failure,
|
||||
is_caller_scoped_auth_failure,
|
||||
is_caller_timeout_408,
|
||||
is_stream_stall,
|
||||
)
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
_REQUEST_SCOPED_STATUS_CODES, # pyright: ignore[reportPrivateUsage] # the fallback path's own rule
|
||||
)
|
||||
|
||||
self._sync_allowed_fails()
|
||||
exception_status: Final[object] = getattr(exception, "status_code", "")
|
||||
if (
|
||||
is_advisor_orchestration_failure(exception)
|
||||
or is_stream_stall(exception)
|
||||
or (generic and cast_exception_status_to_int(exception_status) in _REQUEST_SCOPED_STATUS_CODES) # pyright: ignore[reportArgumentType] # the status as raised
|
||||
or is_caller_timeout_408(model_call_details, exception_status, ended=datetime.now()) # pyright: ignore[reportArgumentType] # noqa: DTZ005 # see logged_call
|
||||
):
|
||||
return None
|
||||
deployment_id: Final[object] = getattr(exception, "failed_deployment_id", None)
|
||||
if not isinstance(deployment_id, str):
|
||||
return None
|
||||
deployment: Final = self._deployment_of(deployment_id)
|
||||
if is_caller_scoped_auth_failure(deployment, exception_status): # pyright: ignore[reportArgumentType] # see above
|
||||
return None
|
||||
return self._handle.report_failure(
|
||||
deployment_id,
|
||||
{
|
||||
**attempt_failure(exception),
|
||||
"status": failure_status(exception_status),
|
||||
"caller_scoped": _is_per_user_provider_request(kwargs) or self.per_user_deployment(deployment_id),
|
||||
},
|
||||
deployment_cooldown=True,
|
||||
count_only=exception_status == "",
|
||||
cooldown_time=None if deployment is None else _deployment_cooldown(deployment.model_dump()),
|
||||
)
|
||||
|
||||
def deployment_callback_on_failure(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
completion_response: object,
|
||||
start_time: object,
|
||||
end_time: object,
|
||||
) -> bool:
|
||||
"""Count the failure and cool the deployment down when the rules say so."""
|
||||
call: Final = self._failure_call(kwargs)
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
return call is not None and self.cooled_down(
|
||||
self._drive_sync_steps(call),
|
||||
_logged_deployment_id(kwargs),
|
||||
getattr(kwargs.get("exception"), "status_code", ""),
|
||||
configured=first_cooldown(as_str_mapping(litellm_params.get("model_info")), litellm_params),
|
||||
error=kwargs.get("exception"),
|
||||
)
|
||||
|
||||
async def report_stream_failure(
|
||||
self, kwargs: Mapping[str, object], callbacks: Mapping[str, bool] | None = None
|
||||
) -> None:
|
||||
"""A stream's failure after its attempt returned it, as an async call's failure callbacks
|
||||
record it: counted, cooled down by the rules, and one more request this minute
|
||||
(``async_deployment_callback_on_failure``). No cooldown event: the stream logs its failure
|
||||
through ``dispatch_failure_handlers``, whose sync callbacks run in an executor thread, where
|
||||
``set_cooldown_deployments`` has no loop to schedule one on."""
|
||||
call: Final = self._failure_call(kwargs, counts_request=True, callbacks=callbacks)
|
||||
if call is not None:
|
||||
_ = await self._drive_cache_steps(call)
|
||||
|
||||
def report_stream_failure_sync(
|
||||
self, kwargs: Mapping[str, object], callbacks: Mapping[str, bool] | None = None
|
||||
) -> None:
|
||||
"""``report_stream_failure`` for a sync stream: a sync call's failure callback counts no request.
|
||||
No cooldown event: the stream logs its failure in a ``threading.Thread`` of its own, which
|
||||
has no running loop whatever the consumer's thread has."""
|
||||
call: Final = self._failure_call(kwargs, callbacks=callbacks)
|
||||
if call is not None:
|
||||
_ = self._drive_sync_steps(call)
|
||||
|
||||
def _failure_call(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
counts_request: bool = False,
|
||||
callbacks: Mapping[str, bool] | None = None,
|
||||
) -> "RouterCall[_FailureReportSteps] | None":
|
||||
"""The Rust call recording a logged failure, or None when it records nothing: one the Python
|
||||
router does not hold against the deployment is still a request of the minute
|
||||
(``async_deployment_callback_on_failure`` makes none of the failure callback's exits)."""
|
||||
from litellm.litellm_core_utils import core_helpers
|
||||
from litellm.router_rust.core import attempt_failure
|
||||
|
||||
self._sync_allowed_fails()
|
||||
exception: Final = kwargs.get("exception")
|
||||
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
|
||||
model_info: Final = litellm_params.get("model_info", {})
|
||||
info: Final = as_str_mapping(model_info)
|
||||
deployment_id: Final = info.get("id") if info is not None else None
|
||||
if deployment_id is None:
|
||||
return None
|
||||
exits: Final = self.failure_exits(exception, kwargs, deployment_id)
|
||||
if exits["outside_deployment"] and not counts_request:
|
||||
return None
|
||||
exception_status: Final[object] = getattr(exception, "status_code", "")
|
||||
described: Final[Mapping[str, object]] = (
|
||||
attempt_failure(exception) if isinstance(exception, BaseException) else {"ok": False, "kind": "APIError"}
|
||||
)
|
||||
failure: Final = {**described, **exits, "status": failure_status(exception_status)}
|
||||
metadata: Final[object] = core_helpers.get_litellm_metadata_from_kwargs(dict(kwargs)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # an untyped helper's dict
|
||||
group: Final = (as_str_mapping(metadata) or {}).get("model_group") # pyright: ignore[reportUnknownArgumentType] # read as a mapping
|
||||
usage_model: Final = (as_str_mapping(metadata) or {}).get("deployment") # pyright: ignore[reportUnknownArgumentType] # see above
|
||||
return self._handle.report_failure(
|
||||
str(deployment_id),
|
||||
failure,
|
||||
deployment_cooldown=True,
|
||||
count_only=exception_status == "",
|
||||
requested_group=group if isinstance(group, str) else None,
|
||||
cooldown_time=_configured_cooldown(model_info, litellm_params),
|
||||
counts_request=counts_request,
|
||||
usage_model=usage_model if isinstance(usage_model, str) else None,
|
||||
callbacks=callbacks,
|
||||
)
|
||||
|
||||
def sync_deployment_callback_on_success(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
completion_response: object,
|
||||
start_time: object,
|
||||
end_time: object,
|
||||
) -> str | None:
|
||||
"""Count a success this minute; returns the counter's key."""
|
||||
return self.report_stream_success_sync(kwargs)
|
||||
|
||||
def report_stream_success_sync(
|
||||
self, kwargs: Mapping[str, object], callbacks: Mapping[str, bool] | None = None
|
||||
) -> str | None:
|
||||
"""``sync_deployment_callback_on_success``, for a sync stream that ended."""
|
||||
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
|
||||
|
||||
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||
return None
|
||||
target: Final = _success_target(kwargs)
|
||||
if target is None:
|
||||
return None
|
||||
key: Final = self._drive_sync_steps(self._handle.count_success(target[0], callbacks)).get("key")
|
||||
return key if isinstance(key, str) else None
|
||||
|
||||
async def repin_after_success(self, kwargs: Mapping[str, object], completion_response: object) -> None:
|
||||
"""The session re-pin ``deployment_callback_on_success`` runs for a success logged against a
|
||||
deployment the router serves (``_repin_session_after_success``). The Rust core runs it from the
|
||||
attempt itself, or from the stream wrapper when a stream ends."""
|
||||
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
|
||||
|
||||
call_type: Final = kwargs.get("call_type")
|
||||
if call_type == "_arealtime" or is_batch_retrieve_call_type(call_type):
|
||||
return
|
||||
target: Final = _success_target(kwargs)
|
||||
# RustCore sets these as attributes in its constructor.
|
||||
hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # RustCore's own attribute
|
||||
if target is None or hosted._core_get_deployment(model_id=target[0]) is None: # pyright: ignore[reportPrivateUsage] # the host registry's id lookup
|
||||
return
|
||||
checks: Final[OptionalChecks] = self.__dict__["_checks"] # pyright: ignore[reportAny] # see above
|
||||
cooldowns: Final[CooldownCache] = self.__dict__["_cooldowns"] # pyright: ignore[reportAny] # see above
|
||||
await checks.repin_session_after_success(
|
||||
as_str_mapping(kwargs.get("litellm_params")) or {}, completion_response, cooldowns
|
||||
)
|
||||
|
||||
async def deployment_callback_on_success(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
completion_response: object,
|
||||
start_time: object,
|
||||
end_time: object,
|
||||
) -> None:
|
||||
"""Count a success and the call's tokens and request against a tracked deployment."""
|
||||
await self.report_stream_success(kwargs, completion_response)
|
||||
|
||||
async def report_stream_success(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
completion_response: object,
|
||||
callbacks: Mapping[str, bool] | None = None,
|
||||
) -> None:
|
||||
"""``deployment_callback_on_success``, for a stream that ended. The re-pin is the host's half
|
||||
of that callback, so it reads ``callbacks`` here."""
|
||||
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import get_counted_usage_tokens
|
||||
|
||||
call_type: Final = kwargs.get("call_type")
|
||||
if call_type == "_arealtime" or is_batch_retrieve_call_type(call_type):
|
||||
return
|
||||
logged: Final = as_str_mapping(kwargs.get("standard_logging_object"))
|
||||
if logged is None:
|
||||
return
|
||||
target: Final = _success_target(kwargs)
|
||||
if target is None:
|
||||
return
|
||||
if (callbacks or {}).get("success", True):
|
||||
await self.repin_after_success(kwargs, completion_response)
|
||||
tokens: Final[object] = logged.get("total_tokens", 0)
|
||||
total: Final = float(tokens) if isinstance(tokens, int | float) else 0.0
|
||||
counted: Final = get_counted_usage_tokens(as_str_mapping(kwargs.get("litellm_params")) or {})
|
||||
usage_model: Final = target[1].get("deployment")
|
||||
await self._drive_cache_steps(
|
||||
self._handle.report_success(
|
||||
target[0],
|
||||
total if counted is None else max(0.0, total - counted),
|
||||
1.0 if counted is None else 0.0,
|
||||
usage_model if isinstance(usage_model, str) else None,
|
||||
self.deployment_tracks_usage(target[0], target[1].get("model_group")),
|
||||
callbacks,
|
||||
)
|
||||
)
|
||||
|
||||
async def count_realtime_session_success(self, deployment_id: str) -> None:
|
||||
"""``increment_deployment_successes_for_current_minute`` for a realtime session that ended
|
||||
cleanly: a success alone, as ``deployment_callback_on_success`` skips ``_arealtime``. A
|
||||
counter that cannot be written does not fail the served session."""
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
||||
try:
|
||||
await self._drive_cache_steps(self._handle.count_success(deployment_id))
|
||||
except Exception as e: # noqa: BLE001 # as the Rust call loop treats a success it could not record
|
||||
verbose_router_logger.debug("count_realtime_session_success: not recorded. Got - %s", e)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"RustHealthCallbacks",
|
||||
"failure_status",
|
||||
"logged_call",
|
||||
]
|
||||
160
litellm/router_rust/surface/health_routing.py
Normal file
160
litellm/router_rust/surface/health_routing.py
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
"""
|
||||
Health-check routing, allowed-fails policies and the failover flags on the host side of the Rust
|
||||
core: how the constructor resolves them, the static document router-core reads them from, and the
|
||||
``health()`` members that read the registry's live model list or the background health checks'
|
||||
results. router-core makes every routing decision (the health-check filter, the cooldown rules,
|
||||
weighted failover); this module only reads Python values into its documents and back.
|
||||
|
||||
The rules are the Python router core's, ported (never imported).
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
from litellm.types.router import AllowedFailsPolicy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_config import RouterConfig
|
||||
from litellm.types.router import Deployment, DeploymentTypedDict
|
||||
|
||||
HEALTH_STATE_KEY: Final = DeploymentHealthCache.CACHE_KEY
|
||||
# The failover flags the Python core keeps as plain router attributes, with their defaults.
|
||||
FAILOVER_FLAGS: Final = ("enable_weighted_failover", "fallback_on_unsupported_params")
|
||||
|
||||
|
||||
def resolved_allowed_fails_policy(value: object) -> AllowedFailsPolicy | None:
|
||||
"""The router's ``allowed_fails_policy`` as the constructor keeps it: a dict is read into an
|
||||
``AllowedFailsPolicy`` (a bad field raises the model's own error), a policy is kept, anything
|
||||
else is no policy."""
|
||||
if isinstance(value, AllowedFailsPolicy):
|
||||
return value
|
||||
if isinstance(value, dict):
|
||||
return AllowedFailsPolicy.model_validate(value)
|
||||
return None
|
||||
|
||||
|
||||
def _staleness(threshold: object) -> float:
|
||||
"""``health_check_staleness_threshold or DEFAULT_HEALTH_CHECK_INTERVAL * multiplier``."""
|
||||
from litellm.constants import DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER
|
||||
|
||||
given: Final = threshold if isinstance(threshold, int | float) and not isinstance(threshold, bool) else None
|
||||
return float(given or (DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER))
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class HealthRouting:
|
||||
"""The health-check routing settings, as the constructor resolves them."""
|
||||
|
||||
enable_health_check_routing: bool = False
|
||||
health_check_staleness: float = 600.0
|
||||
health_check_ignore_transient_errors: bool = False
|
||||
background_health_check_model_groups: frozenset[str] | None = None
|
||||
|
||||
@classmethod
|
||||
def of(cls, config: "RouterConfig") -> "HealthRouting":
|
||||
groups: Final = config.get("background_health_check_model_groups")
|
||||
return cls(
|
||||
enable_health_check_routing=bool(config.get("enable_health_check_routing")),
|
||||
health_check_staleness=_staleness(config.get("health_check_staleness_threshold")),
|
||||
health_check_ignore_transient_errors=bool(config.get("health_check_ignore_transient_errors")),
|
||||
background_health_check_model_groups=frozenset(groups) if groups is not None else None, # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # any iterable of names, as the Python core takes it
|
||||
)
|
||||
|
||||
def static_document(self) -> Mapping[str, object]:
|
||||
groups: Final = self.background_health_check_model_groups
|
||||
return {
|
||||
"enable_health_check_routing": True if self.enable_health_check_routing else None,
|
||||
"health_check_staleness_threshold": self.health_check_staleness,
|
||||
# Only a string can name a model group, so anything else in scope matches nothing.
|
||||
"background_health_check_model_groups": sorted(name for name in groups if isinstance(name, str)) # pyright: ignore[reportUnnecessaryIsInstance] # `of` keeps any iterable the constructor was given
|
||||
if groups is not None
|
||||
else None,
|
||||
}
|
||||
|
||||
def health_checks(self, cache: "DualCache") -> DeploymentHealthCache:
|
||||
return DeploymentHealthCache(cache=cache, staleness_threshold=self.health_check_staleness)
|
||||
|
||||
|
||||
def failover_document(plain: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The failover flags router-core reads, from the router's plain attributes (read for
|
||||
truthiness, as the Python core's checks read them)."""
|
||||
return {
|
||||
"enable_weighted_failover": True if plain.get("enable_weighted_failover") else None,
|
||||
"fallback_on_unsupported_params": False if not plain.get("fallback_on_unsupported_params", True) else None,
|
||||
}
|
||||
|
||||
|
||||
def _model_info(deployment: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return as_str_mapping(deployment.get("model_info")) or {}
|
||||
|
||||
|
||||
def is_blocked(deployment: "Deployment") -> bool:
|
||||
"""``_core_is_deployment_blocked``."""
|
||||
model_info: Final[object | None] = getattr(deployment, "model_info", None)
|
||||
if model_info is None:
|
||||
return False
|
||||
return getattr(model_info, "blocked", None) is True
|
||||
|
||||
|
||||
def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool:
|
||||
"""``_core_are_all_deployments_blocked``."""
|
||||
return len(deployments) > 0 and all(
|
||||
_model_info(deployment).get("blocked") is True # pyright: ignore[reportArgumentType] # a deployment dict
|
||||
for deployment in deployments
|
||||
)
|
||||
|
||||
|
||||
def fully_blocked_model_names(
|
||||
model_list: Sequence[Mapping[str, object]],
|
||||
) -> set[str]: # mutable-ok: the core's type, a fresh set
|
||||
"""``_core_get_fully_blocked_model_names``: groups whose every deployment is blocked."""
|
||||
blocked_by_name: Final[dict[str, bool]] = {} # mutable-ok: built here
|
||||
for deployment in model_list:
|
||||
name = deployment.get("model_name") or ""
|
||||
if not isinstance(name, str) or not name:
|
||||
continue
|
||||
blocked = _model_info(deployment).get("blocked") is True
|
||||
blocked_by_name[name] = blocked_by_name.get(name, True) and blocked
|
||||
return {name for name, fully in blocked_by_name.items() if fully}
|
||||
|
||||
|
||||
async def fully_unhealthy_model_names(
|
||||
model_list: Sequence[Mapping[str, object]], health_checks: DeploymentHealthCache, policy: object
|
||||
) -> set[str]: # mutable-ok: the core's type, a fresh set
|
||||
"""``_core_async_get_fully_unhealthy_model_names``: groups (and team public names) whose
|
||||
every deployment failed its health check; nothing under an allowed_fails_policy."""
|
||||
if policy is not None:
|
||||
return set()
|
||||
unhealthy_ids: Final = await health_checks.async_get_unhealthy_deployment_ids()
|
||||
if not unhealthy_ids:
|
||||
return set()
|
||||
unhealthy_by_name: Final[dict[str, bool]] = {} # mutable-ok: built here
|
||||
for deployment in model_list:
|
||||
model_info = _model_info(deployment)
|
||||
names = [deployment.get("model_name") or ""]
|
||||
team_public_model_name = model_info.get("team_public_model_name")
|
||||
if team_public_model_name:
|
||||
names.append(team_public_model_name)
|
||||
unhealthy = model_info.get("id") in unhealthy_ids
|
||||
for name in names:
|
||||
if not isinstance(name, str) or not name:
|
||||
continue
|
||||
unhealthy_by_name[name] = unhealthy_by_name.get(name, True) and unhealthy
|
||||
return {name for name, fully in unhealthy_by_name.items() if fully}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FAILOVER_FLAGS",
|
||||
"HEALTH_STATE_KEY",
|
||||
"HealthRouting",
|
||||
"all_blocked",
|
||||
"failover_document",
|
||||
"fully_blocked_model_names",
|
||||
"fully_unhealthy_model_names",
|
||||
"is_blocked",
|
||||
"resolved_allowed_fails_policy",
|
||||
]
|
||||
514
litellm/router_rust/surface/indexes.py
Normal file
514
litellm/router_rust/surface/indexes.py
Normal file
|
|
@ -0,0 +1,514 @@
|
|||
# pyright: reportIncompatibleMethodOverride=false
|
||||
# (each view overrides its builtin's mutators with the same behaviour, as a changeset)
|
||||
"""
|
||||
``HostedIndexes``: the registry indexes the Rust core hands out, read back from router-core.
|
||||
|
||||
router-core keeps the indexes the Python router keeps by history (``host_index.rs``): the id,
|
||||
group and team position maps, the group names a ``set_model_list`` was given, the deployment
|
||||
names that only ever grow, the team public names, the provider default ids and the wildcard
|
||||
pattern routers. ``registry().index()``, ``names(kind=...)``, ``patterns()`` and the deprecated
|
||||
attributes return the views here: a ``dict``, ``set`` or ``list`` (so ``isinstance`` and ``json``
|
||||
keep working), one object per router, whose contents are router-core's, read again when an entry
|
||||
point hands a view out (or before a write to it) once the registry moved, not on every registry
|
||||
write. A write to any of them, a nested one included (``index("model_name")["gpt"].append(3)``),
|
||||
is a ``host_index`` changeset that replaces that index in router-core, and the view reads it back.
|
||||
The lookups over these indexes run in router-core (``lookup.rs``), not over the views.
|
||||
"""
|
||||
|
||||
import copy
|
||||
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from functools import partial
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Literal,
|
||||
SupportsIndex,
|
||||
TypeAlias,
|
||||
cast, # noqa: TID251 # any value a caller wrote into a view
|
||||
)
|
||||
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.rust_bridge._native import RouterCoreHandle
|
||||
|
||||
IndexField: TypeAlias = Literal[
|
||||
"id_map",
|
||||
"model_name",
|
||||
"team_model",
|
||||
"model_names",
|
||||
"deployment_names",
|
||||
"team_public_model_names",
|
||||
"provider_default_ids",
|
||||
"patterns",
|
||||
"team_patterns",
|
||||
]
|
||||
_Commit: TypeAlias = Callable[[object], None]
|
||||
# Deployment copies a pattern router routes to, fresh for the caller.
|
||||
Routed: TypeAlias = list[dict[str, object]]
|
||||
_Pairs: TypeAlias = list[tuple[str, list[object]]]
|
||||
|
||||
|
||||
def _stay() -> None:
|
||||
"""A view with no contents of its own to read again before a write."""
|
||||
|
||||
|
||||
def _as_set(value: object) -> set[object]: # mutable-ok: set's own in-place operand
|
||||
return cast(set[object], value) # cast-ok: an in-place set operand, checked by set itself
|
||||
|
||||
|
||||
def _as_mapping(value: object) -> Mapping[object, object]:
|
||||
if not isinstance(value, Mapping):
|
||||
return {}
|
||||
return cast(Mapping[object, object], value) # cast-ok: any mapping a caller wrote
|
||||
|
||||
|
||||
def _json(value: object, depth: int = 0) -> object:
|
||||
"""A view's contents as router-core takes them: plain JSON, a set as a list. A caller wrote them,
|
||||
so a cycle is possible: anything nested past ``DEFAULT_MAX_RECURSE_DEPTH`` is refused."""
|
||||
if depth > DEFAULT_MAX_RECURSE_DEPTH:
|
||||
raise ValueError(f"a router index is nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels")
|
||||
mapping: Final = _as_mapping(value)
|
||||
if mapping or isinstance(value, Mapping):
|
||||
return {str(key): _json(item, depth + 1) for key, item in mapping.items()}
|
||||
if isinstance(value, list | tuple | set | frozenset):
|
||||
items: Final = cast(Iterable[object], value) # cast-ok: any collection a caller wrote
|
||||
return [_json(item, depth + 1) for item in items]
|
||||
return value
|
||||
|
||||
|
||||
class _Positions(list[object]):
|
||||
"""One index's position list (or a pattern's deployments): a write commits the whole list."""
|
||||
|
||||
def __init__(self, items: Iterable[object], commit: _Commit, fresh: Callable[[], None] = _stay) -> None:
|
||||
super().__init__(items)
|
||||
self._commit: Final = commit
|
||||
self._fresh: Final = fresh
|
||||
|
||||
def _load(self, items: Iterable[object]) -> None:
|
||||
"""Show router-core's contents, without a write."""
|
||||
super().__setitem__(slice(None), list(items))
|
||||
|
||||
def _write(self, change: Callable[[list[object]], object]) -> object: # mutable-ok: the copy it changes
|
||||
self._fresh() # a view held across a write elsewhere changes router-core's contents, not stale ones
|
||||
changed: Final = list(self)
|
||||
result: Final = change(changed)
|
||||
self._commit(changed)
|
||||
return result
|
||||
|
||||
def append(self, value: object) -> None:
|
||||
_ = self._write(lambda items: items.append(value))
|
||||
|
||||
def extend(self, values: Iterable[object]) -> None:
|
||||
added: Final = tuple(values)
|
||||
_ = self._write(lambda items: items.extend(added))
|
||||
|
||||
def insert(self, index: SupportsIndex, value: object) -> None:
|
||||
_ = self._write(lambda items: items.insert(index, value))
|
||||
|
||||
def remove(self, value: object) -> None:
|
||||
_ = self._write(lambda items: items.remove(value))
|
||||
|
||||
def pop(self, index: SupportsIndex = -1) -> object:
|
||||
return self._write(lambda items: items.pop(index))
|
||||
|
||||
def clear(self) -> None:
|
||||
_ = self._write(lambda items: items.clear())
|
||||
|
||||
def sort(self, *args: object, **kwargs: object) -> None: # kwargs-ok: list.sort's own arguments
|
||||
_ = self._write(lambda items: items.sort(*args, **kwargs)) # pyright: ignore[reportCallIssue, reportArgumentType] # list.sort's own arguments
|
||||
|
||||
def reverse(self) -> None:
|
||||
_ = self._write(lambda items: items.reverse()) # pyright: ignore[reportUnknownLambdaType] # list.reverse returns None
|
||||
|
||||
def __setitem__(self, index: object, value: object) -> None:
|
||||
_ = self._write(lambda items: items.__setitem__(index, value)) # pyright: ignore[reportCallIssue, reportArgumentType, reportUnknownLambdaType] # list's own item and slice assignment
|
||||
|
||||
def __delitem__(self, index: object) -> None:
|
||||
_ = self._write(lambda items: items.__delitem__(index)) # pyright: ignore[reportArgumentType] # list's own item and slice deletion
|
||||
|
||||
def __iadd__(self, values: Iterable[object]) -> "_Positions":
|
||||
self.extend(values)
|
||||
return self
|
||||
|
||||
def __imul__(self, count: SupportsIndex) -> "_Positions":
|
||||
_ = self._write(lambda items: items.__imul__(count))
|
||||
return self
|
||||
|
||||
def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]:
|
||||
return (list, (list(self),))
|
||||
|
||||
|
||||
class _IndexDict(dict[object, object]):
|
||||
"""A key-to-positions index: a write commits the whole index."""
|
||||
|
||||
def __init__(self, commit: _Commit, fresh: Callable[[], None] = _stay) -> None:
|
||||
super().__init__()
|
||||
self._commit: Final = commit
|
||||
self._fresh: Final = fresh
|
||||
|
||||
def _load(self, items: Mapping[object, object]) -> None:
|
||||
"""Show router-core's contents, without a write."""
|
||||
super().clear()
|
||||
super().update(items)
|
||||
|
||||
def _write(self, change: Callable[[dict[object, object]], object]) -> object: # mutable-ok: the copy it changes
|
||||
self._fresh() # a view held across a write elsewhere changes router-core's contents, not stale ones
|
||||
changed: Final = dict(self)
|
||||
result: Final = change(changed)
|
||||
self._commit(changed)
|
||||
return result
|
||||
|
||||
def __setitem__(self, key: object, value: object) -> None:
|
||||
_ = self._write(lambda items: items.__setitem__(key, value))
|
||||
|
||||
def __delitem__(self, key: object) -> None:
|
||||
_ = self._write(lambda items: items.__delitem__(key))
|
||||
|
||||
def pop(self, key: object, *default: object) -> object:
|
||||
return self._write(lambda items: items.pop(key, *default))
|
||||
|
||||
def popitem(self) -> tuple[object, object]:
|
||||
return cast(tuple[object, object], self._write(lambda items: items.popitem())) # cast-ok: dict.popitem's pair
|
||||
|
||||
def clear(self) -> None:
|
||||
_ = self._write(lambda items: items.clear())
|
||||
|
||||
def update(self, *args: object, **kwargs: object) -> None: # kwargs-ok: dict.update's own arguments
|
||||
_ = self._write(lambda items: items.update(*args, **kwargs)) # pyright: ignore[reportCallIssue, reportArgumentType] # dict.update's own arguments
|
||||
|
||||
def setdefault(self, key: object, default: object = None) -> object:
|
||||
if key in self:
|
||||
return self[key]
|
||||
self[key] = default
|
||||
return self[key]
|
||||
|
||||
def __ior__(self, other: object) -> "_IndexDict":
|
||||
self.update(other)
|
||||
return self
|
||||
|
||||
def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]:
|
||||
return (dict, (dict(self),))
|
||||
|
||||
|
||||
class _NameSet(set[object]):
|
||||
"""``model_names``: a write commits the whole set."""
|
||||
|
||||
def __init__(self, commit: _Commit, fresh: Callable[[], None] = _stay) -> None:
|
||||
super().__init__()
|
||||
self._commit: Final = commit
|
||||
self._fresh: Final = fresh
|
||||
|
||||
def _load(self, items: Iterable[object]) -> None:
|
||||
"""Show router-core's contents, without a write."""
|
||||
super().clear()
|
||||
super().update(items)
|
||||
|
||||
def _write(self, change: Callable[[set[object]], object]) -> object: # mutable-ok: the copy it changes
|
||||
self._fresh() # a view held across a write elsewhere changes router-core's contents, not stale ones
|
||||
changed: Final = set(self)
|
||||
result: Final = change(changed)
|
||||
self._commit(changed)
|
||||
return result
|
||||
|
||||
def add(self, value: object) -> None:
|
||||
_ = self._write(lambda items: items.add(value))
|
||||
|
||||
def discard(self, value: object) -> None:
|
||||
_ = self._write(lambda items: items.discard(value))
|
||||
|
||||
def remove(self, value: object) -> None:
|
||||
_ = self._write(lambda items: items.remove(value))
|
||||
|
||||
def pop(self) -> object:
|
||||
return self._write(lambda items: items.pop())
|
||||
|
||||
def clear(self) -> None:
|
||||
_ = self._write(lambda items: items.clear())
|
||||
|
||||
def update(self, *others: Iterable[object]) -> None:
|
||||
_ = self._write(lambda items: items.update(*others))
|
||||
|
||||
def difference_update(self, *others: Iterable[object]) -> None:
|
||||
_ = self._write(lambda items: items.difference_update(*others))
|
||||
|
||||
def intersection_update(self, *others: Iterable[object]) -> None:
|
||||
_ = self._write(lambda items: items.intersection_update(*others))
|
||||
|
||||
def symmetric_difference_update(self, other: Iterable[object]) -> None:
|
||||
_ = self._write(lambda items: items.symmetric_difference_update(other))
|
||||
|
||||
def __ior__(self, other: object) -> "_NameSet":
|
||||
_ = self._write(lambda items: items.__ior__(_as_set(other)))
|
||||
return self
|
||||
|
||||
def __iand__(self, other: object) -> "_NameSet":
|
||||
_ = self._write(lambda items: items.__iand__(_as_set(other)))
|
||||
return self
|
||||
|
||||
def __isub__(self, other: object) -> "_NameSet":
|
||||
_ = self._write(lambda items: items.__isub__(_as_set(other)))
|
||||
return self
|
||||
|
||||
def __ixor__(self, other: object) -> "_NameSet":
|
||||
_ = self._write(lambda items: items.__ixor__(_as_set(other)))
|
||||
return self
|
||||
|
||||
def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]:
|
||||
return (set, (set(self),))
|
||||
|
||||
|
||||
class HostedPatternRouter(PatternMatchRouter):
|
||||
"""A wildcard pattern router whose patterns are router-core's (``pattern_router``, or one team's
|
||||
router): routing runs in router-core, and ``add_pattern`` / ``remove_deployment`` / a write to
|
||||
``patterns`` are changesets."""
|
||||
|
||||
def __init__(self, owner: "HostedIndexes", team: str | None) -> None:
|
||||
super().__init__() # its empty `patterns` write is ignored: router-core holds them
|
||||
self._owner: Final = owner
|
||||
self._team: Final = team
|
||||
self._view: Final = _IndexDict(self._commit, owner.refresh)
|
||||
|
||||
@property
|
||||
def patterns(self) -> dict[str, list[object]]: # pyright: ignore[reportIncompatibleVariableOverride] # the base class's attribute, read from router-core # mutable-ok: the live patterns, writes are changesets
|
||||
self._owner.refresh()
|
||||
return cast(dict[str, list[object]], self._view) # cast-ok: regex keys, deployment lists
|
||||
|
||||
@patterns.setter
|
||||
def patterns(self, value: Mapping[str, object]) -> None: # pyright: ignore[reportIncompatibleVariableOverride] # see the getter
|
||||
if not hasattr(self, "_owner"):
|
||||
return # the base constructor's empty dict: router-core holds this router's patterns
|
||||
self._commit(value)
|
||||
|
||||
def _show(self, pairs: Iterable[tuple[str, Sequence[object]]]) -> None:
|
||||
shown: Final = {
|
||||
regex: _Positions(deployments, partial(self._write_pattern, regex)) for regex, deployments in pairs
|
||||
}
|
||||
self._view._load(cast(Mapping[object, object], shown)) # pyright: ignore[reportPrivateUsage] # cast-ok: regex to deployments
|
||||
|
||||
def _write_pattern(self, regex: str, deployments: object) -> None:
|
||||
self._commit({**self.patterns, regex: deployments})
|
||||
|
||||
def _commit(self, value: object) -> None:
|
||||
pairs: Final = [
|
||||
[str(regex), _json(deployments)] for regex, deployments in (as_str_mapping(value) or {}).items()
|
||||
]
|
||||
self._owner.write_router(self._team, pairs)
|
||||
|
||||
def route(self, request: str | None, filtered_model_names: Sequence[str] | None = None) -> Routed | None: # pyright: ignore[reportIncompatibleMethodOverride] # the base's list types, read only here
|
||||
if not isinstance(request, str):
|
||||
# The base class matches its regexes against it, which fails and routes nowhere.
|
||||
return super().route(request, None if filtered_model_names is None else list(filtered_model_names)) # pyright: ignore[reportReturnType, reportUnknownMemberType, reportUnknownVariableType] # None, or the base's copies of the deployments
|
||||
return self._owner.route(self._team, request, filtered_model_names)
|
||||
|
||||
def get_models_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> frozenset[str]:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
|
||||
|
||||
if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model goes to the base class
|
||||
return super().get_models_by_pattern(model, custom_llm_provider)
|
||||
provider: Final = custom_llm_provider or inferred_provider(model)
|
||||
return self._owner.models_by_pattern(self._team, [model, f"{provider}/{model}"])
|
||||
|
||||
|
||||
class _TeamRouters(dict[str, HostedPatternRouter]):
|
||||
"""``team_pattern_routers``: team id to its router; removing a team's router is a changeset."""
|
||||
|
||||
def __init__(self, owner: "HostedIndexes") -> None:
|
||||
super().__init__()
|
||||
self._owner: Final = owner
|
||||
|
||||
def _load(self, routers: Mapping[str, HostedPatternRouter]) -> None:
|
||||
super().clear()
|
||||
super().update(routers)
|
||||
|
||||
def __delitem__(self, team: str) -> None:
|
||||
self._owner.drop_team_router(team)
|
||||
|
||||
def pop(self, team: str, *default: object) -> object:
|
||||
if team not in self:
|
||||
if default:
|
||||
return default[0]
|
||||
raise KeyError(team)
|
||||
router: Final = self[team]
|
||||
self._owner.drop_team_router(team)
|
||||
return router
|
||||
|
||||
def clear(self) -> None:
|
||||
self._owner.write({"team_patterns": []})
|
||||
|
||||
def setdefault(self, team: str, default: HostedPatternRouter | None = None) -> HostedPatternRouter:
|
||||
if team not in self:
|
||||
self._owner.write_router(team, [])
|
||||
return self[team]
|
||||
|
||||
|
||||
class HostedIndexes:
|
||||
"""Router-core's host indexes, as the live objects the registry hands out."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
handle: Callable[[], "RouterCoreHandle | None"],
|
||||
apply: Callable[[Mapping[str, object]], None],
|
||||
) -> None:
|
||||
self._handle: Final = handle
|
||||
self._apply: Final = apply
|
||||
self._read_at: tuple[object, int] | None = None
|
||||
self._document: Mapping[str, object] = {}
|
||||
# The views are read again when an entry point hands them out, and before a write to them.
|
||||
self.id_map: Final = _IndexDict(lambda value: self.write({"id_map": value}), self.refresh)
|
||||
self.model_name: Final = _IndexDict(lambda value: self.write({"model_name": value}), self.refresh)
|
||||
self.team_model: Final = _IndexDict(lambda value: self._write_team_model(value), self.refresh)
|
||||
self.model_names: Final = _NameSet(lambda value: self.write({"model_names": value}), self.refresh)
|
||||
self.deployment_names: Final = _Positions(
|
||||
(), lambda value: self.write({"deployment_names": value}), self.refresh
|
||||
)
|
||||
self.provider_default_ids: Final = _Positions(
|
||||
(), lambda value: self.write({"provider_default_ids": value}), self.refresh
|
||||
)
|
||||
self.pattern_router: Final = HostedPatternRouter(self, None)
|
||||
self.team_pattern_routers: Final = _TeamRouters(self)
|
||||
self._team_routers: dict[str, HostedPatternRouter] = {} # mutable-ok: one router object per team
|
||||
self._team_public: frozenset[str] = frozenset()
|
||||
|
||||
# -- reading ----------------------------------------------------------------------------------
|
||||
|
||||
def refresh(self) -> None:
|
||||
"""Read router-core's indexes again when its registry moved since the last read."""
|
||||
handle: Final = self._handle()
|
||||
if handle is None:
|
||||
return
|
||||
version: Final = handle.registry_version()
|
||||
if self._read_at is not None and self._read_at[0] is handle and self._read_at[1] == version:
|
||||
return
|
||||
self._read_at = (handle, version)
|
||||
self._document = as_str_mapping(handle.host_indexes()) or {}
|
||||
self._show()
|
||||
|
||||
def _field(self, name: str) -> object:
|
||||
return self._document.get(name)
|
||||
|
||||
def _show(self) -> None:
|
||||
ids: Final = _as_mapping(self._field("id_map"))
|
||||
self.id_map._load(ids) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
named: Final[dict[object, object]] = {} # mutable-ok: built here
|
||||
for name, positions in (as_str_mapping(self._field("model_name")) or {}).items():
|
||||
named[name] = self._positions(positions, partial(self._write_named, name))
|
||||
self.model_name._load(named) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
teams: Final[dict[object, object]] = {} # mutable-ok: built here
|
||||
for team, public, positions in _triples(self._field("team_model")):
|
||||
teams[(team, public)] = self._positions(positions, partial(self._write_team_key, (team, public)))
|
||||
self.team_model._load(teams) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
self.model_names._load(_items(self._field("model_names"))) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
self.deployment_names._load(_items(self._field("deployment_names"))) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
self.provider_default_ids._load(_items(self._field("provider_default_ids"))) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
self._team_public = frozenset(str(name) for name in _items(self._field("team_public_model_names")))
|
||||
self.pattern_router._show(_pattern_pairs(self._field("patterns"))) # pyright: ignore[reportPrivateUsage] # the router this object owns
|
||||
routers: Final[dict[str, HostedPatternRouter]] = {} # mutable-ok: rebuilt each read
|
||||
for team, pairs in _team_pairs(self._field("team_patterns")):
|
||||
router = self._team_routers.get(team) or HostedPatternRouter(self, team)
|
||||
router._show(pairs) # pyright: ignore[reportPrivateUsage] # the router this object owns
|
||||
routers[team] = router
|
||||
self._team_routers = routers
|
||||
self.team_pattern_routers._load(routers) # pyright: ignore[reportPrivateUsage] # a view this object owns
|
||||
|
||||
def _write_named(self, name: str, items: object) -> None:
|
||||
self.write({"model_name": {**self.model_name, name: items}})
|
||||
|
||||
def _write_team_key(self, key: tuple[str, str], items: object) -> None:
|
||||
self._write_team_model({**self.team_model, key: items})
|
||||
|
||||
def _positions(self, positions: object, commit: _Commit) -> object:
|
||||
"""A position list as a write-through list; anything else a caller wrote, as written."""
|
||||
items: Final = _items(positions)
|
||||
return _Positions(items, commit) if isinstance(positions, list) else positions
|
||||
|
||||
@property
|
||||
def team_public_model_names(self) -> frozenset[str]:
|
||||
self.refresh()
|
||||
return self._team_public
|
||||
|
||||
# -- writing ----------------------------------------------------------------------------------
|
||||
|
||||
def write(self, fields: Mapping[str, object]) -> None:
|
||||
"""Replace the named indexes in router-core, then read them back."""
|
||||
self._apply({"host_index": {name: _json(value) for name, value in fields.items()}})
|
||||
self.refresh()
|
||||
|
||||
def assign(self, field: IndexField, value: object) -> None:
|
||||
"""A deprecated attribute assignment of one index (``router.model_names = {...}``)."""
|
||||
if field == "team_model":
|
||||
self._write_team_model(value)
|
||||
else:
|
||||
self.write({field: value})
|
||||
|
||||
def _write_team_model(self, value: object) -> None:
|
||||
mapping: Final = _as_mapping(value)
|
||||
triples: Final[list[list[object]]] = [] # mutable-ok: JSON triples
|
||||
for key, positions in mapping.items():
|
||||
parts = cast(tuple[object, ...], key) if isinstance(key, tuple) else () # cast-ok: a key a caller wrote
|
||||
if len(parts) == 2:
|
||||
triples.append([str(parts[0]), str(parts[1]), _json(positions)])
|
||||
self.write({"team_model": triples})
|
||||
|
||||
def write_router(self, team: str | None, pairs: list[list[object]]) -> None: # mutable-ok: JSON pairs
|
||||
"""Replace one pattern router's patterns."""
|
||||
if team is None:
|
||||
self.write({"patterns": pairs})
|
||||
return
|
||||
teams: Final = [[t, _json(_pairs_of(router))] for t, router in self._team_routers.items() if t != team]
|
||||
self.write({"team_patterns": [*teams, [team, pairs]]})
|
||||
|
||||
def drop_team_router(self, team: str) -> None:
|
||||
if team not in self._team_routers:
|
||||
raise KeyError(team)
|
||||
self.write({"team_patterns": [[t, _json(_pairs_of(r))] for t, r in self._team_routers.items() if t != team]})
|
||||
|
||||
# -- routing ----------------------------------------------------------------------------------
|
||||
|
||||
def route(self, team: str | None, request: str, filtered: Sequence[str] | None) -> Routed | None:
|
||||
handle: Final = self._handle()
|
||||
routed: Final = (
|
||||
None
|
||||
if handle is None
|
||||
else handle.pattern_route(team, request, None if filtered is None else list(filtered))
|
||||
)
|
||||
return None if routed is None else [dict(deployment) for deployment in routed]
|
||||
|
||||
def models_by_pattern(self, team: str | None, requests: list[str]) -> frozenset[str]: # mutable-ok: read only
|
||||
handle: Final = self._handle()
|
||||
found: Final = None if handle is None else handle.models_by_pattern(team, requests)
|
||||
return frozenset(found or ())
|
||||
|
||||
|
||||
def _items(value: object) -> list[object]: # mutable-ok: a fresh list
|
||||
return list(cast(list[object], value)) if isinstance(value, list) else [] # cast-ok: a JSON list
|
||||
|
||||
|
||||
def _triples(value: object) -> Iterator[tuple[str, str, object]]:
|
||||
for item in _items(value):
|
||||
parts = _items(item)
|
||||
if len(parts) == 3:
|
||||
yield str(parts[0]), str(parts[1]), parts[2]
|
||||
|
||||
|
||||
def _pattern_pairs(value: object) -> _Pairs:
|
||||
pairs: Final[_Pairs] = []
|
||||
for item in _items(value):
|
||||
parts = _items(item)
|
||||
if len(parts) == 2:
|
||||
pairs.append((str(parts[0]), _items(parts[1])))
|
||||
return pairs
|
||||
|
||||
|
||||
def _team_pairs(value: object) -> Iterator[tuple[str, _Pairs]]:
|
||||
for item in _items(value):
|
||||
parts = _items(item)
|
||||
if len(parts) == 2:
|
||||
yield str(parts[0]), _pattern_pairs(parts[1])
|
||||
|
||||
|
||||
def _pairs_of(router: HostedPatternRouter) -> list[list[object]]: # mutable-ok: JSON pairs
|
||||
view: Final[Mapping[object, object]] = router._view # pyright: ignore[reportPrivateUsage] # a router this module owns
|
||||
return [[regex, copy.deepcopy(deployments)] for regex, deployments in view.items()]
|
||||
547
litellm/router_rust/surface/legacy_members.py
Normal file
547
litellm/router_rust/surface/legacy_members.py
Normal file
|
|
@ -0,0 +1,547 @@
|
|||
"""
|
||||
``RustLegacyMembers``: the old public ``Router`` members outside the entry points, on the Rust
|
||||
core: the routing-state lifecycle, the routing callbacks' checks, one provider call's accounting
|
||||
and the pre-routing strategy builders. Each one is the Python router core's, ported (never
|
||||
imported), over the state the Rust core and its host keep.
|
||||
|
||||
Not served, and refused by ``RustCore.__getattr__``: the Python fallback walk
|
||||
(``async_function_with_fallbacks``, ``function_with_fallbacks``, ``async_function_with_retries``,
|
||||
``async_function_with_fallbacks_common_utils``) and ``factory_function``. They run a caller's
|
||||
``original_function``, which picks its own deployment; router-core makes each pick itself and
|
||||
attempts one host op, named by the endpoint's plan.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import traceback
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm._internal_context import service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.add_retry_fallback_headers import response_total_token_count
|
||||
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_parent_otel_span_from_kwargs
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
from litellm.litellm_core_utils.routing_groups import validate_routing_strategy
|
||||
from litellm.litellm_core_utils.token_counter import served_openai_models as _served_openai_models
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.router_rust.invoke.websocket import admit_deployment_turn
|
||||
from litellm.router_rust.steps.callbacks import RustOptionalChecks
|
||||
from litellm.router_rust.steps.prerouting import TIER_PARAMS_NEVER_DROPPED as _TIER_PARAMS_NEVER_DROPPED
|
||||
from litellm.router_rust.surface.members import RustRegistryMembers
|
||||
from litellm.router_strategy.auto_router.auto_router_model_naming import (
|
||||
GatedAutoRouterCapability,
|
||||
capability_limit_violation,
|
||||
count_capability_routers,
|
||||
)
|
||||
from litellm.router_utils.cooldown_handlers import set_cooldown_deployments
|
||||
from litellm.router_utils.endpoint_eligibility import endpoint_selection
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata
|
||||
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
|
||||
from litellm.types.router import RouterCacheEnum
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.router_rust.invoke.endpoints import EndpointCore
|
||||
from litellm.types.router import CustomRoutingStrategyBase, Deployment, RoutingGroup
|
||||
|
||||
# `RoutingArgs.ttl`: a usage counter lives for its minute.
|
||||
_USAGE_TTL: Final = 60
|
||||
|
||||
# What `initialize_assistants_endpoint` binds.
|
||||
_ASSISTANTS_ENDPOINTS: Final = (
|
||||
"acreate_assistants",
|
||||
"adelete_assistant",
|
||||
"aget_assistants",
|
||||
"acreate_thread",
|
||||
"aget_thread",
|
||||
"a_add_message",
|
||||
"aget_messages",
|
||||
"arun_thread",
|
||||
)
|
||||
|
||||
# What `initialize_router_endpoints` binds, with the container file endpoints it reads from
|
||||
# `litellm/containers/endpoints.json`.
|
||||
_ROUTER_ENDPOINTS: Final = frozenset(
|
||||
{
|
||||
"_arealtime",
|
||||
"_aresponses_websocket",
|
||||
"aadapter_generate_content",
|
||||
"aanthropic_messages",
|
||||
"acancel_fine_tuning_job",
|
||||
"acancel_interaction",
|
||||
"acancel_responses",
|
||||
"acompact_responses",
|
||||
"acreate_agent",
|
||||
"acreate_container",
|
||||
"acreate_fine_tuning_job",
|
||||
"acreate_interaction",
|
||||
"acreate_realtime_client_secret",
|
||||
"acreate_realtime_transcription_session",
|
||||
"acreate_realtime_translation_client_secret",
|
||||
"acreate_skill",
|
||||
"adecisions",
|
||||
"adelete_agent",
|
||||
"adelete_container",
|
||||
"adelete_interaction",
|
||||
"adelete_responses",
|
||||
"adelete_skill",
|
||||
"aextraction",
|
||||
"afile_content",
|
||||
"afile_delete",
|
||||
"afile_list",
|
||||
"agenerate_content",
|
||||
"agenerate_content_stream",
|
||||
"aget_agent",
|
||||
"aget_interaction",
|
||||
"aget_responses",
|
||||
"aget_skill",
|
||||
"aimage_edit",
|
||||
"alist_agent_versions",
|
||||
"alist_agents",
|
||||
"alist_containers",
|
||||
"alist_fine_tuning_jobs",
|
||||
"alist_input_items",
|
||||
"alist_skills",
|
||||
"allm_passthrough_route",
|
||||
"amoderation",
|
||||
"anthropic_messages",
|
||||
"aocr",
|
||||
"arealtime_calls",
|
||||
"arealtime_translation_calls",
|
||||
"aresponses",
|
||||
"aretrieve_container",
|
||||
"aretrieve_fine_tuning_job",
|
||||
"asearch",
|
||||
"asystemone",
|
||||
"avector_store_delete",
|
||||
"avector_store_file_content",
|
||||
"avector_store_file_create",
|
||||
"avector_store_file_delete",
|
||||
"avector_store_file_list",
|
||||
"avector_store_file_retrieve",
|
||||
"avector_store_file_update",
|
||||
"avector_store_list",
|
||||
"avector_store_retrieve",
|
||||
"avector_store_search",
|
||||
"avector_store_update",
|
||||
"avideo_content",
|
||||
"avideo_create_character",
|
||||
"avideo_edit",
|
||||
"avideo_extension",
|
||||
"avideo_generation",
|
||||
"avideo_get_character",
|
||||
"avideo_list",
|
||||
"avideo_remix",
|
||||
"avideo_status",
|
||||
"cancel_interaction",
|
||||
"create_agent",
|
||||
"create_container",
|
||||
"create_interaction",
|
||||
"decisions",
|
||||
"delete_agent",
|
||||
"delete_container",
|
||||
"delete_interaction",
|
||||
"extraction",
|
||||
"generate_content",
|
||||
"generate_content_stream",
|
||||
"get_agent",
|
||||
"get_interaction",
|
||||
"list_agent_versions",
|
||||
"list_agents",
|
||||
"list_containers",
|
||||
"ocr",
|
||||
"responses",
|
||||
"retrieve_container",
|
||||
"search",
|
||||
"systemone",
|
||||
"vector_store_create",
|
||||
"vector_store_delete",
|
||||
"vector_store_file_content",
|
||||
"vector_store_file_create",
|
||||
"vector_store_file_delete",
|
||||
"vector_store_file_list",
|
||||
"vector_store_file_retrieve",
|
||||
"vector_store_file_update",
|
||||
"vector_store_list",
|
||||
"vector_store_retrieve",
|
||||
"vector_store_search",
|
||||
"vector_store_update",
|
||||
"video_content",
|
||||
"video_create_character",
|
||||
"video_edit",
|
||||
"video_extension",
|
||||
"video_generation",
|
||||
"video_get_character",
|
||||
"video_list",
|
||||
"video_remix",
|
||||
"video_status",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _dispatch_failure(logging_obj: "Logging | None", error: Exception) -> None:
|
||||
"""The ``## LOG FAILURE EVENT`` of a routing callback that raised."""
|
||||
if logging_obj is not None:
|
||||
asyncio.create_task(
|
||||
logging_obj.dispatch_failure_handlers(
|
||||
exception=error,
|
||||
traceback_exception=traceback.format_exc(),
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _FailedCallParams(TypedDict):
|
||||
"""The ``litellm_params`` of a failed call, as the failure callback reads them."""
|
||||
|
||||
metadata: ReadOnly[Mapping[str, object]]
|
||||
model_info: ReadOnly[NotRequired[Mapping[str, object] | None]]
|
||||
|
||||
|
||||
async def _deployments_filtered_by(
|
||||
callback: CustomLogger,
|
||||
model: str,
|
||||
deployments: object,
|
||||
messages: object,
|
||||
request_kwargs: dict[str, object] | None, # mutable-ok: CustomLogger takes the request dict
|
||||
parent_otel_span: "Span | None",
|
||||
) -> object:
|
||||
return await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts
|
||||
model=model,
|
||||
healthy_deployments=deployments, # pyright: ignore[reportArgumentType] # what the last callback returned
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given
|
||||
request_kwargs=request_kwargs,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
|
||||
def _metadata_model_id(model_info: object) -> object:
|
||||
return model_info.get("id") if isinstance(model_info, dict) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the metadata's model_info dict
|
||||
|
||||
|
||||
class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks):
|
||||
"""Mixed into ``RustCore``, which supplies its host registry, selectors and checks."""
|
||||
|
||||
# RustCore sets these in its constructor.
|
||||
_cache: "DualCache"
|
||||
_cooldown_given: object
|
||||
_custom_strategy: "CustomRoutingStrategyBase | None"
|
||||
# The constructor's argument, or the attribute as assigned.
|
||||
auto_router_capability_limit: "Callable[[], int | None] | None"
|
||||
|
||||
TIER_PARAMS_NEVER_DROPPED: Final = _TIER_PARAMS_NEVER_DROPPED
|
||||
|
||||
# -- registry lookups --------------------------------------------------------------------------
|
||||
|
||||
def get_routing_group(self, model_name: str) -> "RoutingGroup | None":
|
||||
"""The routing group callable as ``model_name`` (``get_routing_group``)."""
|
||||
return self._registry_state.get_routing_group(model_name)
|
||||
|
||||
def team_public_model_names_of(self, team_id: str) -> frozenset[str]:
|
||||
"""The public names ``team_id`` serves its own deployments under (``team_public_model_names_of``)."""
|
||||
return self._registry_state.team_public_model_names_of(team_id)
|
||||
|
||||
def served_openai_models(self, model_name: str) -> tuple[str | None, ...]:
|
||||
"""The OpenAI models ``model_name`` serves (``_served_openai_models``). The Python core keeps
|
||||
an ``lru_cache`` over it, emptied with every model list change."""
|
||||
return _served_openai_models(self._registry_state._core_deployments_named(model_name) or ()) # pyright: ignore[reportPrivateUsage] # the core reads its own host registry
|
||||
|
||||
def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None:
|
||||
"""Why one more router claiming ``capability`` cannot join this router, or None when it can
|
||||
(``auto_router_capability_violation``); the limit is resolved on every call."""
|
||||
resolver: Final = self.auto_router_capability_limit
|
||||
limit: Final = resolver() if resolver is not None else None
|
||||
others: Final = count_capability_routers(
|
||||
(deployment for deployment in self._registry_state._core_model_list if isinstance(deployment, Mapping)), # pyright: ignore[reportPrivateUsage, reportUnnecessaryIsInstance] # the live list
|
||||
capability=capability,
|
||||
)
|
||||
return capability_limit_violation(capability=capability, held=others + 1, limit=limit)
|
||||
|
||||
# -- routing state -----------------------------------------------------------------------------
|
||||
|
||||
def retire_routing_state(self) -> None:
|
||||
"""Stop the background work of the selectors and the budget limiter (``retire_routing_state``)."""
|
||||
self._opt.retire_state((*self._sel.callback_selectors(), *self._sel.override_selectors()))
|
||||
|
||||
def routing_strategy_init(self, routing_strategy: object, routing_strategy_args: Mapping[str, object]) -> None:
|
||||
"""``routing_strategy_init``: the default selector only. As on the Python core the router's
|
||||
``routing_strategy`` setting stays, so picks keep following it; ``apply`` changes it."""
|
||||
verbose_router_logger.info("Routing strategy: %s", routing_strategy)
|
||||
validate_routing_strategy(routing_strategy) # pyright: ignore[reportArgumentType] # the strategy as given
|
||||
# `_reset_custom_routing_strategy`.
|
||||
self._custom_strategy = None
|
||||
_ = self.__dict__.pop("get_available_deployment", None)
|
||||
_ = self.__dict__.pop("_core_async_get_available_deployment", None)
|
||||
self._sel.init_default(routing_strategy, routing_strategy_args, configured=False)
|
||||
|
||||
def set_optional_pre_call_checks(self, optional_pre_call_checks: Iterable[object] | None) -> None:
|
||||
"""``set_optional_pre_call_checks``: drop the toggleable checks the list leaves out, then add it."""
|
||||
self._opt.set(None if optional_pre_call_checks is None else tuple(optional_pre_call_checks))
|
||||
|
||||
def initialize_assistants_endpoint(self) -> None:
|
||||
"""``initialize_assistants_endpoint`` binds the endpoints; this core answers them from
|
||||
``__getattr__``, so what is left to do is to drop an assigned replacement."""
|
||||
for name in _ASSISTANTS_ENDPOINTS:
|
||||
_ = self.__dict__.pop(name, None)
|
||||
|
||||
def initialize_router_endpoints(self) -> None:
|
||||
"""``initialize_router_endpoints``, as ``initialize_assistants_endpoint`` above."""
|
||||
from litellm.containers.endpoint_factory import (
|
||||
_generated_endpoints, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # the names the Python core registers
|
||||
)
|
||||
|
||||
for name in (*_ROUTER_ENDPOINTS, *_generated_endpoints):
|
||||
_ = self.__dict__.pop(name, None)
|
||||
|
||||
# -- pre-routing strategies --------------------------------------------------------------------
|
||||
|
||||
def init_auto_router_deployment(self, deployment: "Deployment") -> None:
|
||||
"""Register the deployment's auto-router (``init_auto_router_deployment``)."""
|
||||
self._registry_state.strategies._init_auto_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # RustCore installs its strategies in its constructor
|
||||
|
||||
def init_complexity_router_deployment(self, deployment: "Deployment") -> None:
|
||||
"""Register the deployment's complexity-router (``init_complexity_router_deployment``)."""
|
||||
self._registry_state.strategies._init_complexity_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # see above
|
||||
|
||||
def init_quality_router_deployment(self, deployment: "Deployment") -> None:
|
||||
"""Register the deployment's quality-router (``init_quality_router_deployment``)."""
|
||||
self._registry_state.strategies._init_quality_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # see above
|
||||
|
||||
def init_adaptive_router_deployment(self, deployment: "Deployment") -> None:
|
||||
"""Register the deployment's adaptive-router and its post-call hook (``init_adaptive_router_deployment``)."""
|
||||
self._registry_state.strategies._init_adaptive_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # see above
|
||||
|
||||
# -- routing callbacks -------------------------------------------------------------------------
|
||||
|
||||
def routing_strategy_pre_call_checks(
|
||||
self,
|
||||
deployment: dict[str, object], # mutable-ok: CustomLogger takes the deployment dict
|
||||
) -> None:
|
||||
"""``routing_strategy_pre_call_checks``: the sync check of each ``CustomLogger`` callback."""
|
||||
for callback in self._opt.routing_callbacks():
|
||||
if isinstance(callback, CustomLogger):
|
||||
callback.pre_call_check(deployment) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict
|
||||
|
||||
async def async_routing_strategy_pre_call_checks(
|
||||
self,
|
||||
deployment: dict[str, object], # mutable-ok: CustomLogger takes the deployment dict
|
||||
parent_otel_span: "Span | None",
|
||||
logging_obj: "Logging | None" = None,
|
||||
) -> None:
|
||||
"""``async_routing_strategy_pre_call_checks``: a check's rate limit cools the deployment down
|
||||
under the Python rules (``set_cooldown_deployments``, cooldown event included)."""
|
||||
from litellm import exceptions as litellm_exceptions
|
||||
|
||||
for callback in self._opt.routing_callbacks():
|
||||
if isinstance(callback, CustomLogger):
|
||||
try:
|
||||
await callback.async_pre_call_check(deployment, parent_otel_span) # pyright: ignore[reportUnknownMemberType] # see above
|
||||
except litellm_exceptions.RateLimitError as e:
|
||||
_dispatch_failure(logging_obj, e)
|
||||
_ = set_cooldown_deployments(
|
||||
ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext
|
||||
exception_status=e.status_code,
|
||||
original_exception=e,
|
||||
deployment=deployment["model_info"]["id"], # pyright: ignore[reportIndexIssue, reportUnknownArgumentType] # read as the Python core reads it
|
||||
time_to_cooldown=self._cooldown_given, # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
_dispatch_failure(logging_obj, e)
|
||||
raise
|
||||
|
||||
async def async_callback_filter_deployments(
|
||||
self,
|
||||
model: str,
|
||||
healthy_deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract
|
||||
messages: object,
|
||||
parent_otel_span: "Span | None",
|
||||
request_kwargs: dict[str, object] | None = None, # mutable-ok: CustomLogger takes the request dict
|
||||
logging_obj: "Logging | None" = None,
|
||||
) -> object:
|
||||
"""``async_callback_filter_deployments``: each routing callback narrows the deployments in turn."""
|
||||
remaining: object = healthy_deployments # rebind-ok: narrowed by each callback
|
||||
for callback in self._opt.routing_callbacks():
|
||||
if isinstance(callback, CustomLogger):
|
||||
try:
|
||||
remaining = await _deployments_filtered_by(
|
||||
callback, model, remaining, messages, request_kwargs, parent_otel_span
|
||||
)
|
||||
except Exception as e:
|
||||
_dispatch_failure(logging_obj, e)
|
||||
raise
|
||||
return remaining
|
||||
|
||||
async def async_admit_deployment_turn(
|
||||
self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
"""Admit a later turn of an open websocket (``async_admit_deployment_turn``)."""
|
||||
port: Final[EndpointCore] = self.__dict__["_port"] # pyright: ignore[reportAny] # RustCore's own attribute
|
||||
return await admit_deployment_turn(port, model_group, opened_on, request_kwargs)
|
||||
|
||||
# -- one provider call -------------------------------------------------------------------------
|
||||
|
||||
async def make_call(
|
||||
self,
|
||||
original_function: Callable[..., object],
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: the endpoint's request
|
||||
) -> object:
|
||||
"""``make_call``: call the endpoint, count its usage and stamp the response's headers. A
|
||||
failed call hands back the io tokens it reserved."""
|
||||
model_group: Final = kwargs.get("model")
|
||||
generic_function: Final = kwargs.get("original_generic_function")
|
||||
endpoint_function: Final = generic_function if callable(generic_function) else original_function
|
||||
try:
|
||||
with endpoint_selection(
|
||||
getattr(endpoint_function, "__name__", None),
|
||||
getattr(endpoint_function, "__module__", None),
|
||||
model=model_group,
|
||||
):
|
||||
response: object = original_function(*args, **kwargs) # rebind-ok: awaited below
|
||||
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
|
||||
awaitable: Final[Awaitable[object]] = response # pyright: ignore[reportAssignmentType] # the endpoint's awaitable
|
||||
response = await awaitable # rebind-ok: awaited here
|
||||
except Exception:
|
||||
with contextlib.suppress(Exception):
|
||||
await async_io_token_refund_failure(
|
||||
self._cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
|
||||
)
|
||||
raise
|
||||
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
|
||||
return await self.set_response_headers(
|
||||
response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs
|
||||
)
|
||||
|
||||
async def increment_deployment_usage_for_response(
|
||||
self, response: object, request_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
"""Count a response's tokens and request against its deployment, once the call returned, and
|
||||
note them on the request for the success callback (``increment_deployment_usage_for_response``)."""
|
||||
if response is None:
|
||||
return
|
||||
try:
|
||||
deployment_metadata: Final = find_deployment_metadata(request_kwargs)
|
||||
model_group: Final = request_kwargs.get("model")
|
||||
if deployment_metadata is None or not isinstance(model_group, str):
|
||||
return
|
||||
deployment_id: Final = _metadata_model_id(deployment_metadata["model_info"])
|
||||
if deployment_id is None:
|
||||
return
|
||||
total_tokens: Final = response_total_token_count(response)
|
||||
deployment_name: Final = deployment_metadata.get("deployment")
|
||||
deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens
|
||||
try:
|
||||
await self._increment_deployment_usage(
|
||||
deployment_id=str(deployment_id),
|
||||
deployment_name=deployment_name if isinstance(deployment_name, str) else None,
|
||||
model_group=model_group,
|
||||
total_tokens=total_tokens,
|
||||
parent_otel_span=get_parent_otel_span_from_kwargs(dict(request_kwargs)),
|
||||
)
|
||||
except Exception:
|
||||
_ = deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None)
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # as the Python core: usage accounting never fails a served call
|
||||
verbose_router_logger.debug(
|
||||
"litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e
|
||||
)
|
||||
|
||||
async def _increment_deployment_usage(
|
||||
self,
|
||||
*,
|
||||
deployment_id: str,
|
||||
deployment_name: str | None,
|
||||
model_group: str,
|
||||
total_tokens: float,
|
||||
parent_otel_span: "Span | None",
|
||||
) -> None:
|
||||
"""``_increment_deployment_usage`` with its one request, in one pipeline; this worker's memory
|
||||
then holds what Redis answered."""
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
|
||||
if not self._registry_state.cached_deployment_tracks_usage(deployment_id, model_group):
|
||||
return
|
||||
current_minute: Final = get_utc_datetime().strftime("%H-%M")
|
||||
tpm_key: Final = RouterCacheEnum.TPM.value.format(
|
||||
id=deployment_id, current_minute=current_minute, model=deployment_name
|
||||
)
|
||||
rpm_key: Final = RouterCacheEnum.RPM.value.format(
|
||||
id=deployment_id, current_minute=current_minute, model=deployment_name
|
||||
)
|
||||
operations: Final[Sequence[RedisPipelineIncrementOperation]] = [
|
||||
RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=_USAGE_TTL)
|
||||
for key, increment_value in ((tpm_key, total_tokens), (rpm_key, 1))
|
||||
]
|
||||
values: Final = await self._cache.async_increment_cache_pipeline(
|
||||
increment_list=list(operations), parent_otel_span=parent_otel_span
|
||||
)
|
||||
if values is not None and self._cache.redis_cache is not None:
|
||||
for operation, value in zip(operations, values):
|
||||
await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=_USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
|
||||
|
||||
async def async_deployment_callback_on_failure(
|
||||
self, kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
"""Count a failed call as one more request of the minute (``async_deployment_callback_on_failure``)."""
|
||||
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||
return
|
||||
litellm_params: Final[_FailedCallParams] = kwargs["litellm_params"] # pyright: ignore[reportAssignmentType] # read as the Python core reads it, a missing key included
|
||||
deployment_name: Final = litellm_params["metadata"].get("deployment", None)
|
||||
model_group: Final = litellm_params["metadata"].get("model_group", None)
|
||||
model_info: Final = litellm_params.get("model_info", {}) or {}
|
||||
deployment_id: Final = model_info.get("id", None)
|
||||
if model_group is None or deployment_id is None:
|
||||
return
|
||||
rpm_key: Final = RouterCacheEnum.RPM.value.format(
|
||||
id=deployment_id,
|
||||
current_minute=get_utc_datetime().strftime("%H-%M"),
|
||||
model=deployment_name,
|
||||
)
|
||||
with service_target(ROUTER_USAGE_TARGET):
|
||||
await self._cache.async_increment_cache( # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
|
||||
key=rpm_key,
|
||||
value=1,
|
||||
parent_otel_span=get_parent_otel_span_from_kwargs(dict(kwargs)),
|
||||
ttl=_USAGE_TTL,
|
||||
)
|
||||
|
||||
# -- the plugin context ------------------------------------------------------------------------
|
||||
|
||||
def _update_kwargs_before_fallbacks(
|
||||
self,
|
||||
model: str,
|
||||
kwargs: dict[str, object], # mutable-ok: the core updates the request kwargs in place
|
||||
metadata_variable_name: str | None = "metadata",
|
||||
) -> None:
|
||||
"""``_update_kwargs_before_fallbacks``: a trace id and the requested group, in place."""
|
||||
_ = kwargs.setdefault("litellm_trace_id", str(uuid.uuid4()))
|
||||
alias: Final = model if self._registry_state._core_get_model_from_alias(model=model) else None # pyright: ignore[reportPrivateUsage] # the core reads its own host registry
|
||||
kwargs.setdefault(metadata_variable_name, {}).update({"model_group": model, "model_group_alias": alias}) # pyright: ignore[reportArgumentType, reportAttributeAccessIssue, reportUnknownMemberType] # the bucket, as the Python core writes it
|
||||
|
||||
def _params_sent_to_deployment(
|
||||
self, request_kwargs: Mapping[str, object], deployment: "Deployment | None"
|
||||
) -> Mapping[str, object]:
|
||||
"""``_params_sent_to_deployment`` of a ``Deployment``, which ``RoutingContext.deployment_params`` hands in."""
|
||||
params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {}
|
||||
configured: Final = as_str_mapping(RoutingContext(self).settings().default_litellm_params) # pyright: ignore[reportArgumentType, reportUnknownMemberType, reportUnknownArgumentType] # the core's RoutingContext
|
||||
defaults: Final = {k: v for k, v in (configured or {}).items() if v is not None}
|
||||
tool_choice: Final = {"tool_choice": params["tool_choice"]} if params.get("tool_choice") is not None else {}
|
||||
return MappingProxyType({**params, **defaults, **tool_choice, **request_kwargs})
|
||||
|
||||
|
||||
__all__ = ["RustLegacyMembers"]
|
||||
163
litellm/router_rust/surface/members.py
Normal file
163
litellm/router_rust/surface/members.py
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
"""
|
||||
``RustRegistryMembers``: the Rust core's public registry members outside ``registry()`` and
|
||||
``apply``: ``add_deployment``, ``set_model_list``, the ``model_list`` attribute, the environment
|
||||
gate, the deprecated id hash, ``print_deployment`` and the pre-call checks the plugin context reads.
|
||||
|
||||
Each one works on the core's host registry (``litellm/router_rust/surface/registry.py``), which hands every
|
||||
change to the Rust core. The rules are the Python router core's, ported (never imported).
|
||||
"""
|
||||
|
||||
import copy
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.router_api import deployment_id
|
||||
from litellm.router_rust.steps.precall import count_tokens, pre_call_checks
|
||||
from litellm.router_rust.surface.registry import HostRegistry, model_info_is_active_for_environment
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_python_host import PythonHost
|
||||
from litellm.router_rust.core import RustCore
|
||||
|
||||
|
||||
class _CallCounter:
|
||||
"""``total_calls``, ``success_calls`` or ``fail_calls``: plain attributes on the Python core,
|
||||
which a write replaces and a ``del`` removes; every call then reads what is there."""
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
self._name: Final = name
|
||||
|
||||
def __get__(self, core: "RustCore | None", owner: type | None = None) -> object:
|
||||
if core is None:
|
||||
return self
|
||||
try:
|
||||
return getattr(core._endpoints.counters, self._name) # pyright: ignore[reportPrivateUsage] # the core's own counters
|
||||
except AttributeError:
|
||||
raise core._no_attribute(self._name) from None # pyright: ignore[reportPrivateUsage] # see above
|
||||
|
||||
def __set__(self, core: "RustCore", value: object) -> None:
|
||||
setattr(core._endpoints.counters, self._name, value) # pyright: ignore[reportPrivateUsage] # see __get__
|
||||
|
||||
def __delete__(self, core: "RustCore") -> None:
|
||||
try:
|
||||
delattr(core._endpoints.counters, self._name) # pyright: ignore[reportPrivateUsage] # see __get__
|
||||
except AttributeError:
|
||||
raise core._no_attribute(self._name) from None # pyright: ignore[reportPrivateUsage] # see __get__
|
||||
|
||||
|
||||
class RustRegistryMembers:
|
||||
"""Mixed into ``RustCore``, which supplies its host registry and host."""
|
||||
|
||||
total_calls: Final = _CallCounter("total_calls")
|
||||
success_calls: Final = _CallCounter("success_calls")
|
||||
fail_calls: Final = _CallCounter("fail_calls")
|
||||
|
||||
@property
|
||||
def _registry_state(self) -> HostRegistry:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def _python_host(self) -> "PythonHost":
|
||||
raise NotImplementedError
|
||||
|
||||
def add_deployment(self, deployment: Deployment) -> Deployment | None:
|
||||
"""Add a deployment unless its id exists (``add_deployment``)."""
|
||||
return self._registry_state.add_deployment(deployment)
|
||||
|
||||
def set_model_list(self, model_list: Sequence[object]) -> None:
|
||||
"""Replace every deployment (``set_model_list``); aliases and wildcard routes stay."""
|
||||
self._registry_state.set_model_list(model_list)
|
||||
|
||||
@property
|
||||
def model_list(self) -> Sequence[object]:
|
||||
"""The live model list (the deprecated ``model_list`` attribute)."""
|
||||
return self._registry_state._core_model_list # pyright: ignore[reportPrivateUsage] # the core reads its own host registry
|
||||
|
||||
def _assign_model_list(self, model_list: Sequence[object]) -> None:
|
||||
"""``router.model_list = [...]``: the list is swapped as given, without reindexing, and the
|
||||
registry version moves once."""
|
||||
self._registry_state.replace_model_list(model_list) # pyright: ignore[reportArgumentType] # the list as assigned, as the core keeps it
|
||||
|
||||
@staticmethod
|
||||
def deployment_is_active_for_environment(deployment: Deployment) -> bool:
|
||||
"""Whether LITELLM_ENVIRONMENT admits the deployment (``deployment_is_active_for_environment``)."""
|
||||
return model_info_is_active_for_environment(model_info=deployment.model_info) # pyright: ignore[reportArgumentType] # ModelInfo reads like a mapping
|
||||
|
||||
@staticmethod
|
||||
def generate_model_id(model_group: str, litellm_params: Mapping[str, object]) -> str:
|
||||
"""The deprecated id hash (``generate_model_id``)."""
|
||||
from litellm.router_deprecated import (
|
||||
_note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice
|
||||
)
|
||||
|
||||
_note("generate_model_id", "litellm.router_api.deployment_id(model_group, litellm_params)")
|
||||
return deployment_id(model_group, litellm_params)
|
||||
|
||||
async def set_response_headers(
|
||||
self, response: object, model_group: str | None = None, request_kwargs: Mapping[str, object] | None = None
|
||||
) -> object:
|
||||
"""Stamp the group's remaining usage and the model id on a response (``set_response_headers``)."""
|
||||
headed: Final[object] = await self._registry_state.set_response_headers( # pyright: ignore[reportAny] # the response as the core hands it back
|
||||
response, model_group, None if request_kwargs is None else dict(request_kwargs)
|
||||
)
|
||||
return headed
|
||||
|
||||
def replay_model_cost_registrations(self) -> None:
|
||||
"""Re-register the served deployments' cost-map entries (``replay_model_cost_registrations``)."""
|
||||
self._registry_state.replay_model_cost_registrations()
|
||||
|
||||
def _count_pre_call_check_tokens(
|
||||
self,
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
input: str | list[object] | None, # mutable-ok: the request field, as the Python core takes it
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
tokenizer_model: str = "",
|
||||
deployments: Sequence[Mapping[str, object]] = (),
|
||||
) -> int:
|
||||
"""``_count_pre_call_check_tokens``, which context compaction counts with."""
|
||||
return count_tokens(self._python_host, messages, input, request_kwargs, tokenizer_model, deployments)
|
||||
|
||||
def _pre_call_checks(
|
||||
self,
|
||||
model: str,
|
||||
healthy_deployments: Sequence[Mapping[str, object]],
|
||||
messages: object = None,
|
||||
input: object = None,
|
||||
request_kwargs: Mapping[str, object] | None = None,
|
||||
input_token_count: int | None = None,
|
||||
skip_inline_token_count: bool = False,
|
||||
request_deployments: Sequence[Mapping[str, object]] | None = None,
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
"""``_pre_call_checks``, which context compaction admits a client-managed request with."""
|
||||
return pre_call_checks(
|
||||
self._python_host,
|
||||
self._registry_state, # pyright: ignore[reportArgumentType] # the host registry serves the reads
|
||||
model,
|
||||
healthy_deployments,
|
||||
messages,
|
||||
input,
|
||||
request_kwargs,
|
||||
input_token_count,
|
||||
skip_inline_token_count,
|
||||
request_deployments,
|
||||
)
|
||||
|
||||
def print_deployment(self, deployment: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""A copy of the deployment, of the type it was given, with its api key masked
|
||||
(``print_deployment``). A ``Deployment`` and its params read and write like dicts."""
|
||||
try:
|
||||
printed: Final = copy.deepcopy(deployment)
|
||||
litellm_params: Final = printed["litellm_params"]
|
||||
if self._python_host.dynamic().redact_user_api_key_info:
|
||||
masker: Final = SensitiveDataMasker(visible_prefix=2, visible_suffix=0)
|
||||
printed["litellm_params"] = masker.mask_dict(litellm_params) # pyright: ignore[reportArgumentType, reportIndexIssue] # the deployment's params, written as the Python core writes them
|
||||
elif "api_key" in litellm_params: # pyright: ignore[reportOperatorIssue] # see the docstring
|
||||
litellm_params["api_key"] = litellm_params["api_key"][:2] + "*" * 10 # pyright: ignore[reportIndexIssue] # see the docstring
|
||||
except Exception as e:
|
||||
verbose_router_logger.debug("Error occurred while printing deployment - %s", e)
|
||||
raise
|
||||
else:
|
||||
return printed
|
||||
871
litellm/router_rust/surface/model_info.py
Normal file
871
litellm/router_rust/surface/model_info.py
Normal file
|
|
@ -0,0 +1,871 @@
|
|||
# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false, reportMissingParameterType=false, reportAny=false, reportExplicitAny=false, reportMissingTypeArgument=false, reportAttributeAccessIssue=false, reportTypedDictNotRequiredAccess=false
|
||||
# (ported from the Python router core as it is written there, over untyped litellm helpers)
|
||||
"""
|
||||
``HostModelInfo``: model info, group info, usage reads and response headers for the Rust core's
|
||||
host registry.
|
||||
|
||||
router-core holds the rules: a deployment's model info (``get_router_model_info``,
|
||||
``get_deployment_model_info``: ``router-core/src/model_info.rs``), a model group's info and the
|
||||
configured-value reads (``_set_model_group_info``, ``get_model_listing_info``, the configured token
|
||||
limits, mode, display name and service tiers: ``router-core/src/group_info.rs``), and the usage
|
||||
counters' keys (``RouterCoreHandle.usage_key``). What lives in this process is answered here as a
|
||||
query asks for it, at query time: the cost map (``get_model_info``, ``litellm.model_cost``),
|
||||
provider resolution, the wildcard pattern routes, what the background refresh discovered, and a
|
||||
group deployment's validated params, supported params, endpoints, fast mode and reasoning efforts.
|
||||
``ModelGroupInfo`` is built here from the group's first deployment (its validation is the model's)
|
||||
and takes what router-core's merge assigned.
|
||||
|
||||
Still the Python router core's rules here: the model group alias a group-info query resolves
|
||||
first (the alias map is the registry's live Python objects), the cache of discovered limits, the
|
||||
configured ``max_input_tokens`` read (a value echoed from the loaded cost map is skipped), and the
|
||||
response headers.
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import dataclasses
|
||||
import math
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast # noqa: TID251 # the core's own signatures, ported unchanged
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import overload
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.add_retry_fallback_headers import (
|
||||
apply_quality_router_decision_headers,
|
||||
apply_remaining_usage_headers,
|
||||
apply_response_model_id,
|
||||
complexity_router_decision_headers,
|
||||
ensure_response_additional_headers,
|
||||
prepare_response_for_header_attachment,
|
||||
)
|
||||
from litellm.litellm_core_utils.batch_utils import deployment_supported_endpoints
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, coerce_token_limit
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.reasoning_effort_capability import (
|
||||
deployment_is_catalog_mapped,
|
||||
intersect_supported_reasoning_efforts,
|
||||
resolve_supported_reasoning_efforts,
|
||||
)
|
||||
from litellm.litellm_core_utils.routing_common_utils import team_may_use_deployment
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
|
||||
from litellm.router_python_host import python_host
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import build_io_token_rate_limit_headers
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata
|
||||
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
|
||||
from litellm.types.router import (
|
||||
Deployment,
|
||||
DeploymentModelListingInfo,
|
||||
DeploymentTypedDict,
|
||||
DiscoveredDeploymentModelInfo,
|
||||
LiteLLM_Params,
|
||||
ModelGroupInfo,
|
||||
)
|
||||
from litellm.types.utils import COST_MAP_LOOKUP_KEY, ModelInfo, echoed_cost_map_fields
|
||||
from litellm.types.utils import ModelInfo as ModelMapInfo
|
||||
from litellm.utils import get_utc_datetime, is_generalized_model_info
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.rust_bridge._native import (
|
||||
RouterCall,
|
||||
RouterCoreHandle,
|
||||
_ModelInfoSteps, # pyright: ignore[reportPrivateUsage] # the stub's step shapes
|
||||
)
|
||||
|
||||
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
# The key router-core's merge reads as an opaque value (``router-core/src/model_info.rs``).
|
||||
_OPAQUE_KEY: Final = "__litellm_python_value__"
|
||||
|
||||
|
||||
class _Wire:
|
||||
"""Values as router-core's model-info merge reads them: plain JSON, with what JSON cannot spell
|
||||
(a datetime, a tuple, a non-finite float) passed as an opaque token the merge moves between
|
||||
layers without reading, and restored to the same object on the way back."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._values: Final[list[object]] = [] # mutable-ok: the tokens of one query
|
||||
|
||||
def to_wire(self, value: object) -> object:
|
||||
if value is None or isinstance(value, (bool, str)):
|
||||
return value
|
||||
if isinstance(value, int):
|
||||
return value if -(2**63) <= value < 2**64 else self._token(value)
|
||||
if isinstance(value, float):
|
||||
return value if math.isfinite(value) else self._token(value)
|
||||
if isinstance(value, list):
|
||||
return [self.to_wire(item) for item in value]
|
||||
if isinstance(value, dict) and all(isinstance(key, str) for key in value):
|
||||
return {str(key): self.to_wire(item) for key, item in value.items()}
|
||||
return self._token(value)
|
||||
|
||||
def from_wire(self, value: object) -> object:
|
||||
if isinstance(value, dict):
|
||||
if len(value) == 1 and _OPAQUE_KEY in value:
|
||||
return self._values[int(value[_OPAQUE_KEY])]
|
||||
return {str(key): self.from_wire(item) for key, item in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [self.from_wire(item) for item in value]
|
||||
return value
|
||||
|
||||
def _token(self, value: object) -> Mapping[str, int]:
|
||||
self._values.append(value)
|
||||
return {_OPAQUE_KEY: len(self._values) - 1}
|
||||
|
||||
|
||||
def _drive_model_info(
|
||||
call: "RouterCall[_ModelInfoSteps]", answer: Callable[[Mapping[str, object]], object]
|
||||
) -> Mapping[str, object]:
|
||||
"""Drive one model-info merge to its end, answering each lookup with ``answer``. A lookup that
|
||||
raises raises out of the query, as it raised in the Python router."""
|
||||
outcome: object = None # rebind-ok: each lookup's answer
|
||||
while True:
|
||||
step = call.next(outcome)
|
||||
if step["kind"] == "done":
|
||||
return step
|
||||
if step["kind"] == "failed":
|
||||
if str(step.get("error_kind")) == "model_info":
|
||||
exception: object = getattr(builtins, str(step.get("error_exception")), ValueError)
|
||||
raise (exception if isinstance(exception, type) else ValueError)(step.get("error_detail"))
|
||||
raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}")
|
||||
outcome = answer(step)
|
||||
|
||||
|
||||
def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {})
|
||||
|
||||
|
||||
def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None:
|
||||
"""``_operator_max_input_tokens``: the deployment's own limit, skipping one a saved
|
||||
``/model/info`` response copied from the cost map."""
|
||||
model_info: Final = _configured_model_info(deployment)
|
||||
catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY)
|
||||
catalog_entry: Final = (
|
||||
GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None
|
||||
)
|
||||
if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry):
|
||||
return None
|
||||
return coerce_token_limit(model_info.get("max_input_tokens"))
|
||||
|
||||
|
||||
# The ModelGroupInfo fields router-core's group merge compares against, as the model validated them.
|
||||
_GROUP_COMPARED_FIELDS: Final = frozenset(
|
||||
{
|
||||
"providers",
|
||||
"max_input_tokens",
|
||||
"max_output_tokens",
|
||||
"input_cost_per_token",
|
||||
"output_cost_per_token",
|
||||
"supports_fast_mode",
|
||||
"supported_endpoints",
|
||||
"supported_reasoning_efforts",
|
||||
}
|
||||
)
|
||||
# The ModelGroupInfo fields the Python router assigns tuples to.
|
||||
_GROUP_TUPLE_FIELDS: Final = frozenset({"supported_endpoints", "supported_reasoning_efforts"})
|
||||
|
||||
|
||||
def _group_field(field: str, value: object) -> object:
|
||||
"""A merged group field as the Python router assigns it: a tuple field takes the list as a tuple."""
|
||||
if field in _GROUP_TUPLE_FIELDS and isinstance(value, list):
|
||||
return tuple(cast(list[object], value)) # cast-ok: the merge's list
|
||||
return value
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _GroupBuild:
|
||||
"""One group-info merge as the host answers it: the deployments it reads, each one's validated
|
||||
params once its first lookup builds them, and the group built from the first deployment."""
|
||||
|
||||
model_list: "Sequence[dict[str, Any]]" # mutable-ok: the stored entries, read only
|
||||
team_id: str | None
|
||||
# Filled as the lookups arrive.
|
||||
params: dict[int, LiteLLM_Params] = dataclasses.field(default_factory=dict) # mutable-ok: see above
|
||||
info: ModelGroupInfo | None = None
|
||||
|
||||
|
||||
class HostModelInfo:
|
||||
"""Mixed into ``HostRegistry``, which holds the model list, the indexes and the host."""
|
||||
|
||||
def _live_handle(self) -> "RouterCoreHandle | None": ... # the registry's: the handle ``RustCore`` set
|
||||
|
||||
def _core_handle(self) -> "RouterCoreHandle":
|
||||
handle: Final = self._live_handle()
|
||||
if handle is None:
|
||||
raise RuntimeError("the host registry reads model info through its Rust core")
|
||||
return handle
|
||||
|
||||
def _core_get_model_list_of_routed_group(
|
||||
self, model_group: str, team_id: str | None = None
|
||||
) -> list: # mutable-ok: a fresh list, as the core returns it
|
||||
if team_id is not None and self._team_deployments_named(team_id, model_group):
|
||||
return self._core_get_all_deployments(model_name=model_group, team_id=team_id)
|
||||
named: Final = [
|
||||
*self._core_get_all_deployments(model_name=model_group),
|
||||
*self._core_get_model_list_from_routing_groups(model_name=model_group),
|
||||
]
|
||||
return named or self._get_wildcard_deployments(model_name=model_group)
|
||||
|
||||
def _get_wildcard_deployments(
|
||||
self, model_name: str, team_id: str | None = None
|
||||
) -> list[dict[str, object]]: # mutable-ok: a fresh list
|
||||
"""The deployments a wildcard pattern routes ``model_name`` to, renamed to it (the routing
|
||||
runs in router-core, through the registry's pattern router views)."""
|
||||
team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None
|
||||
matches: Final = [
|
||||
*self._core_pattern_router.get_deployments_by_pattern(model=model_name),
|
||||
*(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()),
|
||||
]
|
||||
return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches] # pyright: ignore[reportCallIssue] # a stored entry
|
||||
|
||||
def _core_get_model_group_info(
|
||||
self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None
|
||||
) -> ModelGroupInfo | None:
|
||||
"""
|
||||
For a given model group name, return the combined model info
|
||||
|
||||
Returns:
|
||||
- ModelGroupInfo if able to construct a model group
|
||||
- None if error constructing model group info or hidden model group (unless include_hidden)
|
||||
"""
|
||||
## Check if model group alias
|
||||
if model_group in self._core_model_group_alias:
|
||||
item: Final = self._core_model_group_alias[model_group]
|
||||
if isinstance(item, str):
|
||||
_router_model_group = item # rebind-ok: ported from the Python core unchanged
|
||||
elif isinstance(item, dict):
|
||||
if item["hidden"] is True and not include_hidden:
|
||||
return None
|
||||
else:
|
||||
_router_model_group = item["model"] # rebind-ok: ported from the Python core unchanged
|
||||
else:
|
||||
return None
|
||||
|
||||
return self._set_model_group_info(
|
||||
model_group=_router_model_group,
|
||||
user_facing_model_group_name=model_group,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
## Check if actual model
|
||||
return self._set_model_group_info(
|
||||
model_group=model_group, user_facing_model_group_name=model_group, team_id=team_id
|
||||
)
|
||||
|
||||
def _set_model_group_info(
|
||||
self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None
|
||||
) -> ModelGroupInfo | None:
|
||||
"""
|
||||
For a given model group name, return the combined model info
|
||||
|
||||
Returns:
|
||||
- ModelGroupInfo if able to construct a model group
|
||||
- None if error constructing model group info
|
||||
|
||||
router-core merges the group (``router-core/src/group_info.rs``); this picks its
|
||||
deployments, answers what needs this process (validated params, model info, provider,
|
||||
capabilities), builds ``ModelGroupInfo`` from the first deployment and sets what the merge
|
||||
assigned on it.
|
||||
"""
|
||||
serves_team_deployments: Final = team_id is not None and bool(
|
||||
self._team_deployments_named(team_id, model_group)
|
||||
)
|
||||
model_list: Final = self._core_get_model_list_of_routed_group(model_group, team_id)
|
||||
description: Final[str | None] = next(
|
||||
(d for d in (_configured_model_info(m).get("description") for m in model_list) if isinstance(d, str)),
|
||||
None,
|
||||
)
|
||||
# The wildcard check is asked only for a deployment that is not the group's by name.
|
||||
pattern_routes: Final = (
|
||||
not serves_team_deployments
|
||||
and any("model_name" in m and m["model_name"] != model_group for m in model_list)
|
||||
and self._core_pattern_router.route(model_group) is not None
|
||||
)
|
||||
wire: Final = _Wire()
|
||||
group: Final = _GroupBuild(model_list=model_list, team_id=team_id)
|
||||
call: Final = self._core_handle().model_group_info(
|
||||
{
|
||||
"model_group": model_group,
|
||||
"user_facing_model_group_name": user_facing_model_group_name,
|
||||
"serves_team_deployments": serves_team_deployments,
|
||||
"pattern_routes": pattern_routes,
|
||||
"description": description,
|
||||
"deployments": [wire.to_wire(model) for model in model_list],
|
||||
}
|
||||
)
|
||||
done: Final = _drive_model_info(call, partial(self._answer_group_lookup, group=group, wire=wire))
|
||||
answer: Final = as_str_mapping(done["value"])
|
||||
if answer is None or group.info is None:
|
||||
return None
|
||||
for field, value in (as_str_mapping(answer["assignments"]) or {}).items():
|
||||
setattr(group.info, field, _group_field(field, wire.from_wire(value)))
|
||||
return group.info
|
||||
|
||||
def _answer_group_lookup(self, step: Mapping[str, object], group: "_GroupBuild", wire: "_Wire") -> object:
|
||||
"""One lookup of a group-info merge: a deployment's facts, or the group's first entry."""
|
||||
match step["lookup"]:
|
||||
case "group_params":
|
||||
index: Final = int(str(step["index"]))
|
||||
litellm_params: Final = LiteLLM_Params(**group.model_list[index]["litellm_params"])
|
||||
group.params[index] = litellm_params # rebind-ok: the merge's state, which each lookup advances
|
||||
return {
|
||||
"model": litellm_params.model,
|
||||
"configurable_clientside_auth_params": wire.to_wire(
|
||||
litellm_params.configurable_clientside_auth_params
|
||||
),
|
||||
}
|
||||
case "group_deployment":
|
||||
return self._group_deployment(group, int(str(step["index"])), wire)
|
||||
case "group_capabilities":
|
||||
deployment: Final = group.model_list[int(str(step["index"]))]
|
||||
merged_info: Final[object] = wire.from_wire(step["model_info"])
|
||||
model_info: Final = cast(Mapping[str, object], merged_info) # cast-ok: the merged model info
|
||||
litellm_model: Final = str(step["litellm_model"])
|
||||
llm_provider: Final = str(step["llm_provider"])
|
||||
efforts: Final = resolve_supported_reasoning_efforts(
|
||||
model_info, deployment_is_mapped=step["mapped"] is True
|
||||
)
|
||||
return {
|
||||
"supported_endpoints": wire.to_wire(
|
||||
list(
|
||||
deployment_supported_endpoints(
|
||||
model_info.get("supported_endpoints"), deployment, llm_provider
|
||||
)
|
||||
) # pyright: ignore[reportArgumentType] # a stored entry, as the core passes it
|
||||
if team_may_use_deployment(_configured_model_info(deployment).get("team_id"), group.team_id) # pyright: ignore[reportArgumentType] # the configured team id
|
||||
else []
|
||||
),
|
||||
"supports_fast_mode": AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider),
|
||||
"reasoning_efforts": None if efforts is None else list(efforts),
|
||||
}
|
||||
case "group_seed":
|
||||
entry: Final[object] = wire.from_wire(step["info"])
|
||||
seed: Final = cast(dict[str, Any], entry) # cast-ok: the first deployment's entry
|
||||
group.info = ModelGroupInfo(**seed) # rebind-ok: the merge's state, built from the first deployment
|
||||
return group.info.model_dump(mode="json", include=set(_GROUP_COMPARED_FIELDS))
|
||||
case "reasoning_intersect":
|
||||
current: Final = cast("Sequence[str] | None", step["current"]) # cast-ok: plain JSON efforts
|
||||
resolved: Final = cast("Sequence[str] | None", step["resolved"]) # cast-ok: plain JSON efforts
|
||||
intersected: Final = intersect_supported_reasoning_efforts(current, resolved)
|
||||
return None if intersected is None else list(intersected)
|
||||
case _:
|
||||
return self._answer_model_info_lookup(step, wire)
|
||||
|
||||
def _group_deployment(self, group: "_GroupBuild", index: int, wire: "_Wire") -> Mapping[str, object]:
|
||||
"""A group deployment's model info (None when it fails), whether the catalog maps it, and its
|
||||
provider; with no model info, the provider's supported params."""
|
||||
model: Final = group.model_list[index]
|
||||
litellm_params: Final = group.params[index]
|
||||
model_info_dict: Final = model.get("model_info", {})
|
||||
model_info: object
|
||||
try:
|
||||
model_id: Final = model_info_dict.get("id", None)
|
||||
model_info = ( # rebind-ok: one branch binds it
|
||||
self._core_get_deployment_model_info(model_id=model_id, model_name=litellm_params.model)
|
||||
if model_id is not None
|
||||
else None
|
||||
)
|
||||
except Exception: # noqa: BLE001 # ported as the core writes it
|
||||
model_info = None # rebind-ok: see above
|
||||
deployment_is_mapped: Final = deployment_is_catalog_mapped(model_info, model_info_dict) # pyright: ignore[reportArgumentType] # the merged model info, or None
|
||||
litellm_model, llm_provider = self._model_group_llm_provider(litellm_params)
|
||||
answer: Final[dict[str, object]] = { # mutable-ok: built here
|
||||
"model_info": wire.to_wire(model_info),
|
||||
"mapped": deployment_is_mapped,
|
||||
"litellm_model": litellm_model,
|
||||
"llm_provider": llm_provider,
|
||||
}
|
||||
if model_info is None:
|
||||
answer["supported_openai_params"] = wire.to_wire(
|
||||
self._host.invoke("get_supported_openai_params")(model=litellm_model, custom_llm_provider=llm_provider)
|
||||
)
|
||||
return answer
|
||||
|
||||
def _cached_get_model_group_info(self, model_group: str) -> ModelGroupInfo | None:
|
||||
"""
|
||||
Cached version of get_model_group_info, uses @lru_cache wrapper
|
||||
|
||||
This is a speed optimization, since set_response_headers makes a call to get_model_group_info on every request
|
||||
"""
|
||||
return self._core_get_model_group_info(model_group)
|
||||
|
||||
def _core_cached_model_group_info(self, model_group: str) -> ModelGroupInfo | None:
|
||||
return self._cached_get_model_group_info(model_group)
|
||||
|
||||
@overload
|
||||
def _core_get_router_model_info(
|
||||
self,
|
||||
deployment: Union[dict, "Deployment"], # mutable-ok: ported from the Python core unchanged
|
||||
received_model_name: str,
|
||||
id: None = None,
|
||||
) -> ModelMapInfo:
|
||||
pass
|
||||
|
||||
@overload
|
||||
def _core_get_router_model_info(self, deployment: None, received_model_name: str, id: str) -> ModelMapInfo:
|
||||
pass
|
||||
|
||||
def _core_get_router_model_info(
|
||||
self,
|
||||
deployment: Union[dict, "Deployment"] | None, # mutable-ok: ported from the Python core unchanged
|
||||
received_model_name: str,
|
||||
id: str | None = None,
|
||||
) -> ModelMapInfo:
|
||||
"""
|
||||
For a given model id, return the model info (max tokens, input cost, output cost, etc.).
|
||||
|
||||
Augment litellm info with additional params set in `model_info`.
|
||||
|
||||
For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set.
|
||||
For every other provider, a base_model naming the same model as `model:` can only widen its context window.
|
||||
|
||||
router-core merges it (``router-core/src/model_info.rs``); this checks the deployment and
|
||||
answers the cost-map lookups the merge asks for while it runs.
|
||||
|
||||
Returns
|
||||
- ModelInfo - If found -> typed dict with max tokens, input cost, etc.
|
||||
|
||||
Raises:
|
||||
- ValueError -> If model is not mapped yet
|
||||
"""
|
||||
if id is not None:
|
||||
_deployment: Final = self._core_get_deployment(model_id=id)
|
||||
if _deployment is not None:
|
||||
deployment = _deployment # rebind-ok: ported from the Python core unchanged
|
||||
|
||||
if deployment is None:
|
||||
raise ValueError("Deployment not found")
|
||||
|
||||
model_info_base_model: Final = (deployment.get("model_info") or {}).get("base_model", None)
|
||||
params_base_model: Final = (
|
||||
(deployment.get("litellm_params") or {}).get("base_model", None) if model_info_base_model is None else None
|
||||
)
|
||||
|
||||
## GET PROVIDER - reuse LiteLLM_Params if already constructed
|
||||
litellm_params_data: Final = deployment.get("litellm_params")
|
||||
litellm_params: LiteLLM_Params
|
||||
if isinstance(litellm_params_data, LiteLLM_Params):
|
||||
litellm_params = litellm_params_data # rebind-ok: ported from the Python core unchanged
|
||||
elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data:
|
||||
litellm_params = LiteLLM_Params.model_validate( # rebind-ok: ported from the Python core unchanged
|
||||
litellm_params_data
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Deployment missing valid litellm_params. "
|
||||
f"Got: {type(litellm_params_data).__name__}, "
|
||||
f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}"
|
||||
)
|
||||
|
||||
wire: Final = _Wire()
|
||||
user_model_info: dict[str, object] = {} # mutable-ok: filled when the merge asks for it # rebind-ok: see above
|
||||
|
||||
def read_user_model_info() -> object:
|
||||
nonlocal user_model_info
|
||||
raw_user_model_info: Final = deployment.get("model_info")
|
||||
user_model_info = dict( # rebind-ok: see above
|
||||
_MODEL_INFO_ADAPTER.validate_python(
|
||||
raw_user_model_info.model_dump(exclude_none=True)
|
||||
if isinstance(raw_user_model_info, BaseModel)
|
||||
else raw_user_model_info or {}
|
||||
)
|
||||
)
|
||||
return wire.to_wire(user_model_info)
|
||||
|
||||
def answer(step: Mapping[str, object]) -> object:
|
||||
if step["lookup"] == "provider":
|
||||
_model, custom_llm_provider, _, _ = self._host.invoke("get_llm_provider")(
|
||||
model=litellm_params.model,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
return [_model, custom_llm_provider]
|
||||
if step["lookup"] == "user_model_info":
|
||||
return read_user_model_info()
|
||||
return self._answer_model_info_lookup(step, wire)
|
||||
|
||||
query: Final = {
|
||||
"deployment_id": wire.to_wire((deployment.get("model_info") or {}).get("id")),
|
||||
"model_info_base_model": wire.to_wire(model_info_base_model),
|
||||
"params_base_model": wire.to_wire(params_base_model),
|
||||
"params_model": litellm_params.model,
|
||||
"extra_base_model": wire.to_wire((litellm_params.model_extra or {}).get("base_model")),
|
||||
"received_model_name": received_model_name,
|
||||
}
|
||||
done: Final = _drive_model_info(self._core_handle().router_model_info(query), answer)
|
||||
merged: Final = wire.from_wire(done["info"])
|
||||
if not isinstance(merged, dict):
|
||||
return merged # pyright: ignore[reportReturnType] # the core returns get_model_info's answer as is
|
||||
# The deployment's own values are its objects, as the Python router merges them.
|
||||
for key in cast(list[str], done["user_keys"]): # cast-ok: router-core's key list
|
||||
merged[key] = user_model_info[key]
|
||||
return merged # pyright: ignore[reportReturnType] # the merged map, as the core builds it
|
||||
|
||||
def _core_get_deployment_model_info(self, model_id: str, model_name: str) -> ModelInfo | None:
|
||||
"""
|
||||
For a given model id, return the model info
|
||||
|
||||
1. Check if model_id is in model info
|
||||
2. If not, check if litellm model name is in model info
|
||||
3. If not, return None
|
||||
|
||||
router-core merges it (``router-core/src/model_info.rs``); this answers its lookups.
|
||||
"""
|
||||
wire: Final = _Wire()
|
||||
call: Final = self._core_handle().deployment_model_info(model_id, wire.to_wire(model_name))
|
||||
done: Final = _drive_model_info(call, partial(self._answer_model_info_lookup, wire=wire))
|
||||
return wire.from_wire(done["info"]) # pyright: ignore[reportReturnType] # the merged map, as the core builds it
|
||||
|
||||
def _answer_model_info_lookup(self, step: Mapping[str, object], wire: "_Wire") -> object:
|
||||
"""One cost-map lookup of a model-info merge, answered from this process."""
|
||||
match step["lookup"]:
|
||||
case "model_info":
|
||||
return self._model_info_row(wire.from_wire(step["model"]), wire, step)
|
||||
case "cost_entry":
|
||||
key: Final = str(step["key"])
|
||||
model_cost: Final = litellm.model_cost if step["process"] else self._host.dynamic().model_cost
|
||||
return {"entry": wire.to_wire(model_cost.get(key))} if key in model_cost else None
|
||||
case "pattern_models":
|
||||
return self._pattern_models(str(step["name"]), wire)
|
||||
case "bedrock_base":
|
||||
return get_bedrock_base_model(str(step["model_id"]))
|
||||
case "discovered":
|
||||
return wire.to_wire(dict(self._core_get_discovered_model_info(wire.from_wire(step["id"])))) # pyright: ignore[reportArgumentType] # the deployment's id, as configured
|
||||
case "log":
|
||||
log: Final = verbose_router_logger.error if step["level"] == "error" else verbose_router_logger.debug
|
||||
log("%s", step["message"])
|
||||
return None
|
||||
case other:
|
||||
raise ValueError(f"router-core asked an unknown model-info lookup {other!r}")
|
||||
|
||||
def _model_info_row(self, model: object, wire: "_Wire", step: Mapping[str, object]) -> object:
|
||||
"""``get_model_info(model=...)`` as a lookup answers it: the row, ``{"failed": True}`` when a
|
||||
tolerated lookup raises, and for an exact lookup whether a generalization rule made it."""
|
||||
# The base model's own window is read through the process host, as the Python router reads it.
|
||||
host: Final = python_host if step["exact"] else self._host
|
||||
try:
|
||||
row: Final = host.invoke("get_model_info")(model=model)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model
|
||||
if step["tolerate"]:
|
||||
return {"failed": True}
|
||||
raise
|
||||
if step["exact"]:
|
||||
return {"row": wire.to_wire(row), "generalized": is_generalized_model_info(row)}
|
||||
return wire.to_wire(row)
|
||||
|
||||
def _pattern_models(self, name: str, wire: "_Wire") -> object:
|
||||
"""The ``{id, model}`` of each deployment a wildcard pattern routes ``name`` to."""
|
||||
potential_models: Final = self._core_pattern_router.route(name)
|
||||
if potential_models is None:
|
||||
return None
|
||||
routed: Final[list[object]] = [] # mutable-ok: built here
|
||||
for potential_model in potential_models:
|
||||
try:
|
||||
routed.append(
|
||||
{
|
||||
"id": wire.to_wire((potential_model.get("model_info") or {}).get("id")),
|
||||
"model": wire.to_wire((potential_model.get("litellm_params") or {}).get("model")),
|
||||
}
|
||||
)
|
||||
except Exception: # noqa: BLE001 # ported as the core writes it
|
||||
verbose_router_logger.debug("model info lookup failed; ignored as the core ignores it")
|
||||
return routed
|
||||
|
||||
def _core_get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]:
|
||||
cached: Final[object] = self._discovered_model_info_cache.get_cache(model_id)
|
||||
if (
|
||||
model_id is not None
|
||||
and isinstance(cached, DiscoveredDeploymentModelInfo)
|
||||
and cached.deployment is self._core_get_model_info(model_id)
|
||||
):
|
||||
configured: Final = TypeAdapter(Mapping[str, object]).validate_python(cached.deployment["model_info"])
|
||||
return MappingProxyType({key: value for key, value in cached.limits.items() if configured.get(key) is None})
|
||||
return MappingProxyType({})
|
||||
|
||||
def _core_get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None:
|
||||
"""
|
||||
Return what the concrete deployments behind model_name contribute to its
|
||||
/v1/models entry: the cost-map keys for their underlying models, plus the widest
|
||||
configured or discovered token limits. Resolved via O(1) index
|
||||
lookup.
|
||||
|
||||
Returns None for wildcard-expanded or unknown names, where the listed name is the
|
||||
real model name and no deployment-specific information exists, and treats a
|
||||
malformed configured limit as absent rather than failing the listing.
|
||||
|
||||
The whole group is read rather than just its first deployment, so a group that
|
||||
mixes models does not advertise a window that depends on config order; the widest
|
||||
one is reported, which is what get_model_group_info already shows the Admin UI.
|
||||
Keys are deduplicated, so the ordinary group of interchangeable deployments of one
|
||||
model still costs the caller a single cost-map lookup. Unlike get_model_group_info,
|
||||
this never triggers pattern matching or deep copies, so it is safe to call per
|
||||
listed model on the /v1/models hot path.
|
||||
|
||||
router-core combines the deployments (``router-core/src/group_info.rs``).
|
||||
"""
|
||||
deployments: Final = self._core_deployments_named(model_name)
|
||||
if not deployments:
|
||||
return None
|
||||
wire: Final = _Wire()
|
||||
call: Final = self._core_handle().model_listing_info([wire.to_wire(deployment) for deployment in deployments])
|
||||
done: Final = _drive_model_info(call, partial(self._answer_model_info_lookup, wire=wire))
|
||||
listing: Final = as_str_mapping(done["value"]) or {}
|
||||
return DeploymentModelListingInfo(
|
||||
cost_map_keys=tuple(cast(list[str], listing["cost_map_keys"])), # cast-ok: router-core's key list
|
||||
max_input_tokens=cast("int | None", listing["max_input_tokens"]), # cast-ok: router-core's limit
|
||||
max_output_tokens=cast("int | None", listing["max_output_tokens"]), # cast-ok: router-core's limit
|
||||
)
|
||||
|
||||
def _core_get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]":
|
||||
"""
|
||||
Return (max_input_tokens, max_output_tokens) configured or discovered for a concrete
|
||||
deployment of model_name, via O(1) index lookup.
|
||||
|
||||
Returns (None, None) for wildcard-expanded or unknown names, and treats a
|
||||
malformed configured value as absent rather than failing the caller.
|
||||
|
||||
Deliberately reads one deployment rather than aggregating the group the way
|
||||
get_model_listing_info does: its caller truncates an embedding input to this
|
||||
value, so the widest window in a mixed group would be the wrong answer there.
|
||||
"""
|
||||
deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name)
|
||||
if deployment is None:
|
||||
return (None, None)
|
||||
wire: Final = _Wire()
|
||||
configured: Final = {
|
||||
key: wire.to_wire(value) for key, value in deployment.model_info.model_dump(exclude_none=True).items()
|
||||
}
|
||||
call: Final = self._core_handle().configured_token_limits(wire.to_wire(deployment.model_info.id), configured)
|
||||
done: Final = _drive_model_info(call, partial(self._answer_model_info_lookup, wire=wire))
|
||||
limits: Final = cast("list[int | None]", done["value"]) # cast-ok: router-core's two limits
|
||||
return (limits[0], limits[1])
|
||||
|
||||
def _core_get_configured_mode(self, model_name: str) -> "str | None":
|
||||
"""Return the mode explicitly configured for a concrete deployment."""
|
||||
deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name)
|
||||
if deployment is None:
|
||||
return None
|
||||
return self._core_handle().configured_text(_Wire().to_wire(deployment.model_info.get("mode")))
|
||||
|
||||
def _core_get_configured_display_name(self, model_name: str) -> "str | None":
|
||||
"""
|
||||
Return the display_name explicitly configured in a concrete deployment's
|
||||
model_info for model_name, via O(1) index lookup.
|
||||
|
||||
Returns None for wildcard-expanded or unknown names, and treats a
|
||||
non-string or empty configured value as absent rather than failing the
|
||||
listing. Like get_configured_token_limits, this never triggers pattern
|
||||
matching or deep copies, so it is safe to call per listed model on the
|
||||
/v1/models hot path.
|
||||
"""
|
||||
deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name)
|
||||
if deployment is None:
|
||||
return None
|
||||
return self._core_handle().configured_text(_Wire().to_wire(deployment.model_info.get("display_name")))
|
||||
|
||||
def _core_get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]:
|
||||
"""
|
||||
Return the service_tiers value each deployment a request from team_id
|
||||
to model_name can route to (see _routable_deployments) configures in
|
||||
its model_info, unvalidated and in model_list order, None for a
|
||||
deployment that sets none; the caller validates the shape it expects
|
||||
and decides how the deployments combine.
|
||||
|
||||
Returns an empty tuple for wildcard-expanded or unknown names.
|
||||
"""
|
||||
wire: Final = _Wire()
|
||||
tiers: Final = self._core_handle().configured_service_tiers(
|
||||
[
|
||||
wire.to_wire(dict(_configured_model_info(deployment)))
|
||||
for deployment in self._routable_deployments(model_name, team_id)
|
||||
]
|
||||
)
|
||||
return tuple(wire.from_wire(tier) for tier in tiers)
|
||||
|
||||
def _core_get_configured_max_input_tokens(
|
||||
self, model_name: str, team_id: str | None = None
|
||||
) -> tuple[int | None, ...]:
|
||||
"""
|
||||
Return the max_input_tokens each deployment a request from team_id to
|
||||
model_name can route to (see _routable_deployments) sets in its own
|
||||
model_info, in model_list order, None for a deployment that sets none
|
||||
or a malformed one. Unlike get_model_listing_info, neither the cost map
|
||||
nor a discovered limit is read, so only an operator's own value counts,
|
||||
and a value a saved /model/info response copied from the cost map is
|
||||
skipped too.
|
||||
|
||||
Returns an empty tuple for wildcard-expanded or unknown names.
|
||||
"""
|
||||
return tuple(
|
||||
_operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id)
|
||||
)
|
||||
|
||||
async def _core_get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]:
|
||||
"""
|
||||
Returns current tpm/rpm usage for model group
|
||||
|
||||
Parameters:
|
||||
- model_group: str - the received model name from the user (can be a wildcard route).
|
||||
|
||||
Returns:
|
||||
- usage: Tuple[tpm, rpm]
|
||||
"""
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock
|
||||
tpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged
|
||||
rpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged
|
||||
|
||||
model_list: Final = self._core_get_model_list(model_name=model_group)
|
||||
if model_list is None: # no matching deployments
|
||||
return None, None
|
||||
|
||||
for model in model_list:
|
||||
id: str | None = model.get("model_info", {}).get("id")
|
||||
litellm_model: str | None = model["litellm_params"].get(
|
||||
"model"
|
||||
) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written.
|
||||
if id is None or litellm_model is None:
|
||||
continue
|
||||
tpm_keys.append(self._core_handle().usage_key(id, litellm_model, "tpm", current_minute))
|
||||
rpm_keys.append(self._core_handle().usage_key(id, litellm_model, "rpm", current_minute))
|
||||
combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys
|
||||
|
||||
with service_target(ROUTER_USAGE_TARGET):
|
||||
combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys)
|
||||
if combined_tpm_rpm_values is None:
|
||||
return None, None
|
||||
|
||||
rpm_start: Final = len(tpm_keys)
|
||||
tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[:rpm_start] # mutable-ok: ported from the Python core unchanged # fmt: skip
|
||||
rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[rpm_start:] # mutable-ok: ported from the Python core unchanged # fmt: skip
|
||||
|
||||
## TPM
|
||||
tpm_usage: int | None = None # rebind-ok: ported from the Python core unchanged
|
||||
if tpm_usage_list is not None:
|
||||
for t in tpm_usage_list:
|
||||
if isinstance(t, int):
|
||||
if tpm_usage is None:
|
||||
tpm_usage = 0
|
||||
tpm_usage += t
|
||||
## RPM
|
||||
rpm_usage: int | None = None # rebind-ok: ported from the Python core unchanged
|
||||
if rpm_usage_list is not None:
|
||||
for t in rpm_usage_list:
|
||||
if isinstance(t, int):
|
||||
if rpm_usage is None:
|
||||
rpm_usage = 0
|
||||
rpm_usage += t
|
||||
return tpm_usage, rpm_usage
|
||||
|
||||
async def get_model_group_io_token_usage(self, model_group: str) -> tuple[int | None, int | None]:
|
||||
"""
|
||||
Returns current ITPM/OTPM usage for a model group (sum across deployments).
|
||||
"""
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
itpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged
|
||||
otpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged
|
||||
|
||||
model_list: Final = self._core_get_model_list(model_name=model_group)
|
||||
if model_list is None:
|
||||
return None, None
|
||||
|
||||
for model in model_list:
|
||||
model_id: str | None = model.get("model_info", {}).get("id")
|
||||
litellm_model: str | None = model["litellm_params"].get("model")
|
||||
if model_id is None or litellm_model is None:
|
||||
continue
|
||||
itpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "itpm", current_minute))
|
||||
otpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "otpm", current_minute))
|
||||
|
||||
with service_target(ROUTER_USAGE_TARGET):
|
||||
combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys)
|
||||
if combined_values is None:
|
||||
return None, None
|
||||
|
||||
itpm_values: Final = combined_values[: len(itpm_keys)]
|
||||
otpm_values: Final = combined_values[len(itpm_keys) :]
|
||||
|
||||
total_itpm: int | None = None # rebind-ok: ported from the Python core unchanged
|
||||
for value in itpm_values:
|
||||
if isinstance(value, int):
|
||||
total_itpm = (total_itpm or 0) + value
|
||||
|
||||
total_otpm: int | None = None # rebind-ok: ported from the Python core unchanged
|
||||
for value in otpm_values:
|
||||
if isinstance(value, int):
|
||||
total_otpm = (total_otpm or 0) + value
|
||||
|
||||
return total_itpm, total_otpm
|
||||
|
||||
async def _core_get_remaining_model_group_usage(
|
||||
self, model_group: str
|
||||
) -> dict[str, int]: # mutable-ok: ported from the Python core unchanged
|
||||
model_group_info: Final = self._cached_get_model_group_info(model_group)
|
||||
|
||||
returned_dict: Final[dict[str, int]] = {} # mutable-ok: ported from the Python core unchanged
|
||||
|
||||
# ITPM/OTPM groups emit input/output token headers, but they may also set
|
||||
# tpm/rpm, so build both sets rather than returning early - clients and
|
||||
# prometheus gauges that read the standard headers still get data.
|
||||
if model_group_info is not None and (model_group_info.itpm is not None or model_group_info.otpm is not None):
|
||||
current_itpm, current_otpm = await self.get_model_group_io_token_usage(model_group)
|
||||
returned_dict.update(
|
||||
build_io_token_rate_limit_headers(
|
||||
itpm_limit=model_group_info.itpm,
|
||||
otpm_limit=model_group_info.otpm,
|
||||
current_itpm=current_itpm,
|
||||
current_otpm=current_otpm,
|
||||
)
|
||||
)
|
||||
|
||||
tpm_limit: Final = model_group_info.tpm if model_group_info is not None else None
|
||||
rpm_limit: Final = model_group_info.rpm if model_group_info is not None else None
|
||||
|
||||
if tpm_limit is None and rpm_limit is None:
|
||||
return returned_dict
|
||||
|
||||
current_tpm, current_rpm = await self._core_get_model_group_usage(model_group)
|
||||
|
||||
if tpm_limit is not None:
|
||||
returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0)
|
||||
returned_dict["x-ratelimit-limit-tokens"] = tpm_limit
|
||||
if rpm_limit is not None:
|
||||
returned_dict["x-ratelimit-remaining-requests"] = rpm_limit - (current_rpm or 0)
|
||||
returned_dict["x-ratelimit-limit-requests"] = rpm_limit
|
||||
|
||||
return returned_dict
|
||||
|
||||
async def set_response_headers(
|
||||
self,
|
||||
response: object,
|
||||
model_group: str | None = None,
|
||||
request_kwargs: dict[str, object] | None = None, # mutable-ok: ported from the Python core unchanged
|
||||
) -> Any: # noqa: ANN401 # the core's signature, ported unchanged
|
||||
"""
|
||||
Add the most accurate rate limit headers for a given model response.
|
||||
|
||||
## TODO: add model group rate limit headers
|
||||
# - if healthy_deployments > 1, return model group rate limit headers
|
||||
# - else return the model's rate limit headers
|
||||
"""
|
||||
response = prepare_response_for_header_attachment(response) # rebind-ok: ported from the Python core unchanged
|
||||
if response is None:
|
||||
return response
|
||||
|
||||
additional_headers: Final = ensure_response_additional_headers(response)
|
||||
apply_response_model_id(
|
||||
response, find_deployment_metadata(request_kwargs) if request_kwargs is not None else None
|
||||
)
|
||||
additional_headers["x-litellm-model-group"] = model_group
|
||||
apply_quality_router_decision_headers(additional_headers, request_kwargs)
|
||||
additional_headers.update(complexity_router_decision_headers(request_kwargs))
|
||||
|
||||
if model_group is not None:
|
||||
remaining_usage: Final = await self._core_get_remaining_model_group_usage(model_group)
|
||||
apply_remaining_usage_headers(additional_headers, remaining_usage)
|
||||
return response
|
||||
375
litellm/router_rust/surface/model_list.py
Normal file
375
litellm/router_rust/surface/model_list.py
Normal file
|
|
@ -0,0 +1,375 @@
|
|||
"""
|
||||
``HostedModelList``: the Rust core's live ``model_list``, read back from router-core.
|
||||
|
||||
router-core holds the deployments. This list is what ``registry().all()`` and the deprecated
|
||||
``model_list`` attribute hand out: a ``list`` (so callers that check for one, serialize it or
|
||||
validate it keep working) whose contents are router-core's deployment records, read through the
|
||||
handle after every write. Every write to it is a changeset instead of a Python-side mutation:
|
||||
|
||||
- ``append`` / ``extend`` / an ``insert`` at the end add at the end (``append``)
|
||||
- removing one deployment whose id no other entry shares is a ``delete``
|
||||
- any other write (an assignment, a slice, a removal of one of several entries sharing an id,
|
||||
``clear``, ``sort``, ...) hands router-core the resulting list (``replace``)
|
||||
|
||||
A build serves the deployments it added once it ends (or fails), and the list reads them back
|
||||
then. After writes that only appended, it reads just the appended records
|
||||
(``RouterCoreHandle.deployments_from``), so adding deployments one at a time costs each one its own
|
||||
record rather than the whole list. These writes are the raw ones a caller makes: like the Python
|
||||
router's list, they reindex nothing.
|
||||
|
||||
An entry read back from router-core is a dict whose writes (its own, or a nested dict's) hand
|
||||
router-core the list again, so an edit in place reaches it as the Python router's list would show
|
||||
it. Each deployment keeps one Python object while router-core's record for it is unchanged, keyed by
|
||||
its id and its place among entries sharing that id, so identity checks on stored entries (the
|
||||
model info caches) mean what they did: the same object until a write changes that deployment.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import datetime
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
SupportsIndex,
|
||||
TypeAlias,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # any value a caller wrote into the list
|
||||
overload,
|
||||
)
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.types.router import ModelAccessWindow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.rust_bridge._native import RouterCoreHandle
|
||||
|
||||
Entry: TypeAlias = dict[str, object]
|
||||
# (id, its place among the entries sharing it) -> (the record it was read as, the entry object)
|
||||
_Stable: TypeAlias = dict[tuple[str, int], tuple[object, Entry]]
|
||||
|
||||
|
||||
_DATETIME: Final = TypeAdapter(datetime.datetime)
|
||||
_Field = TypeVar("_Field") # rebind-ok: a TypeVar cannot be Final
|
||||
|
||||
|
||||
def _dumped_as(adapter: TypeAdapter[_Field]) -> Callable[[object], object]:
|
||||
"""A value validated by ``adapter``, then dumped back to its Python form."""
|
||||
|
||||
def dump(value: object) -> object:
|
||||
return adapter.dump_python(adapter.validate_python(value), exclude_none=True)
|
||||
|
||||
return dump
|
||||
|
||||
|
||||
# The ``ModelInfo`` fields whose Python value is no JSON value, which router-core holds as strings.
|
||||
_TYPED_MODEL_INFO: Final[Mapping[str, Callable[[object], object]]] = MappingProxyType(
|
||||
{
|
||||
"created_at": _dumped_as(_DATETIME),
|
||||
"updated_at": _dumped_as(_DATETIME),
|
||||
"ptu_effective_from": _dumped_as(_DATETIME),
|
||||
"ptu_effective_to": _dumped_as(_DATETIME),
|
||||
"access_windows": _dumped_as(TypeAdapter(tuple[ModelAccessWindow, ...])),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def typed_deployment(record: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""A deployment record as ``Deployment.to_json`` dumped it: its ``model_info`` typed again. One
|
||||
with none of those fields comes back as it is, and so does a value a raw write left unreadable."""
|
||||
info: Final = as_str_mapping(record.get("model_info"))
|
||||
if info is None or info.keys().isdisjoint(_TYPED_MODEL_INFO):
|
||||
return record
|
||||
typed: Final[dict[str, object]] = {} # mutable-ok: filled here
|
||||
for key, dumped in _TYPED_MODEL_INFO.items():
|
||||
value = info.get(key)
|
||||
if value is None:
|
||||
continue
|
||||
try:
|
||||
typed[key] = dumped(value)
|
||||
except ValidationError:
|
||||
continue
|
||||
return {**record, "model_info": {**info, **typed}}
|
||||
|
||||
|
||||
def _record(value: object) -> Mapping[str, object]:
|
||||
return cast(Mapping[str, object], value) # cast-ok: a JSON object router-core wrote, or a dict inside one
|
||||
|
||||
|
||||
def _entry_id(entry: object) -> str | None:
|
||||
info: Final = as_str_mapping((as_str_mapping(entry) or {}).get("model_info")) or {}
|
||||
value: Final = info.get("id")
|
||||
return None if value is None else str(value)
|
||||
|
||||
|
||||
def _keys(
|
||||
entries: Iterable[object], seen: Mapping[str, int] = MappingProxyType({})
|
||||
) -> tuple[list[tuple[str, int] | None], Mapping[str, int]]: # mutable-ok: a fresh list
|
||||
"""Each entry's identity key (its id and how many earlier entries share it, ``seen`` counting
|
||||
the ids of entries before these), and the counts once past them."""
|
||||
counts: Final[dict[str, int]] = dict(seen) # mutable-ok: counted here
|
||||
keys: Final[list[tuple[str, int] | None]] = [] # mutable-ok: filled here
|
||||
for entry in entries:
|
||||
entry_id = _entry_id(entry)
|
||||
if entry_id is None:
|
||||
keys.append(None)
|
||||
continue
|
||||
nth = counts.get(entry_id, 0)
|
||||
counts[entry_id] = nth + 1
|
||||
keys.append((entry_id, nth))
|
||||
return keys, counts
|
||||
|
||||
|
||||
class _HostedEntry(dict[str, object]):
|
||||
"""A deployment entry read back from router-core: a write to it, or to a dict inside it, hands
|
||||
router-core the list again (``replace``), so it learns what the entry now says. A copy is a
|
||||
plain dict."""
|
||||
|
||||
def __init__(self, record: Mapping[str, object], written: Callable[[], None]) -> None:
|
||||
super().__init__(
|
||||
(key, value if (nested := as_str_mapping(value)) is None else _HostedEntry(nested, written))
|
||||
for key, value in record.items()
|
||||
)
|
||||
self._written: Final = written
|
||||
|
||||
def __setitem__(self, key: str, value: object) -> None:
|
||||
super().__setitem__(key, value)
|
||||
self._written()
|
||||
|
||||
def __delitem__(self, key: str) -> None:
|
||||
super().__delitem__(key)
|
||||
self._written()
|
||||
|
||||
def pop(self, key: str, *default: object) -> object:
|
||||
value: Final = super().pop(key, *default)
|
||||
self._written()
|
||||
return value
|
||||
|
||||
def popitem(self) -> tuple[str, object]:
|
||||
item: Final = super().popitem()
|
||||
self._written()
|
||||
return item
|
||||
|
||||
def clear(self) -> None:
|
||||
super().clear()
|
||||
self._written()
|
||||
|
||||
def update(self, *args: object, **kwargs: object) -> None: # kwargs-ok: dict.update's own arguments
|
||||
super().update(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # dict.update's own arguments
|
||||
self._written()
|
||||
|
||||
def setdefault(self, key: str, default: object = None) -> object:
|
||||
if key in self:
|
||||
return self[key]
|
||||
self[key] = default
|
||||
return default
|
||||
|
||||
def __ior__(self, other: object) -> "_HostedEntry": # pyright: ignore[reportIncompatibleMethodOverride] # dict's in-place merge, as a write
|
||||
self.update(other)
|
||||
return self
|
||||
|
||||
def __copy__(self) -> dict[str, object]: # mutable-ok: a plain copy
|
||||
return dict(self)
|
||||
|
||||
def __deepcopy__(self, memo: dict[int, object]) -> dict[str, object]: # mutable-ok: a plain copy
|
||||
return copy.deepcopy(dict(self), memo)
|
||||
|
||||
def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]:
|
||||
return (dict, (copy.deepcopy(dict(self)),))
|
||||
|
||||
|
||||
class HostedModelList(list[Entry]):
|
||||
"""The live ``model_list``: router-core's deployments, every write a changeset."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
handle: Callable[[], "RouterCoreHandle | None"],
|
||||
apply: Callable[[Mapping[str, object]], None],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._handle: Final = handle
|
||||
self._apply: Final = apply
|
||||
# The handle and registry version the contents were last read at.
|
||||
self._read_at: tuple[object, int] | None = None
|
||||
self._stable: _Stable = {}
|
||||
# How many of the shown entries hold each id, for the keys of entries appended next.
|
||||
self._id_counts: Mapping[str, int] = {}
|
||||
# While a build runs: the length, id counts and read to go back to when it ends.
|
||||
self._building: tuple[int, Mapping[str, int], tuple[object, int] | None] | None = None
|
||||
|
||||
# -- reading ----------------------------------------------------------------------------------
|
||||
|
||||
def refresh(self) -> None:
|
||||
"""Read router-core's deployments again when its registry moved since the last read."""
|
||||
handle: Final = self._handle()
|
||||
if handle is None or self._building is not None:
|
||||
# While a build runs, the list shows what it accepted until router-core serves it.
|
||||
return
|
||||
version: Final = handle.registry_version()
|
||||
if self._read_at is not None and self._read_at[0] is handle and self._read_at[1] == version:
|
||||
return
|
||||
read_at: Final = self._read_at
|
||||
self._read_at = (handle, version)
|
||||
if read_at is not None and read_at[0] is handle and handle.registry_rewritten() <= read_at[1]:
|
||||
# Only appends since the last read: the shown entries stand, and the new ones follow.
|
||||
super().extend(self._kept(handle.deployments_from(len(self)), appended=True))
|
||||
return
|
||||
self._show(self._kept(handle.deployments()))
|
||||
|
||||
def _kept(self, records: Sequence[object], *, appended: bool = False) -> list[Entry]: # mutable-ok: a fresh list
|
||||
"""``records`` as entries, each the object it was while its record is unchanged; when
|
||||
``appended``, they follow the shown entries, which keep theirs."""
|
||||
stable: Final[_Stable] = self._stable if appended else {}
|
||||
keys, counts = _keys(records, self._id_counts if appended else {})
|
||||
entries: Final[list[Entry]] = [] # mutable-ok: filled here
|
||||
for key, record in zip(keys, records, strict=True):
|
||||
typed = typed_deployment(_record(record))
|
||||
held = self._stable.get(key) if key is not None else None
|
||||
kept = held is not None and held[0] == typed
|
||||
entry = held[1] if held is not None and kept else _HostedEntry(typed, self._entry_written)
|
||||
if key is not None:
|
||||
stable[key] = (held[0] if held is not None and kept else copy.deepcopy(typed), entry)
|
||||
entries.append(entry)
|
||||
self._stable = stable
|
||||
self._id_counts = counts
|
||||
return entries
|
||||
|
||||
# -- a build ----------------------------------------------------------------------------------
|
||||
|
||||
def begin_build(self, *, replacing: bool) -> None:
|
||||
"""A build starts: router-core serves what it adds once it ends, so until then the list
|
||||
shows each deployment it accepts here (``built``), after the entries it keeps, or alone
|
||||
when it replaces them, as the Python router's list grows while it adds them."""
|
||||
self.refresh()
|
||||
if replacing:
|
||||
self._show([])
|
||||
self._id_counts = {}
|
||||
self._building = (len(self), self._id_counts, self._read_at)
|
||||
|
||||
def built(self, entry: Mapping[str, object]) -> None:
|
||||
"""A deployment the build accepted, shown until the build ends."""
|
||||
if self._building is not None:
|
||||
super().extend(self._kept([entry], appended=True))
|
||||
|
||||
def end_build(self) -> None:
|
||||
"""The build ended (or failed): router-core serves what it added, read back in place of
|
||||
the entries shown meanwhile, each kept as the object it was shown as while unchanged."""
|
||||
if self._building is None:
|
||||
return
|
||||
start, counts, read_at = self._building
|
||||
self._building = None
|
||||
super().__delitem__(slice(start, None))
|
||||
self._id_counts = counts
|
||||
self._read_at = read_at
|
||||
self.refresh()
|
||||
|
||||
def _seed(self, entries: Sequence[object]) -> None:
|
||||
"""Keep ``entries``, the list router-core now holds, as the objects its records for them
|
||||
read back as."""
|
||||
for key, entry in zip(_keys(entries)[0], entries, strict=True):
|
||||
self._seed_one(key, entry)
|
||||
|
||||
def _seed_appended(self, entries: Sequence[object]) -> None:
|
||||
"""Keep ``entries``, appended after the shown ones, as the objects they read back as."""
|
||||
for key, entry in zip(_keys(entries, self._id_counts)[0], entries, strict=True):
|
||||
self._seed_one(key, entry)
|
||||
|
||||
def _seed_one(self, key: tuple[str, int] | None, entry: object) -> None:
|
||||
if key is not None and isinstance(entry, dict):
|
||||
stored: Final = cast(Entry, entry) # cast-ok: a model_list entry is a dict
|
||||
self._stable[key] = (copy.deepcopy(stored), stored)
|
||||
|
||||
def _entry_written(self) -> None:
|
||||
"""A write inside an entry: router-core takes the list as it now reads."""
|
||||
self._replace(list(self))
|
||||
|
||||
def _show(self, entries: list[Entry]) -> None: # mutable-ok: the entries shown
|
||||
super().__setitem__(slice(None), entries)
|
||||
|
||||
# -- writing ----------------------------------------------------------------------------------
|
||||
|
||||
def _replace(self, entries: list[Entry]) -> None: # mutable-ok: the list router-core takes
|
||||
# Seeded first: router-core's records for them read back as these objects however soon
|
||||
# the list is read again.
|
||||
self._seed(entries)
|
||||
self._apply({"replace": entries})
|
||||
self.refresh()
|
||||
|
||||
def replace(self, entries: Iterable[object]) -> None:
|
||||
"""Every deployment at once (``router.model_list = [...]``)."""
|
||||
self._replace(cast(list[Entry], list(entries))) # cast-ok: the entries as assigned, which router-core validates
|
||||
|
||||
def append(self, value: Entry) -> None:
|
||||
self.refresh()
|
||||
self._seed_appended([value])
|
||||
self._apply({"append": [value]})
|
||||
self.refresh()
|
||||
|
||||
def extend(self, values: Iterable[Entry]) -> None:
|
||||
for value in tuple(values):
|
||||
self.append(value)
|
||||
|
||||
def __iadd__(self, values: Iterable[Entry]) -> "HostedModelList": # pyright: ignore[reportIncompatibleMethodOverride] # list's in-place add, as a changeset
|
||||
self.extend(values)
|
||||
return self
|
||||
|
||||
def insert(self, index: SupportsIndex, value: Entry) -> None:
|
||||
if int(index) >= len(self):
|
||||
self.append(value)
|
||||
return
|
||||
changed: Final = list(self)
|
||||
changed.insert(index, value)
|
||||
self._replace(changed)
|
||||
|
||||
@overload
|
||||
def __setitem__(self, index: SupportsIndex, value: Entry) -> None: ...
|
||||
@overload
|
||||
def __setitem__(self, index: slice, value: Iterable[Entry]) -> None: ...
|
||||
def __setitem__(self, index: SupportsIndex | slice, value: object) -> None:
|
||||
changed: Final = list(self)
|
||||
changed[index] = value # pyright: ignore[reportArgumentType, reportCallIssue] # list's own item and slice assignment
|
||||
self._replace(changed)
|
||||
|
||||
def __delitem__(self, index: SupportsIndex | slice) -> None:
|
||||
changed: Final = list(self)
|
||||
if isinstance(index, slice):
|
||||
del changed[index]
|
||||
self._replace(changed)
|
||||
return
|
||||
removed: Final = changed.pop(index)
|
||||
removed_id: Final = _entry_id(removed)
|
||||
if removed_id is None or any(_entry_id(entry) == removed_id for entry in changed):
|
||||
# A delete removes every deployment with the id; the list removes only this one.
|
||||
self._replace(changed)
|
||||
return
|
||||
self._apply({"delete": [removed_id]})
|
||||
self.refresh()
|
||||
|
||||
def pop(self, index: SupportsIndex = -1) -> Entry:
|
||||
entry: Final = self[index]
|
||||
del self[index]
|
||||
return entry
|
||||
|
||||
def remove(self, value: Entry) -> None:
|
||||
del self[self.index(value)]
|
||||
|
||||
def clear(self) -> None:
|
||||
self._replace([])
|
||||
|
||||
def sort(self, *args: object, **kwargs: object) -> None: # kwargs-ok: list.sort's own arguments
|
||||
changed: Final = list(self)
|
||||
changed.sort(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # list.sort's own arguments
|
||||
self._replace(changed)
|
||||
|
||||
def reverse(self) -> None:
|
||||
self._replace(list(reversed(self)))
|
||||
|
||||
def __imul__(self, count: SupportsIndex) -> "HostedModelList": # pyright: ignore[reportIncompatibleMethodOverride] # list's in-place repeat, as a changeset
|
||||
self._replace(list(self) * int(count))
|
||||
return self
|
||||
|
||||
def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]:
|
||||
# A copy or pickle of the live list is a plain list of its entries.
|
||||
return (list, (list(self),))
|
||||
2191
litellm/router_rust/surface/registry.py
Normal file
2191
litellm/router_rust/surface/registry.py
Normal file
File diff suppressed because it is too large
Load diff
30
litellm/router_rust/surface/replay.py
Normal file
30
litellm/router_rust/surface/replay.py
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
"""The Rust core's live routers, whose deployments a price data reload re-registers in
|
||||
``litellm.model_cost`` until they retire (``_live_routers`` / ``_replay_live_router_model_cost``)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import weakref
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.utils import add_live_deployment_replay
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_rust.surface.registry import HostRegistry
|
||||
|
||||
_live_registries: Final[weakref.WeakSet[HostRegistry]] = weakref.WeakSet()
|
||||
|
||||
|
||||
def track_live_registry(registry: HostRegistry) -> None:
|
||||
_live_registries.add(registry)
|
||||
|
||||
|
||||
def untrack_live_registry(registry: HostRegistry) -> None:
|
||||
_live_registries.discard(registry)
|
||||
|
||||
|
||||
def _replay_live_registries() -> None:
|
||||
for registry in tuple(_live_registries):
|
||||
registry.replay_model_cost_registrations()
|
||||
|
||||
|
||||
add_live_deployment_replay(_replay_live_registries)
|
||||
295
litellm/router_rust/surface/retries.py
Normal file
295
litellm/router_rust/surface/retries.py
Normal file
|
|
@ -0,0 +1,295 @@
|
|||
"""
|
||||
The retry members ``RustCore`` serves: ``should_retry_this_error`` and
|
||||
``get_num_retries_from_retry_policy``. router-core decides both; this module only describes the
|
||||
Python exception the way an attempt failure is described and reads the answer back.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import operator
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.types.router import RetryPolicy
|
||||
|
||||
# The most retries a request's `num_retries` asks router-core for (its count is a u32).
|
||||
_MOST_RETRIES: Final = 2**32 - 1
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_rust.core import Resolved
|
||||
from litellm.rust_bridge._native import ( # pyright: ignore[reportPrivateUsage] # the stub's handle
|
||||
RouterCoreHandle,
|
||||
)
|
||||
|
||||
|
||||
def resolved_retry_policy(value: object) -> RetryPolicy | None:
|
||||
"""The router's ``retry_policy`` as the Python router's constructor keeps it: a dict is read
|
||||
into a ``RetryPolicy``, a ``RetryPolicy`` is kept, anything else is no policy."""
|
||||
if isinstance(value, RetryPolicy):
|
||||
return value
|
||||
if isinstance(value, Mapping):
|
||||
fields: Final[Mapping[str, int | None]] = value # pyright: ignore[reportUnknownVariableType] # the policy's fields
|
||||
return RetryPolicy(**fields) # pyright: ignore[reportArgumentType] # validated by the model
|
||||
return None
|
||||
|
||||
|
||||
def _is_policy(entry: object) -> bool:
|
||||
return isinstance(entry, RetryPolicy | Mapping)
|
||||
|
||||
|
||||
def reads_as_policy(entry: object) -> bool:
|
||||
"""Whether ``resolve_retry_policy`` reads ``entry`` without raising: no entry, a ``RetryPolicy``,
|
||||
or a mapping ``RetryPolicy(**entry)`` accepts."""
|
||||
if entry is None or isinstance(entry, RetryPolicy):
|
||||
return True
|
||||
if not isinstance(entry, Mapping):
|
||||
return False
|
||||
try:
|
||||
_ = RetryPolicy(**entry) # pyright: ignore[reportUnknownArgumentType] # validated by the model
|
||||
except Exception: # noqa: BLE001 # any error the lookup would raise
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def is_policy_per_group(policies: object) -> bool:
|
||||
"""Whether ``model_group_retry_policy`` is what router-core's settings hold: a mapping of
|
||||
group names to policies (or nothing)."""
|
||||
if policies is None:
|
||||
return True
|
||||
if not isinstance(policies, Mapping):
|
||||
return False
|
||||
entries: Final[Mapping[object, object]] = policies # pyright: ignore[reportUnknownVariableType] # read as objects
|
||||
return all(isinstance(group, str) and _is_policy(entry) for group, entry in entries.items())
|
||||
|
||||
|
||||
def readable_retry_policy(value: object) -> RetryPolicy | None:
|
||||
"""A router ``retry_policy`` written as a plain attribute, as router-core's settings hold it.
|
||||
The attribute keeps what was written; a value the lookup raises for reaches router-core as
|
||||
``retry_policy_raises`` instead."""
|
||||
return resolved_retry_policy(value) if reads_as_policy(value) else None
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class GroupPolicies:
|
||||
"""A ``model_group_retry_policy`` as router-core reads it, resolved the way
|
||||
``resolve_retry_policy`` reads it for each group the call may evaluate.
|
||||
|
||||
``document`` holds a group's policy where the lookup finds one, and an empty policy (which
|
||||
grants nothing) where it finds ``None``; a group it does not find is absent, so the router's
|
||||
``retry_policy`` applies. ``raising`` names the groups whose lookup raises (an ``in``/``[]``
|
||||
the value does not support, or an entry that is not a policy, whose fields
|
||||
``get_num_retries_from_retry_policy`` then fails to read): the Python retry loop raises that
|
||||
error in place of the attempt's on a hop's first failure.
|
||||
"""
|
||||
|
||||
document: Mapping[str, object]
|
||||
raising: frozenset[str]
|
||||
|
||||
|
||||
def _looked_up(policies: object, group: str) -> tuple[bool, object]:
|
||||
"""``(found, entry)`` as ``resolve_retry_policy`` looks ``group`` up; raises as it raises."""
|
||||
found: Final = group in policies # pyright: ignore[reportOperatorIssue] # any value, as the Python router reads it
|
||||
return (True, policies[group]) if found else (False, None) # pyright: ignore[reportIndexIssue, reportUnknownVariableType] # see above
|
||||
|
||||
|
||||
def _policy_document(entry: object) -> object:
|
||||
"""A readable entry as router-core reads it: no entry grants nothing, a mapping is read the way
|
||||
``RetryPolicy(**entry)`` reads it."""
|
||||
if entry is None:
|
||||
return {}
|
||||
if isinstance(entry, Mapping):
|
||||
return RetryPolicy(**entry) # pyright: ignore[reportUnknownArgumentType] # checked by reads_as_policy
|
||||
return entry
|
||||
|
||||
|
||||
def group_policies(policies: object, groups: Iterable[str]) -> GroupPolicies:
|
||||
"""``policies`` resolved for ``groups`` (a mapping is read by its own keys)."""
|
||||
if policies is None:
|
||||
return GroupPolicies({}, raising=frozenset())
|
||||
if isinstance(policies, Mapping):
|
||||
entries: Final[Mapping[object, object]] = policies # pyright: ignore[reportUnknownVariableType] # read as objects
|
||||
# A key that is no string never matches a group name.
|
||||
named: Final = {group: entry for group, entry in entries.items() if isinstance(group, str)}
|
||||
readable: Final = {group: entry for group, entry in named.items() if reads_as_policy(entry)}
|
||||
return GroupPolicies(
|
||||
{group: _policy_document(entry) for group, entry in readable.items()},
|
||||
raising=frozenset(named) - frozenset(readable),
|
||||
)
|
||||
document: Final[dict[str, object]] = {} # mutable-ok: built here; router-core reads it as plain JSON
|
||||
raising: Final[set[str]] = set() # mutable-ok: the groups whose lookup raises
|
||||
for group in groups:
|
||||
try:
|
||||
found, entry = _looked_up(policies, group)
|
||||
except Exception: # noqa: BLE001 # the Python retry loop raises it on a hop's first failure
|
||||
raising.add(group)
|
||||
continue
|
||||
if not found:
|
||||
continue
|
||||
if not reads_as_policy(entry):
|
||||
raising.add(group)
|
||||
continue
|
||||
document[group] = _policy_document(entry)
|
||||
return GroupPolicies(document, raising=frozenset(raising))
|
||||
|
||||
|
||||
def static_group_policies(policies: object) -> object:
|
||||
"""The router's ``model_group_retry_policy`` as router-core's static settings hold it. A value
|
||||
that is not a mapping of policies is resolved per request instead (``group_policies``)."""
|
||||
return group_policies(policies, ()).document if isinstance(policies, Mapping) else None # pyright: ignore[reportUnknownArgumentType] # read as objects
|
||||
|
||||
|
||||
def raising_groups(request: Mapping[str, object], resolved: GroupPolicies) -> Sequence[str]:
|
||||
"""The groups whose ``model_group_retry_policy`` lookup raises on a hop's first failure, unless
|
||||
the retry loop never reads it: the request opted out with ``num_retries=0``. A streamed Anthropic
|
||||
messages call is no exception: only the stream's own retry rules read such an entry as no policy
|
||||
(``_anthropic_messages_resolved_retry_policy``); a failure before the stream opened reaches the
|
||||
retry loop."""
|
||||
if not resolved.raising or request.get("num_retries") == 0:
|
||||
return ()
|
||||
return sorted(resolved.raising)
|
||||
|
||||
|
||||
def policy_lookup_error(original: BaseException, model_group: str, retry_policy: object, policies: object) -> Exception:
|
||||
"""The error ``get_num_retries_from_retry_policy`` raises for ``model_group`` after ``original``
|
||||
failed (the Python retry loop raises it in the attempt's place, chained to it); a lookup that
|
||||
unexpectedly reads is reported as such."""
|
||||
from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy
|
||||
|
||||
traceback: Final = original.__traceback__
|
||||
try:
|
||||
# Looked up while `original` is being handled, as in the retry loop's `except` block, so
|
||||
# the error chains to it the way Python chains it.
|
||||
try:
|
||||
raise original
|
||||
except BaseException: # noqa: BLE001 # re-raised only to be handled here
|
||||
_ = get_num_retries_from_retry_policy(
|
||||
exception=original, # pyright: ignore[reportArgumentType] # any exception the hop raised
|
||||
retry_policy=retry_policy, # pyright: ignore[reportArgumentType] # the router's resolved policy
|
||||
model_group=model_group,
|
||||
model_group_retry_policy=policies, # pyright: ignore[reportArgumentType] # read as the Python router reads it
|
||||
)
|
||||
except Exception as raised: # noqa: BLE001 # the error the hop ends with
|
||||
return raised
|
||||
finally:
|
||||
original.__traceback__ = traceback # rebind-ok: put back as the hop left it
|
||||
return RuntimeError(f"the model_group_retry_policy lookup for {model_group!r} did not raise")
|
||||
|
||||
|
||||
def router_policy_retries(
|
||||
exception: Exception, retry_policy: object, policies: object, model_group: str | None
|
||||
) -> int | None:
|
||||
"""``get_num_retries_from_retry_policy`` for a router ``model_group_retry_policy`` that is not
|
||||
a mapping of policies, read as the Python router reads it (raising as it raises)."""
|
||||
from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy
|
||||
|
||||
return get_num_retries_from_retry_policy(
|
||||
exception=exception,
|
||||
retry_policy=retry_policy, # pyright: ignore[reportArgumentType] # the router's resolved policy
|
||||
model_group=model_group,
|
||||
model_group_retry_policy=policies, # pyright: ignore[reportArgumentType] # read as the Python router reads it
|
||||
)
|
||||
|
||||
|
||||
def request_retry_fields(value: object) -> Mapping[str, object]:
|
||||
"""A request's ``num_retries`` as router-core reads it. ``async_function_with_retries`` never
|
||||
casts it: it compares it with 0, and counts with ``range`` only once it is positive."""
|
||||
if value is None:
|
||||
return {}
|
||||
if isinstance(value, int):
|
||||
count: Final = int(value)
|
||||
return {"num_retries": min(count, _MOST_RETRIES)} if count >= 0 else {"uncounted_retries": "negative"}
|
||||
try:
|
||||
if value == 0:
|
||||
return {"num_retries": 0}
|
||||
if not value > 0: # pyright: ignore[reportOperatorIssue] # any value, compared as the retry loop compares it
|
||||
return {"uncounted_retries": "not_positive"}
|
||||
return {"num_retries": min(operator.index(value), _MOST_RETRIES)} # pyright: ignore[reportArgumentType] # see above
|
||||
except Exception: # noqa: BLE001 # the retry loop raises it on a hop's first failure
|
||||
return {"uncounted_retries": "raises"}
|
||||
|
||||
|
||||
def retry_count_error(original: BaseException, value: object) -> Exception:
|
||||
"""The error the retry loop raises reading a request's ``num_retries`` after ``original``
|
||||
failed, chained to it; a value that unexpectedly reads is reported as such."""
|
||||
traceback: Final = original.__traceback__
|
||||
try:
|
||||
try:
|
||||
raise original
|
||||
except BaseException: # noqa: BLE001 # re-raised only to be handled here
|
||||
if value > 0: # pyright: ignore[reportOperatorIssue] # `if num_retries > 0:`
|
||||
_ = range(value) # pyright: ignore[reportArgumentType] # `for current_attempt in range(num_retries):`
|
||||
except Exception as raised: # noqa: BLE001 # the error the hop ends with
|
||||
return raised
|
||||
finally:
|
||||
original.__traceback__ = traceback # rebind-ok: put back as the hop left it
|
||||
return RuntimeError(f"the request's num_retries {value!r} did not raise")
|
||||
|
||||
|
||||
def _count(deployments: object) -> int:
|
||||
return len(deployments) if isinstance(deployments, list) else 0 # pyright: ignore[reportUnknownArgumentType] # only the length is read
|
||||
|
||||
|
||||
def _filled(fallbacks: object) -> bool:
|
||||
return isinstance(fallbacks, Sequence) and len(fallbacks) > 0 # pyright: ignore[reportUnknownArgumentType] # only the length is read
|
||||
|
||||
|
||||
class RustRetries:
|
||||
"""Retry members read from the Rust core's handle."""
|
||||
|
||||
_handle: "RouterCoreHandle"
|
||||
_resolved: "Resolved"
|
||||
|
||||
@property
|
||||
def _written_retry_policy(self) -> object:
|
||||
"""The ``retry_policy`` attribute as written (``Resolved`` holds what router-core reads)."""
|
||||
raise NotImplementedError
|
||||
|
||||
def should_retry_this_error(
|
||||
self,
|
||||
error: Exception,
|
||||
healthy_deployments: list[object] | None = None, # mutable-ok: the Python router's signature
|
||||
all_deployments: list[object] | None = None, # mutable-ok: same
|
||||
context_window_fallbacks: list[object] | None = None, # mutable-ok: same
|
||||
content_policy_fallbacks: list[object] | None = None, # mutable-ok: same
|
||||
regular_fallbacks: list[object] | None = None, # mutable-ok: same
|
||||
) -> bool:
|
||||
"""``True`` when the error is retried in the same group; otherwise the error is raised."""
|
||||
from litellm.router_rust.core import attempt_failure
|
||||
|
||||
retried: Final = self._handle.should_retry(
|
||||
dict(attempt_failure(error)),
|
||||
(_count(healthy_deployments), _count(all_deployments)),
|
||||
(context_window_fallbacks is not None, content_policy_fallbacks is not None, _filled(regular_fallbacks)),
|
||||
)
|
||||
if not retried:
|
||||
raise error
|
||||
return True
|
||||
|
||||
def get_num_retries_from_retry_policy(self, exception: Exception, model_group: str | None = None) -> int | None:
|
||||
from litellm.router_rust.core import attempt_failure
|
||||
|
||||
resolved: Final = self._resolved
|
||||
written: Final = self._written_retry_policy
|
||||
if not is_policy_per_group(resolved.model_group_retry_policy) or not reads_as_policy(written):
|
||||
return router_policy_retries(exception, written, resolved.model_group_retry_policy, model_group)
|
||||
# A dict lookup hashes any key and only a string can match; the native call takes strings alone.
|
||||
_ = hash(model_group)
|
||||
return self._handle.retry_policy_retries(
|
||||
dict(attempt_failure(exception)), model_group if isinstance(model_group, str) else None
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GroupPolicies",
|
||||
"RustRetries",
|
||||
"group_policies",
|
||||
"is_policy_per_group",
|
||||
"policy_lookup_error",
|
||||
"raising_groups",
|
||||
"readable_retry_policy",
|
||||
"reads_as_policy",
|
||||
"request_retry_fields",
|
||||
"resolved_retry_policy",
|
||||
"retry_count_error",
|
||||
"router_policy_retries",
|
||||
"static_group_policies",
|
||||
]
|
||||
236
litellm/router_rust/surface/settings.py
Normal file
236
litellm/router_rust/surface/settings.py
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
"""
|
||||
The Rust core's settings semantics on the host side: how ``ChangeSet.settings`` is read (the rules
|
||||
of ``update_settings``), the ``settings().as_dict()`` view (``get_settings``), the router attributes
|
||||
the Python core keeps as plain attributes, deleted settings, and the construction-time side effects
|
||||
of the logging arguments.
|
||||
|
||||
The rules are the Python router core's, ported (never imported). Settings change only through a
|
||||
changeset; a direct attribute write is turned into one by ``RustCore.__setattr__``.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.constants import RUNTIME_UPDATABLE_ROUTER_SETTINGS
|
||||
from litellm.litellm_core_utils.routing_groups import validate_routing_strategy
|
||||
from litellm.types.router import RetryPolicy, RoutingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router_config import RouterConfig
|
||||
from litellm.router_python_host import PythonHost
|
||||
|
||||
MISSING: Final = object()
|
||||
|
||||
# `update_settings` casts these with `int(value)`.
|
||||
_INT_SETTINGS: Final = frozenset({"timeout", "num_retries", "retry_after", "allowed_fails", "cooldown_time"})
|
||||
CONTINUATION_SETTING: Final = "enable_mid_stream_fallback_continuation"
|
||||
|
||||
# Settings the Python core keeps as plain router attributes and `update_settings` sets as given,
|
||||
# with their constructor defaults. ``max_retries`` has none: it exists once a changeset sets it.
|
||||
PLAIN_DEFAULTS: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
"retry_policy": None,
|
||||
"model_group_retry_policy": {},
|
||||
"enable_weighted_failover": False,
|
||||
"fallback_on_unsupported_params": True,
|
||||
}
|
||||
)
|
||||
# `get_settings`: the settings it lists, in its order (`routing_groups` comes last).
|
||||
_VIEW_ORDER: Final = (
|
||||
"routing_strategy_args",
|
||||
"routing_strategy",
|
||||
"allowed_fails",
|
||||
"cooldown_time",
|
||||
"num_retries",
|
||||
"timeout",
|
||||
"max_retries",
|
||||
"retry_after",
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"model_group_retry_policy",
|
||||
"retry_policy",
|
||||
"model_group_alias",
|
||||
"enable_weighted_failover",
|
||||
"enable_mid_stream_fallback_continuation",
|
||||
"enable_tag_filtering",
|
||||
"tag_routing_prefix",
|
||||
"fallback_on_unsupported_params",
|
||||
)
|
||||
|
||||
# Router attributes that are their constructor argument as given, else this default. Arguments
|
||||
# the Rust core declines are only ever at their default here.
|
||||
_ARGUMENT_ATTRIBUTES: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
"set_verbose": False,
|
||||
"debug_level": "INFO",
|
||||
"client_ttl": 3600,
|
||||
"stream_timeout": None,
|
||||
"assistants_config": None,
|
||||
"ignore_invalid_deployments": False,
|
||||
"alerting_config": None,
|
||||
"auto_router_capability_limit": None,
|
||||
"deployment_affinity_ttl_seconds": 3600,
|
||||
"session_affinity_repin_on_failover": False,
|
||||
"model_group_affinity_config": None,
|
||||
}
|
||||
)
|
||||
# Class-level and constant attributes of the Python core.
|
||||
_CONSTANT_ATTRIBUTES: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{"default_cache_time_seconds": 1 * 60 * 60, "tenacity": None, "access_groups": None}
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SettingsUpdate:
|
||||
"""What a ``ChangeSet.settings`` mapping changes, read with ``update_settings``' rules."""
|
||||
|
||||
values: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) # fields and plain settings
|
||||
plain: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) # plain attributes, as stored
|
||||
continuation: object = MISSING # the continuation flag, when given
|
||||
routing_groups: object = MISSING # the routing groups input, when given
|
||||
optional_pre_call_checks: object = MISSING # the optional pre-call checks, when given
|
||||
failure: Exception | None = None # what stopped the reading; the settings before it stay applied
|
||||
|
||||
|
||||
def read_settings_update(settings: Mapping[str, object], routing_strategy: object) -> SettingsUpdate:
|
||||
"""``_core_update_settings``' reading of each setting, in order: names that are not
|
||||
runtime-updatable are ignored, the integer settings are cast with ``int``, a retry policy dict
|
||||
becomes a ``RetryPolicy`` (any other type but ``None`` is dropped), and a strategy other than
|
||||
``routing_strategy``, the router's, is validated. The loop has no rollback: a setting that
|
||||
raises ends it, and the ones before it are applied."""
|
||||
values: Final[dict[str, object]] = {} # mutable-ok: built here, then frozen into the update
|
||||
plain: Final[dict[str, object]] = {} # mutable-ok: same
|
||||
continuation: object = MISSING # rebind-ok: set when the flag is given
|
||||
routing_groups: object = MISSING # rebind-ok: set when routing groups are given
|
||||
optional_checks: object = MISSING # rebind-ok: set when the optional pre-call checks are given
|
||||
failure: Exception | None = None # rebind-ok: set when a setting raises
|
||||
try:
|
||||
for name, value in settings.items():
|
||||
if name not in RUNTIME_UPDATABLE_ROUTER_SETTINGS:
|
||||
verbose_router_logger.debug("Setting %s is not allowed", name)
|
||||
elif name in _INT_SETTINGS:
|
||||
values[name] = int(value) # pyright: ignore[reportArgumentType] # `int(kwargs[var])`, as given: a bad value raises int's own error
|
||||
elif name == "routing_groups":
|
||||
routing_groups = value
|
||||
elif name == "optional_pre_call_checks":
|
||||
# `set_optional_pre_call_checks` reads a list it is given as a set.
|
||||
_ = frozenset[object](value) if value is not None else None # pyright: ignore[reportArgumentType] # as given: a bad value raises frozenset's own error
|
||||
optional_checks = value
|
||||
elif name == CONTINUATION_SETTING:
|
||||
continuation = value
|
||||
elif name == "retry_policy":
|
||||
policy = RetryPolicy.model_validate(value) if isinstance(value, dict) else value
|
||||
if policy is None or isinstance(policy, RetryPolicy):
|
||||
plain[name] = policy
|
||||
elif name in PLAIN_DEFAULTS or name == "max_retries":
|
||||
plain[name] = value
|
||||
else:
|
||||
if name == "routing_strategy" and _normalized(value) != routing_strategy:
|
||||
validate_routing_strategy(value) # pyright: ignore[reportArgumentType] # validated as given
|
||||
values[name] = value
|
||||
except Exception as raised: # noqa: BLE001 # raised by `apply` once the settings before it are applied
|
||||
failure = raised # rebind-ok: None until a setting fails
|
||||
return SettingsUpdate(
|
||||
values=MappingProxyType(values),
|
||||
plain=MappingProxyType(plain),
|
||||
continuation=continuation,
|
||||
routing_groups=routing_groups,
|
||||
optional_pre_call_checks=optional_checks,
|
||||
failure=failure,
|
||||
)
|
||||
|
||||
|
||||
def _normalized(strategy: object) -> object:
|
||||
"""``_normalize_strategy``, which leaves anything but the enum as given."""
|
||||
return strategy.value if isinstance(strategy, RoutingStrategy) else strategy
|
||||
|
||||
|
||||
def initial_plain(config: "RouterConfig") -> dict[str, object]: # mutable-ok: the core's own plain attributes
|
||||
"""The plain attributes as the constructor sets them: a ``retry_policy`` dict is read into a
|
||||
``RetryPolicy`` and any other type but a policy is no policy."""
|
||||
from litellm.router_rust.surface.retries import resolved_retry_policy
|
||||
|
||||
plain: Final = {name: config.get(name, default) for name, default in PLAIN_DEFAULTS.items()}
|
||||
plain["retry_policy"] = resolved_retry_policy(plain["retry_policy"])
|
||||
return plain
|
||||
|
||||
|
||||
def settings_view(
|
||||
values: Mapping[str, object], plain: Mapping[str, object], deleted: frozenset[str]
|
||||
) -> dict[str, object]: # mutable-ok: the settings view is built fresh per call, as `get_settings` builds it
|
||||
"""``get_settings``: the listed settings the router has, then the routing groups. ``values``
|
||||
holds the settings fields and plain attributes the Rust core resolves."""
|
||||
view: Final[dict[str, object]] = {} # mutable-ok: see above
|
||||
for name in _VIEW_ORDER:
|
||||
if name in deleted:
|
||||
continue
|
||||
if name in plain:
|
||||
view[name] = plain[name]
|
||||
elif name in values:
|
||||
view[name] = values[name]
|
||||
view["routing_groups"] = values.get("routing_groups", [])
|
||||
return view
|
||||
|
||||
|
||||
def argument_attribute(config: "RouterConfig", name: str) -> object:
|
||||
"""A router attribute that is its constructor argument (or a constant), else ``MISSING``."""
|
||||
if name in _ARGUMENT_ATTRIBUTES:
|
||||
return config.get(name, _ARGUMENT_ATTRIBUTES[name])
|
||||
if name in _CONSTANT_ATTRIBUTES:
|
||||
return _CONSTANT_ATTRIBUTES[name]
|
||||
return MISSING
|
||||
|
||||
|
||||
def default_litellm_params(config: "RouterConfig") -> dict[str, object]: # mutable-ok: the router's own dict
|
||||
"""``default_litellm_params`` as the constructor completes it: the router timeout, no client
|
||||
retries and the caching groups."""
|
||||
given: Final = _as_dict(config.get("default_litellm_params")) or {}
|
||||
params: Final[dict[str, object]] = dict(given) # mutable-ok: see above
|
||||
params.setdefault("timeout", config.get("timeout"))
|
||||
params.setdefault("max_retries", 0)
|
||||
metadata: Final = _as_dict(params.setdefault("metadata", {}))
|
||||
if metadata is not None:
|
||||
params["metadata"] = {**metadata, "caching_groups": config.get("caching_groups")}
|
||||
return params
|
||||
|
||||
|
||||
def _as_dict(value: object) -> Mapping[str, object] | None:
|
||||
return value if isinstance(value, dict) else None # pyright: ignore[reportUnknownVariableType] # the argument dict, as passed
|
||||
|
||||
|
||||
def request_timeout(config: "RouterConfig") -> float | None:
|
||||
"""The per-attempt ``request_timeout``: only stored when a router timeout is also set."""
|
||||
from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout
|
||||
|
||||
return get_configured_request_timeout() if config.get("timeout") is not None else None
|
||||
|
||||
|
||||
def apply_construction_effects(config: "RouterConfig", host: "PythonHost") -> None:
|
||||
"""What the constructor's logging arguments do: the host stops its debug info, and
|
||||
``set_verbose`` sets the router logger to ``debug_level``."""
|
||||
host.suppress_debug_info()
|
||||
if config.get("set_verbose", False) is True:
|
||||
level: Final = config.get("debug_level", "INFO")
|
||||
if level == "INFO":
|
||||
verbose_router_logger.setLevel(logging.INFO)
|
||||
elif level == "DEBUG":
|
||||
verbose_router_logger.setLevel(logging.DEBUG)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CONTINUATION_SETTING",
|
||||
"MISSING",
|
||||
"PLAIN_DEFAULTS",
|
||||
"SettingsUpdate",
|
||||
"apply_construction_effects",
|
||||
"argument_attribute",
|
||||
"default_litellm_params",
|
||||
"initial_plain",
|
||||
"read_settings_update",
|
||||
"request_timeout",
|
||||
"settings_view",
|
||||
]
|
||||
|
|
@ -125,12 +125,13 @@ _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = (
|
|||
def _first_present(*sources: Mapping[str, Any] | None, key: str) -> int | float | None:
|
||||
"""Return *key* from the first source mapping where it's set, so callers can
|
||||
support a setting living in more than one deployment config location. Sources
|
||||
are checked in order from most to least specific to that setting."""
|
||||
are checked in order from most to least specific to that setting. Only a number
|
||||
counts: a bool or a string is skipped, so the router default applies."""
|
||||
for source in sources:
|
||||
if source is None:
|
||||
continue
|
||||
value = source.get(key)
|
||||
if value is not None:
|
||||
if isinstance(value, int | float) and not isinstance(value, bool):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
|
@ -811,7 +812,7 @@ def _is_allowed_fails_set_on_router(
|
|||
return False
|
||||
|
||||
|
||||
def cast_exception_status_to_int(exception_status: str | int) -> int:
|
||||
def cast_exception_status_to_int(exception_status: object) -> int:
|
||||
if isinstance(exception_status, str):
|
||||
try:
|
||||
exception_status = int(exception_status)
|
||||
|
|
@ -820,6 +821,13 @@ def cast_exception_status_to_int(exception_status: str | int) -> int:
|
|||
"Unable to cast exception status to int %s. Defaulting to status=500.", exception_status
|
||||
)
|
||||
exception_status = 500
|
||||
# A bool or a number no HTTP status can be is read as a 500, as the Rust router core reads it.
|
||||
if (
|
||||
isinstance(exception_status, bool)
|
||||
or not isinstance(exception_status, int)
|
||||
or not 0 <= exception_status <= 65535
|
||||
):
|
||||
return 500
|
||||
return exception_status
|
||||
|
||||
|
||||
|
|
|
|||
90
litellm/router_utils/deployment_numbers.py
Normal file
90
litellm/router_utils/deployment_numbers.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
# pyright: reportUnknownVariableType=false, reportUnknownArgumentType=false
|
||||
"""Numbers a router deployment cannot be served with, refused the same way by both router cores.
|
||||
|
||||
A whole number wider than 64 bits (a `ptu_shares` of `1e20`) does not fit the router's JSON
|
||||
deployment record, and an infinite or NaN cost would read back as null there, so the call
|
||||
would silently be priced from the model cost map. Both are refused when the deployment is
|
||||
built or added, with the error naming the field.
|
||||
"""
|
||||
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
|
||||
# The whole numbers a deployment record holds (serde_json's i64 and u64).
|
||||
SMALLEST_NUMBER: Final = -(2**63)
|
||||
LARGEST_NUMBER: Final = 2**64 - 1
|
||||
# A pricing field: `input_cost_per_token`, `cache_read_input_token_cost`, ...
|
||||
_COST_FIELD: Final = re.compile(r"cost_per|_cost$")
|
||||
|
||||
|
||||
def _reads_as_non_finite(value: object) -> bool:
|
||||
"""Whether a cost, as written, is an infinity or a NaN (``float()`` reads ``"inf"`` and ``"1e400"``)."""
|
||||
if isinstance(value, float):
|
||||
return not math.isfinite(value)
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
try:
|
||||
return not math.isfinite(float(value))
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _is_cost(key: object) -> bool:
|
||||
return isinstance(key, str) and _COST_FIELD.search(key) is not None
|
||||
|
||||
|
||||
def _number_leaves(value: object) -> Iterator[tuple[str, object, bool]]:
|
||||
"""Every value under ``value`` that is no container, in order: its path, and whether it sits in a
|
||||
cost field. Each container is walked once, so an entry that holds itself ends."""
|
||||
pending: Final[list[tuple[object, str, bool]]] = [(value, "", False)] # mutable-ok: the walk's work list
|
||||
walked: Final[set[int]] = set() # mutable-ok: the ids walked so far
|
||||
while pending:
|
||||
current, path, cost = pending.pop()
|
||||
if not isinstance(current, Mapping | list | tuple):
|
||||
yield path, current, cost
|
||||
elif id(current) not in walked:
|
||||
walked.add(id(current))
|
||||
pending.extend(
|
||||
reversed(
|
||||
[(item, f"{path}.{key}".removeprefix("."), cost or _is_cost(key)) for key, item in current.items()]
|
||||
if isinstance(current, Mapping)
|
||||
else [(item, f"{path}[{index}]", cost) for index, item in enumerate(current)]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _refusal(entry: Mapping[str, object], reason: str) -> ValueError:
|
||||
return ValueError(f"Deployment {entry.get('model_name')!r} cannot be served: {reason}")
|
||||
|
||||
|
||||
def refuse_wide_numbers(entry: object) -> None:
|
||||
"""Refuse a deployment holding a whole number wider than 64 bits, by name."""
|
||||
deployment: Final = as_str_mapping(entry)
|
||||
if deployment is None:
|
||||
return
|
||||
for path, value, _ in _number_leaves(deployment):
|
||||
if isinstance(value, int) and not SMALLEST_NUMBER <= value <= LARGEST_NUMBER:
|
||||
raise _refusal(
|
||||
deployment,
|
||||
f"{path} is {value}, outside what the router can hold "
|
||||
f"({SMALLEST_NUMBER} to {LARGEST_NUMBER}). Use a smaller value.",
|
||||
)
|
||||
|
||||
|
||||
def refuse_unservable_numbers(entry: object) -> None:
|
||||
"""``refuse_wide_numbers``, and a deployment with an infinite or NaN cost."""
|
||||
refuse_wide_numbers(entry)
|
||||
deployment: Final = as_str_mapping(entry)
|
||||
if deployment is None:
|
||||
return
|
||||
for path, value, cost in _number_leaves(deployment):
|
||||
if cost and _reads_as_non_finite(value):
|
||||
raise _refusal(
|
||||
deployment,
|
||||
f"{path} is {value!r}, and the router only prices with finite costs. "
|
||||
"Set a finite cost, or remove it to use the model cost map's price.",
|
||||
)
|
||||
|
|
@ -57,10 +57,14 @@ class _EndpointDeclaration(BaseModel):
|
|||
return (value,) if isinstance(value, str) else value
|
||||
|
||||
|
||||
def endpoint_for_call(call_type: object, module: object = None) -> str | None:
|
||||
qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type
|
||||
return CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def endpoint_selection(call_type: object, module: object = None, model: object = None) -> Generator[None]:
|
||||
qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type
|
||||
endpoint: Final = CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None
|
||||
endpoint: Final = endpoint_for_call(call_type, module)
|
||||
token: Final = _REQUEST_SELECTION.set((endpoint, model if isinstance(model, str) else None))
|
||||
try:
|
||||
yield
|
||||
|
|
@ -68,6 +72,10 @@ def endpoint_selection(call_type: object, module: object = None, model: object =
|
|||
_REQUEST_SELECTION.reset(token)
|
||||
|
||||
|
||||
def selected_endpoint() -> str | None:
|
||||
return _REQUEST_SELECTION.get()[0]
|
||||
|
||||
|
||||
def _enforced_endpoints(model_info: object) -> tuple[str, ...] | None:
|
||||
try:
|
||||
declaration: Final = _EndpointDeclaration.model_validate(model_info)
|
||||
|
|
|
|||
|
|
@ -632,6 +632,8 @@ def get_fallback_model_group(
|
|||
## check for specific model group-specific fallbacks
|
||||
for idx, item in enumerate(fallbacks):
|
||||
if isinstance(item, dict):
|
||||
if not item:
|
||||
continue # an empty entry names no group: the next entries still apply
|
||||
fallback_key = next(iter(item))
|
||||
if fallback_key == model_group: # check exact match
|
||||
fallback_model_group = item[model_group]
|
||||
|
|
|
|||
53
litellm/router_utils/global_callback_holders.py
Normal file
53
litellm/router_utils/global_callback_holders.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
"""
|
||||
The routers of this process that hold optional pre-call checks and routing callbacks, whatever core
|
||||
serves them.
|
||||
|
||||
A router that retires hands the checks it vacates in its host's ``callbacks`` list to a surviving
|
||||
router's same-keyed check (``routers_holding_global_callbacks``), and a router that drops a check at
|
||||
runtime keeps it registered while another live router still has one of that type
|
||||
(``live_check_owners``). Both rules look at every router in the process, so the sets live here,
|
||||
outside any one router core.
|
||||
|
||||
A member of ``routers_holding_global_callbacks`` serves ``_global_callback_candidates()``,
|
||||
``_core_adaptive_routers`` and ``_sync_adaptive_router_hooks()``; a member of ``live_check_owners``
|
||||
serves ``optional_callbacks``.
|
||||
"""
|
||||
|
||||
import threading
|
||||
import weakref
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
class GlobalCallbackHolder(Protocol):
|
||||
"""A router whose callbacks a retiring router may hand its vacated checks to."""
|
||||
|
||||
@property
|
||||
def _core_adaptive_routers(self) -> Mapping[str, object]: ...
|
||||
|
||||
def _global_callback_candidates(self) -> "tuple[CustomLogger, ...]": ...
|
||||
|
||||
def _sync_adaptive_router_hooks(self) -> None: ...
|
||||
|
||||
|
||||
class CheckOwner(Protocol):
|
||||
"""A live router with optional pre-call checks."""
|
||||
|
||||
@property
|
||||
def optional_callbacks(self) -> "Sequence[object] | None": ...
|
||||
|
||||
|
||||
routers_holding_global_callbacks: Final["weakref.WeakSet[GlobalCallbackHolder]"] = weakref.WeakSet()
|
||||
routers_holding_global_callbacks_lock: Final = threading.Lock()
|
||||
live_check_owners: Final["weakref.WeakSet[CheckOwner]"] = weakref.WeakSet()
|
||||
|
||||
__all__ = [
|
||||
"CheckOwner",
|
||||
"GlobalCallbackHolder",
|
||||
"live_check_owners",
|
||||
"routers_holding_global_callbacks",
|
||||
"routers_holding_global_callbacks_lock",
|
||||
]
|
||||
|
|
@ -6,6 +6,7 @@ and exposes it for router candidate filtering.
|
|||
"""
|
||||
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -78,7 +79,7 @@ class DeploymentHealthCache:
|
|||
merged: Final = {
|
||||
model_id: state
|
||||
for model_id, state in {**existing, **states}.items()
|
||||
if isinstance(state, dict) and (now - state.get("timestamp", 0)) < expiry_seconds
|
||||
if isinstance(state, dict) and (now - _timestamp_of(state)) < expiry_seconds # pyright: ignore[reportUnknownArgumentType] # a cached record
|
||||
}
|
||||
self.cache.set_cache(
|
||||
key=self.CACHE_KEY,
|
||||
|
|
@ -101,7 +102,7 @@ class DeploymentHealthCache:
|
|||
for model_id, state in raw.items()
|
||||
if isinstance(state, dict)
|
||||
and not state.get("is_healthy", True)
|
||||
and (now - state.get("timestamp", 0)) < self.staleness_threshold
|
||||
and (now - _timestamp_of(state)) < self.staleness_threshold # pyright: ignore[reportUnknownArgumentType] # a cached record
|
||||
}
|
||||
|
||||
@with_service_target(HEALTH_CHECKS_TARGET)
|
||||
|
|
@ -129,3 +130,10 @@ class DeploymentHealthCache:
|
|||
str(e),
|
||||
)
|
||||
return set()
|
||||
|
||||
|
||||
def _timestamp_of(state: Mapping[str, object]) -> float:
|
||||
"""A health record's timestamp; one that is no number reads as 0 (stale), so a bad record
|
||||
drops out alone instead of turning health routing off for every deployment."""
|
||||
timestamp: Final = state.get("timestamp", 0)
|
||||
return float(timestamp) if isinstance(timestamp, int | float) and not isinstance(timestamp, bool) else 0.0
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from asyncio import Future
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Never, final
|
||||
from typing import Generic, Literal, Never, NotRequired, TypedDict, TypeVar, final
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
|
|
@ -218,6 +218,8 @@ __all__ = [
|
|||
"NativeDiagnosticProcessor",
|
||||
"ProcessReservedForForking",
|
||||
"ResponsesWebSocketConnection",
|
||||
"RouterCall",
|
||||
"RouterCoreHandle",
|
||||
"RustBridgeDeclined",
|
||||
"RustUpstreamError",
|
||||
"TokenCounter",
|
||||
|
|
@ -296,3 +298,493 @@ class _SecretManagerRuntime:
|
|||
timeout: float | httpx.Timeout | None = None,
|
||||
primary_secret_name: str | None = None,
|
||||
) -> Future[JsonValue]: ...
|
||||
|
||||
# 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",
|
||||
"BadGatewayError",
|
||||
"GuardrailIntervention",
|
||||
"Declined",
|
||||
"CallerCredentialAuthenticationError",
|
||||
"CallerCredentialRateLimitError",
|
||||
"APIError",
|
||||
]
|
||||
|
||||
# One event a call reported since the last step (router-core `RouterEvent`), carried on the next step.
|
||||
class _Event(TypedDict):
|
||||
kind: str
|
||||
|
||||
# Every step `RouterCall.next` returns carries the events reported since the last one.
|
||||
class _Step(TypedDict):
|
||||
events: list[_Event]
|
||||
|
||||
class _AttemptStep(_Step):
|
||||
kind: Literal["attempt"]
|
||||
deployment_id: str
|
||||
# The deployment as a `model_list` entry, with its id in `model_info`.
|
||||
deployment: dict[str, JsonValue]
|
||||
# The name this attempt serves: the requested model or the fallback target.
|
||||
model: str
|
||||
model_group: str
|
||||
target_order: int | None
|
||||
candidate_ids: list[str]
|
||||
retry: int
|
||||
fallback_depth: int
|
||||
fallback_level: int
|
||||
timeout: float | None
|
||||
# Request params the fallback walk's dict targets merged in.
|
||||
params: dict[str, JsonValue] | None
|
||||
# The walk level whose request metadata the attempt carries.
|
||||
scope: int
|
||||
max_retries: int
|
||||
model_group_size: int
|
||||
|
||||
class _SleepStep(_Step):
|
||||
kind: Literal["sleep"]
|
||||
seconds: float
|
||||
|
||||
class _AllowFallbackStep(_Step):
|
||||
kind: Literal["allow_fallback"]
|
||||
target: str
|
||||
|
||||
class _RefuseFallbacksStep(_Step):
|
||||
kind: Literal["refuse_fallbacks"]
|
||||
# The walk level whose chain the dict targets are.
|
||||
level: int
|
||||
targets: list[dict[str, JsonValue]]
|
||||
|
||||
class _SelectStep(_Step):
|
||||
kind: Literal["select"]
|
||||
stage: Literal[
|
||||
"plugins",
|
||||
"strategy",
|
||||
"tags",
|
||||
"pre_call",
|
||||
"text_only",
|
||||
"callback_filters",
|
||||
"routing_plugins",
|
||||
"provider_request",
|
||||
]
|
||||
strategy: str | None
|
||||
model_group: str
|
||||
# The name the request resolved from, which strategies key their state by.
|
||||
model: str
|
||||
# Whose selector a strategy reads: "default", "group" (with `group`) or "override".
|
||||
scope: str
|
||||
group: str | None
|
||||
# The deployments the name resolved to, before health filters.
|
||||
pool: list[str]
|
||||
candidates: list[str]
|
||||
# `routing_plugins` only: whether the plugins run.
|
||||
run: NotRequired[bool]
|
||||
# `callback_filters` only: the ids the request already failed on.
|
||||
excluded_ids: NotRequired[list[str]]
|
||||
skipped_ids: NotRequired[list[str]]
|
||||
|
||||
# Pre-routing before a pick resolves `model`: answered with the model group to resolve instead
|
||||
# (a string), a custom strategy's deployment document (a mapping), or None.
|
||||
class _PreRouteStep(_Step):
|
||||
kind: Literal["pre_route"]
|
||||
model: str
|
||||
fallback_depth: int
|
||||
retry: int
|
||||
|
||||
# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`,
|
||||
# and `increment_usage`, which set `keys`; `value` is the JSON to store (`set`), the amount
|
||||
# (`increment`) or one amount per key (`increment_usage`); `ttl` is in seconds; `local` reads or
|
||||
# writes this worker's in-memory tier only. 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, delete and
|
||||
# increment_usage -> None.
|
||||
class _CacheStep(_Step):
|
||||
kind: Literal["cache"]
|
||||
op: Literal["get", "set", "increment", "increment_usage", "batch_get", "delete"]
|
||||
key: str | None
|
||||
keys: list[str] | None
|
||||
value: JsonValue
|
||||
ttl: float | None
|
||||
local: bool
|
||||
|
||||
# One lookup of a model-info query (router-core `ModelInfoLookup`): its `lookup` name and that
|
||||
# lookup's arguments, answered as the lookup documents.
|
||||
class _ModelInfoStep(_Step):
|
||||
kind: Literal["model_info"]
|
||||
lookup: Literal[
|
||||
"model_info",
|
||||
"cost_entry",
|
||||
"provider",
|
||||
"pattern_models",
|
||||
"bedrock_base",
|
||||
"discovered",
|
||||
"user_model_info",
|
||||
"log",
|
||||
"group_params",
|
||||
"group_deployment",
|
||||
"group_capabilities",
|
||||
"group_seed",
|
||||
"reasoning_intersect",
|
||||
]
|
||||
model: NotRequired[JsonValue]
|
||||
tolerate: NotRequired[bool]
|
||||
exact: NotRequired[bool]
|
||||
key: NotRequired[str]
|
||||
process: NotRequired[bool]
|
||||
name: NotRequired[str]
|
||||
model_id: NotRequired[str]
|
||||
id: NotRequired[JsonValue]
|
||||
level: NotRequired[Literal["debug", "error"]]
|
||||
message: NotRequired[str]
|
||||
index: NotRequired[int]
|
||||
model_info: NotRequired[JsonValue]
|
||||
litellm_model: NotRequired[JsonValue]
|
||||
llm_provider: NotRequired[JsonValue]
|
||||
mapped: NotRequired[bool]
|
||||
info: NotRequired[dict[str, JsonValue]]
|
||||
current: NotRequired[JsonValue]
|
||||
resolved: NotRequired[JsonValue]
|
||||
|
||||
class _FailedStep(_Step):
|
||||
kind: Literal["failed"]
|
||||
error: str
|
||||
error_kind: Literal[
|
||||
"config",
|
||||
"unknown_model",
|
||||
"no_healthy_deployments",
|
||||
"attempt",
|
||||
"cache",
|
||||
"retired",
|
||||
"paused",
|
||||
"not_pass_through",
|
||||
"no_pass_through_deployments",
|
||||
"only_markers",
|
||||
"ptu_reserved",
|
||||
"context_window_exceeded",
|
||||
"unsupported_params",
|
||||
"routing_plugins_left_nothing",
|
||||
"routing_plugins_need_async",
|
||||
"multiple_teams",
|
||||
"reserved",
|
||||
"invalid_router_weights",
|
||||
"non_numeric_weight",
|
||||
"missing_routing_group",
|
||||
"default_deployment_raises",
|
||||
"index_out_of_range",
|
||||
"model_info",
|
||||
]
|
||||
attempt_failed: bool
|
||||
status: int | None
|
||||
failure_kind: _FailureKind | None
|
||||
error_model: NotRequired[str]
|
||||
error_detail: NotRequired[str]
|
||||
error_params: NotRequired[list[str]]
|
||||
# `model_info` only: the builtin exception the lookup raised.
|
||||
error_exception: NotRequired[str]
|
||||
|
||||
# The end of a `RouterCoreHandle.start` (and the stream recovery calls).
|
||||
class _DoneStep(_Step):
|
||||
kind: Literal["done"]
|
||||
deployment_id: str
|
||||
model_group: str
|
||||
attempted_retries: int
|
||||
max_retries: int
|
||||
attempted_fallbacks: int
|
||||
total_tokens: int
|
||||
|
||||
# The end of a `RouterCoreHandle.select`.
|
||||
class _SelectedStep(_Step):
|
||||
kind: Literal["done"]
|
||||
deployment_id: str
|
||||
model_group: str
|
||||
deployment: dict[str, JsonValue]
|
||||
|
||||
# The end of a `RouterCoreHandle.healthy`.
|
||||
class _HealthyStep(_Step):
|
||||
kind: Literal["done"]
|
||||
# The name was a deployment id.
|
||||
single: bool
|
||||
deployment_ids: list[str]
|
||||
deployments: list[dict[str, JsonValue]]
|
||||
|
||||
# The end of a `RouterCoreHandle.report_failure`.
|
||||
class _ReportedStep(_Step):
|
||||
kind: Literal["done"]
|
||||
cooled: bool
|
||||
# The cooldown applied, in seconds; None when the deployment was not cooled down.
|
||||
cooldown_time: float | None
|
||||
|
||||
# The end of a `RouterCoreHandle.report_success` or `count_success`.
|
||||
class _CountedStep(_Step):
|
||||
kind: Literal["done"]
|
||||
key: str
|
||||
|
||||
# One `os.environ/` reference a `secrets` step asks for: a top-level param (`key`), or an
|
||||
# env-only `dataSources[source].parameters[key]` (`env_only`).
|
||||
class _SecretRef(TypedDict):
|
||||
name: str
|
||||
key: str
|
||||
env_only: bool
|
||||
source: NotRequired[int]
|
||||
|
||||
# The deployment build's `os.environ/` references: answered with one value per ref, in order.
|
||||
class _SecretsStep(_Step):
|
||||
kind: Literal["secrets"]
|
||||
index: int
|
||||
refs: list[_SecretRef]
|
||||
|
||||
# The deployment build's steps for one deployment. `validate` and `added` are answered with
|
||||
# `{"entry": <model_list entry>}`, `{"skip": True}` or `{"invalid": True}`; `provider` with
|
||||
# `{"provider": str | None}` or `{"invalid": True}`.
|
||||
class _BuildStep(_Step):
|
||||
kind: Literal["validate", "provider", "added"]
|
||||
path: Literal["create", "add"]
|
||||
index: int
|
||||
entry: dict[str, JsonValue]
|
||||
declared_id: str | None
|
||||
duplicate_id: bool
|
||||
provider: str | None
|
||||
|
||||
class _BuiltDeployment(TypedDict):
|
||||
index: int
|
||||
id: str | None
|
||||
provider: str | None
|
||||
|
||||
# The end of a `RouterCoreHandle.build`.
|
||||
class _BuiltStep(_Step):
|
||||
kind: Literal["done"]
|
||||
built: list[_BuiltDeployment]
|
||||
|
||||
# The end of a model-info query: `router_model_info` sets `info` (the merged map or None) and
|
||||
# `user_keys`; `deployment_model_info` sets `info`; the others set `value`.
|
||||
class _ModelInfoDoneStep(_Step):
|
||||
kind: Literal["done"]
|
||||
info: NotRequired[dict[str, JsonValue] | None]
|
||||
user_keys: NotRequired[list[str]]
|
||||
value: NotRequired[JsonValue]
|
||||
|
||||
# `RouterCoreHandle.resolve`: what a name resolves to.
|
||||
class _Resolution(TypedDict):
|
||||
model_group: str
|
||||
kind: str
|
||||
ids: list[str]
|
||||
|
||||
# The steps each kind of `RouterCall` returns, by the handle method that starts it.
|
||||
_CallSteps = (
|
||||
_AttemptStep
|
||||
| _SleepStep
|
||||
| _AllowFallbackStep
|
||||
| _RefuseFallbacksStep
|
||||
| _SelectStep
|
||||
| _PreRouteStep
|
||||
| _CacheStep
|
||||
| _DoneStep
|
||||
| _FailedStep
|
||||
)
|
||||
_SelectSteps = _SelectStep | _PreRouteStep | _CacheStep | _SelectedStep | _FailedStep
|
||||
_HealthySteps = _SelectStep | _PreRouteStep | _CacheStep | _HealthyStep | _FailedStep
|
||||
_BuildSteps = _SecretsStep | _BuildStep | _BuiltStep | _FailedStep
|
||||
_ModelInfoSteps = _ModelInfoStep | _ModelInfoDoneStep | _FailedStep
|
||||
_FailureReportSteps = _CacheStep | _ReportedStep | _FailedStep
|
||||
_SuccessReportSteps = _CacheStep | _CountedStep | _FailedStep
|
||||
|
||||
class _AttemptSucceeded(TypedDict):
|
||||
ok: Literal[True]
|
||||
total_tokens: NotRequired[int]
|
||||
blocked: NotRequired[Literal["content_filter", "refusal"] | None]
|
||||
# A stream whose success or failure the host reports when it ends (`report_success`,
|
||||
# `report_failure`): the attempt counts its request and `total_tokens` only.
|
||||
deferred: NotRequired[bool]
|
||||
# Whether the attempt counts against its deployment's usage; unset, its own limits decide.
|
||||
tracks_usage: NotRequired[bool | None]
|
||||
|
||||
class _AttemptFailed(TypedDict):
|
||||
ok: Literal[False]
|
||||
kind: _FailureKind
|
||||
status: NotRequired[int | None]
|
||||
retry_after: NotRequired[float | None]
|
||||
num_retries: NotRequired[int | None]
|
||||
# The deployment the exception names, which a retry skips over the attempt's own.
|
||||
failed_deployment_id: NotRequired[str | None]
|
||||
detail: NotRequired[str | None]
|
||||
agentic_followup: NotRequired[bool]
|
||||
# The exception has no `status_code`, which no cooldown rule reads.
|
||||
status_absent: NotRequired[bool]
|
||||
# The failure counts against no per-minute requests.
|
||||
uncounted_request: NotRequired[bool]
|
||||
|
||||
# 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[_CallSteps]: ...
|
||||
# Driven like `start`; its `done` step is a `_SelectedStep`.
|
||||
def select(self, model: str, request: _Document) -> RouterCall[_SelectSteps]: ...
|
||||
# Driven like `select`; its `done` step is a `_HealthyStep`.
|
||||
def healthy(self, model: str, request: _Document) -> RouterCall[_HealthySteps]: ...
|
||||
# `{"replace" | "append" | "upsert": [deployments], "delete": [ids], "settings": {dynamic document}, "index": {...}}`.
|
||||
def apply(self, changes: _Document) -> int: ...
|
||||
# Build model_list entries through the host, driven like `start`: `_SecretsStep` and
|
||||
# `_BuildStep` steps, then a `_BuiltStep`. `path`: "create" (`set_model_list`) or "add";
|
||||
# `install`: "replace", "append" or "upsert".
|
||||
def build(
|
||||
self, entries: Sequence[object], path: str, resolve_secrets: bool, install: str
|
||||
) -> RouterCall[_BuildSteps]: ...
|
||||
# Keep the providers `other`'s builds resolved for the same deployments.
|
||||
def adopt_providers(self, other: RouterCoreHandle) -> None: ...
|
||||
# `get_router_model_info`, driven like `start`: `model_info` lookup steps (a `lookup` field
|
||||
# and its arguments), then a done step with `info` (the merged map or None) and `user_keys`.
|
||||
def router_model_info(self, query: _Document) -> RouterCall[_ModelInfoSteps]: ...
|
||||
# `get_deployment_model_info`, driven like `router_model_info`; its done step has `info`.
|
||||
def deployment_model_info(self, model_id: str, model_name: object) -> RouterCall[_ModelInfoSteps]: ...
|
||||
# `_set_model_group_info` for a `GroupInfoQuery` document, driven like `router_model_info`; its
|
||||
# done step has `value`: None, or `{"assignments": {field: value}}`.
|
||||
def model_group_info(self, query: _Document) -> RouterCall[_ModelInfoSteps]: ...
|
||||
# `get_configured_token_limits`, driven like `router_model_info`; `value` is `[input, output]`.
|
||||
def configured_token_limits(self, id: object, configured: _Document) -> RouterCall[_ModelInfoSteps]: ...
|
||||
# `get_model_listing_info`, driven like `router_model_info`; `value` has `cost_map_keys`,
|
||||
# `max_input_tokens` and `max_output_tokens`.
|
||||
def model_listing_info(self, deployments: list[object]) -> RouterCall[_ModelInfoSteps]: ...
|
||||
# A configured string that is not blank (`get_configured_mode`, `get_configured_display_name`).
|
||||
@staticmethod
|
||||
def configured_text(value: object) -> str | None: ...
|
||||
# Each model_info's `service_tiers`, unvalidated.
|
||||
@staticmethod
|
||||
def configured_service_tiers(model_infos: list[object]) -> list[object]: ...
|
||||
# The routing-cache key a deployment counts `metric` (tpm, rpm, itpm, otpm) under in `minute`.
|
||||
def usage_key(self, id: str, model: str, metric: str, minute: str) -> str: ...
|
||||
def registry_version(self) -> int: ...
|
||||
# The version the deployments before the newest appends last changed at.
|
||||
def registry_rewritten(self) -> int: ...
|
||||
# A failure the router did not see itself, driven like `start`: cache steps, then a
|
||||
# `_ReportedStep`.
|
||||
def report_failure(
|
||||
self,
|
||||
id: str,
|
||||
failure: Mapping[str, object],
|
||||
deployment_cooldown: bool = False,
|
||||
count_only: bool = False,
|
||||
requested_group: str | None = None,
|
||||
cooldown_time: float | None = None,
|
||||
uncounted: bool = False,
|
||||
counts_request: bool = False,
|
||||
usage_model: str | None = None,
|
||||
callbacks: Mapping[str, bool] | None = None,
|
||||
) -> RouterCall[_FailureReportSteps]: ...
|
||||
# A success a logging callback saw, driven like `report_failure`; `done` carries `key`.
|
||||
def report_success(
|
||||
self,
|
||||
id: str,
|
||||
total_tokens: float,
|
||||
requests: float,
|
||||
usage_model: str | None = None,
|
||||
tracks_usage: bool | None = None,
|
||||
callbacks: Mapping[str, bool] | None = None,
|
||||
) -> RouterCall[_SuccessReportSteps]: ...
|
||||
# Count a success against any id, driven like `report_failure`; `done` carries `key`.
|
||||
def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ...
|
||||
# [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment.
|
||||
def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ...
|
||||
# Every deployment as a `model_list` entry, in config order.
|
||||
def deployments(self) -> list[dict[str, JsonValue]]: ...
|
||||
def deployments_from(self, start: int) -> list[dict[str, JsonValue]]: ...
|
||||
def deployment(self, id: str) -> dict[str, JsonValue] | None: ...
|
||||
# Lookups over the registry indexes, answering what each index holds as written.
|
||||
# Registry lookups (router-core `lookup.rs`): rows the host reads its stored entries at
|
||||
# (`{"at": position, "name"?, "group"?, "member"?}`, `{"pattern": entry, "name"}`, or a
|
||||
# `{"raw": value}` / `{"raw_positions": value}` its own read raises at).
|
||||
def lookup_all(
|
||||
self, model_name: str, team_id: str | None = None, rename: str | None = None
|
||||
) -> list[dict[str, JsonValue]]: ...
|
||||
def lookup_across_teams(self, model: str) -> tuple[list[dict[str, JsonValue]], int]: ...
|
||||
def lookup_config(self) -> list[dict[str, JsonValue]]: ...
|
||||
def lookup_model_ids(
|
||||
self, model_name: str | None = None, exclude_team_models: bool = False
|
||||
) -> tuple[list[str], dict[str, JsonValue] | None]: ...
|
||||
def lookup_model_list(
|
||||
self, model_name: str | None, team_id: str | None, scope: Mapping[str, object]
|
||||
) -> list[dict[str, JsonValue]]: ...
|
||||
def lookup_model_names(
|
||||
self, team_id: str | None, scope: Mapping[str, object]
|
||||
) -> tuple[list[str], dict[str, JsonValue] | None]: ...
|
||||
def lookup_access_groups(
|
||||
self, model_name: str | None, access_group: str | None, team_id: str | None, scope: Mapping[str, object]
|
||||
) -> tuple[list[tuple[str, list[str]]], dict[str, JsonValue] | None]: ...
|
||||
def lookup_name_serving(
|
||||
self, model_id: str, scope: Mapping[str, object]
|
||||
) -> tuple[bool, str | None, dict[str, JsonValue] | None]: ...
|
||||
def has_model_name(self, name: str) -> bool: ...
|
||||
def lookup_team_model_has_alternatives(self, id: str) -> tuple[bool, dict[str, JsonValue] | None]: ...
|
||||
def lookup_group_usable_by_team(
|
||||
self, model_group: str, team_id: str | None
|
||||
) -> tuple[int | None, dict[str, JsonValue] | None]: ...
|
||||
def lookup_team_public_name(
|
||||
self, model_id: str, team_id: str
|
||||
) -> tuple[int | None, dict[str, JsonValue] | None]: ...
|
||||
def lookup_routable(self, model_group: str, team_id: str | None) -> list[dict[str, JsonValue]]: ...
|
||||
def lookup_wildcard_route(
|
||||
self, model: str, team_id: str | None, requests: list[str]
|
||||
) -> list[dict[str, JsonValue]]: ...
|
||||
def lookup_candidate_ids(
|
||||
self, resolved: str, team_id: str | None, members: list[str] | None, requests: list[str]
|
||||
) -> tuple[list[str], dict[str, JsonValue] | None]: ...
|
||||
def has_model_id(self, id: str) -> bool: ...
|
||||
def id_position(self, id: str) -> JsonValue | None: ...
|
||||
def id_positions(self) -> list[tuple[str, JsonValue]]: ...
|
||||
def named_positions(self, name: str) -> JsonValue | None: ...
|
||||
def team_positions(self, team: str, public: str) -> JsonValue | None: ...
|
||||
def team_public_names_of(self, team: str) -> list[str]: ...
|
||||
def litellm_model_positions(self, model: str) -> list[int]: ...
|
||||
# The registry indexes the host hands out (`HostIndexes`), as JSON.
|
||||
def host_indexes(self) -> dict[str, JsonValue]: ...
|
||||
def provider_default_ids(self) -> list[str]: ...
|
||||
# `PatternMatchRouter.route` over the wildcard router, or `team`'s.
|
||||
def pattern_route(
|
||||
self, team: str | None, request: str, filtered: list[str] | None = None
|
||||
) -> list[dict[str, JsonValue]] | None: ...
|
||||
# `get_models_by_pattern` over the wildcard router, or `team`'s.
|
||||
def models_by_pattern(self, team: str | None, requests: list[str]) -> list[str] | None: ...
|
||||
# The provider the host resolved when it built deployment `id` (its `provider` build step).
|
||||
def provider(self, id: str) -> str | None: ...
|
||||
def resolve(self, name: str, specific_deployment: bool = False) -> _Resolution | None: ...
|
||||
def routable(self, model: str, request: _Document) -> list[dict[str, JsonValue]]: ...
|
||||
# `_get_routing_context` for a pick of `model`: (strategy, scope, group).
|
||||
def routing_context(self, model: str, routing_strategy: str | None = None) -> tuple[str, str, str | None]: ...
|
||||
def retire(self) -> None: ...
|
||||
# A stream that failed before content: what may still take over it, and the retry budget for
|
||||
# a failure of `kind`: {"can_retry", "can_fall_back", "refusal_fallback": bool,
|
||||
# "retry_budget": int | None, "policy_applies": bool}.
|
||||
def stream_policy(
|
||||
self, request: _Document, origin: _Document, kind: str | None = None
|
||||
) -> dict[str, bool | int | None]: ...
|
||||
# The same-group retry and the fallback walk for such a stream, driven like `start`.
|
||||
def retry_stream(
|
||||
self, op: str, request: _Document, origin: _Document, failure: Mapping[str, object]
|
||||
) -> RouterCall[_CallSteps]: ...
|
||||
def fallback_stream(
|
||||
self, op: str, request: _Document, origin: _Document, failure: Mapping[str, object]
|
||||
) -> RouterCall[_CallSteps]: ...
|
||||
def retry_policy_retries(self, failure: Mapping[str, object], model_group: str | None = None) -> int | None: ...
|
||||
def allowed_fails_for(self, failure: Mapping[str, object]) -> int | None: ...
|
||||
def should_retry(
|
||||
self, failure: Mapping[str, object], counts: tuple[int, int], fallbacks: tuple[bool, bool, bool]
|
||||
) -> bool: ...
|
||||
|
||||
_S = TypeVar("_S", covariant=True)
|
||||
|
||||
@final
|
||||
class RouterCall(Generic[_S]):
|
||||
# The answer to the last step, which Rust validates against it: attempt: _AttemptSucceeded |
|
||||
# _AttemptFailed; sleep: None; allow_fallback: bool; refuse_fallbacks: bool; select: the remaining
|
||||
# candidate ids, pick first; pre_route: see _PreRouteStep; cache: see _CacheStep; secrets, build and model_info:
|
||||
# see those steps; first call: None.
|
||||
def next(self, outcome: object) -> _S: ...
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ class Route(str, Enum):
|
|||
OCR = "ocr"
|
||||
TOKEN_COUNTER = "token_counter"
|
||||
TOKENIZER = "tokenizer"
|
||||
ROUTER = "router"
|
||||
|
||||
|
||||
class Target(str, Enum):
|
||||
|
|
@ -224,6 +225,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"})),
|
||||
)
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue