From df50c7572a7906350c4ff0a410980ecb29a22873 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Mon, 2 Nov 2026 02:52:28 +0000 Subject: [PATCH] feat(router): add the Rust router core with the host as its only boundary 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; the remaining Rust differences (the router's own outcome callbacks) are listed in rust_parity_pending.txt. --- litellm-rust/Cargo.lock | 20 + litellm-rust/Cargo.toml | 3 + litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/lib.rs | 5 + .../crates/python-bridge/src/router/config.rs | 369 ++ .../crates/python-bridge/src/router/host.rs | 1351 +++++++ .../crates/python-bridge/src/router/mod.rs | 1862 +++++++++ .../python-bridge/src/router/retries.rs | 55 + .../crates/python-bridge/src/router/stream.rs | 94 + litellm-rust/crates/router-core/Cargo.toml | 18 + litellm-rust/crates/router-core/src/access.rs | 236 ++ litellm-rust/crates/router-core/src/build.rs | 382 ++ litellm-rust/crates/router-core/src/cache.rs | 145 + litellm-rust/crates/router-core/src/call.rs | 1286 ++++++ .../crates/router-core/src/call/stream.rs | 413 ++ litellm-rust/crates/router-core/src/config.rs | 361 ++ litellm-rust/crates/router-core/src/error.rs | 419 ++ .../crates/router-core/src/fallback.rs | 310 ++ .../crates/router-core/src/group_info.rs | 639 +++ litellm-rust/crates/router-core/src/health.rs | 651 +++ .../crates/router-core/src/health_routing.rs | 126 + litellm-rust/crates/router-core/src/host.rs | 573 +++ .../crates/router-core/src/host_index.rs | 647 +++ litellm-rust/crates/router-core/src/lib.rs | 624 ++- litellm-rust/crates/router-core/src/lookup.rs | 851 ++++ litellm-rust/crates/router-core/src/mock.rs | 44 + .../crates/router-core/src/model_info.rs | 709 ++++ .../crates/router-core/src/precall.rs | 300 ++ .../crates/router-core/src/prerouting.rs | 53 + litellm-rust/crates/router-core/src/ptu.rs | 96 + .../crates/router-core/src/registry.rs | 1358 +++++++ .../crates/router-core/src/resolve.rs | 411 ++ litellm-rust/crates/router-core/src/retry.rs | 216 + .../crates/router-core/src/routing_groups.rs | 48 + .../crates/router-core/src/routing_plugins.rs | 114 + litellm-rust/crates/router-core/src/select.rs | 673 +++ .../crates/router-core/src/settings.rs | 103 + .../crates/router-core/src/strategy.rs | 110 + .../crates/router-core/src/weights.rs | 81 + .../crates/router-core/tests/call_loop.rs | 411 ++ litellm/router_api.py | 100 +- litellm/router_attempt.py | 208 + litellm/router_core.py | 32 +- litellm/router_deprecated.py | 357 +- litellm/router_python_core.py | 14 +- litellm/router_python_host.py | 91 +- litellm/router_rust/__init__.py | 13 + litellm/router_rust/core.py | 3599 +++++++++++++++++ litellm/router_rust/invoke/README.md | 13 + litellm/router_rust/invoke/__init__.py | 1 + litellm/router_rust/invoke/endpoints.py | 2117 ++++++++++ .../router_rust/invoke/prompt_management.py | 131 + litellm/router_rust/invoke/shaping.py | 634 +++ litellm/router_rust/invoke/slots.py | 108 + litellm/router_rust/invoke/streaming.py | 1794 ++++++++ litellm/router_rust/invoke/walk.py | 292 ++ litellm/router_rust/invoke/websocket.py | 131 + litellm/router_rust/steps/README.md | 12 + litellm/router_rust/steps/__init__.py | 1 + litellm/router_rust/steps/access.py | 85 + litellm/router_rust/steps/callbacks.py | 558 +++ litellm/router_rust/steps/plugins.py | 87 + litellm/router_rust/steps/precall.py | 318 ++ litellm/router_rust/steps/prerouting.py | 1097 +++++ litellm/router_rust/steps/selection.py | 405 ++ litellm/router_rust/surface/README.md | 13 + litellm/router_rust/surface/__init__.py | 1 + litellm/router_rust/surface/alerting.py | 33 + litellm/router_rust/surface/fallbacks.py | 132 + litellm/router_rust/surface/health.py | 220 + litellm/router_rust/surface/health_routing.py | 157 + litellm/router_rust/surface/indexes.py | 499 +++ litellm/router_rust/surface/members.py | 95 + litellm/router_rust/surface/model_info.py | 832 ++++ litellm/router_rust/surface/model_list.py | 319 ++ litellm/router_rust/surface/registry.py | 2037 ++++++++++ litellm/router_rust/surface/replay.py | 30 + litellm/router_rust/surface/retries.py | 243 ++ litellm/router_rust/surface/settings.py | 218 + .../router_utils/global_callback_holders.py | 53 + litellm/rust_bridge/_native.pyi | 313 +- litellm/rust_bridge/catalog.py | 2 + litellm/utils.py | 9 + .../router_litellm_globals_check.py | 1 + .../router_surface_check.py | 8 +- tests/unit/test_router/core/conftest.py | 49 +- .../core/construction_precall_helpers.py | 98 + .../test_router/core/endpoints_helpers.py | 210 + .../test_router/core/fallbacks_helpers.py | 105 + .../test_router/core/gaps_l11100_helpers.py | 112 + .../unit/test_router/core/registry_helpers.py | 47 + .../core/retries_health_helpers.py | 198 + .../test_router/core/rust_parity_pending.txt | 4 + tests/unit/test_router/core/scripted_host.py | 91 +- .../test_router/core/selection_helpers.py | 88 + .../test_router/core/streaming_helpers.py | 469 +++ ...t_construction_precall_characterization.py | 1271 ++++++ ...uction_precall_characterization_precall.py | 663 +++ .../core/test_deprecated_members.py | 665 +++ tests/unit/test_router/core/test_endpoints.py | 20 +- .../test_endpoints_characterization_batch.py | 516 +++ ...test_endpoints_characterization_generic.py | 780 ++++ ...st_endpoints_characterization_inference.py | 648 +++ .../core/test_fallbacks_characterization.py | 584 +++ ...est_fallbacks_characterization_outcomes.py | 675 ++++ ...test_fallbacks_characterization_refusal.py | 101 + .../core/test_gaps_adapter_timeout.py | 34 + .../unit/test_router/core/test_gaps_l0001.py | 445 ++ .../core/test_gaps_l11100_registry.py | 657 +++ .../core/test_gaps_l11100_strategy.py | 438 ++ .../unit/test_router/core/test_gaps_l14000.py | 299 ++ .../core/test_gaps_l14000_selection.py | 491 +++ .../unit/test_router/core/test_gaps_l2800.py | 431 ++ .../core/test_gaps_l2800_streaming.py | 177 + .../unit/test_router/core/test_gaps_l5600.py | 479 +++ .../core/test_gaps_l5600_anthropic.py | 103 + .../core/test_gaps_l8400_precall.py | 178 + .../core/test_gaps_l8400_pricing.py | 240 ++ .../core/test_gaps_l8400_retries.py | 138 + .../core/test_gaps_request_overrides.py | 227 ++ .../core/test_gaps_router_core_crate.py | 283 ++ .../core/test_gaps_router_core_crate_picks.py | 372 ++ .../core/test_health_registry_views.py | 174 + ..._host_settings_changes_characterization.py | 428 ++ .../core/test_registry_changesets.py | 5 + ..._registry_characterization_construction.py | 379 ++ .../test_registry_characterization_lookups.py | 682 ++++ ...st_registry_characterization_model_info.py | 540 +++ .../test_registry_characterization_writes.py | 516 +++ ...tries_health_characterization_cooldowns.py | 845 ++++ ...retries_health_characterization_retries.py | 636 +++ tests/unit/test_router/core/test_selection.py | 39 +- .../test_selection_characterization_groups.py | 747 ++++ ...t_selection_characterization_prerouting.py | 505 +++ ...t_selection_characterization_strategies.py | 394 ++ ...test_selection_characterization_traffic.py | 145 + ...st_streaming_characterization_anthropic.py | 485 +++ .../test_streaming_characterization_chat.py | 669 +++ ...est_streaming_characterization_outcomes.py | 382 ++ ...st_streaming_characterization_responses.py | 403 ++ 140 files changed, 53727 insertions(+), 215 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/router/config.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/host.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/retries.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/stream.rs create mode 100644 litellm-rust/crates/router-core/src/access.rs create mode 100644 litellm-rust/crates/router-core/src/build.rs create mode 100644 litellm-rust/crates/router-core/src/cache.rs create mode 100644 litellm-rust/crates/router-core/src/call.rs create mode 100644 litellm-rust/crates/router-core/src/call/stream.rs create mode 100644 litellm-rust/crates/router-core/src/config.rs create mode 100644 litellm-rust/crates/router-core/src/error.rs create mode 100644 litellm-rust/crates/router-core/src/fallback.rs create mode 100644 litellm-rust/crates/router-core/src/group_info.rs create mode 100644 litellm-rust/crates/router-core/src/health.rs create mode 100644 litellm-rust/crates/router-core/src/health_routing.rs create mode 100644 litellm-rust/crates/router-core/src/host.rs create mode 100644 litellm-rust/crates/router-core/src/host_index.rs create mode 100644 litellm-rust/crates/router-core/src/lookup.rs create mode 100644 litellm-rust/crates/router-core/src/mock.rs create mode 100644 litellm-rust/crates/router-core/src/model_info.rs create mode 100644 litellm-rust/crates/router-core/src/precall.rs create mode 100644 litellm-rust/crates/router-core/src/prerouting.rs create mode 100644 litellm-rust/crates/router-core/src/ptu.rs create mode 100644 litellm-rust/crates/router-core/src/registry.rs create mode 100644 litellm-rust/crates/router-core/src/resolve.rs create mode 100644 litellm-rust/crates/router-core/src/retry.rs create mode 100644 litellm-rust/crates/router-core/src/routing_groups.rs create mode 100644 litellm-rust/crates/router-core/src/routing_plugins.rs create mode 100644 litellm-rust/crates/router-core/src/select.rs create mode 100644 litellm-rust/crates/router-core/src/settings.rs create mode 100644 litellm-rust/crates/router-core/src/strategy.rs create mode 100644 litellm-rust/crates/router-core/src/weights.rs create mode 100644 litellm-rust/crates/router-core/tests/call_loop.rs create mode 100644 litellm/router_attempt.py create mode 100644 litellm/router_rust/__init__.py create mode 100644 litellm/router_rust/core.py create mode 100644 litellm/router_rust/invoke/README.md create mode 100644 litellm/router_rust/invoke/__init__.py create mode 100644 litellm/router_rust/invoke/endpoints.py create mode 100644 litellm/router_rust/invoke/prompt_management.py create mode 100644 litellm/router_rust/invoke/shaping.py create mode 100644 litellm/router_rust/invoke/slots.py create mode 100644 litellm/router_rust/invoke/streaming.py create mode 100644 litellm/router_rust/invoke/walk.py create mode 100644 litellm/router_rust/invoke/websocket.py create mode 100644 litellm/router_rust/steps/README.md create mode 100644 litellm/router_rust/steps/__init__.py create mode 100644 litellm/router_rust/steps/access.py create mode 100644 litellm/router_rust/steps/callbacks.py create mode 100644 litellm/router_rust/steps/plugins.py create mode 100644 litellm/router_rust/steps/precall.py create mode 100644 litellm/router_rust/steps/prerouting.py create mode 100644 litellm/router_rust/steps/selection.py create mode 100644 litellm/router_rust/surface/README.md create mode 100644 litellm/router_rust/surface/__init__.py create mode 100644 litellm/router_rust/surface/alerting.py create mode 100644 litellm/router_rust/surface/fallbacks.py create mode 100644 litellm/router_rust/surface/health.py create mode 100644 litellm/router_rust/surface/health_routing.py create mode 100644 litellm/router_rust/surface/indexes.py create mode 100644 litellm/router_rust/surface/members.py create mode 100644 litellm/router_rust/surface/model_info.py create mode 100644 litellm/router_rust/surface/model_list.py create mode 100644 litellm/router_rust/surface/registry.py create mode 100644 litellm/router_rust/surface/replay.py create mode 100644 litellm/router_rust/surface/retries.py create mode 100644 litellm/router_rust/surface/settings.py create mode 100644 litellm/router_utils/global_callback_holders.py create mode 100644 tests/unit/test_router/core/construction_precall_helpers.py create mode 100644 tests/unit/test_router/core/endpoints_helpers.py create mode 100644 tests/unit/test_router/core/fallbacks_helpers.py create mode 100644 tests/unit/test_router/core/gaps_l11100_helpers.py create mode 100644 tests/unit/test_router/core/registry_helpers.py create mode 100644 tests/unit/test_router/core/retries_health_helpers.py create mode 100644 tests/unit/test_router/core/rust_parity_pending.txt create mode 100644 tests/unit/test_router/core/selection_helpers.py create mode 100644 tests/unit/test_router/core/streaming_helpers.py create mode 100644 tests/unit/test_router/core/test_construction_precall_characterization.py create mode 100644 tests/unit/test_router/core/test_construction_precall_characterization_precall.py create mode 100644 tests/unit/test_router/core/test_deprecated_members.py create mode 100644 tests/unit/test_router/core/test_endpoints_characterization_batch.py create mode 100644 tests/unit/test_router/core/test_endpoints_characterization_generic.py create mode 100644 tests/unit/test_router/core/test_endpoints_characterization_inference.py create mode 100644 tests/unit/test_router/core/test_fallbacks_characterization.py create mode 100644 tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py create mode 100644 tests/unit/test_router/core/test_fallbacks_characterization_refusal.py create mode 100644 tests/unit/test_router/core/test_gaps_adapter_timeout.py create mode 100644 tests/unit/test_router/core/test_gaps_l0001.py create mode 100644 tests/unit/test_router/core/test_gaps_l11100_registry.py create mode 100644 tests/unit/test_router/core/test_gaps_l11100_strategy.py create mode 100644 tests/unit/test_router/core/test_gaps_l14000.py create mode 100644 tests/unit/test_router/core/test_gaps_l14000_selection.py create mode 100644 tests/unit/test_router/core/test_gaps_l2800.py create mode 100644 tests/unit/test_router/core/test_gaps_l2800_streaming.py create mode 100644 tests/unit/test_router/core/test_gaps_l5600.py create mode 100644 tests/unit/test_router/core/test_gaps_l5600_anthropic.py create mode 100644 tests/unit/test_router/core/test_gaps_l8400_precall.py create mode 100644 tests/unit/test_router/core/test_gaps_l8400_pricing.py create mode 100644 tests/unit/test_router/core/test_gaps_l8400_retries.py create mode 100644 tests/unit/test_router/core/test_gaps_request_overrides.py create mode 100644 tests/unit/test_router/core/test_gaps_router_core_crate.py create mode 100644 tests/unit/test_router/core/test_gaps_router_core_crate_picks.py create mode 100644 tests/unit/test_router/core/test_health_registry_views.py create mode 100644 tests/unit/test_router/core/test_host_settings_changes_characterization.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_construction.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_lookups.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_model_info.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_writes.py create mode 100644 tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py create mode 100644 tests/unit/test_router/core/test_retries_health_characterization_retries.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_groups.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_prerouting.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_strategies.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_traffic.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_anthropic.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_chat.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_outcomes.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_responses.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f91deaaa401..12ecafa026e 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4331,6 +4331,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "async-trait", "aws-sdk-secretsmanager", "bytes", "criterion", @@ -4357,6 +4358,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-router-core", "litellm-router-types", "litellm-secret-value", "litellm-secrets", @@ -4439,6 +4441,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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 75cac2a5d63..dda216c4dcf 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" } litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } +litellm-router-core = { path = "crates/router-core" } litellm-tracing = { path = "crates/tracing" } litellm-spend-clickhouse = { path = "crates/spend-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } @@ -78,6 +79,8 @@ litellm-cost = { path = "crates/cost" } litellm-model-catalog = { path = "crates/model-catalog" } litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } +arc-swap = "1.9.2" +async-trait = "0.1.91" askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..5a63ecd0004 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -25,6 +25,7 @@ litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true litellm-storage-clickhouse.workspace = true litellm-host.workspace = true +async-trait.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true @@ -49,6 +50,7 @@ litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "has litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true +litellm-router-core.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index fdcb969a423..7ecbc6907ed 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -12,6 +12,7 @@ mod marshal; mod model_registry; mod preflight; mod python_settings; +mod router; mod routes; mod secrets; mod tokenizer; @@ -36,6 +37,8 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::router::{RouterCall, RouterCoreHandle}; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -120,6 +123,8 @@ mod tests { "bundled_model_cost_map", "TokenCounter", "Tokenizer", + "RouterCoreHandle", + "RouterCall", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs new file mode 100644 index 00000000000..43787d1e1af --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -0,0 +1,369 @@ +//! 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, RouterConfig, RouterDynamicConfig, RouterRequest, + config::{Fallbacks, RetryPolicy}, + 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", + "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, Error> { + match value { + Value::Object(object) => Ok(object), + _ => Err(Error::NotObject { document }), + } +} + +fn shaped(document: Document, object: Map) -> Result { + serde_json::from_value(Value::Object(object)) + .map_err(|source| Error::Shape { document, source }) +} + +/// Drop unset and ignored keys, then reject any key outside `known`. +fn routing_keys( + document: Document, + value: Value, + known: &[&str], +) -> Result, Error> { + let mut object = object(document, value)?; + object.retain(|key, value| { + let empty_unset = EMPTY_MEANS_UNSET.contains(&key.as_str()) + && value.as_array().is_some_and(Vec::is_empty); + !(value.is_null() || empty_unset || IGNORED_KEYS.contains(&key.as_str())) + }); + if let Some(key) = object.keys().find(|key| !known.contains(&key.as_str())) { + return Err(Error::UnknownKey { + document, + key: key.clone(), + }); + } + Ok(object) +} + +pub(crate) fn router_config(value: Value) -> Result { + let object = routing_keys(Document::Static, value, STATIC_KEYS)?; + shaped(Document::Static, object) +} + +pub(crate) fn dynamic_config(value: Value) -> Result { + let object = routing_keys(Document::Dynamic, value, DYNAMIC_KEYS)?; + shaped(Document::Dynamic, object) +} + +/// A call's request as Python spells it. +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestSpec { + model: Option, + team_id: Option, + proxy_admin: Option, + allowed_models: Option>, + admin_model: Option, + specific_deployment: Option, + num_retries: Option, + disable_fallbacks: Option, + router_weights: Option, + pinned_deployment_id: Option, + routing_strategy: Option, + filter_tags: Option, + /// `listed` or `healthy`: which pass-through pick the request makes. + pass_through: Option, + target_order: Option, + excluded_ids: Option>, + retry_skipped_ids: Option>, + pre_routes: Option, + has_image: Option, + unrouted: Option, + model_group_retry_policy: Option>, + retry_policy_raises: Option>, + #[serde(default, deserialize_with = "present")] + fallbacks: Option>, + #[serde(default, deserialize_with = "present")] + context_window_fallbacks: Option>, + #[serde(default, deserialize_with = "present")] + content_policy_fallbacks: Option>, + max_fallbacks: Option, + mock_failure: Option, + cooldown_time: Option, + health_check_probe: Option, + body: Option, +} + +/// A key the request passed, even as `None`: `Some(None)`. +fn present<'de, D, T>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + T: Deserialize<'de>, +{ + Option::::deserialize(deserializer).map(Some) +} + +/// Read a request; `model` replaces the request's own when given. +pub(crate) fn request(value: Value, model: Option<&str>) -> Result { + let spec: RequestSpec = shaped(Document::Request, object(Document::Request, value)?)?; + let model = model + .map(str::to_owned) + .or(spec.model) + .ok_or(Error::MissingModel)?; + Ok(RouterRequest { + model, + body: spec.body.unwrap_or_default(), + team_id: spec.team_id, + 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, + 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), + unrouted: spec.unrouted.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), + 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>, + /// Deployments added at the end, beside any that share their id. + #[serde(default)] + append: Vec, + #[serde(default)] + upsert: Vec, + #[serde(default)] + delete: Vec, + /// A new dynamic document (the same keys `RouterCoreHandle(static, dynamic)` takes). + #[serde(default)] + settings: Option, + /// The registry indexes the host wrote (`IndexOverrides`); `{}` drops them. + #[serde(default)] + index: Option, + /// The deployment at the id map's position for this id, removed with every index shifted. + #[serde(default)] + remove: Option, + /// A write to the indexes the host hands out (`HostIndexes`), field by field. + #[serde(default)] + host_index: Option, +} + +pub(crate) fn changes(value: Value) -> Result { + 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(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 = 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) { + 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!(router_config(json).unwrap_err().declined(), declined); + } + + #[rstest] + fn python_router_defaults_read_as_unset() { + let dynamic = dynamic_config(json!({ + "num_retries": null, "fallbacks": [], "context_window_fallbacks": [], + "content_policy_fallbacks": [], "model_group_alias": {"gpt": "gpt-4"}, + })) + .unwrap(); + assert_eq!(dynamic.num_retries, None); + assert!(dynamic.context_window_fallbacks.is_none()); + assert!(dynamic.content_policy_fallbacks.is_none()); + assert_eq!(dynamic.model_alias_map["gpt"], "gpt-4"); + } + + #[rstest] + #[case::own_model(json!({"model": "gpt", "num_retries": 1}), None, "gpt")] + #[case::override_model(json!({"model": "gpt"}), Some("other"), "other")] + fn a_request_reads_its_model( + #[case] json: Value, + #[case] model: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(request(json, model).unwrap().model, expected); + } + + #[rstest] + #[case::unknown_field(json!({"model": "gpt", "temperature": 1}))] + #[case::no_model(json!({"team_id": "t"}))] + fn a_bad_request_is_a_value_error(#[case] json: Value) { + assert!(!request(json, None).unwrap_err().declined()); + } + + #[rstest] + fn a_change_set_reads_upserts_and_deletes() { + let changes = changes(json!({ + "upsert": [{"model_name": "m", "litellm_params": {"model": "openai/x"}}], + "delete": ["old"], + })) + .unwrap(); + assert_eq!(changes.upserts.len(), 1); + assert_eq!(changes.deletes, ["old"]); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs new file mode 100644 index 00000000000..3d611643f2b --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -0,0 +1,1351 @@ +//! The host behind `RouterCoreHandle`. Rust never awaits Python: each host operation the call +//! loop needs becomes an [`Ask`] parked in the call's [`Exchange`], and the call future stays +//! pending until Python runs it and hands the [`Answer`] back through [`Driver::resume`]. +//! +//! The driver polls the call future with a no-op waker while the call's exchange is bound to the +//! current thread, so a host future finds its exchange without the router passing it through. +//! +//! The routing cache is a host operation too: [`StepCache`] parks every [`RouterCache`] op as a +//! `cache` step, and Python runs it against the router's own `DualCache` (memory plus Redis), so +//! Python and Rust read and write the same cooldowns, fail counters and usage. + +use std::{ + cell::RefCell, + future::{Future, poll_fn}, + pin::Pin, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + task::{Context, Poll, Waker}, + time::Duration, +}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output, + Payload, PreRoute, PreRouted, Provided, RouterCache, RouterEvent, RouterRequest, + RouterResponse, RoutingContext, SecretRef, SelectStage, build::BuiltDeployment, + host::ModelInfoLookup, model_info::MergedModelInfo, precall::PreCallFacts, +}; +use serde_json::{Map, Value}; + +/// The `kind` of a step dict. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum StepKind { + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "sleep")] + Sleep, + #[strum(serialize = "allow_fallback")] + AllowFallback, + #[strum(serialize = "select")] + Select, + #[strum(serialize = "pre_route")] + PreRoute, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "secrets")] + Secrets, + #[strum(serialize = "validate")] + Validate, + #[strum(serialize = "provider")] + Provider, + #[strum(serialize = "added")] + Added, + #[strum(serialize = "model_info")] + ModelInfo, + #[strum(serialize = "done")] + Done, + #[strum(serialize = "failed")] + Failed, +} + +/// Which selection stage Python is asked to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum Stage { + /// Pre-call checks and sub-routers: return the deployments that remain. + Plugins, + /// A strategy router-core does not implement: return the candidates with the pick first. + Strategy(String), + /// Tag routing: return the deployments that remain. + Tags, + /// `enable_pre_call_checks`: return the facts the filter decides by. + PreCall, + /// `_deployment_declares_no_vision`: return the candidates that declare no vision support. + TextOnly, + /// The routing plugins: run them when `run`, then return the candidate models they left. + RoutingPlugins { run: bool }, + /// The host's `CustomLogger` deployment filters: return the ids that remain, or a failure. + CallbackFilters, +} + +/// The `op` of a `cache` step: one per [`RouterCache`] method. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum CacheOp { + /// `async_get_cache(key)`: answered with the stored JSON value, or `None`. + #[strum(serialize = "get")] + Get, + /// `async_set_cache(key, value, ttl)`: answered with `None`. + #[strum(serialize = "set")] + Set, + /// `async_increment(key, value, ttl)`: answered with the new number. The TTL applies when + /// the key is created. + #[strum(serialize = "increment")] + Increment, + /// `async_batch_get_cache(keys)`: answered with one JSON value or `None` per key, in order. + #[strum(serialize = "batch_get")] + BatchGet, + /// `async_delete_cache(key)`: answered with `None`. + #[strum(serialize = "delete")] + Delete, +} + +/// Python's outcome of a `cache` step. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum CacheOutcome { + /// `get`: the stored value (`None` when missing). + Value(Option), + /// `batch_get`: one entry per key. + Values(Vec>), + /// `increment`: the counter after the increment. + Number(f64), + /// `set` and `delete`. + Stored, +} + +impl CacheOutcome { + /// The op this outcome answers (`Stored` answers both `set` and `delete`). + fn answers(&self, op: CacheOp) -> bool { + matches!( + (self, op), + (Self::Value(_), CacheOp::Get) + | (Self::Values(_), CacheOp::BatchGet) + | (Self::Number(_), CacheOp::Increment) + | (Self::Stored, CacheOp::Set | CacheOp::Delete) + ) + } +} + +/// What a call is blocked on, for Python to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum Ask { + Attempt { + deployment_id: String, + /// The picked deployment as configured: `model_name`, `litellm_params`, `model_info` + /// (with its `id`). The host builds the request from it; it never looks the deployment up + /// in a Python router. + deployment: Value, + /// The name the attempt resolved: the requested name or the fallback target. + model: String, + model_group: String, + /// The `order` level of a same-group order fallback. + target_order: Option, + /// Every deployment `model` resolved to for this caller, the picked one included. + candidate_ids: Vec, + retry: u32, + fallback_depth: u32, + fallback_level: u32, + timeout: Option, + /// Request params the fallback walk's dict targets merged in. + params: Option>, + /// The walk level whose request metadata the attempt carries. + scope: u32, + max_retries: u32, + model_group_size: usize, + }, + Sleep { + seconds: f64, + }, + AllowFallback { + target: String, + }, + Select { + stage: Stage, + model_group: String, + /// The name the request resolved from, which strategies key their state by. + name: String, + /// Whose selector a strategy reads: `default`, `group` (with `group`) or `override`. + scope: &'static str, + group: Option, + /// The deployments the name resolved to, before health filters (the tag stage reads them). + pool: Vec, + candidates: Vec, + }, + /// 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`. + PreRoute { + model: String, + fallback_depth: u32, + retry: u32, + }, + /// One [`RouterCache`] op against the host's cache. `key` is set for every op but + /// `batch_get`, which sets `keys`; `value` is set for `set` (the JSON to store) and + /// `increment` (the amount); `ttl` is in seconds. + Cache { + op: CacheOp, + key: Option, + keys: Option>, + value: Option, + ttl: Option, + /// The op reads or writes this worker's in-memory tier only (`local_only=True`). + local: bool, + }, + /// The deployment build's `os.environ/` references: answered with one value per ref. + Secrets { + index: usize, + refs: Vec, + }, + /// The deployment build's `validate`, `provider` and `added` steps for one deployment. + Build { + kind: StepKind, + path: BuildPath, + index: usize, + entry: Value, + declared_id: Option, + duplicate_id: bool, + /// The provider the `provider` step resolved (`added` only). + provider: Option, + }, + /// One lookup of a model-info query ([`ModelInfoLookup`], as its `lookup` field and + /// arguments): answered as the lookup documents. + ModelInfo { + lookup: Value, + }, +} + +impl Ask { + pub(crate) fn kind(&self) -> StepKind { + match self { + Self::Attempt { .. } => StepKind::Attempt, + Self::Sleep { .. } => StepKind::Sleep, + Self::AllowFallback { .. } => StepKind::AllowFallback, + Self::Select { .. } => StepKind::Select, + Self::PreRoute { .. } => StepKind::PreRoute, + Self::Cache { .. } => StepKind::Cache, + Self::Secrets { .. } => StepKind::Secrets, + Self::Build { kind, .. } => *kind, + Self::ModelInfo { .. } => StepKind::ModelInfo, + } + } +} + +/// Python's answer to an [`Ask`]. +#[derive(Debug)] +pub(crate) enum Answer { + /// The attempt's total tokens, or its failure. + Attempt(Result), + /// A successful attempt whose response the provider blocked, or a stream whose outcome the + /// host reports when it ends (`deferred`), with its tokens. + Served { + total_tokens: u64, + blocked: Option, + deferred: bool, + }, + Slept, + AllowFallback(bool), + /// The candidate ids that remain, pick first. + Select(Vec), + /// The model group to resolve instead, a custom strategy's deployment, or `None`. + PreRoute(Option), + /// The facts a `pre_call` select step asked for. + PreCall(PreCallFacts), + /// The candidate models a `routing_plugins` select step reported. + Models(Option>), + /// The ids a `callback_filters` select step left, or the error a filter raised. + Filtered(Result, AttemptError>), + Cache(CacheOutcome), + /// One value per ref of a `secrets` step. + Secrets(Vec), + Validated(Built), + Provided(Provided), + Added(Built), + /// A `model_info` lookup's answer, as the host spelled it. + ModelInfo(Value), +} + +impl Answer { + fn kind(&self) -> StepKind { + match self { + Self::Attempt(_) | Self::Served { .. } => StepKind::Attempt, + Self::Slept => StepKind::Sleep, + Self::AllowFallback(_) => StepKind::AllowFallback, + Self::Select(_) | Self::PreCall(_) | Self::Models(_) | Self::Filtered(_) => { + StepKind::Select + } + Self::PreRoute(_) => StepKind::PreRoute, + Self::Cache(_) => StepKind::Cache, + Self::Secrets(_) => StepKind::Secrets, + Self::Validated(_) => StepKind::Validate, + Self::Provided(_) => StepKind::Provider, + Self::Added(_) => StepKind::Added, + Self::ModelInfo(_) => StepKind::ModelInfo, + } + } + + fn check(&self, ask: &Ask) -> Result<(), DriveError> { + if self.kind() != ask.kind() { + return Err(DriveError::WrongAnswer { + expected: ask.kind(), + got: self.kind(), + }); + } + if let (Self::Select(ids) | Self::Filtered(Ok(ids)), Ask::Select { candidates, .. }) = + (self, ask) + && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) + { + return Err(DriveError::UnknownCandidate { id: id.clone() }); + } + if let (Self::Secrets(values), Ask::Secrets { refs, .. }) = (self, ask) + && values.len() != refs.len() + { + return Err(DriveError::SecretsLength { + expected: refs.len(), + got: values.len(), + }); + } + if let (Self::Cache(outcome), Ask::Cache { op, keys, .. }) = (self, ask) { + if !outcome.answers(*op) { + return Err(DriveError::WrongCacheOutcome { op: *op }); + } + if let (CacheOutcome::Values(values), Some(keys)) = (outcome, keys) + && values.len() != keys.len() + { + return Err(DriveError::BatchLength { + expected: keys.len(), + got: values.len(), + }); + } + } + Ok(()) + } +} + +/// How a driven future ended well. +#[derive(Debug)] +pub(crate) enum Finish { + /// `start`: the call's response. + Call(RouterResponse), + /// `select`: the deployment picked and the model group it serves. + Selected { + deployment_id: String, + model_group: String, + /// The deployment as served: a priority group's member carries its priority as `order`. + deployment: Value, + }, + /// `healthy`: the deployments a pick would choose among, as served. + Healthy { + /// A deployment named by its id, which the listing returns alone rather than as a list. + single: bool, + deployment_ids: Vec, + deployments: Vec, + }, + /// `report_failure`: whether the deployment was cooled down. + Reported { cooled: bool }, + /// `report_success` / `count_success`: the success counter's key, if one was counted. + Counted { key: Option }, + /// `build`: the deployments it installed, in order. + Built(Vec), + /// `router_model_info`: the merged model info (`None` when the lookup answered `None`). + RouterModelInfo(Option), + /// `deployment_model_info`: the merged model info, if any. + DeploymentModelInfo(Option>), + /// A model-info query that ends with one JSON answer (`model_group_info`, + /// `configured_token_limits`, `model_listing_info`). + Value(Value), +} + +/// One step of a call: what it is blocked on, or how it ended. +#[derive(Debug)] +pub(crate) enum Step { + Ask(Ask), + Done(Finish), + Failed(Error), +} + +/// A misuse of the step protocol by the caller. +#[derive(Debug, thiserror::Error)] +pub(crate) enum DriveError { + #[error("the router call already finished")] + Finished, + #[error("the router call has not asked for anything yet, so the outcome must be None")] + Unasked, + #[error("the router call is waiting for an outcome of its `{expected}` step")] + MissingAnswer { expected: StepKind }, + #[error("the router call is waiting for an outcome of its `{expected}` step, not `{got}`")] + WrongAnswer { expected: StepKind, got: StepKind }, + #[error("deployment `{id}` was not one of the select step's candidates")] + UnknownCandidate { id: String }, + #[error("the outcome does not answer a `{op}` cache step")] + WrongCacheOutcome { op: CacheOp }, + #[error("a `batch_get` cache step for {expected} keys was answered with {got} values")] + BatchLength { expected: usize, got: usize }, + #[error("a `secrets` step for {expected} references was answered with {got} values")] + SecretsLength { expected: usize, got: usize }, + #[error("the router call is pending on something other than a step")] + Stalled, +} + +impl DriveError { + /// Whether the caller passed a bad outcome, rather than driving a call that cannot go on. + pub(crate) fn is_request(&self) -> bool { + match self { + Self::Unasked + | Self::MissingAnswer { .. } + | Self::WrongAnswer { .. } + | Self::UnknownCandidate { .. } + | Self::WrongCacheOutcome { .. } + | Self::BatchLength { .. } + | Self::SecretsLength { .. } => true, + Self::Finished | Self::Stalled => false, + } + } +} + +/// The question a call is blocked on and the answer that releases it. +#[derive(Default)] +struct Exchange { + asked: Option, + answer: Option, + /// Events the call reported since the last step, handed to Python with the next one. + events: Vec, +} + +type Shared = Arc>; + +fn lock(exchange: &Mutex) -> MutexGuard<'_, Exchange> { + exchange.lock().unwrap_or_else(PoisonError::into_inner) +} + +thread_local! { + static CURRENT: RefCell> = const { RefCell::new(None) }; +} + +/// Binds an exchange to this thread for one poll and restores the previous one after. +struct Bound(Option); + +impl Bound { + fn new(exchange: &Shared) -> Self { + Self(CURRENT.replace(Some(exchange.clone()))) + } +} + +impl Drop for Bound { + fn drop(&mut self) { + CURRENT.set(self.0.take()); + } +} + +/// Park `question` in the bound exchange and wait for its answer. +async fn ask(question: Ask) -> Answer { + let mut question = Some(question); + poll_fn(move |_| { + CURRENT.with_borrow(|current| { + let Some(exchange) = current else { + return Poll::Pending; + }; + let mut exchange = lock(exchange); + if let Some(question) = question.take() { + exchange.asked = Some(question); + return Poll::Pending; + } + exchange.answer.take().map_or(Poll::Pending, Poll::Ready) + }) + }) + .await +} + +const CHECKED: &str = "the driver checks each answer against its ask"; + +fn ttl_seconds(ttl: Option) -> Option { + ttl.map(|ttl| ttl.as_secs_f64()) +} + +/// An increment amount as Python's counters store it: an int when it is whole. +fn amount_value(amount: f64) -> Value { + if amount.fract() == 0.0 && amount.abs() < i64::MAX as f64 { + Value::from(amount as i64) + } else { + Value::from(amount) + } +} + +/// The host's cache: every op is a `cache` step Python runs against the router's `DualCache`. +pub(crate) struct StepCache; + +async fn cache_step( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, +) -> CacheOutcome { + ask_cache(op, key, keys, value, ttl, false).await +} + +async fn ask_cache( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, + local: bool, +) -> CacheOutcome { + let question = Ask::Cache { + op, + key: key.map(str::to_owned), + keys: keys.map(<[String]>::to_vec), + value, + ttl: ttl_seconds(ttl), + local, + }; + let Answer::Cache(outcome) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + outcome +} + +#[async_trait] +impl RouterCache for StepCache { + async fn async_get_cache(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + cache_step(CacheOp::Get, Some(key), None, None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Set, Some(key), None, Some(value), ttl).await; + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + cache_step(CacheOp::Increment, Some(key), None, value, ttl).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } + + async fn async_batch_get_cache( + &self, + keys: &[String], + ) -> Result>, litellm_cache::Error> { + let CacheOutcome::Values(values) = + cache_step(CacheOp::BatchGet, None, Some(keys), None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Delete, Some(key), None, None, None).await; + Ok(()) + } + + async fn async_get_local(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + ask_cache(CacheOp::Get, Some(key), None, None, None, true).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_increment_local( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + ask_cache(CacheOp::Increment, Some(key), None, value, ttl, true).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } +} + +/// A deployment as Python spells a `model_list` entry, with its id in `model_info`. +pub(crate) fn deployment_record(deployment: &Deployment) -> Value { + deployment.record() +} + +/// The [`Host`] every `RouterCoreHandle` router runs on. +pub(crate) struct StepHost; + +#[async_trait] +impl Host for StepHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + let question = Ask::Attempt { + deployment_id: attempt.deployment.id.clone(), + deployment: deployment_record(&attempt.deployment), + model: attempt.name.clone(), + model_group: attempt.model_group.clone(), + target_order: attempt.target_order, + candidate_ids: attempt.candidate_ids.clone(), + retry: attempt.retry, + fallback_depth: attempt.fallback_depth, + fallback_level: attempt.fallback_level, + timeout: attempt.timeout.map(|timeout| timeout.as_secs_f64()), + params: attempt + .params + .as_ref() + .map(|params| Box::new(Value::Object(Map::clone(params)))), + scope: attempt.scope, + max_retries: attempt.max_retries, + model_group_size: attempt.model_group_size, + }; + let (result, blocked, deferred) = match ask(question).await { + Answer::Attempt(result) => (result, None, false), + Answer::Served { + total_tokens, + blocked, + deferred, + } => (Ok(total_tokens), blocked, deferred), + _ => unreachable!("{CHECKED}"), + }; + result.map(|total_tokens| Output { + value: Payload(Arc::new(())), + total_tokens, + blocked, + deferred, + }) + } + + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let question = Ask::Select { + stage: match stage { + SelectStage::Plugins => Stage::Plugins, + SelectStage::Strategy(strategy) => Stage::Strategy(strategy.to_string()), + SelectStage::Tags => Stage::Tags, + SelectStage::TextOnly => Stage::TextOnly, + }, + pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let Answer::Select(ids) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + Ok(ids + .iter() + .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) + .collect()) + } + + async fn pre_route(&self, pre: &PreRoute<'_>) -> Result, Error> { + let question = Ask::PreRoute { + model: pre.name.to_owned(), + fallback_depth: pre.fallback_depth, + retry: pre.retry, + }; + let Answer::PreRoute(routed) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + Ok(routed) + } + + async fn routing_plugin_models( + &self, + ctx: &RoutingContext<'_>, + run: bool, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let question = Ask::Select { + stage: Stage::RoutingPlugins { run }, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: Vec::new(), + }; + let Answer::Models(models) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + Ok(models) + } + + async fn pre_call_facts( + &self, + ctx: &RoutingContext<'_>, + candidates: &[Arc], + ) -> Result { + let (scope, group) = ctx.scope.wire(); + let question = Ask::Select { + stage: Stage::PreCall, + pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let Answer::PreCall(facts) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + Ok(facts) + } + + async fn callback_filters( + &self, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let question = Ask::Select { + stage: Stage::CallbackFilters, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let Answer::Filtered(result) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + let ids = result.map_err(Error::Attempt)?; + Ok(ids + .iter() + .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) + .collect()) + } + + async fn allow_fallback(&self, _request: &RouterRequest, target: &str) -> bool { + let question = Ask::AllowFallback { + target: target.to_owned(), + }; + let Answer::AllowFallback(allowed) = ask(question).await else { + unreachable!("{CHECKED}"); + }; + allowed + } + + /// A zero sleep is asked too: Python's retry loop awaits `asyncio.sleep(0)` between sibling + /// retries, which yields to the event loop. + async fn sleep(&self, duration: Duration) { + let question = Ask::Sleep { + seconds: duration.as_secs_f64(), + }; + let Answer::Slept = ask(question).await else { + unreachable!("{CHECKED}"); + }; + } + + /// Kept for Python only when it acts on it: an exhausted retry loop is stamped on the + /// error Python raises, and a fallback hop tells it which target a failed walk ended on. + fn observe(&self, event: &RouterEvent) { + if !matches!( + event, + RouterEvent::RetriesExhausted { .. } + | RouterEvent::Fallback { .. } + | RouterEvent::LevelFailed { .. } + | RouterEvent::MockFailed { .. } + | RouterEvent::RetryLogged { .. } + | RouterEvent::RetryPolicyRaised { .. } + | RouterEvent::FallbackDone { .. } + | RouterEvent::Blocked { .. } + ) { + return; + } + CURRENT.with_borrow(|current| { + if let Some(exchange) = current { + lock(exchange).events.push(event.clone()); + } + }); + } + + fn cache(&self) -> Arc { + Arc::new(StepCache) + } + + async fn secrets(&self, index: usize, refs: &[SecretRef]) -> Result, Error> { + let Answer::Secrets(values) = ask(Ask::Secrets { + index, + refs: refs.to_vec(), + }) + .await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn validate(&self, build: &BuildAsk<'_>) -> Result { + let Answer::Validated(built) = ask(build_ask(StepKind::Validate, build, None)).await else { + unreachable!("{CHECKED}"); + }; + Ok(built) + } + + async fn provider(&self, build: &BuildAsk<'_>) -> Result { + let Answer::Provided(provided) = ask(build_ask(StepKind::Provider, build, None)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(provided) + } + + async fn added(&self, build: &BuildAsk<'_>, provider: Option<&str>) -> Result { + let Answer::Added(built) = ask(build_ask(StepKind::Added, build, provider)).await else { + unreachable!("{CHECKED}"); + }; + Ok(built) + } + + async fn model_info(&self, lookup: &ModelInfoLookup) -> Result { + let lookup = serde_json::to_value(lookup).unwrap_or(Value::Null); + let Answer::ModelInfo(answer) = ask(Ask::ModelInfo { lookup }).await else { + unreachable!("{CHECKED}"); + }; + Ok(answer) + } +} + +fn build_ask(kind: StepKind, build: &BuildAsk<'_>, provider: Option<&str>) -> Ask { + Ask::Build { + kind, + path: build.path, + index: build.index, + entry: Value::Object(build.entry.clone()), + declared_id: build.declared_id.map(str::to_owned), + duplicate_id: build.duplicate_id, + provider: provider.map(str::to_owned), + } +} + +type CallFuture = Pin> + Send>>; + +/// Poll `future` once with `exchange` bound. +fn poll_bound(future: Pin<&mut F>, exchange: &Shared) -> Poll { + let _bound = Bound::new(exchange); + future.poll(&mut Context::from_waker(Waker::noop())) +} + +/// Drives one router call a step at a time. +pub(crate) struct Driver { + future: Option, + exchange: Shared, + waiting: Option, +} + +impl Driver { + pub(crate) fn new(call: impl Future> + Send + 'static) -> Self { + Self { + future: Some(Box::pin(call)), + exchange: Shared::default(), + waiting: None, + } + } + + /// The step the call is blocked on, if it asked for one. + pub(crate) fn waiting(&self) -> Option<&Ask> { + self.waiting.as_ref() + } + + /// The events reported since the last call. + pub(crate) fn take_events(&self) -> Vec { + std::mem::take(&mut lock(&self.exchange).events) + } + + /// Hand back the answer to the waiting step (`None` on the first call) and run the call to + /// its next step or its end. + pub(crate) fn resume(&mut self, answer: Option) -> Result { + let Some(future) = self.future.as_mut() else { + return Err(DriveError::Finished); + }; + match (&self.waiting, answer) { + (None, None) => {} + (None, Some(_)) => return Err(DriveError::Unasked), + (Some(ask), None) => { + return Err(DriveError::MissingAnswer { + expected: ask.kind(), + }); + } + (Some(ask), Some(answer)) => { + answer.check(ask)?; + lock(&self.exchange).answer = Some(answer); + } + } + self.waiting = None; + match poll_bound(future.as_mut(), &self.exchange) { + Poll::Ready(result) => { + self.future = None; + Ok(match result { + Ok(finish) => Step::Done(finish), + Err(error) => Step::Failed(error), + }) + } + Poll::Pending => { + let Some(ask) = lock(&self.exchange).asked.take() else { + self.future = None; + return Err(DriveError::Stalled); + }; + self.waiting = Some(ask.clone()); + Ok(Step::Ask(ask)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + time::{SystemTime, UNIX_EPOCH}, + }; + + use litellm_router_core::{ + FailureKind, Op, Router, RouterConfig, RouterDynamicConfig, SeededRandomness, + }; + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + fn router(config: Value, dynamic: Value) -> Router { + let config: RouterConfig = serde_json::from_value(config).unwrap(); + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, Arc::new(StepHost)) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() + } + + fn deployment(group: &str, id: &str) -> Value { + json!({"model_name": group, "litellm_params": {"model": format!("openai/{id}")}, "model_info": {"id": id}}) + } + + fn start(router: &Router, model: &str) -> Driver { + let router = router.clone(); + let request = RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + }; + Driver::new(async move { router.call(Op::Completion, request).await.map(Finish::Call) }) + } + + fn rate_limited() -> Answer { + Answer::Attempt(Err(AttemptError::new(FailureKind::RateLimit, Some(429)))) + } + + /// A dict-backed stand-in for Python's `DualCache` that answers `cache` steps and keeps + /// every one it answered. + #[derive(Default)] + struct Store { + entries: HashMap, + asked: Vec, + } + + impl Store { + fn answer(&mut self, ask: &Ask) -> Answer { + let Ask::Cache { + op, + key, + keys, + value, + .. + } = ask + else { + panic!("not a cache step: {ask:?}"); + }; + self.asked.push(ask.clone()); + let key = key.clone().unwrap_or_default(); + Answer::Cache(match op { + CacheOp::Get => CacheOutcome::Value(self.entries.get(&key).cloned()), + CacheOp::Set => { + self.entries.insert(key, value.clone().unwrap()); + CacheOutcome::Stored + } + CacheOp::Increment => { + let current = self.entries.get(&key).and_then(Value::as_f64); + let next = + current.unwrap_or(0.0) + value.as_ref().and_then(Value::as_f64).unwrap(); + self.entries.insert(key, amount_value(next)); + CacheOutcome::Number(next) + } + CacheOp::BatchGet => CacheOutcome::Values( + keys.iter() + .flatten() + .map(|key| self.entries.get(key).cloned()) + .collect(), + ), + CacheOp::Delete => { + self.entries.remove(&key); + CacheOutcome::Stored + } + }) + } + + /// Resume with `answer`, answering every cache step on the way, up to the next other + /// step. + fn resume(&mut self, driver: &mut Driver, answer: Option) -> Step { + let mut step = driver.resume(answer).unwrap(); + while let Step::Ask(ask @ Ask::Cache { .. }) = &step { + let answer = self.answer(ask); + step = driver.resume(Some(answer)).unwrap(); + } + step + } + } + + #[fixture] + fn pair() -> Router { + router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({}), + ) + } + + #[rstest] + fn a_first_attempt_success_is_done(pair: Router) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let Step::Ask(Ask::Attempt { + deployment_id, + model_group, + retry: 0, + fallback_depth: 0, + .. + }) = store.resume(&mut driver, None) + else { + panic!("the call should start with an attempt"); + }; + assert_eq!(model_group, "primary"); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(Answer::Attempt(Ok(12)))) + else { + panic!("a successful attempt should finish the call"); + }; + assert_eq!(response.deployment_id, deployment_id); + assert_eq!(response.output.total_tokens, 12); + assert_eq!( + (response.attempted_retries, response.attempted_fallbacks), + (0, 0) + ); + assert!(matches!(driver.resume(None), Err(DriveError::Finished))); + } + + #[rstest] + fn cooldown_reads_are_cache_steps(pair: Router) { + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::Ask(Ask::Cache { + op: CacheOp::BatchGet, + key: None, + keys: Some(keys), + value: None, + ttl: None, + local: false, + }) = step + else { + panic!("the call should read cooldowns first, got {step:?}"); + }; + assert_eq!(keys, ["deployment:a1:cooldown", "deployment:a2:cooldown"]); + } + + #[rstest] + #[case::lone_deployment_backs_off(vec![deployment("primary", "a1")], true)] + #[case::sibling_retries_at_once( + vec![deployment("primary", "a1"), deployment("primary", "a2")], + false + )] + fn a_rate_limit_retries_after_the_backoff( + #[case] model_list: Vec, + #[case] backs_off: bool, + ) { + let router = router(json!({"model_list": model_list}), json!({"num_retries": 1})); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let step = store.resume(&mut driver, Some(rate_limited())); + let Step::Ask(Ask::Sleep { seconds }) = step else { + panic!("a retry should sleep first, got {step:?}"); + }; + assert_eq!(seconds > 0.0, backs_off, "slept {seconds}s"); + let step = store.resume(&mut driver, Some(Answer::Slept)); + assert!( + matches!(step, Step::Ask(Ask::Attempt { retry: 1, .. })), + "expected the retry, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(Answer::Attempt(Ok(1)))) + else { + panic!("the retry should finish the call"); + }; + assert_eq!(response.attempted_retries, 1); + } + + #[rstest] + fn a_cooldown_python_wrote_is_read_back_by_the_next_call() { + let pair_without_retries = router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + store.entries.insert( + "deployment:a1:cooldown".into(), + json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}), + ); + for _ in 0..8 { + let mut call = start(&pair_without_retries, "primary"); + let Step::Ask(Ask::Attempt { deployment_id, .. }) = store.resume(&mut call, None) + else { + panic!("the call should attempt the healthy deployment"); + }; + assert_eq!(deployment_id, "a2", "the cooled-down deployment was picked"); + } + } + + /// The host only runs the attempt; the core records its outcome (cooldowns, fail counters, + /// usage) through the host's cache. + #[rstest] + #[case::failure(rate_limited())] + #[case::success(Answer::Attempt(Ok(7)))] + fn attempt_outcomes_are_recorded_by_the_core(#[case] outcome: Answer) { + let tracked = |id: &str| { + json!({ + "model_name": "primary", + "litellm_params": {"model": format!("openai/{id}"), "tpm": 100, "cooldown_time": 7}, + "model_info": {"id": id, "allowed_fails": 3}, + }) + }; + let router = router( + json!({"model_list": [tracked("a1"), tracked("a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let step = store.resume(&mut driver, Some(outcome)); + assert!( + matches!(step, Step::Done(_) | Step::Failed(_)), + "got {step:?}" + ); + let writes: Vec<&Ask> = store + .asked + .iter() + .filter(|ask| { + matches!(ask, Ask::Cache { op, .. } if !matches!(op, CacheOp::Get | CacheOp::BatchGet)) + }) + .collect(); + assert!(!writes.is_empty(), "the core wrote no outcome"); + } + + #[rstest] + fn a_failed_group_falls_back_once_python_allows_it() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let bad_request = + Answer::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + let step = store.resume(&mut driver, Some(bad_request)); + assert!( + matches!(&step, Step::Ask(Ask::AllowFallback { target }) if target == "backup"), + "expected the fallback check, got {step:?}" + ); + let step = store.resume(&mut driver, Some(Answer::AllowFallback(true))); + assert!( + matches!(&step, Step::Ask(Ask::Attempt { deployment_id, fallback_depth: 1, .. }) if deployment_id == "b1"), + "expected the fallback attempt, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(Answer::Attempt(Ok(3)))) + else { + panic!("the fallback should finish the call"); + }; + assert_eq!(response.model_group, "backup"); + assert_eq!(response.attempted_fallbacks, 1); + } + + #[rstest] + fn a_refused_fallback_fails_with_the_attempt_error() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let bad_request = + Answer::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + store.resume(&mut driver, Some(bad_request)); + let Step::Failed(error) = store.resume(&mut driver, Some(Answer::AllowFallback(false))) + else { + panic!("a refused fallback should fail the call"); + }; + assert_eq!( + error.attempt().map(|a| a.kind), + Some(FailureKind::BadRequest) + ); + } + + #[rstest] + #[case::none_while_waiting(None, StepKind::Attempt)] + #[case::wrong_kind(Some(Answer::Slept), StepKind::Attempt)] + #[case::cache_outcome(Some(Answer::Cache(CacheOutcome::Stored)), StepKind::Attempt)] + fn a_mismatched_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] answer: Option, + #[case] expected: StepKind, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + store.resume(&mut driver, None); + let error = driver.resume(answer).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.waiting().map(Ask::kind), Some(expected)); + assert!(matches!( + store.resume(&mut driver, Some(Answer::Attempt(Ok(1)))), + Step::Done(_) + )); + } + + #[rstest] + #[case::wrong_op(Answer::Cache(CacheOutcome::Stored))] + #[case::number_for_batch(Answer::Cache(CacheOutcome::Number(1.0)))] + #[case::short_batch(Answer::Cache(CacheOutcome::Values(vec![None])))] + #[case::not_a_cache_outcome(Answer::Slept)] + fn a_mismatched_cache_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] answer: Answer, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::Ask( + ask @ Ask::Cache { + op: CacheOp::BatchGet, + .. + }, + ) = &step + else { + panic!("expected the cooldown read, got {step:?}"); + }; + let error = driver.resume(Some(answer)).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.waiting(), Some(ask)); + let answer = store.answer(ask); + assert_eq!( + store.resume(&mut driver, Some(answer)).unwrap_kind(), + StepKind::Attempt + ); + } + + #[rstest] + fn a_strategy_python_runs_asks_for_a_selection_and_keeps_the_pick() { + let router = router( + json!({ + "model_list": [deployment("primary", "a1"), deployment("primary", "a2")], + "routing_strategy": "least-busy", + }), + json!({}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + let step = store.resume(&mut driver, None); + let Step::Ask(Ask::Select { + stage: Stage::Strategy(strategy), + candidates, + .. + }) = step + else { + panic!("a host strategy should ask for a selection, got {step:?}"); + }; + assert_eq!(strategy, "least-busy"); + assert_eq!(candidates, ["a1", "a2"]); + let step = store.resume(&mut driver, Some(Answer::Select(vec!["a2".to_owned()]))); + assert!( + matches!(&step, Step::Ask(Ask::Attempt { deployment_id, .. }) if deployment_id == "a2"), + "expected the picked deployment, got {step:?}" + ); + } + + #[rstest] + fn a_selection_is_driven_through_cache_steps(pair: Router) { + let request = RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + }; + let mut store = Store::default(); + let mut driver = Driver::new(async move { + pair.select(&Op::Completion, &request) + .await + .map(|deployment| Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + deployment: deployment_record(&deployment), + }) + }); + let Step::Done(Finish::Selected { + deployment_id, + model_group, + .. + }) = store.resume(&mut driver, None) + else { + panic!("a selection should finish with the pick"); + }; + assert!(["a1", "a2"].contains(&deployment_id.as_str())); + assert_eq!(model_group, "primary"); + assert!(matches!( + store.asked.as_slice(), + [Ask::Cache { + op: CacheOp::BatchGet, + .. + }] + )); + } + + trait StepKindOf { + fn unwrap_kind(self) -> StepKind; + } + + impl StepKindOf for Step { + fn unwrap_kind(self) -> StepKind { + match self { + Step::Ask(ask) => ask.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + } + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs new file mode 100644 index 00000000000..975ea4e7c04 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -0,0 +1,1862 @@ +//! `RouterCoreHandle`: router-core behind a step-driven handle. Python drives each call through +//! [`RouterCall::next`], running every attempt, sleep, fallback check, selection and cache op the +//! call loop asks for and handing the outcome back; Rust never awaits Python. `select` is driven +//! the same way, since reading cooldowns is a cache step. + +mod config; +mod host; +mod retries; +mod stream; + +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, SystemTime}, +}; + +use litellm_error::Report; +use litellm_host_python::{from_py_argument, json_loads, to_py}; +use litellm_router_core::lookup::{ListScope, Lookup, Row}; +use litellm_router_core::{ + AttemptError, Blocked, BuildPath, Built, Deployment, DeploymentSpec, Error, FailureKind, + FailureReport, Op, PreRouted, Provided, RetryCause, Router, RouterEvent, SecretAt, + build::Install, + health::{UsageMetric, current_minute, usage_key}, + registry::{DeploymentFilter, ResolveKind, ResolveScope}, + strategy::routing_context, +}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::{PyBool, PyDict, PyList, PyString}, +}; +use serde_json::Value; + +use self::host::{ + Answer, Ask, CacheOp, CacheOutcome, DriveError, Driver, Finish, Stage, Step, StepHost, StepKind, +}; +use crate::errors::RustBridgeDeclined; + +/// The `error_kind` of a `failed` step. +#[derive(Clone, Copy, Debug, strum::IntoStaticStr)] +enum ErrorKind { + #[strum(serialize = "config")] + Config, + #[strum(serialize = "unknown_model")] + UnknownModel, + #[strum(serialize = "no_healthy_deployments")] + NoHealthyDeployments, + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "retired")] + Retired, + #[strum(serialize = "paused")] + Paused, + #[strum(serialize = "not_pass_through")] + NotPassThrough, + #[strum(serialize = "no_pass_through_deployments")] + NoPassThroughDeployments, + #[strum(serialize = "only_markers")] + OnlyMarkers, + #[strum(serialize = "ptu_reserved")] + PtuReserved, + #[strum(serialize = "context_window_exceeded")] + ContextWindowExceeded, + #[strum(serialize = "unsupported_params")] + UnsupportedParams, + #[strum(serialize = "routing_plugins_left_nothing")] + RoutingPluginsLeftNothing, + #[strum(serialize = "routing_plugins_need_async")] + RoutingPluginsNeedAsync, + #[strum(serialize = "multiple_teams")] + MultipleTeams, + #[strum(serialize = "reserved")] + Reserved, + #[strum(serialize = "invalid_router_weights")] + InvalidRouterWeights, + #[strum(serialize = "non_numeric_weight")] + NonNumericWeight, + #[strum(serialize = "missing_routing_group")] + MissingRoutingGroup, + #[strum(serialize = "default_deployment_raises")] + DefaultDeploymentRaises, + #[strum(serialize = "index_out_of_range")] + IndexOutOfRange, + #[strum(serialize = "model_info")] + ModelInfo, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + match error { + Error::Config(_) => Self::Config, + Error::UnknownModel { .. } => Self::UnknownModel, + Error::NoHealthyDeployments { .. } => Self::NoHealthyDeployments, + Error::Attempt(_) => Self::Attempt, + Error::Cache(_) => Self::Cache, + Error::Retired => Self::Retired, + Error::Paused { .. } => Self::Paused, + Error::NotPassThrough { .. } => Self::NotPassThrough, + Error::NoPassThroughDeployments { .. } => Self::NoPassThroughDeployments, + Error::OnlyMarkers { .. } => Self::OnlyMarkers, + Error::PtuReserved { .. } => Self::PtuReserved, + Error::ContextWindowExceeded { .. } => Self::ContextWindowExceeded, + Error::UnsupportedParams { .. } => Self::UnsupportedParams, + Error::RoutingPluginsLeftNothing { .. } => Self::RoutingPluginsLeftNothing, + Error::RoutingPluginsNeedAsync => Self::RoutingPluginsNeedAsync, + Error::MultipleTeams { .. } => Self::MultipleTeams, + Error::Reserved { .. } => Self::Reserved, + Error::InvalidRouterWeights => Self::InvalidRouterWeights, + Error::NonNumericWeight { .. } => Self::NonNumericWeight, + Error::MissingRoutingGroup { .. } => Self::MissingRoutingGroup, + Error::DefaultDeploymentRaises { .. } => Self::DefaultDeploymentRaises, + Error::IndexOutOfRange { .. } => Self::IndexOutOfRange, + Error::ModelInfo { .. } => Self::ModelInfo, + } + } +} + +fn config_error_to_pyerr(error: config::Error) -> PyErr { + let message = Report::new(&error).to_string(); + if error.declined() { + RustBridgeDeclined::new_err(message) + } else { + PyValueError::new_err(message) + } +} + +fn drive_error_to_pyerr(error: DriveError) -> PyErr { + let message = Report::new(&error).to_string(); + if error.is_request() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) + } +} + +/// A document Python passes as a mapping or as its JSON text. +fn document(value: &Bound<'_, PyAny>) -> PyResult { + match value.cast::() { + Ok(text) => { + let parsed = json_loads(value.py(), text.to_str()?.as_bytes())?; + from_py_argument(parsed.bind(value.py())) + } + Err(_) => from_py_argument(value), + } +} + +/// A router-core router whose host operations Python runs. +#[pyclass(frozen)] +pub(crate) struct RouterCoreHandle { + router: Router, +} + +#[pymethods] +impl RouterCoreHandle { + #[new] + fn new(config: &Bound<'_, PyAny>, dynamic: &Bound<'_, PyAny>) -> PyResult { + let (config, dynamic) = (document(config)?, document(dynamic)?); + let build = || { + let config = config::router_config(config)?; + let dynamic = config::dynamic_config(dynamic)?; + Ok(Router::new(config, dynamic, Arc::new(StepHost))?) + }; + build() + .map(|router| Self { router }) + .map_err(config_error_to_pyerr) + } + + /// Start a call; nothing runs until the first `next(None)`. + fn start(&self, op: &str, request: &Bound<'_, PyAny>) -> PyResult { + let op: Op = op.parse().unwrap_or_else(|_| Op::Other(op.to_owned())); + let request = config::request(document(request)?, None).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router.call(op, request).await.map(Finish::Call) + })) + } + + /// Pick the deployment `model` routes to now, without calling it. Driven like `start`: its + /// steps are the cache reads and selections the pick needs, and it ends with `done` + /// (`deployment_id`, `model_group`) or `failed`. + fn select(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = + config::request(document(request)?, Some(model)).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let deployment = router.select(&Op::Completion, &request).await?; + Ok(Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + deployment: host::deployment_record(&deployment), + }) + })) + } + + /// List the deployments a pick for `model` would choose among, driven like `select`. + fn healthy(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = + config::request(document(request)?, Some(model)).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let candidates = router.healthy(&Op::Completion, &request).await?; + // A deployment named by its id is listed alone, as a dict (`_core_async_get_healthy_deployments`). + let single = candidates.kind.is_single() && candidates.healthy.len() == 1; + let deployments = candidates.healthy; + Ok(Finish::Healthy { + single, + deployment_ids: deployments.iter().map(|d| d.id.clone()).collect(), + deployments: deployments + .iter() + .map(|d| host::deployment_record(d)) + .collect(), + }) + })) + } + + /// Apply `{"upsert": [deployments], "delete": [ids], "settings": {dynamic document}}` and + /// return the new registry version. + fn apply(&self, changes: &Bound<'_, PyAny>) -> PyResult { + let changes = config::changes(document(changes)?).map_err(config_error_to_pyerr)?; + self.router + .apply(changes) + .map_err(|error| config_error_to_pyerr(error.into())) + } + + /// Build `entries` (model_list entries) through the host, driven like `start`: its steps + /// are `secrets`, `validate`, `provider` and `added`, and it ends with `done` (`built`: one + /// `{"index", "id", "provider"}` per deployment installed). `path` is `create` + /// (`set_model_list`) or `add` (`add_deployment`); `install` is `replace`, `append` or + /// `upsert`. + #[pyo3(signature = (entries, path, resolve_secrets, install))] + fn build( + &self, + entries: &Bound<'_, PyAny>, + path: &str, + resolve_secrets: bool, + install: &str, + ) -> PyResult { + let entries: Vec> = + serde_json::from_value(document(entries)?) + .map_err(|error| PyValueError::new_err(format!("bad model_list: {error}")))?; + let path: BuildPath = path + .parse() + .map_err(|_| PyValueError::new_err(format!("`{path}` is not a build path")))?; + let install: Install = install + .parse() + .map_err(|_| PyValueError::new_err(format!("`{install}` is not a build install")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .build(path, entries, resolve_secrets, install) + .await + .map(Finish::Built) + })) + } + + /// `get_router_model_info` for the deployment `query` describes (a + /// `RouterModelInfoQuery` document), driven like `start`: its steps are `model_info` + /// lookups, and it ends with `done` (`info`: the merged map or None, `user_keys`: the keys + /// whose values are the deployment's own). + fn router_model_info(&self, query: &Bound<'_, PyAny>) -> PyResult { + let query: litellm_router_core::model_info::RouterModelInfoQuery = + serde_json::from_value(document(query)?) + .map_err(|error| PyValueError::new_err(format!("bad model-info query: {error}")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .router_model_info(&query) + .await + .map(Finish::RouterModelInfo) + })) + } + + /// `get_deployment_model_info` for deployment `model_id` serving `model_name`, driven like + /// `router_model_info`; it ends with `done` (`info`: the merged map or None). + fn deployment_model_info( + &self, + model_id: &str, + model_name: &Bound<'_, PyAny>, + ) -> PyResult { + let model_id = model_id.to_owned(); + let model_name = from_py_argument::(model_name)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .deployment_model_info(&model_id, &model_name) + .await + .map(Finish::DeploymentModelInfo) + })) + } + + /// `_set_model_group_info` for the deployments `query` names (a `GroupInfoQuery` document), + /// driven like `router_model_info`; it ends with `done` (`value`: None when no deployment + /// matched, else `{"assignments": {field: value}}`, what the merge set on the group the host + /// built from its first deployment). + fn model_group_info(&self, query: &Bound<'_, PyAny>) -> PyResult { + let query: litellm_router_core::group_info::GroupInfoQuery = + serde_json::from_value(document(query)?) + .map_err(|error| PyValueError::new_err(format!("bad group-info query: {error}")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let info = router.model_group_info(&query).await?; + Ok(Finish::Value(info.map_or( + Value::Null, + |info| serde_json::json!({"assignments": info.assignments}), + ))) + })) + } + + /// `get_configured_token_limits` for the deployment with `id` and its `configured` model_info, + /// driven like `router_model_info`; it ends with `done` (`value`: `[max_input, max_output]`). + fn configured_token_limits( + &self, + id: &Bound<'_, PyAny>, + configured: &Bound<'_, PyAny>, + ) -> PyResult { + let id = from_py_argument::(id)?; + let Value::Object(configured) = document(configured)? else { + return Err(PyValueError::new_err( + "the configured model_info must be a dict", + )); + }; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let (input, output) = router.configured_token_limits(id, &configured).await?; + Ok(Finish::Value(serde_json::json!([input, output]))) + })) + } + + /// `get_model_listing_info` over the `deployments` a listed name names, driven like + /// `router_model_info`; it ends with `done` (`value`: `{"cost_map_keys", "max_input_tokens", + /// "max_output_tokens"}`). + fn model_listing_info(&self, deployments: &Bound<'_, PyAny>) -> PyResult { + let deployments: Vec = from_py_argument(deployments)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let listing = router.model_listing_info(&deployments).await?; + Ok(Finish::Value(serde_json::json!({ + "cost_map_keys": listing.cost_map_keys, + "max_input_tokens": listing.max_input_tokens, + "max_output_tokens": listing.max_output_tokens, + }))) + })) + } + + /// `get_configured_mode` / `get_configured_display_name`: `value` when it is a string that is + /// not blank. + #[staticmethod] + fn configured_text(value: &Bound<'_, PyAny>) -> PyResult> { + let value = from_py_argument::(value).ok(); + Ok(litellm_router_core::group_info::configured_text( + value.as_ref(), + )) + } + + /// `get_configured_service_tiers`: each of `model_infos`' `service_tiers`, unvalidated. + #[staticmethod] + fn configured_service_tiers( + py: Python<'_>, + model_infos: &Bound<'_, PyAny>, + ) -> PyResult>> { + let model_infos: Vec> = from_py_argument(model_infos)?; + litellm_router_core::group_info::configured_service_tiers(&model_infos) + .iter() + .map(|tiers| to_py(py, tiers)) + .collect() + } + + /// The routing-cache key deployment `id` serving litellm `model` counts `metric` (`tpm`, + /// `rpm`, `itpm` or `otpm`) under in `minute` (`HH-MM`). + fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult { + let metric: UsageMetric = metric + .parse() + .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; + Ok(usage_key(id, model, metric, minute)) + } + + /// Keep the providers `other`'s build resolved for the deployments this router serves too. + fn adopt_providers(&self, other: &Self) { + self.router.adopt_providers(&other.router); + } + + fn registry_version(&self) -> u64 { + self.router.registry().version() + } + + /// The version the deployments before the newest appends last changed at: a list read at + /// or after it is current once it reads `deployments_from(len(list))`. + fn registry_rewritten(&self) -> u64 { + self.router.registry().rewritten() + } + + /// Record a failure the router did not see itself (a health check, an interceptor, a + /// logging callback), driven like `start`: its steps are cache ops, and it ends with `done` + /// (`cooled`). `counts_request` also counts the failed request against the deployment's + /// per-minute requests, as an async call's failure callback does. + #[pyo3(signature = (id, failure, deployment_cooldown=false, count_only=false, requested_group=None, cooldown_time=None, uncounted=false, counts_request=false))] + #[allow(clippy::too_many_arguments)] + fn report_failure( + &self, + id: &str, + failure: &Bound<'_, PyDict>, + deployment_cooldown: bool, + count_only: bool, + requested_group: Option, + cooldown_time: Option, + uncounted: bool, + counts_request: bool, + ) -> PyResult { + let failure = match attempt_outcome(failure)? { + Err(failure) => failure, + Ok(_) => { + return Err(PyValueError::new_err( + "a reported failure needs `ok: False`", + )); + } + }; + let router = self.router.clone(); + let id = id.to_owned(); + let report = FailureReport { + deployment_cooldown, + cooldown_time, + count_only, + requested_group, + uncounted, + counts_request, + }; + Ok(RouterCall::new(async move { + router + .report_failure(&id, &failure, report) + .await + .map(|cooled| Finish::Reported { cooled }) + })) + } + + /// Record a success a logging callback saw, driven like `report_failure`; it ends with + /// `done` (`key`: the success counter's key, or None for an unknown deployment). + fn report_success(&self, id: &str, total_tokens: f64, requests: f64) -> RouterCall { + let router = self.router.clone(); + let id = id.to_owned(); + RouterCall::new(async move { + router + .report_success(&id, total_tokens, requests) + .await + .map(|key| Finish::Counted { key }) + }) + } + + /// Count a success against `id`, configured or not; ends with `done` (`key`). + fn count_success(&self, id: &str) -> RouterCall { + let router = self.router.clone(); + let id = id.to_owned(); + RouterCall::new(async move { + router + .count_success(&id) + .await + .map(|key| Finish::Counted { key: Some(key) }) + }) + } + + /// The usage keys of this minute for every deployment of `model_group`: + /// `[(id, tpm_key, rpm_key)]`, or None when the group has no deployment. + fn usage_keys(&self, model_group: &str) -> Option> { + let minute = current_minute(SystemTime::now()); + let group = self.router.registry().group(model_group)?; + Some( + group + .iter() + .map(|deployment| { + let model = &deployment.params.model; + ( + deployment.id.clone(), + usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), + ) + }) + .collect(), + ) + } + + /// Every configured deployment, as `model_list` entries with their ids, in config order. + fn deployments<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments(&DeploymentFilter::default()) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployments from position `start` on, as `deployments()` lists them. + fn deployments_from<'py>(&self, py: Python<'py>, start: usize) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments_from(start) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployment with this id, or None. + fn deployment(&self, py: Python<'_>, id: &str) -> PyResult>> { + self.router + .registry() + .deployment(id) + .map(|deployment| to_py(py, &host::deployment_record(&deployment))) + .transpose() + } + + // -- registry lookups (router-core `lookup.rs`): rows the host reads its stored entries at -- + + /// `_core_get_all_deployments`, as rows. + #[pyo3(signature = (model_name, team_id=None, rename=None))] + fn lookup_all<'py>( + &self, + py: Python<'py>, + model_name: &str, + team_id: Option<&str>, + rename: Option<&str>, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).all(model_name, team_id, rename)) + } + + /// `_team_deployments_across_teams`: the rows, and how many teams own them. + fn lookup_across_teams<'py>( + &self, + py: Python<'py>, + model: &str, + ) -> PyResult<(Bound<'py, PyAny>, usize)> { + let registry = self.router.registry(); + let (rows, teams) = Lookup::new(®istry).across_teams(model); + Ok((rows_to_py(py, &rows)?, teams)) + } + + /// `_core_config_deployments`, as rows. + fn lookup_config<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).config()) + } + + /// `_core_get_model_ids`: the ids, or the row the read raises at. + #[pyo3(signature = (model_name=None, exclude_team_models=false))] + fn lookup_model_ids<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + exclude_team_models: bool, + ) -> PyResult<(Vec, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).model_ids(model_name, exclude_team_models) { + Ok(ids) => Ok((ids, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// `_core_get_model_list`, as rows. + #[pyo3(signature = (model_name, team_id, scope))] + fn lookup_model_list<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + rows_to_py( + py, + &Lookup::new(®istry).model_list(model_name, team_id, &scope), + ) + } + + /// `_core_get_model_names`, or the row the read raises at. + fn lookup_model_names<'py>( + &self, + py: Python<'py>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(Vec, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).model_names(team_id, &scope) { + Ok(names) => Ok((names, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// `_core_get_model_access_groups`: `(group, names)` pairs in the order the groups appear, + /// or the row the read raises at. + fn lookup_access_groups<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + access_group: Option<&str>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(AccessGroups, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).access_groups(model_name, access_group, team_id, &scope) { + Ok(groups) => Ok((groups, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// Whether the whole list lists any row, and the name of its first row whose litellm model + /// is `model_id` (or ends in `/model_id` or `:model_id`); or the row the read raises at. + fn lookup_name_serving<'py>( + &self, + py: Python<'py>, + model_id: &str, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(bool, Option, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).name_serving(model_id, &scope) { + Ok((listed, name)) => Ok((listed, name, None)), + Err(raw) => Ok((false, None, Some(row_to_py(py, &raw)?))), + } + } + + /// Whether `model_names` (the group names) holds `name`. + fn has_model_name(&self, name: &str) -> bool { + self.router + .registry() + .host_indexes() + .model_names + .contains(&Value::String(name.to_owned())) + } + + /// `_core_team_model_has_alternatives`, or the row the read raises at. + fn lookup_team_model_has_alternatives<'py>( + &self, + py: Python<'py>, + id: &str, + ) -> PyResult<(bool, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).team_model_has_alternatives(id) { + Ok(found) => Ok((found, None)), + Err(raw) => Ok((false, Some(row_to_py(py, &raw)?))), + } + } + + /// The position of the group's first deployment the team may use, or the row the read + /// raises at. + fn lookup_group_usable_by_team<'py>( + &self, + py: Python<'py>, + model_group: &str, + team_id: Option<&str>, + ) -> PyResult<(Option, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).group_deployment_usable_by_team(model_group, team_id) { + Ok(position) => Ok((position, None)), + Err(raw) => Ok((None, Some(row_to_py(py, &raw)?))), + } + } + + /// The position of the team public name's first deployment, or the row the read raises at. + fn lookup_team_public_name<'py>( + &self, + py: Python<'py>, + model_id: &str, + team_id: &str, + ) -> PyResult<(Option, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).team_public_name_deployment(model_id, team_id) { + Ok(position) => Ok((position, None)), + Err(raw) => Ok((None, Some(row_to_py(py, &raw)?))), + } + } + + /// `_routable_deployments`, as rows. + fn lookup_routable<'py>( + &self, + py: Python<'py>, + model_group: &str, + team_id: Option<&str>, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).routable(model_group, team_id)) + } + + /// `_core_wildcard_route_deployments`: the team router's routes, else the global router's. + fn lookup_wildcard_route<'py>( + &self, + py: Python<'py>, + model: &str, + team_id: Option<&str>, + requests: Vec, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py( + py, + &Lookup::new(®istry).wildcard_route(model, team_id, &requests), + ) + } + + /// `_core_get_candidate_model_ids_for_route` for the resolved name: the members' ids when it + /// names a callable routing group (`members`), else for a name that is no model group the + /// team's or the wildcard route's, else the group's; or the row the read raises at. + fn lookup_candidate_ids<'py>( + &self, + py: Python<'py>, + resolved: &str, + team_id: Option<&str>, + members: Option>, + requests: Vec, + ) -> PyResult<(Vec, Option>)> { + let registry = self.router.registry(); + let lookup = Lookup::new(®istry); + match lookup.candidate_ids(resolved, team_id, members.as_deref(), &requests) { + Ok(ids) => Ok((ids, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// Whether the id map holds `id` (`_has_model_id`). + fn has_model_id(&self, id: &str) -> bool { + self.router.registry().host_indexes().has_id(id) + } + + /// What the id map holds for `id`, as written (a position, or what a caller wrote), or None. + fn id_position<'py>(&self, py: Python<'py>, id: &str) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .id_entry(id) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// The id map's entries, in order: `(id, what it holds)`. + fn id_positions<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let pairs: Vec = registry + .host_indexes() + .id_map + .iter() + .map(|(id, held)| Value::Array(vec![Value::String(id.clone()), held.clone()])) + .collect(); + to_py(py, &Value::Array(pairs)).map(|value| value.into_bound(py)) + } + + /// What the name index holds for `name`, as written, or None when it does not index it. + fn named_positions<'py>( + &self, + py: Python<'py>, + name: &str, + ) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .named(name) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// What the team index holds for `(team, public)`, as written, or None. + fn team_positions<'py>( + &self, + py: Python<'py>, + team: &str, + public: &str, + ) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .team_named(team, public) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// `team_public_model_names_of`: the public names the team index holds for `team`. + fn team_public_names_of(&self, team: &str) -> Vec { + self.router + .registry() + .host_indexes() + .team_public_names_of(team) + } + + /// The positions of the deployments whose `litellm_params.model` is `model`. + fn litellm_model_positions(&self, model: &str) -> Vec { + self.router + .registry() + .deployments(&DeploymentFilter::default()) + .iter() + .enumerate() + .filter(|(_, deployment)| { + deployment.raw_params.get("model").and_then(Value::as_str) == Some(model) + }) + .map(|(position, _)| position) + .collect() + } + + /// `PatternMatchRouter.route` over the wildcard router (or `team`'s): the first pattern's + /// deployments, each with the model the wildcards matched, or None. + #[pyo3(signature = (team, request, filtered=None))] + fn pattern_route<'py>( + &self, + py: Python<'py>, + team: Option<&str>, + request: &str, + filtered: Option>, + ) -> PyResult>> { + let registry = self.router.registry(); + let indexes = registry.host_indexes(); + let router = match team { + None => Some(&indexes.patterns), + Some(team) => indexes + .team_patterns + .iter() + .find(|(held, _)| held == team) + .map(|(_, router)| router), + }; + router + .and_then(|router| router.route(request, filtered.as_deref())) + .map(|routed| to_py(py, &Value::Array(routed)).map(|value| value.into_bound(py))) + .transpose() + } + + /// `get_models_by_pattern`: the models the first matching (request, pattern) pair routes to. + fn models_by_pattern(&self, team: Option<&str>, requests: Vec) -> Option> { + let registry = self.router.registry(); + let indexes = registry.host_indexes(); + let router = match team { + None => Some(&indexes.patterns), + Some(team) => indexes + .team_patterns + .iter() + .find(|(held, _)| held == team) + .map(|(_, router)| router), + }; + router.and_then(|router| router.models_by_pattern(&requests)) + } + + /// The indexes the host hands out (`HostIndexes`): id map, name and team indexes, group and + /// deployment names, team public names, provider default ids and the pattern routers. + fn host_indexes<'py>(&self, py: Python<'py>) -> PyResult> { + let indexes = serde_json::to_value(self.router.registry().host_indexes()) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + to_py(py, &indexes).map(|value| value.into_bound(py)) + } + + /// `provider_default_deployment_ids` as the indexes hold it now: the ids of the wildcard + /// deployments the cooldown check skips. + fn provider_default_ids(&self) -> Vec { + self.router + .registry() + .host_indexes() + .provider_default_ids + .clone() + } + + /// The provider the host resolved when it built the deployment with this id (its + /// `provider` build step), or None when there is no such deployment or it never went + /// through a host build. + fn provider(&self, id: &str) -> Option { + self.router + .registry() + .deployment(id) + .and_then(|deployment| deployment.resolved_provider.clone()) + } + + /// What `name` resolves to: `{"model_group", "kind", "ids"}`, or None for an unknown name. + #[pyo3(signature = (name, specific_deployment=false))] + fn resolve<'py>( + &self, + py: Python<'py>, + name: &str, + specific_deployment: bool, + ) -> PyResult>> { + let registry = self.router.registry(); + let Some(resolution) = registry.resolve( + name, + ResolveScope { + specific_deployment, + }, + ) else { + return Ok(None); + }; + let dict = PyDict::new(py); + dict.set_item("model_group", resolution.group)?; + dict.set_item("kind", format!("{:?}", resolution.kind))?; + dict.set_item( + "ids", + PyList::new(py, resolution.deployments.iter().map(|d| d.id.clone()))?, + )?; + Ok(Some(dict)) + } + + /// The deployments a pick for `model` would resolve to for `request`, before any team, + /// cooldown or strategy filter (`_core_get_routable_deployments`); empty when it refuses. + fn routable<'py>( + &self, + py: Python<'py>, + model: &str, + request: &Bound<'py, PyAny>, + ) -> PyResult> { + let request = + config::request(document(request)?, Some(model)).map_err(config_error_to_pyerr)?; + let records = litellm_router_core::resolve::routable( + &self.router.registry(), + &self.router.settings().dynamic, + &request, + ) + .iter() + .map(|d| to_py(py, &host::deployment_record(d))) + .collect::>>()?; + PyList::new(py, records) + } + + /// `_get_routing_context` for a pick of `model`: `(strategy, scope, group)`, spelled as a + /// strategy select step spells them. A name that resolves to nothing reads its routing group + /// membership as is. + #[pyo3(signature = (model, routing_strategy=None))] + fn routing_context( + &self, + model: &str, + routing_strategy: Option<&str>, + ) -> (String, &'static str, Option) { + let registry = self.router.registry(); + let settings = self.router.settings(); + let (kind, group) = registry + .resolve(model, ResolveScope::default()) + .map_or((ResolveKind::Group, model.to_owned()), |resolution| { + (resolution.kind, resolution.group) + }); + let (strategy, scope) = routing_context( + ®istry, + &settings.routing_strategy, + routing_strategy, + &kind, + &group, + ); + let (scope, group) = scope.wire(); + (strategy.to_string(), scope, group.map(str::to_owned)) + } + + fn retire(&self) { + self.router.retire(); + } + + /// What may still take over a stream from `origin` that fails before content, and the + /// retry budget for a failure of `kind` (see `stream::policy`). + #[pyo3(signature = (request, origin, kind=None))] + fn stream_policy<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + origin: &Bound<'py, PyAny>, + kind: Option<&str>, + ) -> PyResult> { + stream::policy(py, &self.router, request, origin, kind) + } + + /// Retry a stream from `origin` that failed before content within its group, driven like + /// `start`: `done` when a retry returned a response, `failed` with the newest failure. + fn retry_stream( + &self, + op: &str, + request: &Bound<'_, PyAny>, + origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, + ) -> PyResult { + stream::recover( + &self.router, + stream::Recovery::Retry, + op, + request, + origin, + failure, + ) + } + + /// The fallback walk for a stream from `origin` that failed before content, driven like + /// `start`. + fn fallback_stream( + &self, + op: &str, + request: &Bound<'_, PyAny>, + origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, + ) -> PyResult { + stream::recover( + &self.router, + stream::Recovery::Fallback, + op, + request, + origin, + failure, + ) + } + + /// `get_num_retries_from_retry_policy` for a failure shaped like an attempt outcome. + #[pyo3(signature = (failure, model_group=None))] + fn retry_policy_retries( + &self, + failure: &Bound<'_, PyDict>, + model_group: Option<&str>, + ) -> PyResult> { + retries::policy_retries(&self.router, failure, model_group) + } + + /// `_core_get_allowed_fails_from_policy` for a failure shaped like an attempt outcome. + fn allowed_fails_for(&self, failure: &Bound<'_, PyDict>) -> PyResult> { + retries::allowed_fails(&self.router, failure) + } + + /// `should_retry_this_error` as a predicate, for a failure shaped like an attempt outcome. + fn should_retry( + &self, + failure: &Bound<'_, PyDict>, + counts: (usize, usize), + fallbacks: (bool, bool, bool), + ) -> PyResult { + retries::should_retry(failure, counts, fallbacks) + } +} + +/// One router call, driven a step at a time by Python. +#[pyclass(frozen)] +pub(crate) struct RouterCall { + driver: Mutex, +} + +impl RouterCall { + fn new( + future: impl std::future::Future> + Send + 'static, + ) -> Self { + Self { + driver: Mutex::new(Driver::new(future)), + } + } +} + +#[pymethods] +impl RouterCall { + /// Hand back the outcome of the last step (`None` on the first call) and return the next + /// step, or the call's end. + #[pyo3(signature = (outcome))] + fn next<'py>( + &self, + py: Python<'py>, + outcome: Option>, + ) -> PyResult> { + let mut driver = self.driver.lock().unwrap_or_else(PoisonError::into_inner); + let answer = match driver.waiting() { + None => match outcome { + None => None, + Some(_) => return Err(drive_error_to_pyerr(DriveError::Unasked)), + }, + Some(ask) => Some(answer(ask, outcome)?), + }; + let step = driver.resume(answer).map_err(drive_error_to_pyerr)?; + let dict = step_dict(py, step)?; + dict.set_item("events", events_list(py, driver.take_events())?)?; + Ok(dict) + } +} + +/// `(group, names)` pairs, in the order the groups appear. +type AccessGroups = Vec<(String, Vec)>; + +/// Lookup rows as the host reads them. +fn rows_to_py<'py>(py: Python<'py>, rows: &[Row]) -> PyResult> { + let rows = Value::Array(rows.iter().map(Row::wire).collect()); + to_py(py, &rows).map(|value| value.into_bound(py)) +} + +fn row_to_py<'py>(py: Python<'py>, row: &Row) -> PyResult> { + to_py(py, &row.wire()).map(|value| value.into_bound(py)) +} + +/// `{"aliases": [{alias, target}], "groups": [{name, members}], "wildcard": [requests]}`. +fn list_scope(scope: &Bound<'_, PyAny>) -> PyResult { + serde_json::from_value(document(scope)?) + .map_err(|error| PyValueError::new_err(format!("bad lookup scope: {error}"))) +} + +/// Read Python's outcome for `ask`. +fn answer(ask: &Ask, outcome: Option>) -> PyResult { + let missing = || { + drive_error_to_pyerr(DriveError::MissingAnswer { + expected: ask.kind(), + }) + }; + match ask { + Ask::Sleep { .. } => Ok(Answer::Slept), + Ask::AllowFallback { .. } => Ok(Answer::AllowFallback( + outcome.ok_or_else(missing)?.extract()?, + )), + Ask::Select { + stage: Stage::PreCall, + .. + } => { + let facts = from_py_argument(&outcome.ok_or_else(missing)?)?; + serde_json::from_value(facts) + .map(Answer::PreCall) + .map_err(|error| PyValueError::new_err(format!("bad pre-call facts: {error}"))) + } + Ask::Select { + stage: Stage::RoutingPlugins { .. }, + .. + } => { + // `{"models": [...] | None}`: a bare None would read as no outcome. + let outcome = outcome.ok_or_else(missing)?; + let models = outcome.cast::()?.get_item("models")?; + Ok(Answer::Models(match models { + Some(models) if !models.is_none() => Some(models.extract()?), + _ => None, + })) + } + Ask::Select { + stage: Stage::CallbackFilters, + .. + } => { + // The remaining ids, or a failure mapping as an attempt reports one. + let outcome = outcome.ok_or_else(missing)?; + if let Ok(failure) = outcome.cast::() { + return match attempt_outcome(failure)? { + Err(error) => Ok(Answer::Filtered(Err(error))), + Ok(_) => Err(PyValueError::new_err( + "a `callback_filters` outcome is the remaining ids or a failure", + )), + }; + } + Ok(Answer::Filtered(Ok(outcome.extract()?))) + } + Ask::Select { .. } => Ok(Answer::Select(outcome.ok_or_else(missing)?.extract()?)), + Ask::PreRoute { .. } => pre_route_outcome(outcome).map(Answer::PreRoute), + Ask::Attempt { .. } => { + let outcome = outcome.ok_or_else(missing)?; + let outcome = outcome.cast::()?; + let blocked = match outcome.get_item("blocked")? { + Some(value) if !value.is_none() => { + let name: String = value.extract()?; + Some(name.parse::().map_err(|_| { + PyValueError::new_err(format!( + "`{name}` is not a way a response is blocked" + )) + })?) + } + Some(_) | None => None, + }; + let deferred = outcome + .get_item("deferred")? + .is_some_and(|flag| flag.is_truthy().unwrap_or(false)); + attempt_outcome(outcome).map(|result| match result { + Ok(total_tokens) if blocked.is_some() || deferred => Answer::Served { + total_tokens, + blocked, + deferred, + }, + result => Answer::Attempt(result), + }) + } + Ask::Cache { op, .. } => cache_outcome(*op, outcome).map(Answer::Cache), + Ask::Secrets { .. } => Ok(Answer::Secrets(from_py_argument( + &outcome.ok_or_else(missing)?, + )?)), + // `None` is an answer here (no row, no entry), so a missing outcome reads as one. + Ask::ModelInfo { .. } => Ok(Answer::ModelInfo(match outcome { + Some(outcome) if !outcome.is_none() => from_py_argument(&outcome)?, + _ => Value::Null, + })), + Ask::Build { kind, .. } => { + let outcome = outcome.ok_or_else(missing)?; + let outcome = outcome.cast::()?; + let invalid = outcome + .get_item("invalid")? + .is_some_and(|flag| flag.is_truthy().unwrap_or(false)); + match kind { + StepKind::Provider => Ok(Answer::Provided(if invalid { + Provided::Invalid + } else { + Provided::Provider(match outcome.get_item("provider")? { + Some(provider) if !provider.is_none() => Some(provider.extract()?), + _ => None, + }) + })), + _ => { + let built = if invalid { + Built::Invalid + } else { + match outcome.get_item("entry")? { + Some(entry) if !entry.is_none() => { + Built::Entry(from_py_argument(&entry)?) + } + _ => Built::Skip, + } + }; + Ok(if *kind == StepKind::Validate { + Answer::Validated(built) + } else { + Answer::Added(built) + }) + } + } + } + } +} + +/// A pre-route answer: `None`, a model group name, or a custom strategy's deployment document. +fn pre_route_outcome(outcome: Option>) -> PyResult> { + let Some(outcome) = outcome.filter(|outcome| !outcome.is_none()) else { + return Ok(None); + }; + if let Ok(model) = outcome.cast::() { + return Ok(Some(PreRouted::Model(model.to_str()?.to_owned()))); + } + let spec: DeploymentSpec = serde_json::from_value(document(&outcome)?) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + let deployment = Deployment::from_spec(spec) + .map_err(|error| PyValueError::new_err(Report::new(&error).to_string()))?; + Ok(Some(PreRouted::Deployment(Arc::new(deployment)))) +} + +/// `get`: the JSON value or None; `batch_get`: a list of JSON values or None, one per key; +/// `increment`: the new int or float; `set` and `delete`: None. +fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult { + let outcome = outcome.filter(|outcome| !outcome.is_none()); + let json = |value: Option| value.filter(|value| !value.is_null()); + match op { + CacheOp::Get => outcome + .map(|outcome| from_py_argument::(&outcome)) + .transpose() + .map(|value| CacheOutcome::Value(json(value))), + CacheOp::BatchGet => { + let outcome = outcome.ok_or_else(|| { + PyValueError::new_err("a `batch_get` cache outcome must be a list, not None") + })?; + if !outcome.is_instance_of::() { + return Err(PyValueError::new_err( + "a `batch_get` cache outcome must be a list with one entry per key", + )); + } + let values: Vec> = from_py_argument(&outcome)?; + Ok(CacheOutcome::Values(values.into_iter().map(json).collect())) + } + CacheOp::Increment => outcome + .filter(|outcome| !outcome.is_instance_of::()) + .and_then(|outcome| outcome.extract::().ok()) + .map(CacheOutcome::Number) + .ok_or_else(|| { + PyValueError::new_err("an `increment` cache outcome must be the new int or float") + }), + CacheOp::Set | CacheOp::Delete => match outcome { + None => Ok(CacheOutcome::Stored), + Some(_) => Err(PyValueError::new_err(format!( + "a `{op}` cache step has no outcome, so pass None" + ))), + }, + } +} + +/// `{"ok": True, "total_tokens": int}` (an attempt's may add `"blocked": str` and +/// `"deferred": bool`) or `{"ok": False, "kind": str, "status": int | None, +/// "retry_after": float | None, "num_retries": int | None, "detail": str | None}`. +pub(crate) fn attempt_outcome(outcome: &Bound<'_, PyDict>) -> PyResult> { + fn item<'py, T: FromPyObjectOwned<'py>>( + outcome: &Bound<'py, PyDict>, + key: &str, + ) -> PyResult> { + match outcome.get_item(key)? { + Some(value) if !value.is_none() => Ok(Some(value.extract().map_err(Into::into)?)), + Some(_) | None => Ok(None), + } + } + let ok: bool = item(outcome, "ok")? + .ok_or_else(|| PyValueError::new_err("an attempt outcome needs `ok`"))?; + if ok { + return Ok(Ok(item(outcome, "total_tokens")?.unwrap_or(0))); + } + let kind: String = item(outcome, "kind")? + .ok_or_else(|| PyValueError::new_err("a failed attempt outcome needs `kind`"))?; + let kind: FailureKind = kind + .parse() + .map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind")))?; + let retry_after = item::(outcome, "retry_after")? + .map(|seconds| { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("`retry_after` must be a non-negative number")) + }) + .transpose()?; + let mut failure = AttemptError::new(kind, item(outcome, "status")?); + failure.retry_after = retry_after; + failure.num_retries = item(outcome, "num_retries")?; + failure.detail = item::(outcome, "detail")?.map(Into::into); + failure.recorded = item::(outcome, "recorded")?.unwrap_or(false); + failure.names_provider = item::(outcome, "llm_provider")?.unwrap_or(false); + failure.classes = item::>(outcome, "classes")? + .unwrap_or_default() + .into_iter() + .map(Into::into) + .collect(); + Ok(Err(failure)) +} + +/// `[{"kind": "retries_exhausted", "model_group": str, "attempted": int, "max": int} +/// | {"kind": "fallback", "to": str, "order": int | None, "level": int}]`. +/// The `cause` of a `retry_logged` or `retry_policy_raised` event, with what names it. +fn cause_items(dict: &Bound<'_, PyDict>, cause: RetryCause) -> PyResult<()> { + match cause { + RetryCause::Attempt => dict.set_item("cause", "attempt"), + RetryCause::Level => dict.set_item("cause", "level"), + RetryCause::Other => dict.set_item("cause", "other"), + RetryCause::Pick(picked) => { + dict.set_item("cause", "pick")?; + pick_items(dict, &picked.0) + } + } +} + +/// A pick's error, as the `failed` step names it (`error_kind` and what names it), with the +/// model its message names (`failed_model`) when the error carries one. +fn pick_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { + error_items(dict, error)?; + match error { + Error::NoHealthyDeployments { model } | Error::UnknownModel { model } => { + dict.set_item("failed_model", model) + } + _ => Ok(()), + } +} + +/// `error_kind` and the fields that name the error. +fn error_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { + dict.set_item("error", Report::new(error).to_string())?; + dict.set_item("error_kind", <&'static str>::from(ErrorKind::from(error)))?; + match error { + Error::OnlyMarkers { model } + | Error::PtuReserved { model } + | Error::MultipleTeams { model } + | Error::DefaultDeploymentRaises { model } + | Error::IndexOutOfRange { model } => dict.set_item("error_model", model), + Error::ContextWindowExceeded { model, detail } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", detail) + } + Error::UnsupportedParams { model, params } => { + dict.set_item("error_model", model)?; + dict.set_item("error_params", PyList::new(dict.py(), params)?) + } + Error::RoutingPluginsLeftNothing { models } => { + dict.set_item("error_params", PyList::new(dict.py(), models)?) + } + Error::Reserved { + model, + until, + timezone, + } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", format!("{until} {timezone}")) + } + Error::NonNumericWeight { metric } => dict.set_item("error_detail", metric), + Error::MissingRoutingGroup { group } => dict.set_item("error_model", group), + Error::ModelInfo { exception, detail } => { + dict.set_item("error_exception", *exception)?; + dict.set_item("error_detail", detail) + } + _ => Ok(()), + } +} + +fn events_list(py: Python<'_>, events: Vec) -> PyResult> { + let list = PyList::empty(py); + for event in events { + let dict = PyDict::new(py); + match event { + RouterEvent::RetriesExhausted { + model_group, + attempted, + max, + level, + picked, + } => { + dict.set_item("kind", "retries_exhausted")?; + dict.set_item("model_group", model_group)?; + dict.set_item("attempted", attempted)?; + dict.set_item("max", max)?; + dict.set_item("level", level)?; + dict.set_item("picked", picked)?; + } + RouterEvent::Fallback { + from, + to, + order, + depth, + level, + scope, + parent, + max_fallbacks, + } => { + dict.set_item("kind", "fallback")?; + dict.set_item("from", from)?; + dict.set_item("to", to)?; + dict.set_item("order", order)?; + dict.set_item("depth", depth)?; + dict.set_item("level", level)?; + dict.set_item("scope", scope)?; + dict.set_item("parent", parent)?; + dict.set_item("max_fallbacks", max_fallbacks)?; + } + RouterEvent::Blocked { + level, + model, + blocked, + } => { + dict.set_item("kind", "blocked")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("blocked", <&'static str>::from(blocked))?; + } + RouterEvent::FallbackDone { + from, + to, + level, + scope, + parent, + succeeded, + } => { + dict.set_item("kind", "fallback_done")?; + dict.set_item("from", from)?; + dict.set_item("to", to)?; + dict.set_item("level", level)?; + dict.set_item("scope", scope)?; + dict.set_item("parent", parent)?; + dict.set_item("succeeded", succeeded)?; + } + RouterEvent::RetryLogged { + scope, + level, + model, + attempted_retries, + cause, + } => { + dict.set_item("kind", "retry_logged")?; + dict.set_item("scope", scope)?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("attempted_retries", attempted_retries)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetryPolicyRaised { + level, + model, + cause, + } => { + dict.set_item("kind", "retry_policy_raised")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + cause_items(&dict, cause)?; + } + RouterEvent::LevelFailed { + level, + model, + error, + handling, + } => { + dict.set_item("kind", "level_failed")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("handling", handling)?; + pick_items(&dict, &error.0)?; + } + RouterEvent::MockFailed { + mock, + level, + model, + num_retries, + } => { + dict.set_item("kind", "mock_failed")?; + dict.set_item("mock", <&'static str>::from(mock))?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("num_retries", num_retries)?; + } + _ => continue, + } + list.append(dict)?; + } + Ok(list) +} + +fn step_dict(py: Python<'_>, step: Step) -> PyResult> { + let dict = PyDict::new(py); + let kind = match &step { + Step::Ask(ask) => ask.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + }; + dict.set_item("kind", <&'static str>::from(kind))?; + match step { + Step::Ask(ask) => ask_items(py, &dict, ask)?, + Step::Done(finish) => done_items(&dict, finish)?, + Step::Failed(error) => { + let attempt = error.attempt(); + error_items(&dict, &error)?; + dict.set_item("attempt_failed", attempt.is_some())?; + dict.set_item("status", attempt.and_then(|failure| failure.status))?; + dict.set_item( + "failure_kind", + attempt.map(|failure| <&'static str>::from(failure.kind)), + )?; + } + } + Ok(dict) +} + +fn ask_items(py: Python<'_>, dict: &Bound<'_, PyDict>, ask: Ask) -> PyResult<()> { + match ask { + Ask::Attempt { + deployment_id, + deployment, + model, + model_group, + target_order, + candidate_ids, + retry, + fallback_depth, + fallback_level, + timeout, + params, + scope, + max_retries, + model_group_size, + } => { + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + dict.set_item("model_group_size", model_group_size)?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("deployment", to_py(py, &deployment)?)?; + dict.set_item("model", model)?; + dict.set_item("model_group", model_group)?; + dict.set_item("target_order", target_order)?; + dict.set_item("candidate_ids", PyList::new(py, candidate_ids)?)?; + dict.set_item("retry", retry)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("fallback_level", fallback_level)?; + dict.set_item("timeout", timeout)?; + dict.set_item("params", params.map(|p| to_py(py, &p)).transpose()?) + } + Ask::Sleep { seconds } => dict.set_item("seconds", seconds), + Ask::AllowFallback { target } => dict.set_item("target", target), + Ask::Select { + stage, + model_group, + name, + scope, + group, + pool, + candidates, + } => { + let (stage, strategy) = match stage { + Stage::Plugins => ("plugins", None), + Stage::Strategy(strategy) => ("strategy", Some(strategy)), + Stage::Tags => ("tags", None), + Stage::PreCall => ("pre_call", None), + Stage::TextOnly => ("text_only", None), + Stage::CallbackFilters => ("callback_filters", None), + Stage::RoutingPlugins { run } => { + dict.set_item("run", run)?; + ("routing_plugins", None) + } + }; + dict.set_item("pool", PyList::new(py, pool)?)?; + dict.set_item("stage", stage)?; + dict.set_item("strategy", strategy)?; + dict.set_item("model_group", model_group)?; + dict.set_item("model", name)?; + dict.set_item("scope", scope)?; + dict.set_item("group", group)?; + dict.set_item("candidates", PyList::new(py, candidates)?) + } + Ask::PreRoute { + model, + fallback_depth, + retry, + } => { + dict.set_item("model", model)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("retry", retry) + } + Ask::Cache { + op, + key, + keys, + value, + ttl, + local, + } => { + dict.set_item("op", <&'static str>::from(op))?; + dict.set_item("key", key)?; + dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?; + dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?; + dict.set_item("ttl", ttl)?; + dict.set_item("local", local) + } + Ask::Secrets { index, refs } => { + dict.set_item("index", index)?; + let refs = refs + .into_iter() + .map(|secret| { + let item = PyDict::new(py); + item.set_item("name", secret.name)?; + match secret.at { + SecretAt::Param(key) => { + item.set_item("key", key)?; + item.set_item("env_only", false)?; + } + SecretAt::DataSource { source, key } => { + item.set_item("source", source)?; + item.set_item("key", key)?; + item.set_item("env_only", true)?; + } + } + Ok(item) + }) + .collect::>>()?; + dict.set_item("refs", PyList::new(py, refs)?) + } + Ask::Build { + kind: _, + path, + index, + entry, + declared_id, + duplicate_id, + provider, + } => { + dict.set_item("path", <&'static str>::from(path))?; + dict.set_item("index", index)?; + dict.set_item("entry", to_py(py, &entry)?)?; + dict.set_item("declared_id", declared_id)?; + dict.set_item("duplicate_id", duplicate_id)?; + dict.set_item("provider", provider) + } + Ask::ModelInfo { lookup } => { + if let Value::Object(lookup) = lookup { + for (key, value) in lookup { + dict.set_item(key, to_py(py, &value)?)?; + } + } + Ok(()) + } + } +} + +fn done_items(dict: &Bound<'_, PyDict>, finish: Finish) -> PyResult<()> { + match finish { + Finish::Call(response) => { + dict.set_item("deployment_id", response.deployment_id)?; + dict.set_item("model_group", response.model_group)?; + dict.set_item("attempted_retries", response.attempted_retries)?; + dict.set_item("max_retries", response.max_retries)?; + dict.set_item("attempted_fallbacks", response.attempted_fallbacks)?; + dict.set_item("total_tokens", response.output.total_tokens) + } + Finish::Selected { + deployment_id, + model_group, + deployment, + } => { + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("model_group", model_group)?; + dict.set_item("deployment", to_py(dict.py(), &deployment)?) + } + Finish::Reported { cooled } => dict.set_item("cooled", cooled), + Finish::Built(built) => { + let py = dict.py(); + let records = built + .into_iter() + .map(|deployment| { + let item = PyDict::new(py); + item.set_item("index", deployment.index)?; + item.set_item( + "id", + deployment.spec.model_info.get("id").and_then(Value::as_str), + )?; + item.set_item("provider", deployment.spec.provider)?; + Ok(item) + }) + .collect::>>()?; + dict.set_item("built", PyList::new(py, records)?) + } + Finish::Counted { key } => dict.set_item("key", key), + Finish::RouterModelInfo(merged) => { + let py = dict.py(); + match merged { + Some(merged) => { + dict.set_item("info", to_py(py, &Value::Object(merged.info))?)?; + dict.set_item("user_keys", merged.user_keys) + } + None => dict.set_item("info", py.None()), + } + } + Finish::DeploymentModelInfo(info) => { + let py = dict.py(); + match info { + Some(info) => dict.set_item("info", to_py(py, &Value::Object(info))?), + None => dict.set_item("info", py.None()), + } + } + Finish::Value(value) => dict.set_item("value", to_py(dict.py(), &value)?), + Finish::Healthy { + single, + deployment_ids, + deployments, + } => { + dict.set_item("single", single)?; + dict.set_item("deployment_ids", deployment_ids)?; + let records = deployments + .iter() + .map(|d| to_py(dict.py(), d)) + .collect::>>()?; + dict.set_item("deployments", PyList::new(dict.py(), records)?) + } + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ffi::c_str, prelude::*, types::PyDict}; + use rstest::rstest; + + const CONFIG: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}} + ]}"#; + + const PAIR: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}} + ]}"#; + + #[rstest] + fn python_drives_calls_and_selections_through_the_handle() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("native", crate::native_module(py)).unwrap(); + locals.set_item("config", CONFIG).unwrap(); + locals.set_item("pair", PAIR).unwrap(); + py.run( + c_str!( + r#" +import time +store = {} +asked = [] + +def answer_cache(step): + asked.append(step) + op, key = step["op"], step["key"] + if op == "get": + return store.get(key) + if op == "batch_get": + return [store.get(k) for k in step["keys"]] + if op == "set": + store[key] = step["value"] + return None + if op == "increment": + store[key] = store.get(key, 0) + step["value"] + return store[key] + if op == "delete": + store.pop(key, None) + return None + raise AssertionError(op) + +def drive(call, outcome): + step = call.next(outcome) + while step["kind"] == "cache": + step = call.next(answer_cache(step)) + return step + +handle = native.RouterCoreHandle(config, {"num_retries": 1}) +picked = drive(handle.select("primary", {}), None) +assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked +assert [step["op"] for step in asked] == ["batch_get"] +assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "local", "events"], list(asked[0]) +call = handle.start("acompletion", '{"model": "primary"}') +version = handle.registry_version() +assert handle.apply({"delete": ["gone"]}) == version + 1 +steps = [drive(call, None)] +steps.append(drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})) +steps.append(drive(call, None)) +steps.append(drive(call, {"ok": True, "total_tokens": 5})) +kinds = [step["kind"] for step in steps] + +cooled = native.RouterCoreHandle(pair, {"num_retries": 0}) +call = cooled.start("acompletion", {"model": "primary"}) +attempt = drive(call, None) +tried = attempt["deployment_id"] +assert attempt["deployment"]["model_info"]["id"] == tried, attempt +assert attempt["deployment"]["litellm_params"]["model"].startswith("openai/"), attempt +assert [d["model_info"]["id"] for d in cooled.deployments()] == ["a1", "a2"] +assert cooled.deployment("a2")["model_name"] == "primary" +assert cooled.resolve("primary")["ids"] == ["a1", "a2"] +assert cooled.resolve("nope") is None +assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed" +# The core records the failure itself: the 429 cools the deployment down in the host's store. +assert f"deployment:{tried}:cooldown" in store, store +other = drive(cooled.select("primary", {}), None) +assert other["kind"] == "done" and other["deployment_id"] != tried, other + +call = handle.select("primary", {}) +step = call.next(None) +assert step["op"] == "batch_get" +for bad in ([], [None, None], "nope", None): + try: + call.next(bad) + except ValueError: + pass + else: + raise AssertionError(f"{bad!r} should not answer a batch_get of one key") +assert drive(call, [None])["kind"] == "done" +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let kinds: Vec = locals + .get_item("kinds") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!(kinds, ["attempt", "sleep", "attempt", "done"]); + }); + } + + #[rstest] + #[case::unknown_key(r#"{"routing_strategy_args": {}}"#, true)] + #[case::bad_json("{", false)] + #[case::not_an_object("[]", false)] + fn construction_failures_raise_the_class_python_reads( + #[case] config: &str, + #[case] declined: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let error = crate::native_module(py) + .getattr("RouterCoreHandle") + .and_then(|handle| handle.call1((config, "{}"))) + .unwrap_err(); + assert_eq!( + error.is_instance_of::(py), + declined, + "{error}" + ); + assert_eq!( + error.is_instance_of::(py), + !declined, + "{error}" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/retries.rs b/litellm-rust/crates/python-bridge/src/router/retries.rs new file mode 100644 index 00000000000..949aae96b34 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/retries.rs @@ -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 { + 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> { + 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> { + 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 { + 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, + }, + )) +} diff --git a/litellm-rust/crates/python-bridge/src/router/stream.rs b/litellm-rust/crates/python-bridge/src/router/stream.rs new file mode 100644 index 00000000000..9767182fd1e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/stream.rs @@ -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 { + serde_json::from_value(document(value)?) + .map_err(|error| PyValueError::new_err(format!("a stream origin: {error}"))) +} + +fn kind(kind: Option<&str>) -> PyResult> { + 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> { + let request = config::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 { + let op: Op = op.parse().unwrap_or_else(|_| Op::Other(op.to_owned())); + let request = config::request(document(request)?, None).map_err(config_error_to_pyerr)?; + let 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) + })) +} diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml index af97bd52c57..ff5e2e8161f 100644 --- a/litellm-rust/crates/router-core/Cargo.toml +++ b/litellm-rust/crates/router-core/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/router-core/src/access.rs b/litellm-rust/crates/router-core/src/access.rs new file mode 100644 index 00000000000..1b91fa417da --- /dev/null +++ b/litellm-rust/crates/router-core/src/access.rs @@ -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, +} + +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 { + 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 { + 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 { + 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 { + let active: Vec = 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>, + team_id: Option<&str>, + now: Timestamp, +) -> (Vec>, Option) { + 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>, + allowed_models: &[String], + model_groups: impl FnOnce() -> BTreeSet, +) -> Vec> { + 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 = 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) { + 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())); + } +} diff --git a/litellm-rust/crates/router-core/src/build.rs b/litellm-rust/crates/router-core/src/build.rs new file mode 100644 index 00000000000..d1bb29466f8 --- /dev/null +++ b/litellm-rust/crates/router-core/src/build.rs @@ -0,0 +1,382 @@ +//! 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>, + resolve_secrets: bool, + mut install: impl FnMut(&BuiltDeployment) -> Result<(), Error> + Send, +) -> Result, Error> { + let declared: Vec> = 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, + declared_id: Option<&str>, + duplicate_id: bool, +) -> Result, 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) -> Vec { + 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) -> Vec { + 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) -> Option { + 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>, + } + + #[async_trait] + impl Host for Recording { + async fn invoke(&self, _attempt: &Attempt) -> Result { + unreachable!("the build runs no attempt") + } + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + async fn allow_fallback(&self, _request: &RouterRequest, _target: &str) -> bool { + true + } + async fn sleep(&self, _duration: std::time::Duration) {} + fn observe(&self, _event: &RouterEvent) {} + fn cache(&self) -> Arc { + Arc::new(InMemoryRouterCache::default()) + } + async fn secrets(&self, _index: usize, refs: &[SecretRef]) -> Result, 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 { + 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 { + 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 { + self.steps + .lock() + .unwrap() + .push(format!("added {} {provider:?}", ask.index)); + Ok(Built::Entry(ask.entry.clone())) + } + } + + fn entries(value: Value) -> Vec> { + 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"); + } +} diff --git a/litellm-rust/crates/router-core/src/cache.rs b/litellm-rust/crates/router-core/src/cache.rs new file mode 100644 index 00000000000..b841e9184b0 --- /dev/null +++ b/litellm-rust/crates/router-core/src/cache.rs @@ -0,0 +1,145 @@ +//! [`RouterCache`]: the DualCache-style store all routing state lives in, under the keys, +//! encodings and TTLs the Python router uses, so Python and Rust workers share it. Method names +//! follow `litellm-cache`'s `BaseCache`, `BatchCache`, `CounterCache` and `DeleteCache`. + +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +use async_trait::async_trait; +use litellm_cache::Error; +use parking_lot::Mutex; +use serde_json::Value; + +#[async_trait] +pub trait RouterCache: Send + Sync { + async fn async_get_cache(&self, key: &str) -> Result, Error>; + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error>; + + /// Add `amount` to the number at `key` (0 when missing) and return the new value. The TTL + /// applies when the key is created, as Redis `INCRBYFLOAT` + `EXPIRE NX` does. + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result; + + async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error>; + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error>; + + /// Read this worker's in-memory tier only (`local_only=True`). + async fn async_get_local(&self, key: &str) -> Result, 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, + ) -> Result { + self.async_increment(key, amount, ttl).await + } +} + +#[derive(Default)] +struct Entry { + value: Value, + expires: Option, +} + +/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`. +#[derive(Clone, Default)] +pub struct InMemoryRouterCache { + entries: Arc>>, +} + +impl InMemoryRouterCache { + fn live(entries: &mut HashMap, key: &str) -> Option { + let expired = entries + .get(key) + .is_some_and(|entry| entry.expires.is_some_and(|at| at <= Instant::now())); + if expired { + entries.remove(key); + } + entries.get(key).map(|entry| entry.value.clone()) + } + + pub fn get(&self, key: &str) -> Option { + Self::live(&mut self.entries.lock(), key) + } + + pub fn increment(&self, key: &str, amount: f64, ttl: Option) -> f64 { + let mut entries = self.entries.lock(); + let current = Self::live(&mut entries, key).and_then(|value| value.as_f64()); + let next = current.unwrap_or(0.0) + amount; + let value = serde_json::Number::from_f64(next) + .map(|number| { + if next.fract() == 0.0 && next.abs() < i64::MAX as f64 { + Value::from(next as i64) + } else { + Value::Number(number) + } + }) + .unwrap_or(Value::Null); + let expires = match (current, entries.get(key)) { + (Some(_), Some(entry)) => entry.expires, + _ => ttl.map(|ttl| Instant::now() + ttl), + }; + entries.insert(key.to_owned(), Entry { value, expires }); + next + } +} + +#[async_trait] +impl RouterCache for InMemoryRouterCache { + async fn async_get_cache(&self, key: &str) -> Result, Error> { + Ok(self.get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error> { + let expires = ttl.map(|ttl| Instant::now() + ttl); + self.entries + .lock() + .insert(key.to_owned(), Entry { value, expires }); + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + Ok(self.increment(key, amount, ttl)) + } + + async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error> { + let mut entries = self.entries.lock(); + Ok(keys + .iter() + .map(|key| Self::live(&mut entries, key)) + .collect()) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + self.entries.lock().remove(key); + Ok(()) + } +} diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs new file mode 100644 index 00000000000..97b12dd958e --- /dev/null +++ b/litellm-rust/crates/router-core/src/call.rs @@ -0,0 +1,1286 @@ +//! The call loop: select, run the attempt, then retry in the same group or fall back to another +//! until it succeeds or the policy is exhausted. Mirrors `Router.async_function_with_fallbacks`, +//! `async_function_with_fallbacks_common_utils`, `async_function_with_retries`, +//! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and +//! `run_async_fallback` in `router_utils/fallback_event_handlers.py`. + +use std::{borrow::Cow, future::Future, pin::Pin, sync::Arc, time::Duration}; + +use serde_json::{Map, Value}; + +use tracing::warn; + +use crate::{ + Core, Snapshot, + config::Op, + error::{AttemptError, ConfigError, Error, FailureKind, Unsupported}, + fallback::{FallbackKeys, FallbackTarget, Fallbacks}, + health::{GroupShape, should_retry_status}, + host::{Attempt, Blocked, Output, PickError, RetryCause, RouterEvent, RouterRequest}, + mock::MockFailure, + prerouting::{PreRoute, PreRouted}, + registry::{ResolveKind, ResolveScope}, + retry::{RetryContext, resolve_retry_policy, should_retry}, + select::{Candidates, SelectInputs}, +}; + +mod stream; + +pub use self::stream::{StreamOrigin, StreamPolicy}; + +/// `INITIAL_RETRY_DELAY`, `MAX_RETRY_DELAY` and `JITTER` in `litellm/constants.py`. +const INITIAL_RETRY_DELAY: f64 = 0.5; +const MAX_RETRY_DELAY: f64 = 8.0; +const JITTER: f64 = 0.75; +/// `calculate_retry_after` obeys a provider's retry-after up to a minute. +const MAX_HONOURED_RETRY_AFTER: Duration = Duration::from_secs(60); + +#[derive(Clone, Debug)] +pub struct RouterResponse { + pub output: Output, + pub deployment_id: String, + pub model_group: String, + pub attempted_retries: u32, + /// The retry budget of the loop a retry succeeded in (`x-litellm-max-retries`); `None` + /// when the first attempt succeeded. + pub max_retries: Option, + pub attempted_fallbacks: u32, +} + +/// A failed attempt with the deployment it ran on, when one was picked. +struct Failed { + error: Error, + deployment_id: Option, +} + +impl From for Failed { + fn from(error: Error) -> Self { + Self { + error, + deployment_id: None, + } + } +} + +/// Where a level of the fallback walk sits: `depth` is `run_async_fallback`'s counter, one more +/// for every target tried, `level` how many fallbacks deep the walk is, and `scope` the level's +/// own copy of the request metadata (every hop copies its parent's). +#[derive(Clone, Copy, Debug, Default)] +struct Hop { + depth: u32, + level: u32, + scope: u32, +} + +/// What the whole walk shares: the attempts made (`attempted_targets`) and the scopes handed out. +#[derive(Debug, Default)] +struct Walk { + attempted: Vec, + scopes: u32, +} + +/// What one level of the walk serves: a group, or the same group at a later `order`, with the +/// request params the walk's dict targets merged so far (`kwargs.update(target)`). +#[derive(Clone, Debug)] +struct Target { + group: String, + order: Option, + params: Option>>, + /// `_excluded_deployment_ids` a weighted failover hop put on the request: never picked. + excluded: Arc<[String]>, + /// `_failover_excluded_ids` in the level's metadata: the deployments weighted failover + /// already excluded on the way here. + failover_excluded: Arc<[String]>, +} + +/// A step of a level's fallback chain: a configured target, or an order level of the group. +#[derive(Clone, Debug)] +enum Next { + Fallback(FallbackTarget), + Order(String, i64), + /// A weighted failover hop: the same group without the deployments already excluded. + Failover(String, Arc<[String]>), +} + +/// An attempt the walk already made (`fallback_attempt_key`). +#[derive(Clone, Debug, PartialEq)] +enum Tried { + Target(FallbackTarget), + Order(String, i64), + Failover(String, Vec), +} + +/// The fallback lists a call reads: the request's own where it passed them, else the router's. +struct Lists<'a> { + fallbacks: Option<&'a Fallbacks>, + context_window: Option<&'a Fallbacks>, + content_policy: Option<&'a Fallbacks>, + keys: &'a FallbackKeys, +} + +impl<'a> Lists<'a> { + fn of(snapshot: &'a Snapshot, request: &'a RouterRequest) -> Self { + let dynamic = &snapshot.settings.dynamic; + let pick = |own: &'a Option>, router: Option<&'a Fallbacks>| match own { + Some(own) => own.as_ref(), + None => router, + }; + Self { + fallbacks: pick(&request.fallbacks, Some(&dynamic.fallbacks)), + context_window: pick( + &request.context_window_fallbacks, + dynamic.context_window_fallbacks.as_ref(), + ), + content_policy: pick( + &request.content_policy_fallbacks, + dynamic.content_policy_fallbacks.as_ref(), + ), + keys: &dynamic.fallback_keys, + } + } + + fn has_regular(&self) -> bool { + self.fallbacks.is_some_and(|list| !list.is_empty()) + } +} + +/// `has_unattempted_fallback_target` over the chain `list` resolves for the level's lookup +/// groups: before the first hop nothing counts as tried, and a dict target with its own params +/// is never skipped. +fn unattempted_chain( + list: &Fallbacks, + lists: &Lists<'_>, + request: &RouterRequest, + target: &Target, + at: Hop, + attempted: &[Tried], +) -> bool { + let mut lookup = vec![target.group.as_str()]; + if request.model != target.group { + lookup.push(request.model.as_str()); + } + list.lookup_groups(&lookup, lists.keys) + .is_some_and(|targets| { + at.level == 0 + || targets.iter().any(|next| match next.identity() { + FallbackTarget::Params(_) => true, + identity => !attempted.contains(&Tried::Target(identity)), + }) + }) +} + +/// The deployments `_core_get_model_list(model_name=group)` lists: the group's own, or a routing +/// group's members. +fn group_members(snapshot: &Snapshot, group: &str) -> Vec> { + snapshot + .registry + .group(group) + .or_else(|| { + snapshot + .registry + .resolve(group, ResolveScope::default()) + .filter(|resolution| resolution.kind == ResolveKind::RoutingGroup) + .map(|resolution| resolution.deployments) + }) + .unwrap_or_default() +} + +/// Whether the op's attempt stamps `failed_deployment_id` on its error (`_completion`, +/// `_acompletion` and the generic API calls), which weighted failover needs. +fn stamps_failed_deployment(op: &Op) -> bool { + match op { + Op::Embedding + | Op::ImageGeneration + | Op::Transcription + | Op::Speech + | Op::TextCompletion => false, + Op::Other(name) => !matches!( + name.as_str(), + "arerank" | "embedding" | "image_generation" | "text_completion" + ), + Op::Completion | Op::Responses | Op::Messages | Op::Ocr => true, + } +} + +/// `is_bad_request_without_fallback` with `fallback_on_unsupported_params` off: an unsupported +/// params error, or a provider's plain 400 bad request. +fn is_bad_request_without_fallback(failure: &AttemptError) -> bool { + if !failure.is_a("BadRequestError") { + return false; + } + if failure.is_a("UnsupportedParamsError") { + return true; + } + if [ + "ContextWindowExceededError", + "ContentPolicyViolationError", + "UnsupportedEndpointError", + ] + .iter() + .any(|class| failure.is_a(class)) + { + return false; + } + failure.status == Some(400) && failure.names_provider +} + +type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// `KeyError`'s classes, most specific first. +/// The failure a retry loop decides by: an attempt's, or the exception a failed pick raised +/// (Python's loop retries a pick's error like any other: no deployment available, a stale +/// routing group's `KeyError`, ...). +fn retried_failure(error: &Error) -> Option> { + match error { + Error::Attempt(failure) => Some(Cow::Borrowed(failure)), + _ => error.pick_failure().map(Cow::Owned), + } +} + +impl Core { + pub(crate) async fn call( + &self, + op: Op, + request: Arc, + ) -> Result { + let snapshot = self.snapshot.load_full(); + let mut walk = Walk::default(); + let target = Target { + group: request.model.clone(), + order: None, + params: None, + excluded: Arc::from([]), + failover_excluded: Arc::from([]), + }; + self.with_fallbacks( + &snapshot, + &op, + &request, + target, + Hop::default(), + &mut walk, + None, + ) + .await + } + + /// One level of `async_function_with_fallbacks`: the target's retry loop, then its + /// fallbacks. A level that fails raises its own error, whatever its fallbacks raised. + /// `resumed` is a failure the level already had (a stream that failed after its attempt + /// returned): its fallbacks run without the retry loop. + #[allow(clippy::too_many_arguments)] + fn with_fallbacks<'a>( + &'a self, + snapshot: &'a Snapshot, + op: &'a Op, + request: &'a Arc, + target: Target, + at: Hop, + walk: &'a mut Walk, + resumed: Option, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mock = request + .mock_failure + .filter(|mock| mock.before_retries() && at.level == 0); + let (error, failed_id) = match (resumed, mock) { + (Some(error), _) => (error, None), + (None, Some(mock)) => (self.mock_failed(mock, &target, at, None), None), + (None, None) => match self + .with_retries(snapshot, op, request, &target, at, &walk.attempted) + .await + { + Ok(response) => return Ok(response), + Err(failed) => (failed.error, failed.deployment_id), + }, + }; + self.level_failed(at, &target, &error, false); + if self.raises_without_fallback(snapshot, request, &error) { + return Err(error); + } + // `async_function_with_fallbacks_common_utils` lists the group's deployments + // (`_core_get_model_list`) for its order fallbacks: a written index position past the + // model list raises there, before any fallback. + if let Err(raised) = crate::resolve::all_deployments( + &snapshot.registry, + &target.group, + request.team_id.as_deref(), + ) { + self.level_failed(at, &target, &raised, true); + return Err(raised); + } + let (targets, order_path) = self.fallback_targets(snapshot, request, &target, &error); + let mut target = target; + // `_maybe_run_weighted_failover`: its own `run_async_fallback`, before the regular one. + let failover = if order_path { + None + } else { + match self + .weighted_failover(snapshot, op, request, &target, &error, failed_id.as_deref()) + .await + { + Ok(failover) => failover, + Err(raised) => { + // `_get_routing_context` raised inside the level's fallback handling: the + // level raises that error, before any regular fallback. + self.level_failed(at, &target, &raised, true); + return Err(raised); + } + } + }; + if let Some(failover) = failover { + let excluded = match &failover { + Next::Failover(_, excluded) => excluded.clone(), + _ => Arc::from([]), + }; + target.failover_excluded = excluded; + if let Ok(response) = self + .run_fallbacks(snapshot, op, request, &target, at, walk, vec![failover]) + .await + { + return Ok(response); + } + } + let Some(targets) = targets else { + return Err(error); + }; + match self + .run_fallbacks(snapshot, op, request, &target, at, walk, targets) + .await + { + Ok(response) => Ok(response), + Err(()) => Err(error), + } + }) + } + + /// `run_async_fallback` over one chain: the depth is checked once on entry, then every + /// target tried counts one more. `Err(())` when the chain served nothing. + #[allow(clippy::too_many_arguments)] + fn run_fallbacks<'a>( + &'a self, + snapshot: &'a Snapshot, + op: &'a Op, + request: &'a Arc, + target: &'a Target, + at: Hop, + walk: &'a mut Walk, + targets: Vec, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let max_fallbacks = request + .max_fallbacks + .unwrap_or(snapshot.settings.max_fallbacks); + if at.depth >= max_fallbacks { + return Err(()); + } + let failed = Tried::Target(FallbackTarget::Name(target.group.clone())); + if !walk.attempted.contains(&failed) { + walk.attempted.push(failed); + } + let mut params = target.params.clone(); + let mut hop = at.depth; + for next in targets { + let (tried, group, order) = match &next { + Next::Order(group, order) => ( + Tried::Order(group.clone(), *order), + group.clone(), + Some(*order), + ), + Next::Failover(group, excluded) => { + let mut key = excluded.to_vec(); + key.sort(); + (Tried::Failover(group.clone(), key), group.clone(), None) + } + Next::Fallback(FallbackTarget::Name(name)) if *name == target.group => continue, + Next::Fallback(fallback) => { + if let Some(model) = fallback.model() + && model != target.group + && !self.host.allow_fallback(request, model).await + { + continue; + } + let group = fallback.model().unwrap_or(&target.group).to_owned(); + (Tried::Target(fallback.identity()), group, None) + } + }; + if walk.attempted.contains(&tried) { + continue; + } + walk.attempted.push(tried); + if let Next::Fallback(fallback) = &next + && let Some(merged) = fallback.params() + { + let mut combined = params.as_deref().cloned().unwrap_or_default(); + combined.extend(merged.iter().map(|(k, v)| (k.clone(), v.clone()))); + params = Some(Arc::new(combined)); + } + hop += 1; + self.host.observe(&RouterEvent::RetryLogged { + scope: at.scope, + level: at.level, + model: target.group.clone(), + attempted_retries: None, + cause: RetryCause::Level, + }); + walk.scopes += 1; + self.host.observe(&RouterEvent::Fallback { + from: target.group.clone(), + to: group.clone(), + order, + depth: hop, + level: at.level + 1, + scope: walk.scopes, + parent: at.scope, + max_fallbacks, + }); + let excluded = match &next { + Next::Failover(_, excluded) => excluded.clone(), + _ => target.excluded.clone(), + }; + let next = Target { + group: group.clone(), + order, + params: params.clone(), + excluded, + failover_excluded: target.failover_excluded.clone(), + }; + let parent = at; + let at = Hop { + depth: hop, + level: at.level + 1, + scope: walk.scopes, + }; + let result = self + .with_fallbacks(snapshot, op, request, next, at, walk, None) + .await; + self.host.observe(&RouterEvent::FallbackDone { + from: target.group.clone(), + to: group, + level: parent.level, + scope: at.scope, + parent: parent.scope, + succeeded: result.is_ok(), + }); + if let Ok(mut response) = result { + response.attempted_fallbacks = hop; + return Ok(response); + } + } + Err(()) + }) + } + + /// `_maybe_run_weighted_failover`'s hop for a failed deployment, when it applies: the + /// simple-shuffle group re-picked without every deployment already excluded, while one + /// that is not cooling down remains. + async fn weighted_failover( + &self, + snapshot: &Snapshot, + op: &Op, + request: &RouterRequest, + target: &Target, + error: &Error, + failed_id: Option<&str>, + ) -> Result, Error> { + if !snapshot.settings.enable_weighted_failover { + return Ok(None); + } + if retried_failure(error).is_some_and(|failure| { + failure.is_a("ContextWindowExceededError") + || failure.is_a("ContentPolicyViolationError") + }) { + return Ok(None); + } + // `_get_routing_context` runs first, whatever failed: a stale routing-group member map + // raises its `KeyError` here. + snapshot.missing_routing_group(request, &target.group)?; + if !stamps_failed_deployment(op) || error.attempt().is_none() { + return Ok(None); + } + let kind = snapshot + .registry + .resolve(&target.group, ResolveScope::default()) + .map_or(ResolveKind::Group, |resolution| resolution.kind); + let (strategy, _) = crate::strategy::routing_context( + &snapshot.registry, + &snapshot.settings.routing_strategy, + request.routing_strategy.as_deref(), + &kind, + &target.group, + ); + if strategy != crate::config::RoutingStrategy::SimpleShuffle { + return Ok(None); + } + let Some(failed_id) = failed_id else { + return Ok(None); + }; + let mut excluded: Vec = target.failover_excluded.to_vec(); + if !excluded.iter().any(|id| id == failed_id) { + excluded.push(failed_id.to_owned()); + } + let ids: Vec = group_members(snapshot, &target.group) + .iter() + .filter(|d| d.usable_by_team(request.team_id.as_deref())) + .map(|d| d.id.clone()) + .collect(); + let cooling = self.health.active_cooldowns(&ids).await.unwrap_or_default(); + let remaining = ids + .iter() + .any(|id| !excluded.contains(id) && !cooling.iter().any(|(cooled, _)| cooled == id)); + Ok(remaining.then(|| Next::Failover(target.group.clone(), excluded.into()))) + } + + /// `async_function_with_fallbacks_common_utils` raises before looking for any fallback: + /// fallbacks are off, the error is a guardrail's, a bad request without + /// `fallback_on_unsupported_params`, or no fallback serves the error at all. + fn raises_without_fallback( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + error: &Error, + ) -> bool { + if request.disable_fallbacks { + return true; + } + match error { + Error::Attempt(failure) => { + failure.kind == FailureKind::GuardrailIntervention + || (!snapshot.settings.fallback_on_unsupported_params + && is_bad_request_without_fallback(failure)) + } + Error::NoHealthyDeployments { .. } + | Error::UnknownModel { .. } + | Error::RoutingPluginsLeftNothing { .. } + | Error::MissingRoutingGroup { .. } + | Error::DefaultDeploymentRaises { .. } + | Error::IndexOutOfRange { .. } + | Error::ContextWindowExceeded { .. } => false, + Error::UnsupportedParams { .. } + | Error::RoutingPluginsNeedAsync + | Error::InvalidRouterWeights + | Error::NonNumericWeight { .. } + | Error::Config(_) + | Error::Cache(_) + | Error::Retired + | Error::Paused { .. } + | Error::NotPassThrough { .. } + | Error::NoPassThroughDeployments { .. } + | Error::OnlyMarkers { .. } + | Error::PtuReserved { .. } + | Error::MultipleTeams { .. } + | Error::Reserved { .. } + | Error::ModelInfo { .. } => true, + } + } + + /// The fallback targets for a failure, or `None` when it must be raised as is + /// (`async_function_with_fallbacks_common_utils`). + fn fallback_targets( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + target: &Target, + error: &Error, + ) -> (Option>, bool) { + // A plain `ValueError` (routing plugins left nothing) takes the generic fallbacks, as a + // missing group does; a pre-call context window refusal takes the context window ones. + // Every other error was raised before (`raises_without_fallback`). + let failure = match error { + Error::Attempt(failure) => Some(failure), + _ => None, + }; + let context_window = matches!(error, Error::ContextWindowExceeded { .. }); + let is_a = |name: &str| { + failure.is_some_and(|f| f.is_a(name)) + || (context_window && name == "ContextWindowExceededError") + }; + let lists = Lists::of(snapshot, request); + // `fallback_lookup_groups`: the group being served, then the group the request was for. + let mut lookup = vec![target.group.as_str()]; + if request.model != target.group { + lookup.push(request.model.as_str()); + } + let chain = |list: &Fallbacks| { + if list.is_client_side(lists.keys) { + Some(list.as_chain()) + } else { + list.lookup_groups(&lookup, lists.keys) + } + }; + let fallbacks_of = |targets: Vec| { + targets.into_iter().map(Next::Fallback).collect::>() + }; + let skip_order = is_a("ContextWindowExceededError") || is_a("ContentPolicyViolationError"); + if !skip_order { + // A routing group's members, with a priority group's priorities as their order. + let mut orders: Vec = group_members(snapshot, &target.group) + .iter() + .filter(|d| d.usable_by_team(request.team_id.as_deref())) + .filter_map(|d| d.order()) + .collect(); + orders.sort_unstable(); + orders.dedup(); + if orders.len() > 1 { + let current = target.order.unwrap_or(orders[0]); + let mut targets: Vec = orders + .into_iter() + .filter(|order| *order > current) + .map(|order| Next::Order(target.group.clone(), order)) + .collect(); + if let Some(list) = lists.fallbacks { + targets.extend(fallbacks_of(chain(list).unwrap_or_default())); + } + if !targets.is_empty() { + return (Some(targets), true); + } + } + } + if let Some(list) = lists.fallbacks + && list.is_client_side(lists.keys) + { + return (Some(fallbacks_of(list.as_chain())), false); + } + let specific = if is_a("ContextWindowExceededError") { + lists.context_window + } else if is_a("ContentPolicyViolationError") { + lists.content_policy + } else { + None + }; + if let Some(list) = specific { + return ( + list.lookup_groups(&lookup, lists.keys).map(fallbacks_of), + false, + ); + } + ( + lists + .fallbacks + .and_then(|list| list.lookup_groups(&lookup, lists.keys)) + .map(fallbacks_of), + false, + ) + } + + async fn with_retries( + &self, + snapshot: &Snapshot, + op: &Op, + request: &Arc, + target: &Target, + at: Hop, + attempted: &[Tried], + ) -> Result { + let walk = (at, attempted); + let mut inputs = SelectInputs { + target_order: target.order, + excluded_ids: target.excluded.clone(), + ..SelectInputs::default() + }; + let counting = SelectInputs { + count_only: true, + ..SelectInputs::default() + }; + let settings = &snapshot.settings; + let first_budget = request.num_retries.unwrap_or(settings.num_retries); + let first = match request.mock_failure { + Some(MockFailure::RateLimit) => { + // `_handle_mock_testing_rate_limit_error`: a lone deployment's `num_retries` + // rides on the mock error, as on a real one. + let lone = snapshot + .registry + .group(&target.group) + .filter(|group| group.len() == 1) + .and_then(|group| group[0].raw_params.get("num_retries")?.as_u64()) + .and_then(|retries| u32::try_from(retries).ok()); + Err(self + .mock_failed(MockFailure::RateLimit, target, at, lone) + .into()) + } + _ => { + self.attempt( + snapshot, + op, + request, + target, + &inputs, + (0, first_budget), + walk, + ) + .await + } + }; + let mut failed = match first { + Ok(response) => return Ok(response), + Err(failed) => failed, + }; + let guardrail = matches!( + &failed.error, + Error::Attempt(failure) if failure.kind == FailureKind::GuardrailIntervention + ); + // The first failure looks the group's retry policy up; a lookup that raises replaces the + // failure with its own error, before any retry (`async_function_with_retries`). + if !guardrail && request.retry_policy_raises.contains(&target.group) { + self.host.observe(&RouterEvent::RetryPolicyRaised { + level: at.level, + model: target.group.clone(), + cause: RetryCause::of(&failed.error), + }); + return Err(Failed { + error: Error::Attempt(AttemptError::host_raised(&["Exception"])), + deployment_id: None, + }); + } + if guardrail { + return Err(failed); + } + let Some(first_failure) = retried_failure(&failed.error) else { + return Err(failed); + }; + let failure: &AttemptError = &first_failure; + let mut num_retries = request + .num_retries + .or(failure.num_retries) + .unwrap_or(settings.num_retries); + let group_policies = request + .model_group_retry_policy + .as_ref() + .unwrap_or(&settings.model_group_retry_policy); + let policy_retries = (request.num_retries != Some(0)) + .then(|| { + resolve_retry_policy( + settings.retry_policy.as_ref(), + &target.group, + Some(group_policies), + )? + .retries_for(failure) + }) + .flatten(); + let policy_applies = policy_retries.is_some(); + num_retries = policy_retries.unwrap_or(num_retries); + drop(first_failure); + // `_async_get_healthy_deployments` finds no group for a deployment id, so a call by id + // counts no deployments when deciding whether to retry. + let counted = |c: &Candidates, n: usize| if c.kind.is_single() { 0 } else { n }; + let all = self + .candidates(snapshot, op, request, &target.group, &counting) + .await + .map(|c| counted(&c, c.all.len())) + .unwrap_or(0); + + for retry in 0..=num_retries { + if retry > 0 { + self.log_retry(target, at, retry, &failed.error); + } + let healthy_ids: Vec = self + .candidates(snapshot, op, request, &target.group, &counting) + .await + .map(|c| { + let ids = c.healthy.iter().map(|d| d.id.clone()); + ids.take(counted(&c, c.healthy.len())).collect() + }) + .unwrap_or_default(); + let healthy = healthy_ids.len(); + // A pick's error is retried as Python's loop retries any exception: a retry policy + // keeps going, else `should_retry_this_error` decides (it raises when nothing is + // healthy). + let Some(failure) = retried_failure(&failed.error) else { + return Err(failed); + }; + if !policy_applies && !self.should_retry(snapshot, request, &failure, healthy, all) { + return Err(failed); + } + if num_retries == 0 { + return Err(failed); + } + if retry == 0 { + self.log_retry(target, at, 0, &failed.error); + } + if let (Some(status), Some(id)) = (failure.status, &failed.deployment_id) + && !should_retry_status(status) + && !inputs.skipped_ids.contains(id) + { + inputs.skipped_ids.push(id.clone()); + } + // An endpoint a deployment does not serve is retried on the others, until every + // one refused (`_every_deployment_refused`). + if failure.is_a("UnsupportedEndpointError") + && healthy_ids.iter().all(|id| inputs.skipped_ids.contains(id)) + { + return Err(failed); + } + let delay = self.time_to_sleep( + snapshot, + failure.retry_after, + num_retries - retry, + num_retries, + healthy, + all, + ); + self.host.sleep(delay).await; + if retry == num_retries { + self.host.observe(&RouterEvent::RetriesExhausted { + model_group: target.group.clone(), + attempted: num_retries, + max: num_retries, + level: at.level, + picked: failed.error.attempt().is_none(), + }); + break; + } + self.host.observe(&RouterEvent::Retry { + model_group: target.group.clone(), + retry: retry + 1, + delay, + }); + match self + .attempt( + snapshot, + op, + request, + target, + &inputs, + (retry + 1, num_retries), + walk, + ) + .await + { + Ok(mut response) => { + response.max_retries = Some(num_retries); + return Ok(response); + } + Err(next) => failed = next, + } + } + Err(failed) + } + + /// The level serving `target` raises a pick's `error` ([`RouterEvent::LevelFailed`]); + /// `handling` when its fallback handling raised it. + fn level_failed(&self, at: Hop, target: &Target, error: &Error, handling: bool) { + if let Some(picked) = PickError::of(error) { + self.host.observe(&RouterEvent::LevelFailed { + level: at.level, + model: target.group.clone(), + error: picked, + handling, + }); + } + } + + /// `_core_log_retry` for a failure in the level's retry loop. + fn log_retry(&self, target: &Target, at: Hop, retry: u32, error: &Error) { + let cause = match RetryCause::of(error) { + RetryCause::Other => RetryCause::Attempt, + cause => cause, + }; + self.host.observe(&RouterEvent::RetryLogged { + scope: at.scope, + level: at.level, + model: target.group.clone(), + attempted_retries: Some(retry), + cause, + }); + } + + /// A `mock_testing_*` failure for the level's group, reported to the host, which raises + /// the matching exception. + fn mock_failed( + &self, + mock: MockFailure, + target: &Target, + at: Hop, + num_retries: Option, + ) -> Error { + self.host.observe(&RouterEvent::MockFailed { + mock, + level: at.level, + model: target.group.clone(), + num_retries, + }); + Error::Attempt(mock.failure(num_retries)) + } + + fn should_retry( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + failure: &AttemptError, + healthy: usize, + all: usize, + ) -> bool { + let lists = Lists::of(snapshot, request); + should_retry( + failure, + RetryContext { + healthy, + all, + has_context_window_fallbacks: lists.context_window.is_some(), + has_content_policy_fallbacks: lists.content_policy.is_some(), + has_regular_fallbacks: lists.has_regular(), + }, + ) + } + + /// `_time_to_sleep_before_retry` and `calculate_retry_after`. + fn time_to_sleep( + &self, + snapshot: &Snapshot, + retry_after: Option, + remaining_retries: u32, + num_retries: u32, + healthy: usize, + all: usize, + ) -> Duration { + if all != 1 && healthy > 0 { + return Duration::ZERO; + } + let jitter = JITTER * self.rng.next_f64(); + if let Some(retry_after) = retry_after + && !retry_after.is_zero() + && retry_after <= MAX_HONOURED_RETRY_AFTER + { + return retry_after + Duration::from_secs_f64(jitter); + } + let exponent = num_retries.saturating_sub(remaining_retries); + let backoff = (INITIAL_RETRY_DELAY * 2f64.powi(exponent as i32)) + .max(snapshot.settings.retry_after.as_secs_f64()) + .min(MAX_RETRY_DELAY); + Duration::from_secs_f64(backoff + jitter) + } + + /// One attempt: select a deployment, run it through the host, record the outcome. + #[allow(clippy::too_many_arguments)] + async fn attempt( + &self, + snapshot: &Snapshot, + op: &Op, + request: &Arc, + target: &Target, + inputs: &SelectInputs, + (retry, max_retries): (u32, u32), + (at, attempted): (Hop, &[Tried]), + ) -> Result { + if request.unrouted { + return self + .unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at) + .await; + } + crate::routing_plugins::refuse_sync_pick( + snapshot.settings.routing_plugins, + &snapshot.settings.routing_strategy, + request, + )?; + let routed = self + .pre_routed(PreRoute { + op, + request, + name: &target.group, + fallback_depth: at.depth, + retry, + }) + .await?; + let (deployment, candidates) = match routed { + PreRouted::Model(name) => { + snapshot.missing_routing_group(request, &name)?; + let candidates = self + .candidates(snapshot, op, request, &name, inputs) + .await?; + if candidates.healthy.is_empty() { + return Err(Error::NoHealthyDeployments { + model: candidates.model_group, + } + .into()); + } + let deployment = self.pick(snapshot, op, request, &candidates).await?; + (deployment, candidates) + } + // A custom strategy's pick is served as is, for the requested name. + PreRouted::Deployment(deployment) => { + let candidates = Candidates { + name: target.group.clone(), + model_group: target.group.clone(), + kind: ResolveKind::Group, + all: vec![deployment.clone()], + healthy: vec![deployment.clone()], + }; + (deployment, candidates) + } + }; + self.host.observe(&RouterEvent::Selected { + model_group: candidates.model_group.clone(), + deployment_id: deployment.id.clone(), + }); + if !self.host.supports(op, deployment.provider()) { + return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op { + op: op.clone(), + provider: deployment.provider().to_owned(), + })) + .into()); + } + // A name no deployment serves is routed to the first default fallback + // (`_common_checks_available_deployment`), still stamped as the name. + let defaulted = candidates.model_group != target.group + && matches!( + crate::resolve::resolve_request( + &snapshot.registry, + &target.group, + request.specific_deployment, + crate::resolve::Caller { + team_id: request.team_id.as_deref(), + proxy_admin: request.proxy_admin, + allowed_models: &request.allowed_models, + }, + None, + snapshot.settings.dynamic.default_ref(), + jiff::Timestamp::now(), + ), + Ok(None) + ); + let attempt = Attempt { + op: op.clone(), + request: request.clone(), + deployment: deployment.clone(), + name: target.group.clone(), + model_group: if defaulted { + target.group.clone() + } else { + candidates.model_group.clone() + }, + target_order: target.order, + candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(), + retry, + fallback_depth: at.depth, + fallback_level: at.level, + timeout: snapshot.settings.timeout, + params: target.params.clone(), + scope: at.scope, + max_retries, + model_group_size: snapshot.registry.model_list_len(&target.group), + }; + // Every attempt runs through the host; the core records every outcome itself, except a + // deferred stream's: its request counts now (`increment_deployment_usage_for_response`), + // its success or failure when the host reports the stream's end. + match self.host.invoke(&attempt).await { + Ok(output) => { + let recorded = if output.deferred { + self.health + .add_usage(&deployment, output.total_tokens as f64, 1.0) + .await + } else { + self.health + .record_success(&deployment, output.total_tokens) + .await + }; + if let Err(error) = recorded { + warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded"); + } + self.host.observe(&RouterEvent::Succeeded { + deployment_id: deployment.id.clone(), + }); + // `_should_raise_content_policy_error`: a blocked response is raised only when + // a content-policy fallback would serve it, and counts as no failure. + let raised = + match output.blocked { + Some(Blocked::ContentFilter) => self + .has_content_policy_fallback(snapshot, request, target, at, attempted), + Some(Blocked::Refusal) => self + .refusal_fallback_available(snapshot, request, target, at, attempted), + None => false, + }; + if let Some(blocked) = output.blocked.filter(|_| raised) { + self.host.observe(&RouterEvent::Blocked { + level: at.level, + model: target.group.clone(), + blocked, + }); + let mut failure = + AttemptError::new(FailureKind::ContentPolicyViolation, Some(400)); + failure.classes = crate::retry::content_policy_classes(); + return Err(Failed { + error: Error::Attempt(failure), + deployment_id: Some(deployment.id.clone()), + }); + } + Ok(RouterResponse { + output, + deployment_id: deployment.id.clone(), + model_group: candidates.model_group, + attempted_retries: retry, + max_retries: None, + attempted_fallbacks: at.depth, + }) + } + Err(failure) => { + if !failure.recorded { + self.record_failure(snapshot, request, &candidates, &deployment, &failure) + .await; + } + Err(Failed { + error: Error::Attempt(failure), + deployment_id: Some(deployment.id.clone()), + }) + } + } + } + + /// An attempt on a name that is no model group: the host serves it under that name, with no + /// deployment picked and no outcome recorded. + async fn unrouted_attempt( + &self, + snapshot: &Snapshot, + op: &Op, + request: &Arc, + target: &Target, + (retry, max_retries): (u32, u32), + at: Hop, + ) -> Result { + let deployment = + Arc::new(crate::registry::Deployment::unrouted(&target.group).map_err(Error::from)?); + let attempt = Attempt { + op: op.clone(), + request: request.clone(), + deployment, + name: target.group.clone(), + model_group: target.group.clone(), + target_order: target.order, + candidate_ids: Vec::new(), + retry, + fallback_depth: at.depth, + fallback_level: at.level, + timeout: snapshot.settings.timeout, + params: target.params.clone(), + scope: at.scope, + max_retries, + model_group_size: snapshot.registry.model_list_len(&target.group), + }; + match self.host.invoke(&attempt).await { + Ok(output) => Ok(RouterResponse { + output, + deployment_id: String::new(), + model_group: target.group.clone(), + attempted_retries: retry, + max_retries: None, + attempted_fallbacks: at.depth, + }), + Err(failure) => Err(Error::Attempt(failure).into()), + } + } + + /// `_has_content_policy_fallback`: the content-policy chain resolves for the level's + /// lookup groups and holds a target the walk has not tried; with no content-policy list, + /// any router `"*"` entry. + fn has_content_policy_fallback( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + target: &Target, + at: Hop, + attempted: &[Tried], + ) -> bool { + let lists = Lists::of(snapshot, request); + match lists.content_policy { + Some(list) => unattempted_chain(list, &lists, request, target, at, attempted), + None => snapshot.settings.dynamic.fallbacks.has_star(), + } + } + + /// `_refusal_fallback_available`: the content-policy list decides when there is one; + /// else a router `"*"` entry, else a generic chain with a target the walk has not tried. + fn refusal_fallback_available( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + target: &Target, + at: Hop, + attempted: &[Tried], + ) -> bool { + if request.disable_fallbacks { + return false; + } + let lists = Lists::of(snapshot, request); + if lists.content_policy.is_some() { + return self.has_content_policy_fallback(snapshot, request, target, at, attempted); + } + if snapshot.settings.dynamic.fallbacks.has_star() { + return true; + } + lists + .fallbacks + .is_some_and(|list| unattempted_chain(list, &lists, request, target, at, attempted)) + } + + async fn record_failure( + &self, + snapshot: &Snapshot, + request: &RouterRequest, + candidates: &Candidates, + deployment: &crate::registry::Deployment, + failure: &AttemptError, + ) { + self.host.observe(&RouterEvent::AttemptFailed { + deployment_id: deployment.id.clone(), + kind: failure.kind, + status: failure.status, + }); + // The default deployment's copy is no configured deployment: Python's cooldown check finds + // no group for its id and records nothing, unless the id names a configured deployment, + // which then takes the failure. + let configured; + let deployment = if candidates.kind == ResolveKind::Default { + match snapshot.registry.deployment(&deployment.id) { + Some(found) => { + configured = found; + &*configured + } + None => { + if let Err(error) = self.health.count_failed_request(deployment).await { + warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded"); + } + return; + } + } + } else { + deployment + }; + let shape = GroupShape { + // The deployment's own group decides, also when the call named the deployment by id. + single_deployment: snapshot + .registry + .group(&deployment.model_name) + .map_or(candidates.all.len() == 1, |group| group.len() == 1) + // A routing group call with alternatives cools its member down so selection + // moves on (`routing_group_alternatives`). + && !(candidates.kind == ResolveKind::RoutingGroup && candidates.all.len() > 1), + has_fallback_group: snapshot + .settings + .dynamic + .fallback_targets(&[candidates.model_group.as_str()]) + .is_some(), + }; + if let Err(error) = self.health.count_failed_request(deployment).await { + warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded"); + } + match self + .health + .record_failure( + &snapshot.settings, + deployment, + shape, + request.cooldown_time, + failure, + ) + .await + { + Ok(Some(cooldown)) => self.host.observe(&RouterEvent::CooldownSet { + deployment_id: deployment.id.clone(), + cooldown, + }), + Ok(None) => {} + Err(error) => { + warn!(deployment_id = %deployment.id, error = %error, "cooldown was not recorded"); + } + } + } +} diff --git a/litellm-rust/crates/router-core/src/call/stream.rs b/litellm-rust/crates/router-core/src/call/stream.rs new file mode 100644 index 00000000000..c4046f74054 --- /dev/null +++ b/litellm-rust/crates/router-core/src/call/stream.rs @@ -0,0 +1,413 @@ +//! 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::{Hop, Target, Tried, Walk}; +use crate::{ + Router, Snapshot, + config::{Op, RouterDynamicConfig}, + error::{AttemptError, Error, FailureKind}, + fallback::{FallbackTarget, Fallbacks}, + health::should_retry_status, + host::RouterRequest, + 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, + /// `fallback_depth` of the attempt. + pub depth: u32, + pub level: u32, + /// The deployment the attempt ran on. + pub deployment_id: Option, + /// 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, +} + +/// 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>) -> 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>) -> Option> { + chain.map(|targets| { + targets + .iter() + .filter_map(|target| target.model().map(str::to_owned)) + .collect() + }) +} + +/// The regular fallback targets for the lookup groups, then `default_fallbacks`, which Python +/// folds into `fallbacks` as `{"*": [...]}`. +fn regular_targets(dynamic: &RouterDynamicConfig, lookup: &[&str]) -> Option> { + names( + dynamic + .fallbacks + .lookup_groups(lookup, &dynamic.fallback_keys), + ) + .or_else(|| (!dynamic.default_fallbacks.is_empty()).then(|| dynamic.default_fallbacks.clone())) +} + +fn has_default_fallbacks(fallbacks: &Fallbacks, default_fallbacks: &[String]) -> bool { + fallbacks.has_star() || !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, + ) -> 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 walk 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 { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + let request = Arc::new(request); + let mut walk = Walk::default(); + if request.model != origin.group { + walk.attempted + .push(Tried::Target(FallbackTarget::Name(request.model.clone()))); + } + self.core + .with_fallbacks( + &snapshot, + &op, + &request, + origin.target(), + Hop { + depth: origin.depth, + level: origin.level, + scope: 0, + }, + &mut walk, + Some(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 { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + let request = Arc::new(request); + let target = origin.target(); + let at = Hop { + 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, + ..SelectInputs::default() + }; + let mut last = failure; + for attempt in origin.attempted_retries..budget { + 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; + 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: &Snapshot, request: &RouterRequest, origin: &StreamOrigin) -> u32 { + if let Some(retries) = request.num_retries { + return retries; + } + 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`. +fn retry_policy<'a>( + snapshot: &'a Snapshot, + request: &'a RouterRequest, + origin: &StreamOrigin, +) -> Option<&'a crate::config::RetryPolicy> { + if request.num_retries == Some(0) { + 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: &Snapshot, + 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: &Snapshot, 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: &Snapshot, + request: &RouterRequest, + origin: &StreamOrigin, +) -> bool { + if request.disable_fallbacks { + return false; + } + if snapshot.settings.enable_weighted_failover { + return true; + } + let mut orders: Vec = snapshot + .registry + .group(&origin.group) + .unwrap_or_default() + .iter() + .filter(|d| d.usable_by_team(request.team_id.as_deref())) + .filter_map(|d| d.order()) + .collect(); + orders.sort_unstable(); + orders.dedup(); + if orders.len() > 1 { + let current = origin.order.unwrap_or(orders[0]); + if orders.iter().any(|order| *order > current) { + return true; + } + } + let dynamic = &snapshot.settings.dynamic; + let lookup = origin.lookup(request); + if let Some(content_policy) = &dynamic.content_policy_fallbacks + && origin.has_unattempted( + request, + names(content_policy.lookup_groups(&lookup, &dynamic.fallback_keys)), + ) + { + return true; + } + if !dynamic.has_regular_fallbacks() { + return false; + } + origin.has_unattempted(request, regular_targets(dynamic, &lookup)) +} + +fn refusal_fallback_available( + snapshot: &Snapshot, + request: &RouterRequest, + origin: &StreamOrigin, +) -> bool { + if request.disable_fallbacks { + return false; + } + let dynamic = &snapshot.settings.dynamic; + let lookup = origin.lookup(request); + if let Some(content_policy) = &dynamic.content_policy_fallbacks { + return origin.has_unattempted( + request, + names(content_policy.lookup_groups(&lookup, &dynamic.fallback_keys)), + ); + } + if has_default_fallbacks(&dynamic.fallbacks, &dynamic.default_fallbacks) { + return true; + } + if !dynamic.has_regular_fallbacks() { + return false; + } + origin.has_unattempted(request, regular_targets(dynamic, &lookup)) +} diff --git a/litellm-rust/crates/router-core/src/config.rs b/litellm-rust/crates/router-core/src/config.rs new file mode 100644 index 00000000000..a05fdc53261 --- /dev/null +++ b/litellm-rust/crates/router-core/src/config.rs @@ -0,0 +1,361 @@ +//! 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, spelled as the Python router's endpoint method. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + Hash, + strum::EnumString, + strum::Display, + serde_with::DeserializeFromStr, +)] +pub enum Op { + #[strum(serialize = "acompletion")] + Completion, + #[strum(serialize = "atext_completion")] + TextCompletion, + #[strum(serialize = "aembedding")] + Embedding, + #[strum(serialize = "aresponses")] + Responses, + #[strum(serialize = "aanthropic_messages")] + Messages, + #[strum(serialize = "aimage_generation")] + ImageGeneration, + #[strum(serialize = "atranscription")] + Transcription, + #[strum(serialize = "aspeech")] + Speech, + #[strum(serialize = "aocr")] + Ocr, + #[strum(default, transparent)] + Other(String), +} + +/// `routing_strategy`. Only simple-shuffle is native; the rest go through `Host::select`. +#[derive( + Clone, + Debug, + Default, + PartialEq, + Eq, + strum::EnumString, + strum::Display, + serde_with::DeserializeFromStr, +)] +pub enum RoutingStrategy { + #[default] + #[strum(serialize = "simple-shuffle")] + SimpleShuffle, + #[strum(serialize = "least-busy")] + LeastBusy, + #[strum(serialize = "usage-based-routing")] + UsageBased, + #[strum(serialize = "usage-based-routing-v2")] + UsageBasedV2, + #[strum(serialize = "latency-based-routing")] + LatencyBased, + #[strum(serialize = "cost-based-routing")] + CostBased, + /// `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, + pub AuthenticationErrorRetries: Option, + pub TimeoutErrorRetries: Option, + pub RateLimitErrorRetries: Option, + pub ContentPolicyViolationErrorRetries: Option, + pub InternalServerErrorRetries: Option, + pub ServiceUnavailableErrorRetries: Option, + pub NotFoundErrorRetries: Option, + pub DefaultRetries: Option, +} + +/// `AllowedFailsPolicy` in `litellm/types/router.py`: allowed fails per minute per exception type. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[allow(non_snake_case)] +pub struct AllowedFailsPolicy { + pub BadRequestErrorAllowedFails: Option, + pub AuthenticationErrorAllowedFails: Option, + pub TimeoutErrorAllowedFails: Option, + pub RateLimitErrorAllowedFails: Option, + pub ContentPolicyViolationErrorAllowedFails: Option, + pub InternalServerErrorAllowedFails: Option, + pub ServiceUnavailableErrorAllowedFails: Option, + pub BadGatewayErrorAllowedFails: Option, + pub NotFoundErrorAllowedFails: Option, +} + +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 { + [ + ("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, + pub routing_strategy: RoutingStrategy, + pub routing_groups: Vec, + /// `_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>, + pub plugins: Vec, + /// Router default `cooldown_time`; `None` is `DEFAULT_COOLDOWN_TIME_SECONDS`. + pub cooldown_time: Option, + pub disable_cooldowns: bool, + pub allowed_fails_policy: Option, + pub retry_policy: Option, + /// `model_group_retry_policy`: a group's own policy, which wins over `retry_policy`. + pub model_group_retry_policy: BTreeMap, + /// `retry_after`: the minimum backoff between retries, in seconds. + pub retry_after: f64, + pub timeout: Option, + pub request_timeout: Option, + pub drop_params: Option, + pub enable_pre_call_checks: bool, + /// `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, + /// `background_health_check_model_groups`: when set, only these groups are health-filtered. + pub background_health_check_model_groups: Option>, + /// `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, + /// Deployment ids that are a provider's default and never cool down + /// (`provider_default_deployment_ids`). + pub provider_default_deployment_ids: Vec, + /// The operations this router must serve: `Host::supports` must accept each one for every + /// configured provider, or the router fails to build. + pub served_ops: Vec, +} + +/// Dynamic config: a versioned snapshot read once per call. +#[derive(Clone, Debug, Deserialize)] +#[serde(default)] +pub struct RouterDynamicConfig { + pub version: u64, + /// `model_group_alias` as Python spells it: each alias maps to a group name, or to + /// `{"model": , "hidden": }`. + #[serde(alias = "model_group_alias", deserialize_with = "alias_targets")] + pub model_alias_map: BTreeMap, + pub fallbacks: Fallbacks, + pub default_fallbacks: Vec, + pub context_window_fallbacks: Option, + pub content_policy_fallbacks: Option, + pub num_retries: Option, + pub max_fallbacks: u32, + /// `allowed_fails`. `Some` means it was set explicitly, which switches cooldowns from the + /// failure-ratio rule to the allowed-fails counter (`_is_allowed_fails_set_on_router`). + pub allowed_fails: Option, + /// Other spellings a fallback key matches. + pub fallback_keys: FallbackKeys, + /// `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, + /// 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": , ...}`; any other value rewrites nothing. +fn alias_targets<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let aliases = BTreeMap::::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, + 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> { + 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> { + 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()) + } + + pub fn has_regular_fallbacks(&self) -> bool { + !self.fallbacks.is_empty() || !self.default_fallbacks.is_empty() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::group_name(json!({"model_group_alias": {"gpt": "gpt-4"}}))] + #[case::hidden_entry(json!({"model_group_alias": {"gpt": {"model": "gpt-4", "hidden": true}}}))] + #[case::native_name(json!({"model_alias_map": {"gpt": "gpt-4"}}))] + fn model_group_alias_reads_both_python_spellings(#[case] wire: serde_json::Value) { + let dynamic: RouterDynamicConfig = serde_json::from_value(wire).unwrap(); + assert_eq!( + dynamic.model_alias_map, + BTreeMap::from([("gpt".to_owned(), "gpt-4".to_owned())]) + ); + } + + #[rstest] + #[case::native("simple-shuffle", RoutingStrategy::SimpleShuffle)] + #[case::unknown("my-strategy", RoutingStrategy::Other("my-strategy".to_owned()))] + fn routing_strategy_parses_its_wire_name( + #[case] wire: &str, + #[case] expected: RoutingStrategy, + ) { + assert_eq!(wire.parse::().unwrap(), expected); + assert_eq!(expected.to_string(), wire); + } +} diff --git a/litellm-rust/crates/router-core/src/error.rs b/litellm-rust/crates/router-core/src/error.rs new file mode 100644 index 00000000000..f1f3156df4e --- /dev/null +++ b/litellm-rust/crates/router-core/src/error.rs @@ -0,0 +1,419 @@ +//! 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 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 }, + /// The routing plugins left none of the healthy deployments' models. + #[error("no deployments left after routing-plugin filtering: {models:?}")] + RoutingPluginsLeftNothing { models: Vec }, + /// 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::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 { + const BAD_REQUEST: &[&str] = &["BadRequestError"]; + let (classes, status): (&[&str], Option) = match self { + Self::UnknownModel { .. } + | Self::OnlyMarkers { .. } + | Self::PtuReserved { .. } + | Self::NotPassThrough { .. } + | Self::NoPassThroughDeployments { .. } + | Self::MultipleTeams { .. } + | Self::Reserved { .. } => (BAD_REQUEST, 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 { + 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::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 credential failed, not the deployment's. + #[strum(serialize = "CallerCredentialError")] + CallerCredential, + #[strum(serialize = "APIError")] + Other, +} + +/// An opaque value the host hands through the router unchanged, such as the original Python +/// exception or response object. +#[derive(Clone)] +pub struct Payload(pub Arc); + +impl fmt::Debug for Payload { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Payload(..)") + } +} + +/// One failed attempt, as typed facts the retry, fallback and cooldown rules decide by. +#[derive(Clone, Debug, thiserror::Error)] +#[error("the deployment call failed with {kind}")] +pub struct AttemptError { + pub kind: FailureKind, + pub status: Option, + /// The provider's `retry-after` (or `retry-after-ms`) header. + pub retry_after: Option, + /// A deployment-level `num_retries` the failure carries (Python's `exception.num_retries`). + pub num_retries: Option, + /// The provider's error text, kept only for the cooldown record Python writes + /// (`exception_received`, masked past 50 characters). + pub detail: Option>, + /// 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>, + /// 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), 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, + pub payload: Option, +} + +impl AttemptError { + pub fn new(kind: FailureKind, status: Option) -> Self { + Self { + kind, + status, + retry_after: None, + num_retries: None, + detail: None, + classes: Vec::new(), + recorded: false, + names_provider: 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, )`. + 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("CallerCredentialError", FailureKind::CallerCredential)] + #[case::other("APIError", FailureKind::Other)] + fn failure_kind_round_trips_its_exception_name(#[case] wire: &str, #[case] kind: FailureKind) { + assert_eq!(wire.parse::().unwrap(), kind); + assert_eq!(<&'static str>::from(kind), wire); + } + + #[rstest] + #[case::empty("")] + #[case::lowercase("ratelimiterror")] + fn failure_kind_rejects_unknown_names(#[case] wire: &str) { + assert!(wire.parse::().is_err()); + } +} diff --git a/litellm-rust/crates/router-core/src/fallback.rs b/litellm-rust/crates/router-core/src/fallback.rs new file mode 100644 index 00000000000..a1071fd6a0a --- /dev/null +++ b/litellm-rust/crates/router-core/src/fallback.rs @@ -0,0 +1,310 @@ +//! 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), + /// 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")] +pub struct Fallbacks(Vec); + +impl From> for Fallbacks { + fn from(entries: Vec) -> 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), +} + +impl FallbackTarget { + fn of(value: &Value) -> Option { + 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), + } + } + + /// `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> { + 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, + pub provider_prefixed: BTreeMap, + pub param_keys: Vec, +} + +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, + 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 { + 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 { + 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 { + 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> { + let results: Vec = 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::*; + + fn names(targets: Option>) -> Option> { + 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>, + ) { + 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); + } +} diff --git a/litellm-rust/crates/router-core/src/group_info.rs b/litellm-rust/crates/router-core/src/group_info.rs new file mode 100644 index 00000000000..6f3385511dc --- /dev/null +++ b/litellm-rust/crates/router-core/src/group_info.rs @@ -0,0 +1,639 @@ +//! 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, +} + +/// 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, +} + +/// The type name Python reports for a value, in its error messages. +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 { + 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) -> Result { + 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 { + 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 { + 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) -> 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, 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, 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, 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, + assignments: Map, +} + +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) -> 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) { + 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) { + 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 = 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, Error> { + let mut merged: Option = None; + let mut totals: [(&str, Option); 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_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_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 { + 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, +) -> Result, 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, +) -> Result<(Option, Option), 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 { + 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]) -> Vec { + 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, + pub max_input_tokens: Option, + pub max_output_tokens: Option, +} + +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 { + 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 { + 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")); + } +} diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs new file mode 100644 index 00000000000..ad6b73f5ad5 --- /dev/null +++ b/litellm-rust/crates/router-core/src/health.rs @@ -0,0 +1,651 @@ +//! [`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}, + registry::Deployment, + settings::Settings, +}; + +/// `RoutingArgs.ttl`: per-minute usage keys expire after a minute. +pub const USAGE_TTL: Duration = Duration::from_secs(60); +/// `DEFAULT_FAILURE_THRESHOLD_PERCENT`. +pub const FAILURE_THRESHOLD_PERCENT: f64 = 0.5; +/// `DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS`. +pub const FAILURE_THRESHOLD_MINIMUM_REQUESTS: u64 = 5; +/// `SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD`. +pub const SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: u64 = 1000; +/// `exception_masker` in `CooldownCache`: the visible prefix of `exception_received`. +const EXCEPTION_VISIBLE_PREFIX: usize = 50; + +pub trait Clock: Send + Sync { + fn now(&self) -> SystemTime; +} + +pub struct SystemClock; + +impl Clock for SystemClock { + fn now(&self) -> SystemTime { + SystemTime::now() + } +} + +fn unix_seconds(time: SystemTime) -> f64 { + time.duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs_f64()) + .unwrap_or_default() +} + +/// `get_utc_datetime().strftime("%H-%M")`. +pub fn current_minute(time: SystemTime) -> String { + let seconds = time + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs()) + .unwrap_or_default() + % 86_400; + format!("{:02}-{:02}", seconds / 3600, (seconds % 3600) / 60) +} + +pub fn cooldown_key(id: &str) -> String { + format!("deployment:{id}:cooldown") +} + +pub fn allowed_fails_key(id: &str, suffix: Option<&str>) -> String { + match suffix { + Some(suffix) => format!("deployment:{id}:allowed_fails:{suffix}"), + None => format!("deployment:{id}:allowed_fails"), + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, 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(value: &f64, serializer: S) -> Result { + if value.fract() == 0.0 && value.abs() < i64::MAX as f64 { + serializer.serialize_i64(*value as i64) + } else { + serializer.serialize_f64(*value) + } +} + +/// `litellm.should_retry`. +pub fn should_retry_status(status: u16) -> bool { + matches!(status, 408 | 409 | 429) || status >= 500 +} + +/// `_is_cooldown_required`: a 4xx cools down only for 429, 401, 402, 404 and 408, every other +/// status does, and a connection error never does. +pub fn is_cooldown_required(failure: &AttemptError) -> bool { + if failure.kind == FailureKind::ApiConnection { + return false; + } + match failure.cooldown_status() { + status @ 400..500 => matches!(status, 429 | 401 | 402 | 408 | 404), + _ => true, + } +} + +/// `SensitiveDataMasker(visible_prefix=50, visible_suffix=0, mask_short_values=False)`. +fn mask_exception(text: &str) -> String { + let length = text.chars().count(); + if length <= EXCEPTION_VISIBLE_PREFIX { + return text.to_owned(); + } + let mut masked: String = text.chars().take(EXCEPTION_VISIBLE_PREFIX).collect(); + masked.extend(std::iter::repeat_n('*', length - EXCEPTION_VISIBLE_PREFIX)); + masked +} + +/// `_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>, + failure: &AttemptError, +) -> Option { + 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, +} + +/// 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, + clock: Arc, +} + +impl Health { + pub fn new(cache: Arc, clock: Arc) -> Self { + Self { cache, clock } + } + + pub fn cache(&self) -> &Arc { + &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, Error> { + let keys: Vec = ids.iter().map(|id| cooldown_key(id)).collect(); + let values = self.cache.async_batch_get_cache(&keys).await?; + let now = self.now(); + Ok(ids + .iter() + .zip(values) + .filter_map(|(id, value)| { + let value: CooldownValue = serde_json::from_value(value?).ok()?; + (value.timestamp + value.cooldown_time - now > 0.0).then(|| (id.clone(), value)) + }) + .collect()) + } + + /// `add_deployment_to_cooldown`. + pub async fn add_cooldown( + &self, + id: &str, + failure: &AttemptError, + cooldown: Duration, + ) -> Result<(), Error> { + let value = CooldownValue { + exception_received: mask_exception(failure.detail.as_deref().unwrap_or_default()), + status_code: failure.cooldown_status().to_string(), + timestamp: self.now(), + cooldown_time: cooldown.as_secs_f64(), + }; + let value = serde_json::to_value(value).unwrap_or_default(); + self.cache + .async_set_cache(&cooldown_key(id), value, Some(cooldown)) + .await?; + Ok(()) + } + + async fn local_count(&self, key: &str) -> Result { + 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 { + 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, + 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 { + 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, + failure: &AttemptError, + ) -> Result, Error> { + 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, 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::CallerCredential + ) && !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 { + 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.is_none() { + 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 { + 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 it has limits to track. + pub async fn record_success( + &self, + deployment: &Deployment, + total_tokens: u64, + ) -> Result<(), Error> { + self.count_success(&deployment.id).await?; + self.add_usage(deployment, total_tokens as f64, 1.0).await + } + + /// `_increment_deployment_usage`: tokens and requests against a deployment that tracks + /// usage, unless there is nothing to add. + pub async fn add_usage( + &self, + deployment: &Deployment, + total_tokens: f64, + requests: f64, + ) -> Result<(), Error> { + if (total_tokens <= 0.0 && requests <= 0.0) || !deployment.tracks_usage() { + return Ok(()); + } + let minute = current_minute(self.clock.now()); + let model = &deployment.params.model; + self.cache + .async_increment( + &usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + total_tokens, + Some(USAGE_TTL), + ) + .await?; + self.cache + .async_increment( + &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, whether or not it tracks usage. + pub async fn count_failed_request(&self, deployment: &Deployment) -> Result<(), Error> { + let minute = current_minute(self.clock.now()); + self.cache + .async_increment( + &usage_key( + &deployment.id, + &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 rstest::rstest; + + use super::*; + + #[rstest] + #[case::rate_limit(FailureKind::RateLimit, Some(429), true)] + #[case::auth(FailureKind::Authentication, Some(401), true)] + #[case::payment(FailureKind::Other, Some(402), true)] + #[case::not_found(FailureKind::NotFound, Some(404), true)] + #[case::timeout(FailureKind::Timeout, Some(408), true)] + #[case::bad_request(FailureKind::BadRequest, Some(400), false)] + #[case::forbidden(FailureKind::PermissionDenied, Some(403), false)] + #[case::server(FailureKind::InternalServer, Some(500), true)] + #[case::no_status(FailureKind::Other, None, true)] + #[case::connection(FailureKind::ApiConnection, Some(500), false)] + fn cooldown_required_follows_python( + #[case] kind: FailureKind, + #[case] status: Option, + #[case] expected: bool, + ) { + assert_eq!( + is_cooldown_required(&AttemptError::new(kind, status)), + expected + ); + } + + #[rstest] + #[case::midnight(0, "00-00")] + #[case::afternoon(13 * 3600 + 7 * 60 + 59, "13-07")] + fn current_minute_is_utc_hours_and_minutes(#[case] seconds: u64, #[case] expected: &str) { + let time = UNIX_EPOCH + Duration::from_secs(86_400 * 3 + seconds); + assert_eq!(current_minute(time), expected); + } + + #[rstest] + #[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, + ) { + 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 + ); + } + + #[rstest] + fn keys_match_python() { + assert_eq!( + ( + cooldown_key("abc"), + allowed_fails_key("abc", Some("generic")), + usage_key("abc", "openai/gpt-4o", UsageMetric::Tpm, "09-05"), + ), + ( + "deployment:abc:cooldown".to_owned(), + "deployment:abc:allowed_fails:generic".to_owned(), + "global_router:abc:openai/gpt-4o:tpm:09-05".to_owned(), + ) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/health_routing.rs b/litellm-rust/crates/router-core/src/health_routing.rs new file mode 100644 index 00000000000..bc3e43a6e92 --- /dev/null +++ b/litellm-rust/crates/router-core/src/health_routing.rs @@ -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 { + 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>, + when_all_unhealthy: WhenAllUnhealthy, + ) -> Vec> { + 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> = 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::>(), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs new file mode 100644 index 00000000000..a4f7c243bbc --- /dev/null +++ b/litellm-rust/crates/router-core/src/host.rs @@ -0,0 +1,573 @@ +//! [`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.`, +//! 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_json::{Map, Value}; + +use crate::{ + cache::RouterCache, + config::{Op, RetryPolicy, RoutingStrategy}, + error::{AttemptError, Error, FailureKind, Payload}, + fallback::Fallbacks, + health::Health, + mock::MockFailure, + 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`. + pub team_id: Option, + /// 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, + /// 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, + pub specific_deployment: bool, + /// A per-request `num_retries`, which wins over the router's. + pub num_retries: Option, + pub disable_fallbacks: bool, + /// `_router_weights` as the request passed it: model group to deployment id to weight, + /// validated when simple-shuffle picks ([`crate::weights::validate`]). + pub router_weights: Option, + /// 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, + /// A per-request `routing_strategy` override, as Python normalizes it. + pub routing_strategy: Option, + /// 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, + /// `_target_order`: the order level to pick from instead of the lowest. + pub target_order: Option, + /// `_excluded_deployment_ids`: never picked, even when that leaves nothing. + pub excluded_ids: Vec, + /// `_retry_skipped_deployment_ids`: skipped unless that leaves nothing. + pub retry_skipped_ids: Vec, + /// 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, + /// 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, + /// Fallback lists the request passed (`Some(None)` when it passed `None`), which replace + /// the router's. + pub fallbacks: Option>, + pub context_window_fallbacks: Option>, + pub content_policy_fallbacks: Option>, + /// A per-request `max_fallbacks`, which replaces the router's. + pub max_fallbacks: Option, + /// 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, + /// A `mock_testing_*` flag the request set. + pub mock_failure: Option, + /// A per-request `model_group_retry_policy`, which replaces the router's. + pub model_group_retry_policy: Option>, + /// 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, + /// `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, +} + +/// One `os.environ/` reference the deployment build found ([`crate::build`]). +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SecretRef { + /// The value as written, `os.environ/`. + 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, + /// 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), + /// 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), + 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 }, + /// `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, + pub deployment: Arc, + /// The name this attempt resolved: the requested name, or the fallback target. + pub name: String, + /// The model group being served, after fallbacks. + pub model_group: String, + /// The `order` level an order fallback targets in the same group. + pub target_order: Option, + /// Every deployment `name` resolved to for this caller, the picked one included. + pub candidate_ids: Vec, + /// 0 for the first attempt in this group, then the retry number. + pub retry: u32, + pub fallback_depth: u32, + /// How many fallbacks deep the walk is: 0 for the requested name. + pub fallback_level: u32, + pub timeout: Option, + /// Request params the fallback walk's dict targets merged in (`kwargs.update(target)`). + pub params: Option>>, + /// The walk 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, + /// A stream the host reports the outcome of when it ends ([`crate::Router::report_success`], + /// [`crate::Router::report_failure`]): the attempt counts its request and `total_tokens` + /// now, and no success. + pub deferred: 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], +} + +/// 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, +} + +/// Routing events, delivered through [`Host::observe`]. +#[derive(Clone, Debug, PartialEq)] +pub enum RouterEvent { + Selected { + model_group: String, + deployment_id: String, + }, + Succeeded { + deployment_id: String, + }, + AttemptFailed { + deployment_id: String, + kind: FailureKind, + status: Option, + }, + CooldownSet { + deployment_id: String, + cooldown: Duration, + }, + Retry { + model_group: String, + retry: u32, + delay: Duration, + }, + Fallback { + from: String, + to: String, + /// The order level of `to` a same-group hop walks to. + order: Option, + depth: u32, + /// The fallback level `to` is tried at: 1 for the request's own fallbacks. + level: u32, + /// The new level's request metadata scope, copied from `parent`'s at the hop. + scope: u32, + parent: u32, + /// The `max_fallbacks` the walk reads, stamped on the hop's request. + max_fallbacks: u32, + }, + /// A hop of the fallback walk 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, + cause: RetryCause, + }, + /// A level of the fallback walk 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, + }, + /// 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, + }, + /// 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); + +impl PickError { + /// The copy of `error` events carry, when a pick raised it. + pub fn of(error: &Error) -> Option { + 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; + + /// Run a selection stage the core does not implement. + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, 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, Error> { + let _ = ask; + Ok(None) + } + + /// Access and budget checks before a cross-group fallback target. + async fn allow_fallback(&self, request: &RouterRequest, target: &str) -> bool; + + async fn sleep(&self, duration: Duration); + + fn observe(&self, event: &RouterEvent); + + /// The shared store routing state lives in. + fn cache(&self) -> Arc; + + /// Whether [`Self::invoke`] can run `op` on `provider`. Checked for every served op and + /// deployment when the router is built or changed, so an unsupported pair fails the config, + /// never a request. `PythonHost` runs every op. + 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>, 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], + ) -> Result { + Ok(PreCallFacts::default()) + } + + /// `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>, + ) -> Result>, 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, 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 { + 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 { + 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 { + 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 { + let _ = lookup; + Ok(Value::Null) + } +} diff --git a/litellm-rust/crates/router-core/src/host_index.rs b/litellm-rust/crates/router-core/src/host_index.rs new file mode 100644 index 00000000000..ac572cad2a8 --- /dev/null +++ b/litellm-rust/crates/router-core/src/host_index.rs @@ -0,0 +1,647 @@ +//! [`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)>, +} + +/// 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, + /// `model_name_to_deployment_indices`: name to its positions. + pub model_name: Map, + /// `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, + pub deployment_names: Vec, + pub team_public_model_names: Vec, + pub provider_default_ids: Vec, + /// `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>, + pub model_name: Option>, + pub team_model: Option>, + pub model_names: Option>, + pub deployment_names: Option>, + pub team_public_model_names: Option>, + pub provider_default_ids: Option>, + pub patterns: Option, + pub team_patterns: Option>, +} + +/// 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) { + 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, 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 = 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 = 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 { + let mut names: Vec = 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 { + 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> { + let filtered: Option> = + 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 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> { + 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 + } +} + +/// `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> { + 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> { + 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) -> bool { + let Some((literal, later)) = runs.split_first() else { + return true; + }; + let mut ends: Vec = 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) -> Vec { + let mut kept: Vec = 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, 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, 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 serde_json::json; + + use super::*; + + fn entry(name: &str, id: &str, model: &str) -> Map { + 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()); + } +} diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 05d95d02c84..0d1cc760a81 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -43,21 +43,621 @@ //! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the //! policy is exhausted. //! -//! # Execution -//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs -//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A -//! started attempt is never replayed on a different executor. +//! # Hosts: the only boundary +//! The core reaches outside itself only through [`Host`], and every attempt runs through +//! `Host::invoke`: the core picks the deployment, the host runs the call, the core records the +//! outcome. There is no second, "native" attempt path in the core. Which code runs an attempt is +//! the host's business: +//! - `PythonHost` (inside a Python process) builds the request for the picked deployment and calls +//! `litellm.`. It never calls back into a Python router. //! -//! # 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, `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. +//! - [`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 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 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, + }, + prerouting::{PreRoute, PreRouted}, + registry::{ChangeSet, Deployment, DeploymentSpec, Registry, RegistryVersion}, + select::{Randomness, SeededRandomness, ThreadRandomness}, + settings::Settings, +}; +use crate::{health::GroupShape, registry::DeploymentFilter, select::SelectInputs}; + +/// The registry and settings one call reads, swapped atomically by `apply`. +pub(crate) struct Snapshot { + pub registry: Arc, + pub settings: Arc, +} + +impl Snapshot { + /// `_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, + rng: Arc, + health: Health, + config: RouterConfig, + snapshot: ArcSwap, + writer: Mutex<()>, + retired: AtomicBool, +} + +impl Core { + /// Reject what this host cannot serve, before any request. + fn validate(&self, registry: &Registry, settings: &Settings) -> Result<(), ConfigError> { + let can_select = self.host.can_select(); + let strategies = std::iter::once(&settings.routing_strategy).chain( + registry + .routing_groups() + .iter() + .map(|group| &group.routing_strategy), + ); + for strategy in strategies { + if !strategy.is_native() && !can_select { + return Err(Unsupported::Strategy(strategy.clone()).into()); + } + } + if let Some(plugin) = settings.plugins.first() + && !can_select + { + return Err(Unsupported::Plugin(plugin.name.clone()).into()); + } + for deployment in registry.deployments(&DeploymentFilter::default()) { + 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, + /// 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, + /// 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, +} + +/// Builds a [`Router`], with test seams. +pub struct RouterBuilder { + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + rng: Arc, + clock: Arc, +} + +impl RouterBuilder { + #[must_use] + pub fn randomness(self, rng: Arc) -> Self { + Self { rng, ..self } + } + + #[must_use] + pub fn clock(self, clock: Arc) -> Self { + Self { clock, ..self } + } + + pub fn build(self) -> Result { + let registry = Registry::build( + self.config.model_list.clone(), + self.dynamic.model_alias_map.clone(), + self.config.routing_groups.clone(), + )? + .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(Snapshot { + registry: Arc::new(Registry::default()), + settings: Arc::new(settings.clone()), + }), + writer: Mutex::new(()), + retired: AtomicBool::new(false), + }; + core.validate(®istry, &settings)?; + core.snapshot.store(Arc::new(Snapshot { + registry: Arc::new(registry), + settings: Arc::new(settings), + })); + Ok(Router { + core: Arc::new(core), + }) + } +} + +/// The router: the entry points `litellm.Router` delegates to. +#[derive(Clone)] +pub struct Router { + core: Arc, +} + +impl Router { + pub fn new( + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + ) -> Result { + Self::builder(config, dynamic, host).build() + } + + pub fn builder( + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + ) -> RouterBuilder { + RouterBuilder { + config, + dynamic, + host, + rng: Arc::new(ThreadRandomness), + clock: Arc::new(SystemClock), + } + } + + fn live(&self) -> Result<(), Error> { + if self.core.retired.load(Ordering::Acquire) { + return Err(Error::Retired); + } + Ok(()) + } + + /// Select, run the attempt, and retry or fall back until it succeeds or the policy is + /// exhausted. + pub async fn call(&self, op: Op, request: RouterRequest) -> Result { + self.live()?; + self.core.call(op, Arc::new(request)).await + } + + /// Resolve, filter and pick a deployment without calling it. + pub async fn select(&self, op: &Op, request: &RouterRequest) -> Result, Error> { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + 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 { + 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.model_group, + }); + } + Ok(candidates) + } + + pub fn registry(&self) -> Arc { + self.core.snapshot.load().registry.clone() + } + + pub fn settings(&self) -> Arc { + self.core.snapshot.load().settings.clone() + } + + pub fn health(&self) -> &Health { + &self.core.health + } + + /// Record an outcome the router did not see itself, such as a health check. Returns whether + /// the deployment was cooled down. + pub async fn report(&self, id: &str, outcome: Outcome) -> Result { + match outcome { + Outcome::Success { total_tokens } => self + .report_success(id, total_tokens as f64, 1.0) + .await + .map(|_| false), + Outcome::Failure(failure) => { + self.report_failure(id, &failure, FailureReport::default()) + .await + } + } + } + + /// 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, + ) -> Result, 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?; + self.core + .health + .add_usage(&deployment, 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) -> Result { + self.core.health.count_success(id).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`). + pub async fn report_failure( + &self, + id: &str, + failure: &AttemptError, + report: FailureReport, + ) -> Result { + if !report.uncounted { + self.core.health.count_failure(id).await?; + } + let snapshot = self.core.snapshot.load_full(); + let Some(deployment) = snapshot.registry.deployment(id) else { + return Ok(false); + }; + if report.counts_request { + self.core.health.count_failed_request(&deployment).await?; + } + // 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(false); + } + 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), + 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 + .map(|cooldown| cooldown.is_some()) + } + + /// Apply one change set atomically and return the new registry version. + pub fn apply(&self, changes: ChangeSet) -> Result { + let _writer = self.core.writer.lock(); + let current = self.core.snapshot.load_full(); + let aliases = changes + .settings + .as_ref() + .map(|dynamic| dynamic.model_alias_map.clone()); + let registry = current.registry.apply(&changes, aliases)?; + let settings = match changes.settings { + Some(dynamic) => Settings::new(&self.core.config, Arc::new(dynamic)), + None => Settings::clone(¤t.settings), + }; + self.core.validate(®istry, &settings)?; + let version = registry.version(); + self.core.snapshot.store(Arc::new(Snapshot { + registry: Arc::new(registry), + settings: Arc::new(settings), + })); + Ok(version) + } + + /// `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, 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, 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, + ) -> Result<(Option, Option), 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::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>, 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>, + resolve_secrets: bool, + install: build::Install, + ) -> Result, Error> { + let model_names: Vec = 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 = 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, + built: Registry, + installed: Vec, + 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(Snapshot { + 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(Snapshot { + 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); + } +} diff --git a/litellm-rust/crates/router-core/src/lookup.rs b/litellm-rust/crates/router-core/src/lookup.rs new file mode 100644 index 00000000000..e0b8bf3ca4e --- /dev/null +++ b/litellm-rust/crates/router-core/src/lookup.rs @@ -0,0 +1,851 @@ +//! 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, + 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, +} + +/// 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, + pub groups: Vec, + /// `get_deployments_by_pattern`'s requests: the name, then the name under its provider. + pub wildcard: Vec, + /// 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, +} + +/// 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, + params: &'a Map, + strips_access_groups: bool, +} + +/// The lookups over one registry version. +pub struct Lookup<'a> { + deployments: &'a [Arc], + 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 { + 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, 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> + 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 { + let rows = |positions: Vec| -> Vec { + 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, 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 { + (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, 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 { + 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) -> Result, 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 { + 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 { + let route = |router: &crate::host_index::PatternRouter| -> Vec { + 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() + } + + /// `_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, 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)>, Row> { + let mut groups: Vec<(String, Vec)> = 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 = 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), 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 { + 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, 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, 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 { + 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, Row> { + let mut ids: Vec = 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, Row> { + if let Some(members) = members { + let mut ids: Vec = 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 { + let route = |router: &crate::host_index::PatternRouter| -> Vec { + 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 { + static EMPTY: std::sync::OnceLock> = 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 + } + + 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()])] + ); + } +} diff --git a/litellm-rust/crates/router-core/src/mock.rs b/litellm-rust/crates/router-core/src/mock.rs new file mode 100644 index 00000000000..e063577defc --- /dev/null +++ b/litellm-rust/crates/router-core/src/mock.rs @@ -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 replaces a level's walk (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) -> 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 + } +} diff --git a/litellm-rust/crates/router-core/src/model_info.rs b/litellm-rust/crates/router-core/src/model_info.rs new file mode 100644 index 00000000000..e47eb4521e7 --- /dev/null +++ b/litellm-rust/crates/router-core/src/model_info.rs @@ -0,0 +1,709 @@ +//! 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, + pub user_keys: Vec, +} + +pub(crate) fn raised(exception: &'static str, detail: impl Into) -> 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 { + 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>, 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, 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, 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 { + 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, 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, + base_model: Option<&str>, + provider: &str, +) -> Result, 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, 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", + "the deployment's base_model has no attribute 'startswith'", + )); + }; + 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 { + let chars: Vec = 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 { + 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::().ok().filter(|n| n.is_finite()) +} + +pub(crate) fn python_int(text: &str) -> Option { + 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::() + .map(Number::from) + .ok() + .or_else(|| text.parse::().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, new: &Map) { + 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 { + 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>, Error> { + let mut custom: Option> = 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 = 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, + provider: (String, String), + user_model_info: Value, + lookups: Mutex>, + } + + #[async_trait] + impl Host for Catalog { + async fn invoke(&self, _attempt: &Attempt) -> Result { + unreachable!("a model-info query runs no attempt") + } + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + async fn allow_fallback(&self, _request: &RouterRequest, _target: &str) -> bool { + true + } + async fn sleep(&self, _duration: std::time::Duration) {} + fn observe(&self, _event: &RouterEvent) {} + fn cache(&self) -> Arc { + Arc::new(InMemoryRouterCache::default()) + } + async fn model_info(&self, lookup: &ModelInfoLookup) -> Result { + 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::, "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::}), + ..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}) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/precall.rs b/litellm-rust/crates/router-core/src/precall.rs new file mode 100644 index 00000000000..04dd22f1c90 --- /dev/null +++ b/litellm-rust/crates/router-core/src/precall.rs @@ -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, + /// `_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, + /// The request's `allowed_model_region`. + pub allowed_model_region: Option, + /// The requested params the check filters on, sorted (`_pre_call_check_requested_chat_params`). + pub requested_params: Vec, + /// Facts by deployment id, for the candidates and the group's deployments. + pub deployments: BTreeMap, +} + +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>, + pool: &[Arc], + facts: &PreCallFacts, +) -> Result>, 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 = (0..candidates.len()).filter(|i| !invalid[*i]).collect(); + let lacking: Vec = 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 = 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 { + 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::>(), + ["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::>(), + ["new"] + ); + } +} diff --git a/litellm-rust/crates/router-core/src/prerouting.rs b/litellm-rust/crates/router-core/src/prerouting.rs new file mode 100644 index 00000000000..e5622512cc6 --- /dev/null +++ b/litellm-rust/crates/router-core/src/prerouting.rs @@ -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 walk 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), +} + +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 { + 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()))) + } +} diff --git a/litellm-rust/crates/router-core/src/ptu.rs b/litellm-rust/crates/router-core/src/ptu.rs new file mode 100644 index 00000000000..9110ac6a67f --- /dev/null +++ b/litellm-rust/crates/router-core/src/ptu.rs @@ -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) -> Option> { + 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>, + team_id: Option<&str>, +) -> (Vec>, bool) { + let team = team_id.unwrap_or(""); + let checks: Vec<(Arc, Option>)> = 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> = 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>) { + 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()) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/registry.rs b/litellm-rust/crates/router-core/src/registry.rs new file mode 100644 index 00000000000..1af16d2c55f --- /dev/null +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -0,0 +1,1358 @@ +//! The configured deployments, indexed by id, model group, alias, wildcard pattern, routing +//! group and team. A [`Registry`] is immutable; [`Registry::apply`] builds the next version. + +use std::{ + collections::{BTreeMap, BTreeSet, HashMap}, + fmt::Write as _, + sync::Arc, +}; + +use litellm_router_types::{LitellmParams, Spelled}; +use serde::Deserialize; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; + +use crate::{ + error::ConfigError, + host_index::{HostIndexes, IndexWrite}, +}; + +/// A registry's version; every change produces a higher one. +pub type RegistryVersion = u64; + +/// A deployment as a config or an `apply` spells it. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +pub struct DeploymentSpec { + pub model_name: String, + pub litellm_params: Map, + #[serde(default)] + pub model_info: Map, + /// The entry's other top-level keys (a config's `tpm`, ...), kept so the host's view of the + /// deployment list reads back the entries it stored. + #[serde(flatten)] + pub extra: Map, + /// The provider the host resolved when it built the deployment ([`crate::build`]); never + /// read from a config or an `apply`. + #[serde(skip)] + pub provider: Option, +} + +/// `routing_groups`: a named subset of model groups sharing a routing strategy. +#[derive(Clone, Debug, PartialEq, Deserialize)] +pub struct RoutingGroupSpec { + pub group_name: String, + pub models: Vec, + pub routing_strategy: crate::config::RoutingStrategy, + /// A priority group's member priorities, served as each member's `order`. + #[serde(default)] + pub model_priorities: Option>, +} + +/// The `model_info` fields the router reads; every other key stays in [`Self::raw`]. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ModelInfo { + pub team_id: Option, + pub team_public_model_name: Option, + pub access_groups: Vec, + pub blocked: bool, + pub cooldown_time: Option, + pub allowed_fails: Option, + /// `model_info.allowed_fails_policy` when it is a dict, read as Python reads it: field name + /// to allowed fails, values as given. + pub allowed_fails_policy: Option>, + pub raw: Map, +} + +impl ModelInfo { + fn from_raw(raw: Map) -> Self { + let text = |key: &str| raw.get(key).and_then(Value::as_str).map(str::to_owned); + Self { + team_id: text("team_id"), + team_public_model_name: text("team_public_model_name"), + access_groups: raw + .get("access_groups") + .and_then(Value::as_array) + .map(|groups| { + groups + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default(), + blocked: raw.get("blocked").and_then(Value::as_bool) == Some(true), + cooldown_time: raw.get("cooldown_time").and_then(Value::as_f64), + allowed_fails: raw + .get("allowed_fails") + .and_then(Value::as_u64) + .and_then(|fails| u32::try_from(fails).ok()), + allowed_fails_policy: raw + .get("allowed_fails_policy") + .and_then(Value::as_object) + .cloned(), + raw, + } + } +} + +/// One configured deployment. +#[derive(Clone, Debug, PartialEq)] +pub struct Deployment { + pub id: String, + pub model_name: String, + pub params: LitellmParams, + /// `litellm_params` as written, in order: what the id hashes and the host receives. + pub raw_params: Map, + pub model_info: ModelInfo, + /// The entry's other top-level keys, as written. + pub extra: Map, + /// The provider the host resolved for it, when it was built through the host. + pub resolved_provider: Option, +} + +impl Deployment { + pub fn from_spec(spec: DeploymentSpec) -> Result { + let params: LitellmParams = serde_json::from_value(Value::Object( + spec.litellm_params.clone(), + )) + .map_err(|source| ConfigError::InvalidDeployment { + model_name: spec.model_name.clone(), + source, + })?; + let id = match spec.model_info.get("id") { + Some(Value::String(id)) => id.clone(), + Some(Value::Number(id)) => id.to_string(), + _ => deployment_id(&spec.model_name, &spec.litellm_params), + }; + Ok(Self { + id, + model_name: spec.model_name, + params, + raw_params: spec.litellm_params, + model_info: ModelInfo::from_raw(spec.model_info), + extra: spec.extra, + resolved_provider: spec.provider, + }) + } + + /// The deployment as a `model_list` entry: its other top-level keys, then `model_name`, + /// `litellm_params` and `model_info` (with its id). + pub fn record(&self) -> Value { + let mut model_info = self.model_info.raw.clone(); + model_info.insert("id".to_owned(), Value::String(self.id.clone())); + let mut record = self.extra.clone(); + record.insert( + "model_name".to_owned(), + Value::String(self.model_name.clone()), + ); + record.insert( + "litellm_params".to_owned(), + Value::Object(self.raw_params.clone()), + ); + record.insert("model_info".to_owned(), Value::Object(model_info)); + Value::Object(record) + } + + /// The stand-in an unrouted attempt carries: the requested name as its group and model, no id. + pub fn unrouted(name: &str) -> Result { + let mut params = Map::new(); + params.insert("model".to_owned(), Value::String(name.to_owned())); + let mut info = Map::new(); + info.insert("id".to_owned(), Value::String(String::new())); + Self::from_spec(DeploymentSpec { + model_name: name.to_owned(), + litellm_params: params, + model_info: info, + extra: Map::new(), + provider: None, + }) + } + + /// The provider: the one the host resolved when it built the deployment, else + /// `custom_llm_provider`, else the `model` prefix before `/`. + pub fn provider(&self) -> &str { + self.resolved_provider + .as_deref() + .or(self.params.custom_llm_provider.as_deref()) + .or_else(|| { + self.params + .model + .split_once('/') + .map(|(provider, _)| provider) + }) + .unwrap_or("openai") + } + + /// A numeric `litellm_params` value, as `simple_shuffle._metric_weight` reads it. + pub fn metric(&self, key: &str) -> Option { + let spelled = |value: &Option>| match value { + Some(Spelled::Value(value)) => Some(*value), + Some(Spelled::Text(_)) | None => None, + }; + match key { + "rpm" => spelled(&self.params.rpm), + "tpm" => spelled(&self.params.tpm), + _ => self.raw_params.get(key).and_then(Value::as_f64), + } + } + + /// `simple_shuffle._metric_weight`: 0 when unset, a number (a bool counts as 0 or 1), or + /// `None` when the value is not numeric. A `rpm`/`tpm` spelled as text reads as unset. + pub fn metric_weight(&self, key: &str) -> Option { + if key != "weight" { + return Some(self.metric(key).unwrap_or(0.0)); + } + match self.raw_params.get(key) { + None | Some(Value::Null) => Some(0.0), + Some(Value::Bool(flag)) => Some(f64::from(u8::from(*flag))), + Some(Value::Number(number)) => number.as_f64(), + Some(_) => None, + } + } + + /// `litellm_params.order` (`get_deployment_order`). + pub fn order(&self) -> Option { + self.raw_params.get("order").and_then(Value::as_i64) + } + + /// Whether `team_id` may use this deployment (`deployment_usable_by_team`). + pub fn usable_by_team(&self, team_id: Option<&str>) -> bool { + self.model_info + .team_id + .as_deref() + .is_none_or(|owner| Some(owner) == team_id) + } + + /// Whether the router counts per-minute usage for this deployment + /// (`_deployment_tracks_usage`, without the model-catalog limits). + pub fn tracks_usage(&self) -> bool { + self.params.tpm.is_some() + || self.params.rpm.is_some() + || self.params.itpm.is_some() + || self.params.otpm.is_some() + } +} + +/// The connection params of one deployment (`credentials(id)`). +#[derive(Clone, Debug, PartialEq)] +pub struct Credentials { + pub provider: String, + pub api_base: Option, + pub api_version: Option, + pub credential_name: Option, + pub params: LitellmParams, +} + +/// What a name resolved through. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ResolveKind { + SpecificDeployment, + Id, + /// A name nothing else serves, routed to the router's `default_deployment` (a copy whose + /// `litellm_params.model` is the name). Served like a deployment named by its id. + Default, + Alias, + RoutingGroup, + Group, + Wildcard, +} + +impl ResolveKind { + /// The name resolved to one deployment that is served as is: a deployment id, or the + /// default deployment (`_common_checks_available_deployment` returns a dict for both). + pub fn is_single(&self) -> bool { + matches!(self, Self::Id | Self::Default) + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Resolution { + pub kind: ResolveKind, + /// The model group the deployments are routed as. + pub group: String, + pub deployments: Vec>, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct ResolveScope { + /// The caller named a deployment's `litellm_params.model` (`specific_deployment=True`). + pub specific_deployment: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct DeploymentFilter<'a> { + pub model_name: Option<&'a str>, + pub team_id: Option<&'a str>, + pub provider: Option<&'a str>, +} + +#[derive(Clone, Debug, Default)] +pub struct NameFilter { + pub include_aliases: bool, +} + +/// The registry indexes as the host wrote them over the ones the deployments build (the +/// deprecated `model_names`, `model_name_to_deployment_indices` and +/// `team_model_to_deployment_indices` writes), which resolution reads instead from then on. +/// Indices point into the deployment list; one past its end is skipped. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct IndexOverrides { + /// `_core_model_names`: the names `name in model_names` finds. + pub model_names: Option>, + /// `_core_model_name_to_deployment_indices`: a model name's deployments. + pub model_name: Option>>, + /// `_core_team_model_to_deployment_indices`, as `(team_id, public name, indices)` entries. + pub team_model: Option)>>, +} + +impl IndexOverrides { + fn is_empty(&self) -> bool { + self.model_names.is_none() && self.model_name.is_none() && self.team_model.is_none() + } +} + +/// One atomic change: optionally a whole new deployment list, then upserts, deletes and, +/// optionally, a new dynamic config. +#[derive(Clone, Debug, Default)] +pub struct ChangeSet { + /// Every deployment, in order, replacing the current list (the host registry's `model_list` + /// after a write; it may hold one id more than once, as the Python router's list can). + pub replace: Option>, + /// Deployments added at the end, beside any that share their id (a deployment the host + /// retries after its provider resolved). + pub appends: Vec, + pub upserts: Vec, + pub deletes: Vec, + pub settings: Option, + /// The indexes the host wrote ([`IndexOverrides`]), replacing the ones held; an empty one + /// goes back to the indexes the deployments build. + pub index: Option, + /// The deployment at the position the id map holds for this id, removed with every index + /// shifted (`_delete_deployment`, and an upsert's removal of the deployment it edits). + pub remove: Option, + /// A host write to the indexes it hands out ([`HostIndexes`]). + pub host_index: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct Registry { + version: RegistryVersion, + /// The version the deployments before the newest appends last changed at: a host that read + /// the list at or after it reads only what was appended since ([`Self::deployments_from`]). + rewritten: RegistryVersion, + deployments: Vec>, + by_id: HashMap, + by_group: BTreeMap>, + by_team: HashMap>, + wildcards: Vec<(String, Vec)>, + aliases: BTreeMap, + routing_groups: Vec, + overrides: Option>, + /// The indexes the host hands out, kept by the Python router's history rules. + host: Arc, + /// `_model_to_group` as Python last built it, when a raw `routing_groups` write replaced the + /// groups since: member calls keep routing by it (see [`Self::member_routing_group`]). + member_groups: Option>>, +} + +impl Registry { + pub fn build( + specs: Vec, + aliases: BTreeMap, + routing_groups: Vec, + ) -> Result { + // One id may name several deployments (a per-organization fan-out, or two entries + // declaring one id): each is kept, and the id finds the last, as the Python router keeps them. + let deployments = specs + .into_iter() + .map(|spec| Deployment::from_spec(spec).map(Arc::new)) + .collect::, _>>()?; + Ok(Self::index(0, deployments, aliases, routing_groups)) + } + + fn index( + version: RegistryVersion, + deployments: Vec>, + aliases: BTreeMap, + routing_groups: Vec, + ) -> Self { + let mut by_id = HashMap::new(); + let mut by_group: BTreeMap> = BTreeMap::new(); + let mut by_team: HashMap> = HashMap::new(); + for (index, deployment) in deployments.iter().enumerate() { + by_id.insert(deployment.id.clone(), index); + by_group + .entry(deployment.model_name.clone()) + .or_default() + .push(index); + if let Some(team) = &deployment.model_info.team_id { + by_team.entry(team.clone()).or_default().push(index); + } + } + // In the order each pattern first appears, then most specific first, as + // `PatternUtils.sorted_patterns` keeps them: a stable sort, longest regex first, + // then the most regex syntax. + let mut wildcards: Vec<(String, Vec)> = Vec::new(); + for deployment in &deployments { + let name = &deployment.model_name; + if name.contains('*') && !wildcards.iter().any(|(pattern, _)| pattern == name) { + wildcards.push(( + name.clone(), + by_group.get(name).cloned().unwrap_or_default(), + )); + } + } + wildcards.sort_by_key(|(pattern, _)| std::cmp::Reverse(pattern_specificity(pattern))); + Self { + version, + rewritten: version, + deployments, + by_id, + by_group, + by_team, + wildcards, + aliases, + routing_groups, + overrides: None, + host: Arc::default(), + member_groups: None, + } + } + + /// This registry with the stale routing-group member map (`None` derives it from the groups). + #[must_use] + pub fn with_member_groups(self, member_groups: Option>) -> Self { + Self { + member_groups: member_groups.map(Arc::new), + ..self + } + } + + pub fn version(&self) -> RegistryVersion { + self.version + } + + /// The version the deployments before the newest appends last changed at (see `rewritten`). + pub fn rewritten(&self) -> RegistryVersion { + self.rewritten + } + + /// The deployments from position `start` on, in list order. + pub fn deployments_from(&self, start: usize) -> &[Arc] { + self.deployments.get(start..).unwrap_or_default() + } + + /// The next registry: upserts, then deletes (`apply_changes`). An upserted deployment whose + /// id exists is removed and appended at the end, as the Python router does on an edit. + pub fn apply( + &self, + changes: &ChangeSet, + aliases: Option>, + ) -> Result { + let mut deployments = match &changes.replace { + Some(specs) => specs + .iter() + .map(|spec| self.carried(spec.clone()).map(Arc::new)) + .collect::, _>>()?, + None => self.deployments.clone(), + }; + for spec in &changes.appends { + deployments.push(Arc::new(self.carried(spec.clone())?)); + } + for spec in &changes.upserts { + let deployment = self.carried(spec.clone())?; + deployments.retain(|existing| existing.id != deployment.id); + deployments.push(Arc::new(deployment)); + } + deployments.retain(|deployment| !changes.deletes.contains(&deployment.id)); + let mut host = HostIndexes::clone(&self.host); + if let Some(id) = &changes.remove + && let Some(position) = host.position_of(id) + && position < deployments.len() + { + deployments.remove(position); + host.removed(id, position) + .map_err(|error| ConfigError::Removal(error.to_string()))?; + } + if let Some(write) = &changes.host_index { + host.write(write.clone()); + } + let overrides = match &changes.index { + Some(index) => (!index.is_empty()).then(|| Arc::new(index.clone())), + None => self.overrides.clone(), + }; + // Appends alone leave every deployment before them as it was. + let appended_only = changes.replace.is_none() + && changes.upserts.is_empty() + && changes.deletes.is_empty() + && changes.remove.is_none(); + let next = Self::index( + self.version + 1, + deployments, + aliases.unwrap_or_else(|| self.aliases.clone()), + self.routing_groups.clone(), + ); + Ok(Self { + overrides, + host: Arc::new(host), + member_groups: self.member_groups.clone(), + rewritten: if appended_only { + self.rewritten + } else { + next.rewritten + }, + ..next + }) + } + + /// The next registry for a build starting over (`set_model_list`): no deployments, and the + /// indexes [`HostIndexes::reset`] clears cleared. + #[must_use] + pub fn emptied(&self) -> Self { + let mut host = HostIndexes::clone(&self.host); + host.reset(); + Self { + overrides: self.overrides.clone(), + host: Arc::new(host), + member_groups: self.member_groups.clone(), + ..Self::index( + self.version + 1, + Vec::new(), + self.aliases.clone(), + self.routing_groups.clone(), + ) + } + } + + /// Add one deployment a build accepted at the end and index it in place (`add` is + /// `add_deployment`, which also names its group). A build installs into a registry of its + /// own this way, so a list of n deployments builds in O(n) rather than reindexing per + /// deployment. + pub fn push_built(&mut self, spec: DeploymentSpec, add: bool) -> Result<(), ConfigError> { + let deployment = Arc::new(self.carried(spec)?); + let index = self.deployments.len(); + let record = deployment.record(); + Arc::make_mut(&mut self.host).added( + record.as_object().expect("a record is an object"), + index, + add, + ); + self.by_id.insert(deployment.id.clone(), index); + let name = &deployment.model_name; + self.by_group.entry(name.clone()).or_default().push(index); + if let Some(team) = &deployment.model_info.team_id { + self.by_team.entry(team.clone()).or_default().push(index); + } + if name.contains('*') { + match self + .wildcards + .iter_mut() + .find(|(pattern, _)| pattern == name) + { + Some((_, indices)) => indices.push(index), + None => { + // A new pattern joins after every one it ties with, as `index` orders them: + // first appearance, then a stable sort by specificity. + self.wildcards.push((name.clone(), vec![index])); + self.wildcards.sort_by_key(|(pattern, _)| { + std::cmp::Reverse(pattern_specificity(pattern)) + }); + } + } + } + self.deployments.push(deployment); + self.version += 1; + Ok(()) + } + + /// The last deployment added. + pub fn last(&self) -> Option<&Arc> { + self.deployments.last() + } + + /// Take `set_model_list`'s `model_names`: the names it was given. + pub fn set_model_names(&mut self, names: Vec) { + Arc::make_mut(&mut self.host).set_model_names(names); + self.version += 1; + } + + /// The indexes the host hands out. + pub fn host_indexes(&self) -> &HostIndexes { + &self.host + } + + /// The deployment `spec` describes, keeping the provider the host resolved for the + /// deployment it replaces (same id, same params) when the spec carries none: a host that + /// hands back its whole list after a write does not resolve every provider again. + fn carried(&self, mut spec: DeploymentSpec) -> Result { + let deployment = Deployment::from_spec(spec.clone())?; + if spec.provider.is_none() + && let Some(previous) = self + .by_id + .get(&deployment.id) + .and_then(|index| self.deployments.get(*index)) + && previous.raw_params == deployment.raw_params + && previous.model_name == deployment.model_name + { + spec.provider.clone_from(&previous.resolved_provider); + return Deployment::from_spec(spec); + } + Ok(deployment) + } + + /// This registry with the providers `from` resolved for the same deployments (same id, + /// same params) and `from`'s host indexes: a router rebuilt for new static settings over the + /// same list keeps what its host resolved and the indexes their history built. + #[must_use] + pub fn with_providers_of(&self, from: &Self) -> Self { + let deployments = self + .deployments + .iter() + .map(|deployment| { + let previous = from + .by_id + .get(&deployment.id) + .and_then(|index| from.deployments.get(*index)) + .filter(|previous| { + previous.raw_params == deployment.raw_params + && previous.model_name == deployment.model_name + }); + match previous { + Some(previous) if deployment.resolved_provider.is_none() => { + Arc::new(Deployment { + resolved_provider: previous.resolved_provider.clone(), + ..Deployment::clone(deployment) + }) + } + _ => deployment.clone(), + } + }) + .collect(); + Self { + deployments, + host: from.host.clone(), + ..self.clone() + } + } + + fn at(&self, indices: &[usize]) -> Vec> { + indices + .iter() + .filter_map(|index| self.deployments.get(*index).cloned()) + .collect() + } + + /// `_core_deployments_named(name)`: the deployments the model-name index holds for `name` + /// (whatever their own name), or `None` when it does not index the name. + pub(crate) fn named(&self, name: &str) -> Option>> { + match self.overrides.as_ref().and_then(|o| o.model_name.as_ref()) { + Some(index) => index.get(name).map(|indices| self.at(indices)), + None => self.by_group.get(name).map(|indices| self.at(indices)), + } + } + + /// Whether the model-name index holds `name` (`name in model_name_to_deployment_indices`). + fn indexes_name(&self, name: &str) -> bool { + match self.overrides.as_ref().and_then(|o| o.model_name.as_ref()) { + Some(index) => index.contains_key(name), + None => self.by_group.contains_key(name), + } + } + + /// `_team_deployments_named(team_id, name)`: the deployments the team index holds for the + /// pair, or `None` when it does not index it. + pub(crate) fn team_named(&self, team_id: &str, name: &str) -> Option>> { + match self.overrides.as_ref().and_then(|o| o.team_model.as_ref()) { + Some(entries) => entries + .iter() + .find(|(team, public, _)| team == team_id && public == name) + .map(|(_, _, indices)| self.at(indices)), + None => { + let own: Vec> = self + .deployments + .iter() + .filter(|d| { + d.model_info.team_id.as_deref() == Some(team_id) + && d.model_info.team_public_model_name.as_deref() == Some(name) + }) + .cloned() + .collect(); + (!own.is_empty()).then_some(own) + } + } + } + + /// Every deployment the team index holds under public name `name`, any team's + /// (`_team_deployments_across_teams`). + pub(crate) fn team_named_across_teams(&self, name: &str) -> Vec> { + match self.overrides.as_ref().and_then(|o| o.team_model.as_ref()) { + Some(entries) => entries + .iter() + .filter(|(_, public, _)| public == name) + .flat_map(|(_, _, indices)| self.at(indices)) + .collect(), + None => self + .deployments + .iter() + .filter(|d| { + d.model_info.team_id.is_some() + && d.model_info.team_public_model_name.as_deref() == Some(name) + }) + .cloned() + .collect(), + } + } + + /// Whether a written index entry names a position past the model list. + fn overflows(&self, indices: &[usize]) -> bool { + indices.iter().any(|index| *index >= self.deployments.len()) + } + + /// The written model-name index entry for `name` names a position past the model list. + pub(crate) fn name_index_overflows(&self, name: &str) -> bool { + self.overrides + .as_ref() + .and_then(|o| o.model_name.as_ref()) + .and_then(|index| index.get(name)) + .is_some_and(|indices| self.overflows(indices)) + } + + /// The written team index entry for `(team_id, name)` names a position past the model list. + pub(crate) fn team_index_overflows(&self, team_id: &str, name: &str) -> bool { + self.overrides + .as_ref() + .and_then(|o| o.team_model.as_ref()) + .and_then(|entries| { + entries + .iter() + .find(|(team, public, _)| team == team_id && public == name) + }) + .is_some_and(|(_, _, indices)| self.overflows(indices)) + } + + /// A written team index entry under public name `name` names a position past the list. + pub(crate) fn across_teams_index_overflows(&self, name: &str) -> bool { + self.overrides + .as_ref() + .and_then(|o| o.team_model.as_ref()) + .is_some_and(|entries| { + entries + .iter() + .any(|(_, public, indices)| public == name && self.overflows(indices)) + }) + } + + /// The group a pick of `name` reads its strategy from, when the stale member map names a + /// group that is no longer configured (`_get_routing_context` raises `KeyError`): `name` + /// after its alias, unless it is a callable routing group. + pub fn missing_member_group(&self, name: &str) -> Option<&str> { + let members = self.member_groups.as_ref()?; + let resolved = self.alias_of(name).unwrap_or(name); + if self.routing_group_named(resolved).is_some() { + return None; + } + let group = members.get(resolved)?; + (!self.routing_groups.iter().any(|g| g.group_name == *group)).then_some(group.as_str()) + } + + /// The routing group a member call of `model_group` routes by: the stale member map's + /// group when there is one (`_model_to_group`), else the non-priority group listing it. + pub(crate) fn member_routing_group(&self, model_group: &str) -> Option<&RoutingGroupSpec> { + match &self.member_groups { + Some(members) => members + .get(model_group) + .and_then(|name| self.routing_groups.iter().find(|g| g.group_name == *name)), + None => self + .routing_groups + .iter() + .filter(|group| !crate::routing_groups::is_priority(group)) + .find(|group| group.models.iter().any(|model| model == model_group)), + } + } + + pub fn deployment(&self, id: &str) -> Option> { + self.by_id + .get(id) + .map(|index| self.deployments[*index].clone()) + } + + pub fn deployments(&self, filter: &DeploymentFilter<'_>) -> Vec> { + let candidates = match (filter.model_name, filter.team_id) { + (Some(name), _) => self + .by_group + .get(name) + .map(|i| self.at(i)) + .unwrap_or_default(), + (None, Some(team)) => self + .by_team + .get(team) + .map(|i| self.at(i)) + .unwrap_or_default(), + (None, None) => self.deployments.clone(), + }; + candidates + .into_iter() + .filter(|d| { + filter + .team_id + .is_none_or(|team| d.model_info.team_id.as_deref() == Some(team)) + }) + .filter(|d| { + filter + .provider + .is_none_or(|provider| d.provider() == provider) + }) + .collect() + } + + /// The model group names, in name order, with aliases when asked. + pub fn names(&self, filter: &NameFilter) -> Vec { + let mut names: Vec = self.by_group.keys().cloned().collect(); + if filter.include_aliases { + names.extend(self.aliases.keys().cloned()); + } + names + } + + pub fn group(&self, name: &str) -> Option>> { + if self + .overrides + .as_ref() + .is_some_and(|o| o.model_name.is_some()) + { + return self.named(name).map(|deployments| { + deployments + .into_iter() + .filter(|d| d.model_name == name) + .collect() + }); + } + self.by_group.get(name).map(|indices| self.at(indices)) + } + + /// How many rows `_core_get_model_list(model_name=name)` lists: the group's deployments, + /// an alias's target group, a routing group's members, and only when none of those, the + /// deployments of the wildcard the name matches. + pub fn model_list_len(&self, name: &str) -> usize { + let count = |group: &str| self.group(group).map_or(0, |deployments| deployments.len()); + let listed = count(name) + + self.aliases.get(name).map_or(0, |target| count(target)) + + self + .routing_groups + .iter() + .find(|group| group.group_name == name) + .map_or(0, |group| { + group.models.iter().map(|model| count(model)).sum() + }); + if listed > 0 { + return listed; + } + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .map_or(0, |(_, indices)| indices.len()) + } + + /// The routing group a model group belongs to, if any. + pub fn routing_group_of(&self, model_name: &str) -> Option<&RoutingGroupSpec> { + self.routing_groups + .iter() + .find(|group| group.models.iter().any(|model| model == model_name)) + } + + /// The strategy a routing group named `group_name` declares. + pub fn resolve_routing_group_strategy( + &self, + group_name: &str, + ) -> Option { + self.routing_groups + .iter() + .find(|group| group.group_name == group_name) + .map(|group| group.routing_strategy.clone()) + } + + pub(crate) fn routing_groups(&self) -> &[RoutingGroupSpec] { + &self.routing_groups + } + + /// Every deployment, in configured order. + pub(crate) fn all(&self) -> &[Arc] { + &self.deployments + } + + /// Whether `name` is a deployment `model_name` (`name in _core_model_names`). + pub(crate) fn has_name(&self, name: &str) -> bool { + match self.overrides.as_ref().and_then(|o| o.model_names.as_ref()) { + Some(names) => names.contains(name), + None => self.by_group.contains_key(name), + } + } + + /// The group `model_group_alias` rewrites `name` to. + pub(crate) fn alias_of(&self, name: &str) -> Option<&str> { + self.aliases.get(name).map(String::as_str) + } + + /// The callable routing group named `name`: a model group of that name shadows it. + pub(crate) fn routing_group_named(&self, name: &str) -> Option<&RoutingGroupSpec> { + (!self.indexes_name(name)) + .then(|| self.routing_groups.iter().find(|g| g.group_name == name)) + .flatten() + } + + /// The most specific wildcard pattern matching `name`, with its deployments. + pub(crate) fn wildcard_for(&self, name: &str) -> Option<(String, Vec>)> { + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .map(|(pattern, indices)| (pattern.clone(), self.at(indices))) + } + + /// Resolve a requested name, in the order of `_common_checks_available_deployment`: + /// a specific deployment, an id, an alias, a routing group, a model group, a wildcard. + pub fn resolve(&self, name: &str, scope: ResolveScope) -> Option { + let resolution = |kind, group: &str, deployments: Vec>| { + (!deployments.is_empty()).then(|| Resolution { + kind, + group: group.to_owned(), + deployments, + }) + }; + if scope.specific_deployment { + let matching = self + .deployments + .iter() + .filter(|d| d.params.model == name) + .cloned() + .collect(); + // An empty match still resolves: the pick then finds no deployment rather than + // falling back to a default (`_common_checks_available_deployment` returns `[]`). + return Some(Resolution { + kind: ResolveKind::SpecificDeployment, + group: name.to_owned(), + deployments: matching, + }); + } + if !self.by_group.contains_key(name) + && let Some(deployment) = self.deployment(name) + { + let group = deployment.model_name.clone(); + return resolution(ResolveKind::Id, &group, vec![deployment]); + } + let routing_group_named = |name: &str| { + // A model group of the same name shadows a routing group (`get_routing_group`). + (!self.by_group.contains_key(name)) + .then(|| self.routing_groups.iter().find(|g| g.group_name == name)) + .flatten() + }; + if let Some(target) = self.aliases.get(name) { + if let Some(routing_group) = routing_group_named(target) { + let deployments = crate::routing_groups::members(routing_group, |model| { + self.by_group + .get(model) + .map(|indices| self.at(indices)) + .unwrap_or_default() + }); + return resolution(ResolveKind::RoutingGroup, target, deployments); + } + return self + .group(target) + .and_then(|deployments| resolution(ResolveKind::Alias, target, deployments)); + } + if let Some(routing_group) = routing_group_named(name) { + let deployments = crate::routing_groups::members(routing_group, |model| { + self.by_group + .get(model) + .map(|indices| self.at(indices)) + .unwrap_or_default() + }); + return resolution(ResolveKind::RoutingGroup, name, deployments); + } + if let Some(indices) = self.by_group.get(name) { + return resolution(ResolveKind::Group, name, self.at(indices)); + } + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .and_then(|(pattern, indices)| { + resolution(ResolveKind::Wildcard, pattern, self.at(indices)) + }) + } + + /// Access group name to the model groups in it. + pub fn access_groups(&self) -> BTreeMap> { + let mut groups: BTreeMap> = BTreeMap::new(); + for deployment in &self.deployments { + for group in &deployment.model_info.access_groups { + groups + .entry(group.clone()) + .or_default() + .insert(deployment.model_name.clone()); + } + } + groups + } + + pub fn credentials(&self, id: &str) -> Option { + self.deployment(id).map(|deployment| Credentials { + provider: deployment.provider().to_owned(), + api_base: deployment.params.api_base.clone(), + api_version: deployment.params.api_version.clone(), + credential_name: deployment.params.litellm_credential_name.clone(), + params: deployment.params.clone(), + }) + } +} + +/// `PatternUtils.calculate_pattern_specificity` over the pattern's regex +/// (`re.escape(pattern).replace(r"\*", "(.*)")`): its length, then how much regex syntax it has. +fn pattern_specificity(pattern: &str) -> (usize, usize) { + let mut regex = String::new(); + for c in pattern.chars() { + match c { + '*' => regex.push_str("(.*)"), + '(' | ')' | '[' | ']' | '{' | '}' | '?' | '+' | '-' | '|' | '^' | '$' | '\\' | '.' + | '&' | '~' | '#' | ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c' => { + regex.push('\\'); + regex.push(c); + } + _ => regex.push(c), + } + } + let complexity = regex + .chars() + .filter(|c| matches!(c, '*' | '+' | '?' | '\\' | '^' | '$' | '|' | '(' | ')')) + .count(); + (regex.chars().count(), complexity) +} + +/// `*` matches any run of characters, as the pattern router's `(.*)` does under `re.match`: +/// anchored at the start only, so the name may run on past the pattern's last literal. +fn wildcard_matches(pattern: &str, name: &str) -> bool { + let mut parts = pattern.split('*'); + let first = parts.next().unwrap_or_default(); + let Some(mut rest) = name.strip_prefix(first) else { + return false; + }; + let parts: Vec<&str> = parts.collect(); + for part in parts { + match rest.find(part) { + Some(at) => rest = &rest[at + part.len()..], + None => return false, + } + } + true +} + +/// `deployment_id` in `litellm/router_api.py`: sha256 over the model group and every +/// `litellm_params` key and value in order. A string is taken as is, a dict as `json.dumps`, and +/// anything else as Python's `str()`. +pub fn deployment_id(model_group: &str, litellm_params: &Map) -> String { + let mut text = String::from(model_group); + for (key, value) in litellm_params { + text.push_str(key); + match value { + Value::String(value) => text.push_str(value), + Value::Object(_) => json_dumps(value, &mut text), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::Array(_) => { + python_str(value, &mut text); + } + } + } + let digest = Sha256::digest(text.as_bytes()); + let mut hex = String::with_capacity(64); + for byte in digest { + let _ = write!(hex, "{byte:02x}"); + } + hex +} + +/// Python's `repr(float)`: the shortest round-trip digits, positional for exponents in +/// `-4..16`, scientific otherwise. +fn python_float(value: f64, out: &mut String) { + if value.is_nan() { + out.push_str("nan"); + return; + } + if value.is_infinite() { + out.push_str(if value > 0.0 { "inf" } else { "-inf" }); + return; + } + let scientific = format!("{value:e}"); + let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); + let exponent: i32 = exponent.parse().unwrap_or(0); + let (sign, mantissa) = mantissa + .strip_prefix('-') + .map_or(("", mantissa), |m| ("-", m)); + let digits: String = mantissa.chars().filter(char::is_ascii_digit).collect(); + out.push_str(sign); + if (-4..16).contains(&exponent) { + if exponent < 0 { + out.push_str("0."); + out.extend(std::iter::repeat_n('0', (-exponent - 1) as usize)); + out.push_str(&digits); + } else { + let point = exponent as usize + 1; + if digits.len() <= point { + out.push_str(&digits); + out.extend(std::iter::repeat_n('0', point - digits.len())); + out.push_str(".0"); + } else { + out.push_str(&digits[..point]); + out.push('.'); + out.push_str(&digits[point..]); + } + } + } else { + out.push_str(&digits[..1]); + if digits.len() > 1 { + out.push('.'); + out.push_str(&digits[1..]); + } + let _ = write!( + out, + "e{}{:02}", + if exponent < 0 { '-' } else { '+' }, + exponent.abs() + ); + } +} + +fn number(value: &serde_json::Number, out: &mut String) { + match value.as_f64() { + Some(float) if !value.is_i64() && !value.is_u64() => python_float(float, out), + Some(_) | None => out.push_str(&value.to_string()), + } +} + +/// `json.dumps` with its defaults: `", "` and `": "` separators and `ensure_ascii`. +fn json_dumps(value: &Value, out: &mut String) { + match value { + Value::Null => out.push_str("null"), + Value::Bool(flag) => out.push_str(if *flag { "true" } else { "false" }), + Value::Number(value) => number(value, out), + Value::String(text) => json_string(text, out), + Value::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + json_dumps(item, out); + } + out.push(']'); + } + Value::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + json_string(key, out); + out.push_str(": "); + json_dumps(item, out); + } + out.push('}'); + } + } +} + +fn json_string(text: &str, out: &mut String) { + out.push('"'); + for ch in text.chars() { + match ch { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + '\u{08}' => out.push_str("\\b"), + '\u{0c}' => out.push_str("\\f"), + ch if ch.is_ascii() && !ch.is_ascii_control() => out.push(ch), + ch => { + let mut units = [0u16; 2]; + for unit in ch.encode_utf16(&mut units) { + let _ = write!(out, "\\u{unit:04x}"); + } + } + } + } + out.push('"'); +} + +/// Python's `str()` of a JSON value: `repr` for the items of a list or dict. +fn python_str(value: &Value, out: &mut String) { + match value { + Value::Null => out.push_str("None"), + Value::Bool(flag) => out.push_str(if *flag { "True" } else { "False" }), + Value::Number(value) => number(value, out), + Value::String(text) => python_repr_str(text, out), + Value::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + python_str(item, out); + } + out.push(']'); + } + Value::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + python_repr_str(key, out); + out.push_str(": "); + python_str(item, out); + } + out.push('}'); + } + } +} + +fn python_repr_str(text: &str, out: &mut String) { + let quote = if text.contains('\'') && !text.contains('"') { + '"' + } else { + '\'' + }; + out.push(quote); + for ch in text.chars() { + match ch { + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + ch if ch == quote => { + out.push('\\'); + out.push(ch); + } + ch if (ch as u32) < 0x20 || ch as u32 == 0x7f => { + let _ = write!(out, "\\x{:02x}", ch as u32); + } + ch => out.push(ch), + } + } + out.push(quote); +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn spec(model_name: &str, params: Value, info: Value) -> DeploymentSpec { + serde_json::from_value(json!({ + "model_name": model_name, "litellm_params": params, "model_info": info, + })) + .unwrap() + } + + /// Expected ids computed with `litellm.router_api.deployment_id`. + #[rstest] + #[case::strings( + "g", + json!({"model": "openai/gpt-4o", "api_key": "sk-1"}), + "30256ba697cc98ce2aff32c3b73c4b90514cb199876aef429ffcd91ae826d025" + )] + #[case::scalars( + "claude", + json!({"model": "anthropic/claude-3", "rpm": 10, "timeout": 1.5, "stream": true, "x": null}), + "645fd1504089c1084316573ccfc5dd1dd9622623b169026657484fb64a2c89a2" + )] + #[case::list_and_dict( + "m", + json!({"model": "azure/x", "tags": ["a", "b'c"], "extra_body": {"k": [1, 2.0, "é"], "z": {"q": false}}}), + "64f11e16dabed321268bd92132de2020290ef193c0cf554b2c74faee8badbb96" + )] + fn deployment_id_matches_python( + #[case] group: &str, + #[case] params: Value, + #[case] expected: &str, + ) { + let Value::Object(params) = params else { + unreachable!() + }; + assert_eq!(deployment_id(group, ¶ms), expected); + } + + #[rstest] + #[case::integral(2.0, "2.0")] + #[case::small(0.00001, "1e-05")] + #[case::large(1e16, "1e+16")] + #[case::fraction(0.1, "0.1")] + fn floats_print_as_python_repr(#[case] value: f64, #[case] expected: &str) { + let mut out = String::new(); + python_float(value, &mut out); + assert_eq!(out, expected); + } + + #[rstest] + #[case::group("gpt-4", ResolveKind::Group, "gpt-4")] + #[case::alias("smart", ResolveKind::Alias, "gpt-4")] + #[case::id("dep-b", ResolveKind::Id, "gpt-4")] + #[case::wildcard("anthropic/claude-x", ResolveKind::Wildcard, "anthropic/*")] + fn resolve_follows_python_order( + #[case] name: &str, + #[case] kind: ResolveKind, + #[case] group: &str, + ) { + let registry = Registry::build( + vec![ + spec( + "gpt-4", + json!({"model": "openai/gpt-4"}), + json!({"id": "dep-a"}), + ), + spec( + "gpt-4", + json!({"model": "azure/gpt-4"}), + json!({"id": "dep-b"}), + ), + spec("anthropic/*", json!({"model": "anthropic/*"}), json!({})), + ], + BTreeMap::from([("smart".to_owned(), "gpt-4".to_owned())]), + Vec::new(), + ) + .unwrap(); + let resolution = registry.resolve(name, ResolveScope::default()).unwrap(); + assert_eq!((resolution.kind, resolution.group.as_str()), (kind, group)); + } + + #[rstest] + fn apply_moves_an_edited_deployment_to_the_end() { + let registry = Registry::build( + vec![ + spec("a", json!({"model": "openai/a"}), json!({"id": "1"})), + spec("a", json!({"model": "openai/b"}), json!({"id": "2"})), + ], + BTreeMap::new(), + Vec::new(), + ) + .unwrap(); + let next = registry + .apply( + &ChangeSet { + upserts: vec![spec("a", json!({"model": "openai/c"}), json!({"id": "1"}))], + ..ChangeSet::default() + }, + None, + ) + .unwrap(); + let ids: Vec = next + .group("a") + .unwrap() + .iter() + .map(|d| d.id.clone()) + .collect(); + assert_eq!( + (next.version(), ids), + (1, vec!["2".to_owned(), "1".to_owned()]) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/resolve.rs b/litellm-rust/crates/router-core/src/resolve.rs new file mode 100644 index 00000000000..369e9995d0a --- /dev/null +++ b/litellm-rust/crates/router-core/src/resolve.rs @@ -0,0 +1,411 @@ +//! 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 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}, + error::Error, + registry::{Deployment, DeploymentSpec, Registry, Resolution, ResolveKind}, + 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], +} + +/// `_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>, Error> { + let keep = |deployments: Vec>| -> Vec> { + 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>, 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 { + 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> { + 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), 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>, + team_id: Option<&str>, + now: Timestamp, +) -> Result>, Error> { + let any = !deployments.is_empty(); + let selectable: Vec> = deployments + .into_iter() + .filter(|d| !is_strategy_marker(d)) + .collect(); + if any && selectable.is_empty() { + return Err(Error::OnlyMarkers { + model: model.to_owned(), + }); + } + let (kept, blocking) = filter_reserved(selectable, 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), + } +} + +/// 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> { + 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> { + let caller = Caller { + team_id: request.team_id.as_deref(), + proxy_admin: request.proxy_admin, + allowed_models: &request.allowed_models, + }; + 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>, + now: Timestamp, +) -> Result, 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, team_id, 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], team_id, 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, team_id, now)?; + return resolved(ResolveKind::Group, model, kept); + } + if let Some((pattern, deployments)) = registry.wildcard_for(model) { + let kept = reserved(model, deployments, team_id, 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], team_id, 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>| { + let allowed = filter_by_access_groups(model, deployments, caller.allowed_models, || { + model_access_groups(registry, model, team_id) + }); + reserved(model, allowed, team_id, 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> = 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(|_| !access_emptied) { + let deployments = narrow(target, all_deployments(registry, target, team_id)?)?; + if !deployments.is_empty() { + return resolved(ResolveKind::Group, target, deployments); + } + } + 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 = serde_json::from_value(entries).unwrap(); + Registry::build(specs, BTreeMap::new(), Vec::new()).unwrap() + } + + fn ids(resolution: Option) -> Vec { + 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")); + } +} diff --git a/litellm-rust/crates/router-core/src/retry.rs b/litellm-rust/crates/router-core/src/retry.rs new file mode 100644 index 00000000000..b0a04ed3c77 --- /dev/null +++ b/litellm-rust/crates/router-core/src/retry.rs @@ -0,0 +1,216 @@ +//! 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>, +) -> 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 { + 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 { + 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> { + [ + "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::GuardrailIntervention + | FailureKind::Declined + | FailureKind::CallerCredential + | FailureKind::Other => &[], + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn failure(kind: FailureKind, status: Option, 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, + ) { + 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, + #[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 + ); + } +} diff --git a/litellm-rust/crates/router-core/src/routing_groups.rs b/litellm-rust/crates/router-core/src/routing_groups.rs new file mode 100644 index 00000000000..fcf6c2febf3 --- /dev/null +++ b/litellm-rust/crates/router-core/src/routing_groups.rs @@ -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>, +) -> Vec> { + 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, order: i64) -> Arc { + let mut prioritized = Deployment::clone(deployment); + prioritized + .raw_params + .insert("order".to_owned(), Value::from(order)); + Arc::new(prioritized) +} diff --git a/litellm-rust/crates/router-core/src/routing_plugins.rs b/litellm-rust/crates/router-core/src/routing_plugins.rs new file mode 100644 index 00000000000..63bbed30cd6 --- /dev/null +++ b/litellm-rust/crates/router-core/src/routing_plugins.rs @@ -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`). +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>, + models: Option<&[String]>, +) -> Result>, Error> { + let Some(models) = models else { + return Ok(healthy); + }; + let kept: Vec> = 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 { + 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::>(), + ["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 + ); + } +} diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs new file mode 100644 index 00000000000..4455810290c --- /dev/null +++ b/litellm-rust/crates/router-core/src/select.rs @@ -0,0 +1,673 @@ +//! Selection: resolve a name, filter its deployments in the order of +//! `Router._core_async_get_healthy_deployments` (team, 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, Snapshot, + config::{Op, RoutingStrategy}, + error::Error, + health_routing::WhenAllUnhealthy, + host::{RouterRequest, RoutingContext, SelectStage}, + registry::{Deployment, 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; +} + +pub struct ThreadRandomness; + +impl Randomness for ThreadRandomness { + fn next_f64(&self) -> f64 { + rand::thread_rng().r#gen() + } +} + +pub struct SeededRandomness(Mutex); + +impl SeededRandomness { + pub fn new(seed: u64) -> Self { + Self(Mutex::new(StdRng::seed_from_u64(seed))) + } +} + +impl Randomness for SeededRandomness { + fn next_f64(&self) -> f64 { + self.0.lock().r#gen() + } +} + +/// 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>, + pub healthy: Vec>, +} + +/// Per-attempt selection inputs the call loop carries. +#[derive(Clone, Debug, Default)] +pub(crate) struct SelectInputs { + pub target_order: Option, + pub skipped_ids: Vec, + /// 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, +} + +/// `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], + model_group: &str, + router_weights: Option<&BTreeMap>>, + rng: &dyn Randomness, +) -> Result>, 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::>() + }); + let metric_sets = ["weight", "rpm", "tpm"].into_iter().map(|metric| { + deployments + .iter() + .map(|d| d.metric_weight(metric)) + .collect::>>() + .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 = 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())); + } + let index = ((rng.next_f64() * deployments.len() as f64) as usize).min(deployments.len() - 1); + Ok(Some(deployments[index].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>, + request: &RouterRequest, +) -> Result>, 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 Snapshot, + 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: &[], + } + } + + /// The eligible deployments for `name`, in filter-chain order. + pub(crate) async fn candidates( + &self, + snapshot: &Snapshot, + op: &Op, + request: &RouterRequest, + name: &str, + inputs: &SelectInputs, + ) -> Result { + let registry = &snapshot.registry; + // 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); + let caller = Caller { + team_id: request.team_id.as_deref(), + proxy_admin: request.proxy_admin, + allowed_models: &request.allowed_models, + }; + 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 { + model: 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(), + }, + }); + } + 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() { + return Err(Error::NoPassThroughDeployments { + model: name.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 { + let ids: Vec = all.iter().map(|d| d.id.clone()).collect(); + self.health.active_cooldowns(&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> = 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()))); + } + // `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 ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + 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, + }); + } + + 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, + }); + } + + 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> = healthy + .iter() + .filter(|d| !text_only.iter().any(|t| t.id == d.id)) + .cloned() + .collect(); + if !vision.is_empty() { + healthy = vision; + } + } + + 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.target_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> = 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, + }) + } + + /// `_pre_call_checks` over `healthy`, with the host's facts. + #[allow(clippy::too_many_arguments)] + async fn pre_call_checks( + &self, + snapshot: &Snapshot, + op: &Op, + request: &RouterRequest, + model_group: &str, + name: &str, + pool: &[Arc], + healthy: Vec>, + ) -> Result>, 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: &Snapshot, + op: &Op, + request: &RouterRequest, + candidates: &Candidates, + ) -> Result, Error> { + let no_deployment = || Error::NoHealthyDeployments { + model: candidates.model_group.clone(), + }; + // 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 { + 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"); + } +} diff --git a/litellm-rust/crates/router-core/src/settings.rs b/litellm-rust/crates/router-core/src/settings.rs new file mode 100644 index 00000000000..d690c466087 --- /dev/null +++ b/litellm-rust/crates/router-core/src/settings.rs @@ -0,0 +1,103 @@ +//! [`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, + pub allowed_fails_policy: Option, + pub retry_policy: Option, + pub model_group_retry_policy: Arc>, + pub num_retries: u32, + pub max_fallbacks: u32, + pub retry_after: Duration, + pub timeout: Option, + pub drop_params: Option, + 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>, + pub enable_weighted_failover: bool, + pub fallback_on_unsupported_params: bool, + pub provider_default_deployment_ids: Arc<[String]>, + pub dynamic: Arc, +} + +/// `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) -> 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 { + 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 + } + } + + /// Python's `allowed_fails` default when set without a value (`litellm.allowed_fails`). + pub fn effective_allowed_fails(&self) -> u32 { + self.allowed_fails.unwrap_or(3) + } +} diff --git a/litellm-rust/crates/router-core/src/strategy.rs b/litellm-rust/crates/router-core/src/strategy.rs new file mode 100644 index 00000000000..0e1f3b49fa8 --- /dev/null +++ b/litellm-rust/crates/router-core/src/strategy.rs @@ -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 { + 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, + ) { + assert_eq!(request_override(raw), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/weights.rs b/litellm-rust/crates/router-core/src/weights.rs new file mode 100644 index 00000000000..3e5d76844ab --- /dev/null +++ b/litellm-rust/crates/router-core/src/weights.rs @@ -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>; + +/// 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 { + 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, 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()); + } +} diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs new file mode 100644 index 00000000000..97797ac0bca --- /dev/null +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -0,0 +1,411 @@ +//! A call through the router end to end: selection, retries, fallbacks and cooldowns against a +//! scripted host. + +use std::{collections::VecDeque, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Deployment, Error, FailureKind, FailureReport, Host, + InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig, + RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage, +}; +use parking_lot::Mutex; +use rstest::{fixture, rstest}; +use serde_json::json; + +#[derive(Default)] +struct ScriptedHost { + /// Failures handed out in order, keyed by model group; a group with none left succeeds. + script: Mutex)>>, + events: Mutex>, + cache: InMemoryRouterCache, + /// Answer every success as a stream whose outcome the host reports when it ends. + deferred: bool, +} + +impl ScriptedHost { + fn failing(group: &str, failures: Vec) -> Self { + Self { + script: Mutex::new(vec![(group.to_owned(), failures.into())]), + ..Self::default() + } + } +} + +#[async_trait] +impl Host for ScriptedHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + let next = self + .script + .lock() + .iter_mut() + .find(|(group, _)| *group == attempt.model_group) + .and_then(|(_, failures)| failures.pop_front()); + match next { + Some(failure) => Err(failure), + None => Ok(Output { + value: Payload(Arc::new(attempt.deployment.id.clone())), + total_tokens: 10, + blocked: None, + deferred: self.deferred, + }), + } + } + + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + + async fn allow_fallback(&self, _request: &RouterRequest, _target: &str) -> bool { + true + } + + async fn sleep(&self, _duration: Duration) {} + + fn observe(&self, event: &RouterEvent) { + self.events.lock().push(event.clone()); + } + + fn cache(&self) -> Arc { + Arc::new(self.cache.clone()) + } +} + +#[fixture] +fn config() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/c"}, "model_info": {"id": "b1"}}, + ] + })) + .unwrap() +} + +fn router(config: RouterConfig, dynamic: serde_json::Value, host: Arc) -> Router { + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, host) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() +} + +fn request() -> RouterRequest { + RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + } +} + +#[rstest] +#[tokio::test] +async fn a_rate_limit_retries_in_the_same_group(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router(config, json!({"num_retries": 2}), host.clone()); + let response = router.call(Op::Completion, request()).await.unwrap(); + assert_eq!( + ( + response.model_group.as_str(), + response.attempted_retries, + response.attempted_fallbacks + ), + ("primary", 1, 0) + ); +} + +#[rstest] +#[tokio::test] +async fn a_bad_request_falls_back_without_retrying(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::BadRequest, Some(400))], + )); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let response = router.call(Op::Completion, request()).await.unwrap(); + let retried = host + .events + .lock() + .iter() + .any(|event| matches!(event, RouterEvent::Retry { .. })); + assert_eq!( + ( + response.deployment_id.as_str(), + response.attempted_fallbacks, + retried + ), + ("b1", 1, false) + ); +} + +#[rstest] +#[tokio::test] +async fn a_rate_limit_cools_the_deployment_down(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router(config, json!({}), host.clone()); + let error = router.call(Op::Completion, request()).await.unwrap_err(); + let cooled = router + .health() + .active_cooldowns(&["a1".to_owned(), "a2".to_owned()]) + .await + .unwrap(); + assert_eq!( + (error.attempt().map(|failure| failure.kind), cooled.len()), + (Some(FailureKind::RateLimit), 1) + ); +} + +#[rstest] +#[case::now(false, (Some(1.0), (10, 1)))] +#[case::deferred(true, (None, (10, 1)))] +#[tokio::test] +async fn a_deferred_attempt_counts_its_request_and_tokens_but_no_success( + #[case] deferred: bool, + #[case] expected: (Option, (u64, u64)), +) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a", "rpm": 100}, "model_info": {"id": "a1"}}, + ] + })) + .unwrap(); + let host = Arc::new(ScriptedHost { + deferred, + ..ScriptedHost::default() + }); + let router = router(config, json!({}), host.clone()); + router.call(Op::Completion, request()).await.unwrap(); + let deployment = router.registry().deployment("a1").unwrap(); + let successes = host + .cache + .get("a1:successes") + .and_then(|count| count.as_f64()); + let usage = router.health().usage(&deployment).await.unwrap(); + assert_eq!((successes, usage), expected); +} + +#[rstest] +#[case::counted(true, 2)] +#[case::uncounted(false, 1)] +#[tokio::test] +async fn a_reported_failure_counts_its_request_when_asked( + #[case] counts_request: bool, + #[case] requests: u64, +) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a", "rpm": 100}, "model_info": {"id": "a1"}}, + ] + })) + .unwrap(); + let host = Arc::new(ScriptedHost { + deferred: true, + ..ScriptedHost::default() + }); + let router = router(config, json!({}), host.clone()); + router.call(Op::Completion, request()).await.unwrap(); + let report = FailureReport { + counts_request, + ..FailureReport::default() + }; + router + .report_failure( + "a1", + &AttemptError::new(FailureKind::RateLimit, Some(429)), + report, + ) + .await + .unwrap(); + let deployment = router.registry().deployment("a1").unwrap(); + let fails = host.cache.get("a1:fails").and_then(|count| count.as_f64()); + let usage = router.health().usage(&deployment).await.unwrap(); + assert_eq!((fails, usage.1), (Some(1.0), requests)); +} + +#[rstest] +#[tokio::test] +async fn a_raising_retry_policy_lookup_ends_the_hop_and_falls_back(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let request = RouterRequest { + retry_policy_raises: vec!["primary".to_owned()], + ..request() + }; + let response = router.call(Op::Completion, request).await.unwrap(); + let events = host.events.lock(); + assert!( + !events + .iter() + .any(|e| matches!(e, RouterEvent::Retry { .. })) + ); + assert!(events.iter().any(|e| matches!( + e, + RouterEvent::RetryPolicyRaised { level: 0, model, .. } if model == "primary" + ))); + assert_eq!(response.model_group, "backup"); +} + +#[rstest] +#[tokio::test] +async fn a_name_nothing_serves_routes_to_the_default_deployment(config: RouterConfig) { + let host = Arc::new(ScriptedHost::default()); + let router = router( + config, + json!({"default_deployment": { + "model_name": "dflt", "litellm_params": {"model": "openai/x"}, "model_info": {"id": "dd"} + }}), + host.clone(), + ); + let request = RouterRequest { + model: "other".to_owned(), + ..RouterRequest::default() + }; + let response = router.call(Op::Completion, request.clone()).await.unwrap(); + assert_eq!( + ( + response.deployment_id.as_str(), + response.model_group.as_str() + ), + ("dd", "other") + ); + let picked = router.select(&Op::Completion, &request).await.unwrap(); + assert_eq!(picked.params.model, "other"); +} + +#[rstest] +#[tokio::test] +async fn a_stale_member_map_naming_a_dropped_group_raises_on_every_pick(mut config: RouterConfig) { + config.routing_group_members = Some([("primary".to_owned(), "gone".to_owned())].into()); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let refused = router.select(&Op::Completion, &request()).await; + assert!(matches!(refused, Err(Error::MissingRoutingGroup { group }) if group == "gone")); + let other = RouterRequest { + model: "backup".to_owned(), + ..RouterRequest::default() + }; + assert!(router.select(&Op::Completion, &other).await.is_ok()); +} + +#[rstest] +#[tokio::test] +async fn written_indexes_replace_the_ones_the_deployments_build(config: RouterConfig) { + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let index: litellm_router_core::registry::IndexOverrides = serde_json::from_value(json!({ + "model_names": ["backup"], + "model_name": {"primary": [1], "backup": [2]}, + })) + .unwrap(); + router + .apply(litellm_router_core::ChangeSet { + index: Some(index), + ..litellm_router_core::ChangeSet::default() + }) + .unwrap(); + for _ in 0..4 { + let picked = router.select(&Op::Completion, &request()).await.unwrap(); + assert_eq!(picked.id, "a2"); + } + router + .apply(litellm_router_core::ChangeSet { + index: Some(litellm_router_core::registry::IndexOverrides::default()), + ..litellm_router_core::ChangeSet::default() + }) + .unwrap(); + assert_eq!(router.registry().group("primary").unwrap().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_retry_policy_retries_a_pick_that_found_no_deployment(config: RouterConfig) { + let host = Arc::new(ScriptedHost::default()); + let mut config = config; + config.retry_policy = serde_json::from_value(json!({"DefaultRetries": 2})).unwrap(); + let router = router(config, json!({"num_retries": 0}), host.clone()); + let request = RouterRequest { + model: "nothing".to_owned(), + ..RouterRequest::default() + }; + let refused = router.call(Op::Completion, request).await; + assert!(matches!(refused, Err(Error::UnknownModel { .. }))); + let events = host.events.lock(); + let retries = events + .iter() + .filter(|event| matches!(event, RouterEvent::Retry { .. })) + .count(); + assert_eq!(retries, 2); + assert!(events.iter().any(|event| matches!( + event, + RouterEvent::RetriesExhausted { + picked: true, + level: 0, + .. + } + ))); +} + +#[rstest] +#[tokio::test] +async fn a_default_deployment_python_cannot_copy_raises_for_a_name_nothing_serves( + config: RouterConfig, +) { + let router = router( + config, + json!({"default_deployment_raises": true}), + Arc::new(ScriptedHost::default()), + ); + let request = RouterRequest { + model: "other".to_owned(), + ..RouterRequest::default() + }; + let refused = router.select(&Op::Completion, &request).await; + assert!(matches!(refused, Err(Error::DefaultDeploymentRaises { model }) if model == "other")); + assert!( + router + .select(&Op::Completion, &self::request()) + .await + .is_ok() + ); +} + +#[rstest] +#[tokio::test] +async fn a_written_index_position_past_the_list_raises_for_that_name(config: RouterConfig) { + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let index: litellm_router_core::registry::IndexOverrides = + serde_json::from_value(json!({"model_name": {"primary": [0, 9], "backup": [2]}})).unwrap(); + router + .apply(litellm_router_core::ChangeSet { + index: Some(index), + ..litellm_router_core::ChangeSet::default() + }) + .unwrap(); + let refused = router.select(&Op::Completion, &request()).await; + assert!(matches!(refused, Err(Error::IndexOutOfRange { model }) if model == "primary")); + let backup = RouterRequest { + model: "backup".to_owned(), + ..RouterRequest::default() + }; + assert!(router.select(&Op::Completion, &backup).await.is_ok()); +} diff --git a/litellm/router_api.py b/litellm/router_api.py index c013a99573f..be05803c677 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -45,6 +45,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 @@ -178,9 +179,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: @@ -200,11 +201,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: @@ -268,9 +335,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: @@ -315,6 +382,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 @@ -531,10 +605,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 @@ -1176,16 +1248,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``).""" diff --git a/litellm/router_attempt.py b/litellm/router_attempt.py new file mode 100644 index 00000000000..543e41ea97c --- /dev/null +++ b/litellm/router_attempt.py @@ -0,0 +1,208 @@ +""" +One router attempt, as the host runs it: the request for the deployment the router core picked. + +The Rust router core picks a deployment and hands the attempt to its host; ``PythonHost`` builds +the keyword arguments for that one deployment here and calls ``litellm.`` with them. Nothing +here selects, retries, falls back or records cooldowns: the core does all of that. Nothing here +reads a Python router either, so this module outlives ``router_python_core.py``. +""" + +import re +import uuid +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter +from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs + +# Request keys the router reads to route the call; the deployment call never receives them. +ROUTING_KEYS: Final = frozenset( + { + "model", + "num_retries", + "disable_fallbacks", + "specific_deployment", + "_router_weights", + "_retry_skipped_deployment_ids", + "_target_order", + "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 for_requested_model(deployment: Mapping[str, object], requested: str | None) -> Mapping[str, object]: + """A wildcard deployment as the router serves ``requested`` through it: its ``model`` with + the wildcards filled in from the requested name (``PatternMatchRouter.route``). Any other + deployment comes back unchanged.""" + model_name: Final = deployment.get("model_name") + params: Final = _as_mapping(deployment.get("litellm_params")) + model: Final = params.get("model") + if requested is None or not isinstance(model_name, str) or "*" not in model_name or not isinstance(model, str): + return deployment + regex: Final = PatternMatchRouter().pattern_to_regex(model_name) + match: Final = re.match(regex, requested) or re.match(regex, f"{inferred_provider(requested)}/{requested}") + if match is None: + return deployment + return { + **deployment, + "litellm_params": {**params, "model": PatternMatchRouter.set_deployment_model_name(match, model)}, + } + + +def attempt_kwargs( + op: str, + request: Mapping[str, object], + deployment: Mapping[str, object], + *, + model_group: str, + default_litellm_params: Mapping[str, object], + timeout: float | None, +) -> Mapping[str, object]: + """The keyword arguments for ``litellm.`` against one deployment + (``_update_kwargs_with_deployment`` and the call sites that merge ``litellm_params``).""" + params: Final = _as_mapping(deployment.get("litellm_params")) + model_info: Final = _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 [] diff --git a/litellm/router_core.py b/litellm/router_core.py index acbe82a0643..1a182108409 100644 --- a/litellm/router_core.py +++ b/litellm/router_core.py @@ -1,14 +1,16 @@ """ ``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one. -``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. The Rust core -(litellm-rust/crates/router-core) implements the same operations. The core is picked once, when -the router is built, from its ``RouterConfig``. +``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. ``RustCore`` +(``litellm/router_rust/core.py`` over litellm-rust/crates/router-core) is a separate implementation +that reaches outside itself only through its host; it does not use ``PythonCore``. The core is +picked once, when the router is built, from its ``RouterConfig``. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable +from litellm._logging import verbose_router_logger from litellm.router_config import RouterConfig if TYPE_CHECKING: @@ -51,6 +53,26 @@ class RouterCore(Protocol): def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore: + """ + The Rust core when the route catalog's router rule allows it and the Rust core serves this + config, else the Python core. Decided once, for the router's life, from the config alone: the + Rust core never builds, wraps or falls back to a Python core. + """ from litellm.router_python_core import PythonCore + from litellm.rust_bridge.catalog import Route, RouteContext, decision + from litellm.rust_bridge.configuration import Decision - return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature + python_class: Final = core_class or PythonCore + chosen: Final = decision(RouteContext(Route.ROUTER)) + if chosen is not Decision.PYTHON: + from litellm.router_rust.core import RustCore + from litellm.rust_bridge._native import RustBridgeDeclined + + try: + return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported + except RustBridgeDeclined as declined: + # A config the Rust core cannot take keeps the router on Python while both cores exist. + if chosen is Decision.RUST_REQUIRED: + raise + verbose_router_logger.info("Router served by the Python core: %s", declined) + return python_class(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 83a1251d3c0..eb7d9edda34 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -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_``) 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_`` object as before. +Settings writes go through ``apply`` / ``aapply`` so the registry version moves; any other old +attribute write sets the core's ``_core_``, 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,7 +19,7 @@ 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. @@ -27,11 +29,11 @@ import inspect import sys import warnings from collections import Counter -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload -from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet +from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet, RoutingContext if TYPE_CHECKING: from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged @@ -39,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 @@ -129,6 +132,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 @@ -164,6 +170,55 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib return _DeprecatedAttribute(replacement, read) +def _async_picker(router: "Router") -> Callable[..., Awaitable[object]]: + """ + 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``. + """ + 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 + + 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: + return await router.select( + model, request_kwargs, messages=messages, input=input, specific_deployment=specific_deployment + ) + + return async_get_available_deployment + + +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: + """ + 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. + """ + 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") + + # (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() @@ -298,109 +353,112 @@ class DeprecatedRouterMembers: 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: 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)", _async_picker) # -- 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( @@ -420,62 +478,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) 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, @@ -486,17 +538,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, @@ -505,92 +561,96 @@ 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) 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 @@ -598,42 +658,42 @@ class DeprecatedRouterMembers: 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, @@ -642,27 +702,29 @@ 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 @@ -678,15 +740,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 @@ -709,31 +771,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( @@ -747,9 +809,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, @@ -757,7 +820,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 @@ -766,12 +829,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 ---------------------------------------------------------------------------------- @@ -808,8 +874,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()") @@ -822,7 +895,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, @@ -833,7 +906,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, @@ -843,7 +916,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=``, diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..f26f72b2a3b 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -353,6 +353,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, @@ -1199,8 +1204,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: @@ -1839,6 +1845,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 @@ -2064,6 +2071,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() @@ -3359,7 +3367,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")): diff --git a/litellm/router_python_host.py b/litellm/router_python_host.py index 8156740144e..8ab2c06b82f 100644 --- a/litellm/router_python_host.py +++ b/litellm/router_python_host.py @@ -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 diff --git a/litellm/router_rust/__init__.py b/litellm/router_rust/__init__.py new file mode 100644 index 00000000000..42cfc56aca4 --- /dev/null +++ b/litellm/router_rust/__init__.py @@ -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. +""" diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py new file mode 100644 index 00000000000..3020a1c29f7 --- /dev/null +++ b/litellm/router_rust/core.py @@ -0,0 +1,3599 @@ +""" +``RustCore``: the router core served by litellm-rust/crates/router-core through ``_native``. + +The Rust core owns the registry, selection, retries, fallbacks and cooldowns. It reaches outside +itself only through its host, and the host never calls back into a Python router: + +- every attempt runs through the host: ``router_attempt`` builds the request for the one + deployment Rust picked and ``PythonHost.invoke(op)`` calls ``litellm.`` with it +- the core records every outcome (cooldowns, allowed-fails counters, usage) itself, in the routing + cache the host builds and owns (``PythonHost.routing_cache``) +- sleeps, fallback access and budget checks, and cache reads and writes are host steps + +``RustCore`` does not build, wrap or read a ``PythonCore``. Settings are resolved from the +``RouterConfig`` and the host's process settings; the views (``registry()``, ``settings()``, +``health()``) are read from the Rust handle. A member the Rust core does not serve raises +``RustCoreUnsupported`` (an ``AttributeError``) instead of reaching a Python core, so deleting +``router_python_core.py`` cannot change what this core does. +""" + +import contextlib +import dataclasses +import inspect +import itertools +import json +import math +from collections.abc import Awaitable, Callable, Mapping, Sequence +from functools import partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._internal_context import service_target +from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_TARGET +from litellm.integrations.custom_guardrail import is_guardrail_intervention +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, + apply_remaining_usage_headers, + apply_response_model_id, + ensure_response_additional_headers, + prepare_response_for_header_attachment, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.routing_common_utils import ( + format_content_policy_fallback_hint, + format_fallback_outcome_message, + format_no_fallback_group_message, + truncate_fallback_error_detail, +) +from litellm.litellm_core_utils.routing_groups import validate_routing_strategy +from litellm.litellm_core_utils.secret_redaction import redact_string +from litellm.router_api import ( + ROUTER_INIT_ARGS, + SETTING_FIELDS, + ApplyResult, + ChangeSet, + RouterHealth, + RouterRegistry, + RouterSettings, + RoutingContext, + RoutingObservers, +) +from litellm.router_attempt import for_requested_model, request_from_call +from litellm.router_config import RouterConfig, RouterDynamicConfig, detached +from litellm.router_python_host import PythonHost +from litellm.router_rust.invoke.endpoints import ( + AttemptCall, + Plan, + RustEndpoints, + attempt_call, + pinned_video_deployment, + plan_for, + stamp_error, + unrouted_fields, +) +from litellm.router_rust.invoke.shaping import Shaping +from litellm.router_rust.invoke.slots import enter_deployment_slot, max_parallel_requests_limit +from litellm.router_rust.invoke.streaming import ( + StreamAttempt, + continuation_flag, + continuation_routing, + open_attempt_stream, + open_attempt_stream_sync, + register_continuation_prefill_check, + shape_attempt_kwargs, +) +from litellm.router_rust.invoke.walk import FallbackOutcomes, WalkMetadata, await_all, log_retry_into +from litellm.router_rust.steps.access import caller_fields, selection_refusal +from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks +from litellm.router_rust.steps.plugins import SYNC_PATH_MESSAGE, candidate_models, run_plugins +from litellm.router_rust.steps.precall import PreCallFacts, declares_no_vision, request_carries_image +from litellm.router_rust.steps.prerouting import PreRouter, PreRoutingStrategies +from litellm.router_rust.steps.selection import ( + SELECTOR_ATTRIBUTES, + StrategySelectors, + group_documents, + group_settings, + member_map, + normalize_strategy, + parse_groups, + pick_async, + pick_sync, + picked_id, + routing_pre_call_checks, +) +from litellm.router_rust.surface.alerting import initialize_alerting +from litellm.router_rust.surface.fallbacks import ( + MOCK_FLAGS, + RustFallbacks, + mock_error, + mock_failure, + refuse_malformed_fallbacks, + rejected_fallback_target, +) +from litellm.router_rust.surface.health import ( + RustHealthCallbacks, + cooled_in_running_loop, + failure_status, +) +from litellm.router_rust.surface.health_routing import ( + HEALTH_STATE_KEY, + HealthRouting, + failover_document, + fully_unhealthy_model_names, + resolved_allowed_fails_policy, +) +from litellm.router_rust.surface.health_routing import all_blocked as deployments_all_blocked +from litellm.router_rust.surface.health_routing import fully_blocked_model_names as blocked_model_names +from litellm.router_rust.surface.health_routing import is_blocked as deployment_is_blocked +from litellm.router_rust.surface.members import RustRegistryMembers +from litellm.router_rust.surface.registry import ( + DEFAULT_DEPLOYMENT_COPIED_ID, + DEFAULT_DEPLOYMENT_FAILING_ID, + INDEX_ATTRIBUTES, + BuildInstall, + HostRegistry, + RustRegistry, + default_deployment_call_error, + default_deployment_copy, + default_deployment_copy_error, + hosted_registry_view, + index_document, + served_default_deployment, +) +from litellm.router_rust.surface.replay import track_live_registry, untrack_live_registry +from litellm.router_rust.surface.retries import ( + GroupPolicies, + RustRetries, + group_policies, + is_policy_per_group, + policy_lookup_error, + raising_groups, + resolved_retry_policy, + static_group_policies, +) +from litellm.router_rust.surface.settings import ( + MISSING, + SettingsUpdate, + apply_construction_effects, + argument_attribute, + default_litellm_params, + initial_plain, + read_settings_update, + request_timeout, + settings_view, +) +from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET +from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery +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, + RouterErrors, + RouterGeneralSettings, + RouterRateLimitError, + RoutingGroup, + RoutingStrategy, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.caching.redis_cache 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 + from litellm.router_api import UntypedDict + from litellm.router_utils.cooldown_cache import CooldownCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + from litellm.rust_bridge._native import ( # pyright: ignore[reportPrivateUsage] # the stub's step shapes + RouterCall, + RouterCoreHandle, + ) + from litellm.types.router import ( + CustomRoutingStrategyBase, + DeploymentTypedDict, + GuardrailTypedDict, + PreRoutingHookResponse, + RoutingPlugin, + ) + +# Async endpoints the Rust core routes: select a deployment, then run the attempt on the host. +RUST_SERVED_OPS: Final[frozenset[str]] = frozenset( + { + "acompletion", + "atext_completion", + "aembedding", + "aresponses", + "aanthropic_messages", + "aimage_generation", + "atranscription", + "aspeech", + "arerank", + "amoderation", + } +) +_FAILURE_KINDS: Final = ( + "RateLimitError", + "ContextWindowExceededError", + "ContentPolicyViolationError", + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "BadRequestError", + "Timeout", + "APIConnectionError", + "InternalServerError", + "ServiceUnavailableError", + "BadGatewayError", + "GuardrailIntervention", + "CallerCredentialError", +) +_DECLINED: Final = frozenset({"MalformedRequestFieldError", "UnsupportedEndpointError"}) +_STEP: Final = TypeAdapter(Mapping[str, object]) +_PROVIDERS: Final = tuple(provider.value for provider in LlmProviders) +_DEFAULT_COOLDOWN_SECONDS: Final = 5.0 + +# Constructor arguments this core reads. Any other argument set to something other than its +# unset value keeps the router on the Python core while both exist (make_core decides that from +# the config alone, before any core is built). +_SERVED_ARGUMENTS: Final = frozenset( + { + "model_list", + "routing_strategy", + "routing_strategy_args", + "routing_groups", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_after", + "timeout", + "default_litellm_params", + "fallback_access_check", + "fallback_budget_check", + "enable_mid_stream_fallback_continuation", + "redis_url", + "redis_host", + "redis_port", + "redis_password", + "redis_db", + "cache_kwargs", + "set_verbose", + "debug_level", + # Read host-side by the endpoint plans (litellm/router_rust/invoke/endpoints.py). + "stream_timeout", + "cache_responses", + "assistants_config", + "search_tools", + "guardrail_list", + "default_priority", + "polling_interval", + "ignore_invalid_deployments", + # Read host-side by the pre-routing strategies (litellm/router_rust/steps/prerouting.py). + "auto_router_capability_limit", + # Decided in router-core (precall.rs); its facts come from litellm/router_rust/steps/precall.py. + "enable_pre_call_checks", + # Run host-side (litellm/router_rust/steps/plugins.py); router-core filters on what they leave. + "plugins", + "retry_policy", + "model_group_retry_policy", + # Health-check routing and the failover flags (litellm/router_rust/surface/health_routing.py). + "enable_health_check_routing", + "health_check_staleness_threshold", + "health_check_ignore_transient_errors", + "background_health_check_model_groups", + "enable_weighted_failover", + "fallback_on_unsupported_params", + # Built host-side as CustomLogger checks (litellm/router_rust/steps/callbacks.py); router-core runs + # their deployment filters through its callback_filters stage. + "optional_pre_call_checks", + "provider_budget_config", + "deployment_affinity_ttl_seconds", + "session_affinity_repin_on_failover", + "model_group_affinity_config", + # Kept as given (litellm/router_rust/surface/settings.py) and read host-side by the endpoint plans + # (default_litellm_params' caching groups) and the slot limits. + "client_ttl", + "caching_groups", + "default_max_parallel_requests", + # Built host-side into a Slack alerting logger (litellm/router_rust/surface/alerting.py). + "alerting_config", + # Kept as given and reported by settings(); nothing in the router reads its fields. + "router_general_settings", + } +) +# Arguments that may be passed at their default without asking for anything. +_UNSET_DEFAULTS: Final[Mapping[str, object]] = MappingProxyType( + { + "cache_responses": False, + "enable_pre_call_checks": False, + "enable_tag_filtering": False, + "retry_policy": None, + "model_group_retry_policy": {}, + } +) +# Settings a changeset may change. Static ones rebuild the Rust core (settings are fixed for a +# core's life); dynamic ones are handed to it as a new dynamic document. +_STATIC_SETTINGS: Final = frozenset( + { + "routing_strategy", + "routing_groups", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_after", + "timeout", + "enable_pre_call_checks", + } +) +_DYNAMIC_SETTINGS: Final = frozenset( + { + "default_deployment", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "model_group_alias", + "default_litellm_params", + "fallback_access_check", + "fallback_budget_check", + "routing_strategy_args", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + } +) +# Settings only the host reads (the endpoint plans, the slot limits, the budget limiter, the +# settings view). A changeset assigns them as given, the Rust core never sees them, and +# `update_settings` ignores them (none is runtime-updatable). +_HOST_SETTINGS: Final = frozenset( + { + "cache_responses", + "search_tools", + "guardrail_list", + "default_max_parallel_requests", + "provider_budget_config", + "health_check_ignore_transient_errors", + "router_general_settings", + "provider_default_deployment_ids", + "kubernetes_pod_discovery", + # Read by router-core's health-check filter: assigning one rebuilds the Rust core. + "enable_health_check_routing", + "background_health_check_model_groups", + } +) +# The constructor's default `router_general_settings`: one object every router built without the +# argument shares, as the Python core's default argument is. +_DEFAULT_GENERAL_SETTINGS: Final = RouterGeneralSettings() + + +# Plain settings `update_settings` sets that the Rust core's retry loop reads. +_RETRY_POLICY_SETTINGS: Final = ("retry_policy", "model_group_retry_policy") + + +# Settings converted by `Resolved._selection_setting`. +_SELECTION_SETTINGS: Final = frozenset( + { + "routing_strategy", + "routing_strategy_args", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + } +) +_NUMBER_SETTINGS: Final = frozenset( + {"num_retries", "max_fallbacks", "allowed_fails", "cooldown_time", "retry_after", "timeout"} +) + + +class RustCoreUnsupported(AttributeError): + """A member or request the Rust core does not serve. It is never forwarded to a Python core.""" + + +def _cache_target(key: str) -> str: + """The service target a routing cache read or write is attributed to, as the Python core's.""" + if key == HEALTH_STATE_KEY: + return HEALTH_CHECKS_TARGET + return ROUTER_COOLDOWNS_TARGET if key.endswith(":cooldown") else ROUTER_USAGE_TARGET + + +# --- failures ----------------------------------------------------------------------------------- + + +def failure_kind(error: BaseException) -> str: + names: Final = tuple(cls.__name__ for cls in type(error).__mro__) + if not _DECLINED.isdisjoint(names): + return "Declined" + if is_guardrail_intervention(error): # pyright: ignore[reportArgumentType] # every attempt error is an Exception + return "GuardrailIntervention" + return next((name for name in names if name in _FAILURE_KINDS), "APIError") + + +def _retry_after_seconds(error: BaseException) -> float | None: + response: Final = getattr(error, "response", None) + headers: Final[object] = getattr(error, "litellm_response_headers", None) or getattr(response, "headers", None) + get: Final[object] = getattr(headers, "get", None) + if headers is None or not callable(get): + return None + milliseconds: Final[object] = get("retry-after-ms") + seconds: Final[object] = get("retry-after") + try: + value: Final = float(str(milliseconds)) / 1000 if milliseconds is not None else float(str(seconds)) + except (TypeError, ValueError): + return None + return value if math.isfinite(value) and value >= 0 else None + + +def attempt_failure(error: BaseException) -> Mapping[str, object]: + status: Final = getattr(error, "status_code", None) + num_retries: Final = getattr(error, "num_retries", None) + return { + "ok": False, + "kind": failure_kind(error), + "status": status if isinstance(status, int) and 100 <= status <= 599 else None, + "retry_after": _retry_after_seconds(error), + "num_retries": num_retries if isinstance(num_retries, int) and num_retries >= 0 else None, + "detail": str(error)[:500], + "classes": [cls.__name__ for cls in type(error).__mro__], + "llm_provider": bool(getattr(error, "llm_provider", "")), + } + + +def total_tokens(response: object) -> int: + usage: Final = getattr(response, "usage", None) + tokens: Final = getattr(usage, "total_tokens", None) + return tokens if isinstance(tokens, int) else 0 + + +# --- settings ----------------------------------------------------------------------------------- + + +def _json_default(value: object) -> object: + dump: Final = getattr(value, "model_dump", None) + return dump(exclude_none=True) if callable(dump) else str(value) + + +def _plain(value: object) -> object: + """``value`` as plain JSON values, the shape the Rust core reads.""" + loaded: Final[object] = json.loads(json.dumps(value, default=_json_default)) # pyright: ignore[reportAny] # plain JSON + return loaded + + +def _document(value: object) -> Mapping[str, object]: + """A mapping as the plain JSON document the Rust core reads.""" + return as_str_mapping(_plain(value)) or {} + + +def _as_list(value: object) -> Sequence[object]: + return value if isinstance(value, list) else () # pyright: ignore[reportUnknownVariableType] # a list's items are read as objects + + +@dataclasses.dataclass(frozen=True, slots=True) +class Resolved: + """The router's settings, resolved from its arguments and the process settings the way the + Python router's constructor resolves them. Replaced, never changed, by ``apply``.""" + + routing_strategy: str | None + routing_strategy_args: Mapping[str, object] + num_retries: int + max_fallbacks: int + allowed_fails: int + explicit_allowed_fails: int | None + cooldown_time: float + disable_cooldowns: bool | None + allowed_fails_policy: object + retry_after: float + timeout: float | None + fallbacks: list[object] | None # mutable-ok: the setting as Python spells it + context_window_fallbacks: list[object] | None # mutable-ok: same + content_policy_fallbacks: list[object] | None # mutable-ok: same + model_group_alias: Mapping[str, object] + default_litellm_params: Mapping[str, object] + fallback_access_check: object + fallback_budget_check: object + enable_mid_stream_fallback_continuation: bool = False + routing_groups: "tuple[RoutingGroup, ...]" = () + enable_tag_filtering: bool = False + tag_filtering_match_any: bool = True + tag_routing_prefix: str = "" + enable_pre_call_checks: bool = False + routing_plugins: "tuple[RoutingPlugin, ...]" = () + retry_policy: object = None + model_group_retry_policy: object = None + health_routing: HealthRouting = dataclasses.field(default_factory=HealthRouting) + # `_core_default_deployment`, as assigned (detached). + default_deployment: object = None + # `_model_to_group` as the Python core last built it, kept once a raw `routing_groups` write + # replaced the groups without rebuilding it. + routing_group_members: Mapping[str, str] | None = None + # `_routing_groups_input`: the groups last passed to the constructor or `update_settings`, which + # a strategy change through `update_settings` re-inits the routing groups from. + routing_groups_input: object = None + assigned: frozenset[str] = frozenset() + + @classmethod + def of(cls, config: RouterConfig, dynamic: RouterDynamicConfig) -> "Resolved": + get: Final = config.get + # The process list is the router's own, as in the Python router: default_fallbacks are + # appended to it in place. + configured: Final[list[object] | None] = ( # mutable-ok: the router's live list, as Python keeps it + list(_as_list(get("fallbacks"))) or (dynamic.fallbacks if dynamic.fallbacks else None) # pyright: ignore[reportAssignmentType] # the process list as given + ) + refuse_malformed_fallbacks(configured) + default_fallbacks: Final = get("default_fallbacks") or dynamic.default_fallbacks + if isinstance(default_fallbacks, str): + raise TypeError( + f"default_fallbacks must be a list of model names, got '{default_fallbacks}'. Use ['{default_fallbacks}']" + ) + if default_fallbacks is not None and configured is not None: + configured.append({"*": default_fallbacks}) + fallbacks: Final[list[object] | None] = ( # mutable-ok: see above + [{"*": default_fallbacks}] if default_fallbacks is not None and configured is None else configured + ) + content_policy: Final = list(_as_list(get("content_policy_fallbacks"))) or ( + list(dynamic.content_policy_fallbacks) if dynamic.content_policy_fallbacks else None + ) + refuse_malformed_fallbacks(content_policy) + num_retries: Final = get("num_retries") + max_fallbacks: Final = get("max_fallbacks") + allowed_fails: Final = get("allowed_fails") + timeout: Final = get("timeout") + strategy: Final = get("routing_strategy") + validate_routing_strategy(strategy) # pyright: ignore[reportArgumentType] # validated as given + return cls( + # An explicit `routing_strategy=None` is stored as None (no strategy picks). + routing_strategy=None + if strategy is None and "routing_strategy" in config.arguments + else normalize_strategy(strategy) or "simple-shuffle", + routing_strategy_args=MappingProxyType(dict(as_str_mapping(get("routing_strategy_args")) or {})), + num_retries=int(str(num_retries)) + if num_retries is not None + else dynamic.num_retries + if dynamic.num_retries is not None + else 2, + max_fallbacks=int(str(max_fallbacks)) + if max_fallbacks is not None + else dynamic.max_fallbacks + if dynamic.max_fallbacks is not None + else dynamic.router_max_fallbacks, + allowed_fails=int(str(allowed_fails)) if allowed_fails is not None else dynamic.allowed_fails, + explicit_allowed_fails=int(str(allowed_fails)) if allowed_fails is not None else None, + cooldown_time=float(str(get("cooldown_time") or _DEFAULT_COOLDOWN_SECONDS)), + disable_cooldowns=_optional_bool(get("disable_cooldowns")), + allowed_fails_policy=resolved_allowed_fails_policy(get("allowed_fails_policy")), + retry_after=float(str(get("retry_after") or 0)), + timeout=float(str(timeout)) if timeout else dynamic.request_timeout, + fallbacks=fallbacks, + context_window_fallbacks=list(_as_list(get("context_window_fallbacks"))) + or (list(dynamic.context_window_fallbacks) if dynamic.context_window_fallbacks else None), + content_policy_fallbacks=content_policy, + model_group_alias=MappingProxyType(dict(as_str_mapping(get("model_group_alias")) or {})), + default_litellm_params=MappingProxyType(dict(as_str_mapping(get("default_litellm_params")) or {})), + fallback_access_check=get("fallback_access_check"), + fallback_budget_check=get("fallback_budget_check"), + enable_mid_stream_fallback_continuation=continuation_flag(get("enable_mid_stream_fallback_continuation")), + enable_tag_filtering=bool(get("enable_tag_filtering")), + tag_filtering_match_any=_true_unless_given(get("tag_filtering_match_any")), + tag_routing_prefix=str(get("tag_routing_prefix") or ""), + enable_pre_call_checks=bool(get("enable_pre_call_checks")), + routing_plugins=tuple(_as_list(get("plugins"))), # pyright: ignore[reportArgumentType] # the plugins as passed + routing_groups_input=get("routing_groups"), + routing_groups=parse_groups( + get("routing_groups"), + frozenset( + str(name) + for entry in _as_list(get("model_list")) + if (name := (as_str_mapping(entry) or {}).get("model_name")) + ), + frozenset(as_str_mapping(get("model_group_alias")) or {}), + ), + retry_policy=resolved_retry_policy(get("retry_policy")), + model_group_retry_policy=get("model_group_retry_policy"), + health_routing=HealthRouting.of(config), + ) + + def changed(self, values: Mapping[str, object], updated: frozenset[str] = frozenset()) -> "Resolved": + """``values`` applied; ``updated`` names those passed as ``ChangeSet.settings`` (which go + through ``update_settings``: only they may name ``routing_strategy_args``, and only they + validate a strategy).""" + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + unknown: Final = sorted( + set(values) + - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES} + - (updated & {"routing_strategy_args", "routing_groups", "enable_tag_filtering", "tag_routing_prefix"}) + ) + if unknown: + raise ValueError(f"not assignable router settings: {unknown}") + unserved: Final = sorted(set(values) - _STATIC_SETTINGS - _DYNAMIC_SETTINGS) + if unserved: + raise RustCoreUnsupported(f"the Rust router core does not serve the setting(s) {', '.join(unserved)} yet") + updates: Final[dict[str, object]] = {} # mutable-ok: built here, then applied once + for name, value in values.items(): + if name in {"model_group_alias", "default_litellm_params"}: + updates[name] = MappingProxyType(dict(as_str_mapping(value) or {})) + elif name in _NUMBER_SETTINGS: + updates[name] = self._number_setting(name, value) + if name == "allowed_fails": + updates["explicit_allowed_fails"] = updates[name] + elif name in _SELECTION_SETTINGS: + updates[name] = self._selection_setting(name, value, validate=name in updated) + elif name == "disable_cooldowns": + updates[name] = _optional_bool(value) + elif name == "enable_pre_call_checks": + updates[name] = bool(value) + elif name in ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks"): + updates[name] = list(_as_list(value)) or None + elif name == "default_deployment": + updates[name] = detached(value) + else: + updates[name] = value + if "routing_groups" in values: + # `_init_routing_groups` rebuilds the member map from the new groups. + updates["routing_group_members"] = None + return dataclasses.replace(self, **updates, assigned=self.assigned | frozenset(values)) # pyright: ignore[reportArgumentType] # each value converted to its field's type above + + def _number_setting(self, name: str, value: object) -> object: + """A numeric setting converted to its field's type.""" + if name in ("num_retries", "max_fallbacks", "allowed_fails"): + return int(str(value)) + if name == "cooldown_time": + # Stored as given (`_core_cooldown_time`), so a failure cools down for that long: + # 0 for no time at all; None leaves the cooldown cache's own default, set by the caller. + return float(str(value)) if value is not None else self.cooldown_time + if name == "retry_after": + return float(str(value or 0)) + return float(str(value)) if value is not None else None + + def _selection_setting(self, name: str, value: object, *, validate: bool) -> object: + """A selection setting converted to its field's type; a strategy changed through + ``update_settings`` is validated.""" + if name == "routing_strategy": + normalized: Final = normalize_strategy(value) or "simple-shuffle" + if validate and normalized != self.routing_strategy: + validate_routing_strategy(value) # pyright: ignore[reportArgumentType] # validated as given + return normalized + if name == "routing_strategy_args": + return MappingProxyType(dict(as_str_mapping(value) or {})) + if name == "tag_routing_prefix": + return str(value or "") + return bool(value) + + def static_document(self, model_list: Sequence[object]) -> Mapping[str, object]: + values: Final = { + "model_list": _plain(list(model_list)), + "routing_strategy": self.routing_strategy, + "cooldown_time": self.cooldown_time, + "disable_cooldowns": self.disable_cooldowns, + "allowed_fails_policy": _plain(self.allowed_fails_policy), + "retry_after": self.retry_after, + "timeout": self.timeout, + "routing_groups": group_documents(self.routing_groups) if self.routing_groups else None, + "enable_pre_call_checks": True if self.enable_pre_call_checks else None, + "routing_plugins": True if self.routing_plugins else None, + # `async_callback_filter_deployments` runs on every async candidate list. + "callback_filters": True, + "retry_policy": _plain(self.retry_policy), + "model_group_retry_policy": _plain(static_group_policies(self.model_group_retry_policy)) or None, + "routing_group_members": dict(self.routing_group_members) + if self.routing_group_members is not None + else None, + **self.health_routing.static_document(), + } + kept: Final = {name: value for name, value in values.items() if value is not None} + # An explicit `routing_strategy=None` reaches the Rust core as "None": no strategy picks. + return kept if self.routing_strategy is not None else {**kept, "routing_strategy": "None"} + + def dynamic_document(self, version: int, model_names: Sequence[str]) -> Mapping[str, object]: + values: Final = { + "version": version, + "model_group_alias": _plain(dict(self.model_group_alias)), + "fallbacks": _plain(self.fallbacks), + "context_window_fallbacks": _plain(self.context_window_fallbacks), + "content_policy_fallbacks": _plain(self.content_policy_fallbacks), + "num_retries": self.num_retries, + "max_fallbacks": self.max_fallbacks, + "allowed_fails": self._router_allowed_fails(), + "fallback_keys": self._fallback_keys(model_names), + "default_deployment": _plain(served_default_deployment(self.default_deployment)[0]), + "default_deployment_raises": True if served_default_deployment(self.default_deployment)[1] else None, + } + return {name: value for name, value in values.items() if value is not None} + + def _router_allowed_fails(self) -> int | None: + """``allowed_fails`` when it counts as set on the router (``_is_allowed_fails_set_on_router``): + passed explicitly, or different from the process default.""" + import litellm + + if self.explicit_allowed_fails is not None: + return self.explicit_allowed_fails + if self.allowed_fails != litellm.allowed_fails: + return self.allowed_fails + return None + + def _fallback_keys(self, model_names: Sequence[str]) -> Mapping[str, object]: + """The spellings a fallback key matches besides the group name (`get_fallback_model_group`).""" + from litellm.router_utils.fallback_event_handlers import ( + _provider_prefixed_model_group, # pyright: ignore[reportPrivateUsage] # the fallback walk's own spelling rule + ) + + keys: Final = tuple( # comprehension-ok: the first group of each fallback entry, across the three lists + str(next(iter(mapping))) + for entries in (self.fallbacks, self.context_window_fallbacks, self.content_policy_fallbacks) + for entry in _as_list(entries) + if (mapping := as_str_mapping(entry)) + ) + names: Final = {*model_names, *self.model_group_alias} + prefixed: Final = { + name: spelled for name in names if (spelled := _provider_prefixed_model_group(name, keys)) is not None + } + from litellm.types.router import LiteLLMParamsTypedDict + + return { + "providers": _PROVIDERS, + "provider_prefixed": prefixed, + "param_keys": list(LiteLLMParamsTypedDict.__annotations__), + } + + +def _alias_target(target: object) -> str | None: + mapping: Final = as_str_mapping(target) + if mapping is not None: + model: Final = mapping.get("model") + return model if isinstance(model, str) else None + return target if isinstance(target, str) else None + + +def _true_unless_given(value: object) -> bool: + return True if value is None else bool(value) + + +def _group_selectors(groups: "Sequence[RoutingGroup]") -> Sequence[tuple[str, str, Mapping[str, object]]]: + return [ + ( + group.group_name, + group.routing_strategy, + as_str_mapping(group.routing_strategy_args) or {}, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # RoutingGroup's untyped dict + ) + for group in groups + ] + + +def _optional_bool(value: object) -> bool | None: + return None if value is None else bool(value) + + +def unsupported_arguments(config: RouterConfig) -> frozenset[str]: + """Arguments the Rust core does not serve.""" + unserved: Final = frozenset( + name + for name, value in config.arguments.items() + if name not in _SERVED_ARGUMENTS and value not in (None, _UNSET_DEFAULTS.get(name)) + ) + return unserved + + +def _team_id(request: Mapping[str, object]) -> str | None: + metadata: Final = as_str_mapping(request.get("metadata") or request.get("litellm_metadata")) + team_id: Final = metadata.get("user_api_key_team_id") if metadata is not None else None + return team_id if isinstance(team_id, str) else None + + +# Request fallback settings, which replace the router's for the call. +_REQUEST_FALLBACK_KEYS: Final = ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks", "max_fallbacks") + + +def _strings(value: object) -> "tuple[str, ...]": + """Every string in ``value``, a name or a fallback list of names and dict entries.""" + if isinstance(value, str): + return (value,) + if isinstance(value, Mapping): + entries: Final[Mapping[object, object]] = value # pyright: ignore[reportUnknownVariableType] # read as objects + return tuple(itertools.chain.from_iterable(map(_strings, itertools.chain.from_iterable(entries.items())))) + if isinstance(value, list | tuple): + items: Final[Sequence[object]] = value # pyright: ignore[reportUnknownVariableType] # see above + return tuple(itertools.chain.from_iterable(map(_strings, items))) # pyright: ignore[reportUnknownArgumentType] # read as objects + return () + + +def _start_fields(request: Mapping[str, object]) -> Mapping[str, object]: + """The request fields the Rust core routes on.""" + fields: Final = { + key: request[key] + for key in ("specific_deployment", "num_retries", "cooldown_time") + if request.get(key) is not None + } + # `record_disable_fallbacks(kwargs, disable_fallbacks is True)`: only True turns them off. + disabled: Final = {"disable_fallbacks": True} if request.get("disable_fallbacks") is True else {} + lists: Final = {key: _plain(request[key]) for key in _REQUEST_FALLBACK_KEYS if key in request} + weights: Final = request.get("_router_weights") + team_id: Final = _team_id(request) + strategy: Final = request.get("routing_strategy") + override: Final = normalize_strategy(strategy) if isinstance(strategy, str | RoutingStrategy) else None + return { + **fields, + **({"routing_strategy": override} if override is not None else {}), + **disabled, + **({"router_weights": weights} if weights is not None else {}), + **lists, + **({"team_id": team_id} if team_id else {}), + **caller_fields(request), + } + + +# --- views -------------------------------------------------------------------------------------- + + +class RustSettings(RouterSettings): + """``settings()``: the resolved settings as a frozen snapshot.""" + + def guardrail(self, name: str) -> "GuardrailTypedDict": + router: Final = self._router + if not isinstance(router, _SettingsSource): + return super().guardrail(name) + return router.available_guardrail(name) # pyright: ignore[reportReturnType] # the configured guardrail entry + + def as_dict(self) -> Mapping[str, object]: + router: Final = self._router + if not isinstance(router, _SettingsSource): + return super().as_dict() + return router.settings_dict() + + +class _SettingsSource: + def settings_dict(self) -> Mapping[str, object]: + raise NotImplementedError + + def available_guardrail(self, name: str) -> object: + raise NotImplementedError + + +class RustHealth(RouterHealth): + """``health()``: cooldowns, fail counts and usage, read through the host's routing cache and + written through the Rust core.""" + + __slots__ = ("_core",) + + def __init__(self, core: "RustCore") -> None: # pyright: ignore[reportMissingSuperCall] # reads the Rust core, not a Router + self._core = core + + def __getattribute__(self, name: str) -> object: + if name in _HEALTH_SERVED or name.startswith("__") or name == "_core": + found: Final[object] = object.__getattribute__(self, name) # pyright: ignore[reportAny] # the view's own member + return found + raise RustCoreUnsupported(f"health().{name} is not served by the Rust router core", name=name, obj=self) + + @property + def store(self) -> "DualCache": + return self._core._cache # pyright: ignore[reportPrivateUsage] # the health view reads its own core + + @property + def cooldowns(self) -> "CooldownCache": + return self._core._cooldowns # pyright: ignore[reportPrivateUsage] # see above + + def failures_this_minute(self, deployment_id: str) -> int: + count: Final[object] = self.store.get_cache(f"{deployment_id}:fails", local_only=True) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return int(count) if isinstance(count, int | float) else 0 + + def record_failure(self, deployment_id: str, exception: Exception, exception_status: str | int) -> bool: + failure: Final = {**attempt_failure(exception), "status": failure_status(exception_status)} + handle: Final = self._core._handle # pyright: ignore[reportPrivateUsage] # see above + call: Final = handle.report_failure(deployment_id, failure, count_only=exception_status == "") + done: Final = self._core._drive_sync_steps(call) # pyright: ignore[reportPrivateUsage] # see above + return cooled_in_running_loop(bool(done.get("cooled"))) + + def team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": + return self._core._hosted._core_team_ptu_ceiling(team_id, model_group) # pyright: ignore[reportPrivateUsage] # the host registry's ported ceiling + + def split_priority(self, model: str, kwargs: Mapping[str, object]) -> tuple[int | None, Mapping[str, object]]: + """``_core_split_request_priority``, as the endpoints resolve a request's priority.""" + return self._core._endpoints._split_priority(model, kwargs) # pyright: ignore[reportPrivateUsage] # see store + + def priority_header_fields( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + """``priority_header_update``, as the endpoints read a priority header.""" + return self._core._endpoints.priority_header_update(model=model, kwargs=kwargs, header=header) # pyright: ignore[reportPrivateUsage] # see store + + @property + def health_checks(self) -> "DeploymentHealthCache": + return self._core._health_checks # pyright: ignore[reportPrivateUsage] # see store + + @property + def latency(self) -> "UntypedDict": # the live latency map, as the Python core keeps it + return self._core._hosted._core_deployment_latency_map # pyright: ignore[reportPrivateUsage] # see team_ptu_ceiling + + @staticmethod + def is_blocked(deployment: "Deployment") -> bool: + return deployment_is_blocked(deployment) + + @staticmethod + def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool: + return deployments_all_blocked(deployments) + + def fully_blocked_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set + return blocked_model_names(self._core._hosted._core_get_model_list() or []) # pyright: ignore[reportPrivateUsage, reportArgumentType] # see team_ptu_ceiling + + async def fully_unhealthy_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set + core: Final = self._core + return await fully_unhealthy_model_names( + core._hosted._core_get_model_list() or [], # pyright: ignore[reportPrivateUsage, reportArgumentType] # see team_ptu_ceiling + core._health_checks, # pyright: ignore[reportPrivateUsage] # see store + core._resolved.allowed_fails_policy, # pyright: ignore[reportPrivateUsage] # see store + ) + + async def usage( # pyright: ignore[reportIncompatibleMethodOverride] # the same two kinds + self, model_group: str, *, kind: Literal["current", "remaining"] = "current" + ) -> tuple[int | None, int | None] | Mapping[str, int]: + """The group's usage this minute, read from the routing cache over the host registry's + live model list, as ``get_model_group_usage`` / ``get_remaining_model_group_usage`` read it.""" + hosted: Final = self._core._hosted # pyright: ignore[reportPrivateUsage] # see above + if kind == "current": + return await hosted._core_get_model_group_usage(model_group) # pyright: ignore[reportPrivateUsage] # the ported usage read + if kind == "remaining": + return await hosted._core_get_remaining_model_group_usage(model_group) # pyright: ignore[reportPrivateUsage] # see above + raise ValueError(f"unknown usage kind {kind!r}") + + +# Deprecated attribute writes (``router. = ...`` reaches the core as ``_core_``) of +# registry state this core keeps on its host registry, and of routing state it rebinds. +_HOSTED_STATE: Final = frozenset( + { + "_core_model_names", + "_core_deployment_names", + "_core_team_public_model_names", + "_core_model_name_to_deployment_indices", + "_core_team_model_to_deployment_indices", + "_core_cached_deployment_model_info", + "_core_pattern_router", + "_core_team_pattern_routers", + "_core_routing_groups", + "_core_auto_routers", + "_core_complexity_routers", + "_core_quality_routers", + "_core_adaptive_routers", + "_core_routing_plugins", + "_core_deployment_latency_map", + } +) +_REBOUND_STATE: Final = MappingProxyType( + { + "_core_cache": "_cache", + "_core_cooldown_cache": "_cooldowns", + "_core_health_state_cache": "_health_checks", + "_core_slack_alerting_logger": "_slack_alerting_logger", + } +) + + +_HEALTH_SERVED: Final = frozenset( + { + "store", + "cooldowns", + "failures_this_minute", + "record_failure", + "usage", + "team_ptu_ceiling", + "split_priority", + "priority_header_fields", + "health_checks", + "latency", + "is_blocked", + "all_blocked", + "fully_blocked_model_names", + "fully_unhealthy_model_names", + } +) + + +# --- the core ----------------------------------------------------------------------------------- + + +_NO_INPUT: Final = object() +_EndpointImpl: TypeAlias = Callable[..., Awaitable[object]] +# A strategy select step's strategy, selector, keyed name and candidate deployments. +_StrategyInputs: TypeAlias = tuple[str, object, str, list[dict[str, object]]] + + +class RustCore( + _SettingsSource, RustRegistryMembers, RustRetries, RustFallbacks, RustHealthCallbacks, RustOptionalChecks +): + """A ``RouterCore`` served by the Rust router core over a ``PythonHost``. It holds no Python core.""" + + def __init__(self, config: RouterConfig, host: PythonHost, resolved: Resolved) -> None: + from litellm.router_utils.cooldown_cache import CooldownCache + + self._config: Final = config + self._host: Final = host + apply_construction_effects(config, host) + self._resolved: Resolved = resolved + # `cooldown_time` as the Python core stores it (`cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS`, + # then each value assigned or updated, as given): what `settings()` reads back. + self._cooldown_given: object = config.get("cooldown_time") or DEFAULT_COOLDOWN_TIME_SECONDS + # Settings the Python core keeps as plain attributes, and settings deleted with `del`. + self._plain: Final = initial_plain(config) + self._deleted: frozenset[str] = frozenset() + self._request_timeout: Final = request_timeout(config) + legacy: Final[dict[str, object]] = {"deployment_stats": {}} # mutable-ok: the core's own legacy attributes + self._legacy: Final = legacy + self._health_view: RustHealth | None = None + self._cache: DualCache = host.routing_cache(config) + self._cooldowns: CooldownCache = CooldownCache( + cache=self._cache, + default_cooldown_time=self._cooldown_given, # pyright: ignore[reportArgumentType] # the constructor's value, as the Python core passes it + ) + self._health_checks: DeploymentHealthCache = resolved.health_routing.health_checks(self._cache) + self._version: int = 1 + self._retired: bool = False + self._settings: RustSettings | None = None + self._tag_filtering_models: tuple[int, bool] | None = None + # The registry indexes written through the deprecated attributes (`_core_`), and the + # document router-core last received for them. + self._index_written: frozenset[str] = frozenset() + self._index_sent: Mapping[str, object] | None = None + self._selectors: Final = StrategySelectors(host, self._cache) + self._selectors.init_default(resolved.routing_strategy, resolved.routing_strategy_args) + self._selectors.init_groups(_group_selectors(resolved.routing_groups)) + self._endpoints: Final = RustEndpoints(self, config) + # The resolver every picked deployment with `kubernetes_pod_discovery` goes through. + self._pod_discovery: KubernetesPodDiscovery = KubernetesPodDiscovery() + general: Final = config.get("router_general_settings", _DEFAULT_GENERAL_SETTINGS) + self._general_settings: object = general or RouterGeneralSettings() + self._registry_view: RouterRegistry | None = None + self._hosted: Final = HostRegistry( + host, + model_group_alias=dict(resolved.model_group_alias), # pyright: ignore[reportArgumentType] # the configured alias map + ignore_invalid_deployments=bool(config.get("ignore_invalid_deployments")), + on_change=self._sync_registry, + ) + self._hosted._core_cache = self._cache # pyright: ignore[reportPrivateUsage] # usage reads come from the cache the Rust core writes + ctx: Final = RoutingContext(self) # pyright: ignore[reportArgumentType] # the strategies read registry(), settings(), health() and select() + strategies: Final = PreRoutingStrategies( + self._hosted, host, ctx, config.get("auto_router_capability_limit"), self._cache + ) + self._hosted.strategies = strategies + self._prerouter: Final = PreRouter(strategies, self._hosted, host, ctx, self._cache) + # `register_plugin`: a custom strategy that replaces every pick (the host answers the + # Rust core's pre_route step with its deployment). + self._custom_strategy: CustomRoutingStrategyBase | None = None + self._custom_discovery: KubernetesPodDiscovery | None = None + # router-core builds the deployments, asking the host registry about each one: the handle + # starts empty, and is rebuilt over the built list for the settings the list's names feed. + self._handle_default_ids: tuple[str, ...] = () # the provider default ids the handle was built with + self._handle: RouterCoreHandle = self._new_handle([]) + self._hosted.build_runner = self._run_build + self._hosted.changes_runner = self._apply_registry_changes + self._hosted.handle_lookup = self._live_handle + self._hosted.build(_as_list(config.get("model_list")) if config.get("model_list") is not None else None) + self._hosted.set_routing_groups(resolved.routing_groups) + track_live_registry(self._hosted) + self._hosted.routable_lookup = self._routable + self._hosted.max_parallel_requests_lookup = self._max_parallel_requests_client + self._hosted._core_routing_plugins = list(resolved.routing_plugins) # pyright: ignore[reportPrivateUsage] # the core fills its own host registry + self._rebuild() + self._join_callback_holders() + self._checks: Final = OptionalChecks(self, host, self._cache, self._hosted, ctx, config.arguments) + self._hosted.checks = self._checks + strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback + self._checks.build(config.get("optional_pre_call_checks"), config.get("model_list")) + self._slack_alerting_logger: SlackAlerting | None = initialize_alerting(host, config.get("alerting_config")) # pyright: ignore[reportArgumentType] # the argument as passed + if self._slack_alerting_logger is not None: + self._checks.slack_alerting_loggers = (self._slack_alerting_logger,) + if resolved.enable_mid_stream_fallback_continuation: + register_continuation_prefill_check(host) + if resolved.routing_strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(ctx, resolved.routing_strategy_args) # pyright: ignore[reportArgumentType] # the core's RoutingContext + + @property + def _registry_state(self) -> HostRegistry: + return self._hosted + + @property + def _python_host(self) -> PythonHost: + return self._host + + @classmethod + def build(cls, config: RouterConfig, host: PythonHost) -> "RustCore": + from litellm.rust_bridge._native import RustBridgeDeclined + + extra: Final = unsupported_arguments(config) + if extra: + raise RustBridgeDeclined(f"router arguments not served by the Rust core: {sorted(extra)}") + RustEndpoints.prepare_host(config, host) + return cls(config, host, Resolved.of(config, host.dynamic())) + + @property + def enable_mid_stream_fallback_continuation(self) -> bool: + return self._resolved.enable_mid_stream_fallback_continuation + + def _set_continuation(self, enabled: bool) -> None: + """``_set_mid_stream_fallback_continuation``: an enabled flag registers the process-wide prefill check.""" + self._resolved = dataclasses.replace(self._resolved, enable_mid_stream_fallback_continuation=enabled) + self._settings = None + if enabled: + register_continuation_prefill_check(self._host) + + def __getattr__(self, name: str) -> object: + if name in SELECTOR_ATTRIBUTES.values(): + return self._selectors.attribute(name) + # Only reached for members this class does not define: the endpoints served through the + # host, else refused. Never forwarded anywhere. + endpoints: Final[object] = self.__dict__.get("_endpoints") + endpoint: Final = endpoints.member(name) if isinstance(endpoints, RustEndpoints) else None + if endpoint is not None: + return endpoint + attribute: Final = self._settings_attribute(name) if "_resolved" in self.__dict__ else MISSING + if attribute is not MISSING: + return attribute + raise RustCoreUnsupported(f"'{name}' is not served by the Rust router core", name=name, obj=self) + + def _settings_attribute(self, name: str) -> object: + """A router attribute that reads a setting, the constructor's arguments or a constant.""" + r: Final = self._resolved + resolved: Final[Mapping[str, object]] = { + "num_retries": r.num_retries, + "max_fallbacks": r.max_fallbacks, + "timeout": r.timeout, + "retry_after": r.retry_after, + "explicit_allowed_fails": r.explicit_allowed_fails, + "routing_strategy_args": r.routing_strategy_args, + "request_timeout": self._request_timeout, + "default_priority": self._endpoints._default_priority, # pyright: ignore[reportPrivateUsage] # parsed by the endpoint plans + } + if name in resolved: + return resolved[name] + if name in self._plain: + return self._plain[name] + if name in self._legacy: + return self._legacy[name] + if name == "healthy_deployments": + return self._hosted._core_model_list # pyright: ignore[reportPrivateUsage] # the live list, as the core keeps it + if name == "get_model_group_io_token_usage": + return self._hosted.get_model_group_io_token_usage + if name == "scheduler": + return self._endpoints._get_scheduler() # pyright: ignore[reportPrivateUsage] # the endpoint plans' scheduler + if name == "chat": + from litellm.main import Chat + + self._legacy["chat"] = Chat(params=default_litellm_params(self._config), router_obj=self) # pyright: ignore[reportArgumentType] # Chat calls the router's endpoints, which this core serves + return self._legacy["chat"] + return argument_attribute(self._config, name) + + @property + def _shaping(self) -> Shaping: + return self._endpoints.shaping(self._resolved) + + # -- building -------------------------------------------------------------------------------- + + def _new_handle(self, model_list: Sequence[object]) -> "RouterCoreHandle": + from litellm.rust_bridge._native import RouterCoreHandle + + names: Final = [ + str(name) + for entry in model_list + if isinstance(name := (as_str_mapping(entry) or {}).get("model_name"), str) + ] + default_ids: Final = self._provider_default_ids() + self._handle_default_ids = default_ids + static: Final = { + **self._resolved.static_document(model_list), + **{name: value for name, value in failover_document(self._plain).items() if value is not None}, + **({"provider_default_deployment_ids": list(default_ids)} if default_ids else {}), + } + return RouterCoreHandle(static, self._resolved.dynamic_document(self._version, names)) + + def _provider_default_ids(self) -> tuple[str, ...]: + """``provider_default_deployment_ids`` as the cooldown check reads it (``deployment in + settings().provider_default_deployment_ids``): the ids it holds now. router-core holds them + with its static settings.""" + handle: Final = self._live_handle() + return () if handle is None else tuple(handle.provider_default_ids()) + + def _pick_discovery(self, *, sync: bool) -> KubernetesPodDiscovery | None: + """The pod discovery a pick resolves through: the one bound when a custom strategy was + registered (``resolve_pods_after_bound``), else the router's own (``resolve_pods_after``); + none for an async pick assigned through the deprecated attribute, which Python calls bare.""" + if not sync and callable(self.__dict__.get("_core_async_get_available_deployment")): + return None + if self._custom_strategy is not None and self._custom_discovery is not None: + return self._custom_discovery + return self._pod_discovery + + # -- the plugin context (``RoutingContext``) ----------------------------------------------- + + def _core_arm_routing_read_prefetch(self, model: str, request_kwargs: Mapping[str, object] | None = None) -> None: + """``RoutingContext.arm_read_prefetch``: declare the cooldown reads (and, for + usage-based-routing-v2, the usage reads) the pick for ``model`` makes on the request's + Redis batch, so admission's flush carries them. A miss costs nothing.""" + from litellm.router_utils.routing_read_batch import RoutingPrefetch + + try: + usage_selector: Final = self._prefetch_usage_selector(model, request_kwargs) + deployments: Final = self._hosted._core_get_model_list(model_name=model) # pyright: ignore[reportPrivateUsage] # see __init__ + if deployments: + RoutingPrefetch.arm(RoutingContext(self), usage_selector, deployments) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # the prefetch reads registry() and health(); its deployments list is untyped + except Exception as e: # noqa: BLE001 # a prefetch is an optimisation, never a reason to fail the request + from litellm._logging import verbose_router_logger + + verbose_router_logger.debug( + "routing read prefetch not armed for %s: %s", + str(model).replace("\r", "").replace("\n", ""), + str(e).replace("\r", "").replace("\n", ""), + ) + + def _prefetch_usage_selector(self, model: str, request_kwargs: Mapping[str, object] | None) -> object: + """``_get_routing_context``'s selector when it is a usage-based-routing-v2 one: the request's + strategy override, else the routing group ``model`` resolves into, else the router's own.""" + from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 + + strategy, selector = self._prefetch_strategy_and_selector(model, request_kwargs or {}) + if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2): + return selector + return None + + def _prefetch_strategy_and_selector(self, model: str, request: Mapping[str, object]) -> tuple[str | None, object]: + raw: Final = request.get("routing_strategy") + override: Final = normalize_strategy(raw) if isinstance(raw, str | RoutingStrategy) else None + if override is not None and override in {"simple-shuffle", *SELECTOR_ATTRIBUTES}: + return override, self._selectors.selector_for(override, "override", None, request) + groups: Final = self._hosted._core_routing_groups # pyright: ignore[reportPrivateUsage] # see __init__ + resolved: Final = self._hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # see __init__ + group_name: Final = ( + resolved + if self._hosted.get_routing_group(resolved) is not None + else next( + ( + group.group_name + for group in groups.values() + if group.routing_strategy != "priority" and resolved in group.models + ), + None, + ) + ) + if group_name is None: + return normalize_strategy(self._resolved.routing_strategy), self._selectors.default + if groups[group_name].routing_strategy == "priority": + return None, None + group_strategy: Final = normalize_strategy(groups[group_name].routing_strategy) + return group_strategy, self._selectors.selector_for(group_strategy or "", "group", group_name, None) + + async def _core_async_pre_routing_hook( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: the request kwargs the hook stamps + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + ) -> "PreRoutingHookResponse | None": + """``RoutingContext.pre_routing_hook``: the pre-routing strategies' answer for ``model`` + (the tier group and its params), with the routing plugins run first.""" + r: Final = self._resolved + owner: Final[object] = self.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return await self._prerouter._hook( # pyright: ignore[reportPrivateUsage] # the core's own pre-router + model, + request_kwargs, + messages=messages, + input=input, + fallback_depth=0, + enable_tag_filtering=r.enable_tag_filtering, + match_any=r.tag_filtering_match_any, + public_router=public if public is not None else self, + run_plugins=partial(self._run_plugins, request_kwargs, messages) + if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list + else None, + ) + + def _core_log_retry( + self, + kwargs: dict[str, object], # mutable-ok: the request kwargs the record joins + e: Exception, + ) -> dict[str, object]: # mutable-ok: the same kwargs, as Python's returns them + """``RoutingContext.log_retry``.""" + return log_retry_into(kwargs, e, self._deployment_params_base_model) + + def _deployment_params_base_model(self, deployment_id: str) -> str | None: + """``_deployment_params_base_model``: the ``base_model`` a deployment declares under ``litellm_params``.""" + deployment: Final = self._hosted._core_get_deployment(model_id=deployment_id) # pyright: ignore[reportPrivateUsage] # see __init__ + if deployment is None: + return None + base_model: Final = (deployment.litellm_params.model_extra or {}).get("base_model") + return base_model if isinstance(base_model, str) else None + + def _sync_routing_state(self) -> None: + """Before a pick: the plugin stage and the registry indexes written through the deprecated + attributes, as the Python core reads both live on every pick.""" + self._sync_plugin_stage() + self._sync_index() + + def _sync_index(self) -> None: + """Hand router-core the registry indexes written through the deprecated attributes + (``model_names``, ``model_name_to_deployment_indices``, ``team_model_to_deployment_indices``) + as the host registry holds them now (its own deployment changes update the written objects, + as the Python core's do), whenever they differ from what it holds.""" + if not self._index_written: + return + document: Final = index_document(self._hosted, self._index_written) + if document != self._index_sent: + _ = self._handle.apply({"index": document}) + self._index_sent = document + + def _sync_plugin_stage(self) -> None: + """Give router-core the routing-plugin stage when the live plugin list (``routing_plugins``, + which callers may assign or append to) gains its first plugin, and drop it when the list + empties: the Python core checks the list on every pick.""" + live: Final = _as_list(self._hosted._core_routing_plugins) # pyright: ignore[reportPrivateUsage] # the live plugin list + if bool(live) == bool(self._resolved.routing_plugins): + return + self._resolved = dataclasses.replace(self._resolved, routing_plugins=tuple(live)) # pyright: ignore[reportArgumentType] # the plugins as listed + self._rebuild() + + def _rebuild(self) -> None: + """A new Rust core for changed static settings, over the same deployments and the + providers the old one's builds resolved for them.""" + old: Final = self._handle + model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ + self._handle = self._new_handle(model_list) + self._handle.adopt_providers(old) + old.retire() + # The live list and indexes read the new handle's (the same entries, kept as the same objects). + self._hosted.refresh() + # The new handle builds its indexes from the deployments: written ones are handed again. + self._index_sent = None + + def _run_build( + self, + entries: Sequence[object], + path: str, + resolve_secrets: bool, + install: BuildInstall, + answer: Callable[[Mapping[str, object]], object], + ) -> None: + """Drive one router-core deployment build (``RouterCoreHandle.build``), answering each of + its steps with ``answer``; router-core installs each deployment as it accepts it, and the + host registry reads the list and indexes back before answering the next step.""" + call: Final = self._handle.build(_plain(list(entries)), path, resolve_secrets, install) # pyright: ignore[reportArgumentType] # plain JSON entries + outcome: object = None # rebind-ok: each step's answer + while True: + step: Mapping[str, object] = call.next(outcome) # pyright: ignore[reportArgumentType, reportAssignmentType] # the next step + self._hosted.refresh() + if step["kind"] == "done": + break + if step["kind"] == "failed": + raise ValueError(str(step.get("error"))) + outcome = _plain(answer(step)) + + def _live_handle(self) -> "RouterCoreHandle | None": + """The handle the live model_list reads, once the constructor made one.""" + from litellm.rust_bridge._native import RouterCoreHandle + + handle: Final[object] = self.__dict__.get("_handle") + return handle if isinstance(handle, RouterCoreHandle) else None + + def _apply_registry_changes(self, changes: Mapping[str, object]) -> None: + """Apply a host registry write router-core's build did not make (a removal is a delete, a + swapped list a replace) as a changeset.""" + _ = self._handle.apply(_plain(dict(changes))) # pyright: ignore[reportArgumentType] # a plain JSON changeset + + def _sync_registry(self) -> None: + """After a host registry write: router-core already holds it (the build installed what it + added, a removal came as a changeset), so only the settings it feeds move.""" + from litellm.rust_bridge._native import RouterCoreHandle + + self._tag_filtering_models = None + handle: Final[object] = self.__dict__.get("_handle") + # A wildcard deployment added an id the cooldown check skips. + if isinstance(handle, RouterCoreHandle) and self._provider_default_ids() != self._handle_default_ids: + self._rebuild() + + # -- entry points ---------------------------------------------------------------------------- + + def registry(self) -> RouterRegistry: + """The registry view, built once (``registry()`` returns the same object each time).""" + if self._registry_view is None: + self._registry_view = hosted_registry_view(self._hosted) + return self._registry_view + + def settings(self) -> RustSettings: + """A cached snapshot, replaced when a changeset changes a setting.""" + if self._settings is None: + self._settings = self._settings_snapshot() + return self._settings + + def _settings_snapshot(self) -> RustSettings: + r: Final = self._resolved + snapshot: Final = RustSettings( + routing_strategy=r.routing_strategy, + fallbacks=r.fallbacks, + context_window_fallbacks=r.context_window_fallbacks, + content_policy_fallbacks=r.content_policy_fallbacks, + default_deployment=r.default_deployment, + default_litellm_params=dict(r.default_litellm_params) + if "default_litellm_params" in r.assigned + else default_litellm_params(self._config), + default_max_parallel_requests=self._endpoints.default_max_parallel_requests, # pyright: ignore[reportArgumentType] # the setting as given + cooldown_time=self._cooldown_given, # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + allowed_fails=r.allowed_fails, + allowed_fails_policy=r.allowed_fails_policy, # pyright: ignore[reportArgumentType] # the policy as passed + disable_cooldowns=r.disable_cooldowns, + enable_pre_call_checks=r.enable_pre_call_checks, + enable_health_check_routing=r.health_routing.enable_health_check_routing, + enable_tag_filtering=r.enable_tag_filtering, + tag_filtering_match_any=r.tag_filtering_match_any, + tag_routing_prefix=r.tag_routing_prefix, + fallback_access_check=r.fallback_access_check, # pyright: ignore[reportArgumentType] # the check as passed + fallback_budget_check=r.fallback_budget_check, # pyright: ignore[reportArgumentType] # the check as passed + cache_responses=self._endpoints.cache_responses, + search_tools=self._endpoints.search_tools, # pyright: ignore[reportArgumentType] # the configured entries, as passed + guardrail_list=self._endpoints.guardrail_list, # pyright: ignore[reportArgumentType] # same + provider_default_deployment_ids=self._hosted._core_provider_default_deployment_ids, # pyright: ignore[reportPrivateUsage] # the host registry's live list + health_check_ignore_transient_errors=r.health_routing.health_check_ignore_transient_errors, + background_health_check_model_groups=r.health_routing.background_health_check_model_groups, + router_general_settings=self._general_settings, # pyright: ignore[reportArgumentType] # the setting as given + provider_budget_config=self._checks.provider_budget_config, # pyright: ignore[reportArgumentType] # the budget limiter's live config + _router=self, # pyright: ignore[reportArgumentType] # the snapshot reads as_dict from this core + kubernetes_pod_discovery=self._pod_discovery, + ) + if not self._deleted: + return snapshot + # A deleted setting reads as unset. + return dataclasses.replace( + snapshot, + **dict.fromkeys(self._deleted, None), # pyright: ignore[reportArgumentType] # settings fields only + _unset=snapshot._unset | self._deleted, # pyright: ignore[reportPrivateUsage] # the snapshot's own unset set + ) + + def available_guardrail(self, name: str) -> object: + return self._endpoints.available_guardrail(name) + + def settings_dict(self) -> Mapping[str, object]: + r: Final = self._resolved + latency: Final = self._selectors.latency_args() if r.routing_strategy == "latency-based-routing" else None + # `get_settings` reports the latency selector's full routing args. + args: Final = latency if latency is not None else dict(r.routing_strategy_args) + values: Final = { + "routing_strategy_args": args, + "routing_strategy": r.routing_strategy, + "allowed_fails": r.allowed_fails, + "cooldown_time": self._cooldown_given, + "num_retries": r.num_retries, + "timeout": r.timeout, + "retry_after": r.retry_after, + "fallbacks": r.fallbacks, + "context_window_fallbacks": r.context_window_fallbacks, + "model_group_alias": dict(r.model_group_alias), + "enable_mid_stream_fallback_continuation": r.enable_mid_stream_fallback_continuation, + "enable_tag_filtering": r.enable_tag_filtering, + "tag_routing_prefix": r.tag_routing_prefix, + "routing_groups": group_settings(r.routing_groups), + } + return settings_view(values, self._plain, self._deleted) + + def health(self) -> RustHealth: + """The health view, built once.""" + if self._health_view is None: + self._health_view = RustHealth(self) + return self._health_view + + def observe(self, logging_obj: "LiteLLMLogging | None" = None, *, request_scoped: bool = False) -> RoutingObservers: + """The routing observers: the Slack alerting logger ``alerting_config`` built, if any. + ``request_scoped=True`` takes the router's callbacks off the host lists; passing a request's + ``logging_obj`` binds its selectors and policies to it. This core registers no outcome + callbacks, so it binds none either.""" + if request_scoped: + self._serve_through_request_logging() + if logging_obj is not None: + self._checks.bind_to_request_logging(logging_obj) + return RoutingObservers(slack_alerting_logger=self._slack_alerting_logger) + + def register_plugin(self, plugin: "CustomRoutingStrategyBase") -> None: + """``_core_set_custom_routing_strategy``: ``plugin`` picks every deployment from now on, until + the routing strategy changes.""" + self._custom_strategy = plugin + # `resolve_pods_after_bound`: its picks resolve pods through the discovery bound now. + self._custom_discovery = self._pod_discovery + # It replaces an async pick assigned through the deprecated attribute, as Python's does. + if "_core_async_get_available_deployment" in self.__dict__: + del self.__dict__["_core_async_get_available_deployment"] + + def set_custom_routing_strategy(self, CustomRoutingStrategy: "CustomRoutingStrategyBase") -> None: + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the deprecated alias notes its use like Python's + ) + + _note("set_custom_routing_strategy", "register_plugin(strategy)") + self.register_plugin(CustomRoutingStrategy) + + def retire(self) -> None: + self._retired = True + untrack_live_registry(self._hosted) + self._retire_callbacks() + if self._hosted.strategies is not None: + self._hosted.strategies.retire() + self._handle.retire() + + def apply(self, changes: ChangeSet) -> ApplyResult: + """``apply_changes``: assigned settings, then ``update_settings``, upserts, deletes and the + Redis cache, in that order. A refused changeset changes nothing.""" + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + if changes.refresh_model_info: + raise TypeError("ChangeSet(refresh_model_info=True) needs aapply()") + unknown: Final = sorted(set(changes.assign) - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES}) + if unknown: + raise ValueError(f"not assignable router settings: {unknown}") + unserved: Final = sorted(set(changes.assign) - _STATIC_SETTINGS - _DYNAMIC_SETTINGS - _HOST_SETTINGS) + if unserved: + raise RustCoreUnsupported(f"the Rust router core does not serve the setting(s) {', '.join(unserved)} yet") + update: Final = read_settings_update(changes.settings) + # `apply_changes` assigns before it runs `update_settings`: a budget config assigned beside + # `optional_pre_call_checks` builds the budget limiter with the new config. + self._assign_host_settings(changes.assign) + if update.optional_pre_call_checks is not MISSING: + self._checks.set(update.optional_pre_call_checks) + if update.continuation is not MISSING: + self._set_continuation(continuation_flag(update.continuation)) + self._apply_settings(changes.assign, update) + # Upserts first, then deletes (`apply_changes`), through the host registry, which hands + # each change to the Rust core. + upserted: Final = [self._hosted._core_upsert_deployment(deployment=d) for d in changes.upsert] # pyright: ignore[reportPrivateUsage] # see __init__ + deleted: Final = [self._hosted._core_delete_deployment(id=i) for i in changes.delete] # pyright: ignore[reportPrivateUsage] # see __init__ + if changes.redis_cache is not None: + # `_core_update_redis_cache`: the routing state's cache gains the Redis cache. + self._cache.attach_redis_cache(changes.redis_cache) + self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache) + return ApplyResult( + version=self._hosted._registry_version, # pyright: ignore[reportPrivateUsage] # see __init__ + upserted=upserted, + deleted=deleted, # pyright: ignore[reportArgumentType] # the removed model_list entries, as the core returns them + ) + + async def aapply(self, changes: ChangeSet) -> ApplyResult: + """``apply``, then the model info refresh, which needs I/O.""" + result: Final = self.apply(dataclasses.replace(changes, refresh_model_info=False, refresh_client=None)) + if changes.refresh_model_info: + await self._hosted.arefresh_model_info(changes.refresh_client) + return result + + def _apply_settings(self, assign: Mapping[str, object], update: SettingsUpdate) -> None: + policies: Final = {name: update.plain[name] for name in _RETRY_POLICY_SETTINGS if name in update.plain} + values: Final = { + **{name: value for name, value in assign.items() if name not in _HOST_SETTINGS}, + **update.values, + } + given: Final = dict(update.values) # the settings passed through `update_settings` + groups_input: Final[object] = ( + update.routing_groups if update.routing_groups is not MISSING else self._resolved.routing_groups_input + ) + if update.routing_groups is not MISSING or self._changes_strategy(update.values): + # `update_settings` re-inits the routing groups from the groups passed, else (on a + # strategy change) from the groups last passed, whatever a raw write assigned since. + given["routing_groups"] = groups_input + values["routing_groups"] = parse_groups( + groups_input, + frozenset(str(d.get("model_name")) for d in self._handle.deployments()), + frozenset(self._resolved.model_group_alias), + ) + flags_before: Final = failover_document(self._plain) + self._plain.update(update.plain) + if failover_document(self._plain) != flags_before: + self._rebuild() + if policies: + self._set_retry_policies(policies) + if self._deleted & set(values): + self._deleted = self._deleted - set(values) + self._settings = None + if values: + previous: Final = self._resolved + resolved: Final = dataclasses.replace( + previous.changed(values, frozenset(given)), routing_groups_input=groups_input + ) + static_changed: Final = bool(_STATIC_SETTINGS & set(values)) + self._resolved = resolved + if set(values) & set(SETTING_FIELDS): + self._settings = None # the snapshot holds settings fields, not the plain attributes + if "model_group_alias" in values: + self._hosted._core_model_group_alias = dict(resolved.model_group_alias) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # see __init__ + self._apply_strategy(previous, resolved, given) + if "routing_groups" in values: + self._settings = None + self._hosted.set_routing_groups(resolved.routing_groups) + if "cooldown_time" in values: + self._cooldown_given = values["cooldown_time"] + if values["cooldown_time"] is None: + # `add_deployment_to_cooldown(cooldown_time=None)` cools down for the cooldown + # cache's default, which keeps the constructor's value. + self._resolved = dataclasses.replace( + resolved, cooldown_time=float(self._cooldowns.default_cooldown_time) + ) + if static_changed: + self._rebuild() + else: + names: Final = [str(d.get("model_name")) for d in self._handle.deployments()] + self._handle.apply({"settings": resolved.dynamic_document(self._version, names)}) + + def _changes_strategy(self, updated: Mapping[str, object]) -> bool: + """``update_settings`` names a ``routing_strategy`` other than the router's.""" + if "routing_strategy" not in updated: + return False + return (normalize_strategy(updated["routing_strategy"]) or "simple-shuffle") != self._resolved.routing_strategy + + def _assign_host_settings(self, assign: Mapping[str, object]) -> None: + """The assigned settings only the host reads, each stored as given (detached, as + ``apply_changes`` stores it) where its readers find it from the next call on.""" + given: Final = {name: detached(value) for name, value in assign.items() if name in _HOST_SETTINGS} + if not given: + return + for name, value in given.items(): + if name in ("enable_health_check_routing", "background_health_check_model_groups"): + self._resolved = dataclasses.replace( + self._resolved, + health_routing=dataclasses.replace(self._resolved.health_routing, **{name: value}), # pyright: ignore[reportArgumentType] # stored as given, as the Python core stores it + ) + elif name == "provider_default_deployment_ids": + # The list wildcard deployments added later are appended to, as on the Python core. + self._hosted._core_provider_default_deployment_ids = value # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # stored as given, as the Python core stores it + elif name == "cache_responses": + self._endpoints.set_cache_responses(bool(value)) + elif name == "provider_budget_config": + # Read when the budget limiter is built; a limiter built already keeps its config. + self._checks.provider_budget_config = value + elif name == "health_check_ignore_transient_errors": + self._resolved = dataclasses.replace( + self._resolved, + health_routing=dataclasses.replace( + self._resolved.health_routing, + health_check_ignore_transient_errors=value, # pyright: ignore[reportArgumentType] # stored as given, as the Python core stores it + ), + ) + elif name == "router_general_settings": + self._general_settings = value + elif name == "kubernetes_pod_discovery": + self._pod_discovery = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it + else: + self._endpoints.assign_setting(name, value) + self._deleted = self._deleted - set(given) + self._settings = None + if {"enable_health_check_routing", "background_health_check_model_groups"} & set(given) or ( + self._provider_default_ids() != self._handle_default_ids + ): + self._rebuild() + + def _set_retry_policies(self, policies: Mapping[str, object]) -> None: + """``update_settings``' ``retry_policy`` and ``model_group_retry_policy``, which the retry + loop reads from then on: router-core holds them with its static settings.""" + previous: Final = self._resolved + resolved: Final = dataclasses.replace( + previous, + retry_policy=resolved_retry_policy(policies.get("retry_policy", previous.retry_policy)), + model_group_retry_policy=policies.get("model_group_retry_policy", previous.model_group_retry_policy), + ) + if resolved != previous: + self._resolved = resolved + self._rebuild() + + def _apply_strategy(self, previous: Resolved, resolved: Resolved, values: Mapping[str, object]) -> None: + """``_core_update_settings``: a changed strategy re-inits the selectors with the args passed + alongside it (or none); changed args alone rebuild the default selector.""" + if "routing_strategy" in values and resolved.routing_strategy != previous.routing_strategy: + # `routing_strategy_init` drops a custom strategy (`_reset_custom_routing_strategy`), + # and an async pick assigned through the deprecated attribute. + self._custom_strategy = None + if "_core_async_get_available_deployment" in self.__dict__: + del self.__dict__["_core_async_get_available_deployment"] + self._selectors.init_default( + resolved.routing_strategy, as_str_mapping(values.get("routing_strategy_args")) or {} + ) + self._selectors.init_groups(_group_selectors(resolved.routing_groups)) + elif "routing_groups" in values: + self._selectors.init_groups(_group_selectors(resolved.routing_groups)) + elif "routing_strategy_args" in values and resolved.routing_strategy_args != previous.routing_strategy_args: + if resolved.routing_strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(RoutingContext(self), resolved.routing_strategy_args) # pyright: ignore[reportArgumentType] # the core's RoutingContext + return + self._selectors.rebuild_default(resolved.routing_strategy_args) + + def upsert_deployment(self, deployment: Deployment) -> Deployment | None: + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice + ) + + _note("upsert_deployment", "apply(ChangeSet(upsert=[deployment]))") + return self.apply(ChangeSet(upsert=[deployment])).upserted[0] + + def delete_deployment(self, id: str) -> Deployment | None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see above + + _note("delete_deployment", "apply(ChangeSet(delete=[id]))") + return self.apply(ChangeSet(delete=[id])).deleted[0] + + def update_settings(self, **kwargs: object) -> None: # kwargs-ok: the deprecated signature, unchanged + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("update_settings", "apply(ChangeSet(settings={...}))") + _ = self.apply(ChangeSet(settings=kwargs)) + + def update_redis_cache(self, cache: "RedisCache") -> None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("update_redis_cache", "apply(ChangeSet(redis_cache=cache))") + _ = self.apply(ChangeSet(redis_cache=cache)) + + 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 + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("get_model_list", "registry().deployments(model_name, team_id=...)") + return self._hosted._core_get_model_list(model_name=model_name, team_id=team_id) # pyright: ignore[reportPrivateUsage, reportReturnType] # the host registry's listing, as the deprecated member returns it + + @staticmethod + def get_valid_args() -> list[str]: # mutable-ok: the old return type, a fresh list + """The arguments ``Router.__init__`` accepts.""" + return list(ROUTER_INIT_ARGS) + + async def arefresh_model_info(self, *, client: "AsyncHTTPHandler | None" = None) -> None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("arefresh_model_info", "aapply(ChangeSet(refresh_model_info=True))") + _ = await self.aapply(ChangeSet(refresh_model_info=True, refresh_client=client)) + + def flush_cache(self) -> None: + """Drop the host's response cache and the routing state.""" + self._host.set_global_cache(None) + self._cache.flush_cache() + self._checks.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + def reset(self) -> None: + """Empty the host's success and failure callback lists (every entry, not only this core's), + drop the retry policy and flush the caches.""" + for name in ("success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"): + self._host.set_callback_list(name, []) + self._plain["retry_policy"] = None + self._set_retry_policies({"retry_policy": None}) + self.flush_cache() + + def __delattr__(self, name: str) -> None: + # `del router.` unsets it: the settings view reports it unset from then on. + from litellm.router_api import SETTING_FIELDS + + # The deprecated attribute deletes the core's own name (``_core_``). + setting: Final = name.removeprefix("_core_") + if setting not in SETTING_FIELDS: + object.__delattr__(self, name) + return + if setting in self._deleted: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{setting}'") + self._deleted = self._deleted | {setting} + self._settings = None + + def __setattr__(self, name: str, value: object) -> None: + # A setting write is a changeset (`router.num_retries = 2` is `apply(ChangeSet(assign=...))`), + # a `model_list` write swaps the live list, and any other attribute is kept as written. + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + if name in PLAIN_SETTING_ATTRIBUTES or name in SETTING_FIELDS or name in ASSIGNABLE_ALIASES: + _ = self.apply(ChangeSet(assign={name: value})) + elif name == "model_list": + self._assign_model_list(_as_list(value)) + elif name in _HOSTED_STATE: + # A deprecated attribute write of the registry state: it lands where this core keeps it. + setattr(self._hosted, name, value) + if name in INDEX_ATTRIBUTES: + self._index_written = self._index_written | {name} + self._sync_index() + elif name == "_core_routing_groups": + # The Python core routes on the assigned groups from the next call on, while its + # member map and group selectors stay as they were; router-core holds the groups and + # the member map with its static settings. + assigned: Final = as_str_mapping(value) or {} + previous: Final = self._resolved + self._resolved = dataclasses.replace( + previous, + routing_groups=tuple(g for g in assigned.values() if isinstance(g, RoutingGroup)), + routing_group_members=previous.routing_group_members + if previous.routing_group_members is not None + else member_map(previous.routing_groups), + ) + self._settings = None + self._rebuild() + elif name == "_core_routing_plugins": + self._sync_plugin_stage() + elif name == "_core_prompt_caching_pin_cache": + # Read when a prompt caching check is built, by `update_redis_cache` and by `flush_cache`. + self._checks.prompt_caching_pin_cache = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it + elif name in _REBOUND_STATE: + # A deprecated attribute write of the routing state (``router.cache = ...``): this core + # reads the new object from now on, as the Python core reads its rebound attribute. + object.__setattr__(self, _REBOUND_STATE[name], value) + if name == "_core_cache": + self._hosted._core_cache = value # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # usage reads come from the routing cache + elif name == "_core_slack_alerting_logger": + from litellm.integrations.custom_logger import CustomLogger + + # `_slack_alerting_loggers` reads the attribute each time: the new logger is the one bound. + self._checks.slack_alerting_loggers = (value,) if isinstance(value, CustomLogger) else () + else: + # Any other attribute is kept as written; `_core_async_get_available_deployment` (the + # deprecated async pick) is then asked by every async pick (see `_pre_route`). + object.__setattr__(self, name, value) + + async def select( # pyright: ignore[reportIncompatibleMethodOverride] # the RouterCore selection + self, + model: str, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the RouterCore signature + *, + kind: str = "pick", + specific_deployment: bool | None = None, + messages: object = None, + input: object = None, + health_check_probe: bool = False, + **_: object, # kwargs-ok: the RouterCore select options this core ignores + ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + """Pick a deployment for ``model`` without calling it (``pick``), list the deployments a + pick would choose among (``healthy``), or pick one for a pass-through route + (``pass_through``, which picks with the sync strategies, as Python does).""" + if kind not in ("pick", "healthy", "pass_through"): + raise ValueError(f"unknown select kind {kind!r}") + request: Final = request_kwargs if request_kwargs is not None else {} + selection: Final = _Selection( + self, + model, + request, + kind, + specific_deployment, + messages, + input, + probe=kind == "healthy" and health_check_probe, + ) + if kind == "pass_through": + return selection.run_sync() + return await selection.run() + + async def async_get_available_deployment_for_pass_through( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: the legacy signature + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + ) -> Mapping[str, object]: + """The full healthy chain, then only ``use_in_pass_through`` deployments, picked with the + async strategies.""" + selection: Final = _Selection( + self, model, request_kwargs, "async_pass_through", specific_deployment, messages, input + ) + picked: Final = await selection.run() + return picked if isinstance(picked, Mapping) else {} + + def get_available_deployment( + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the legacy signature + ) -> Mapping[str, object]: + """The sync pick: no pre-routing, the sync strategies.""" + selection: Final = _Selection( + self, + model, + request_kwargs if request_kwargs is not None else {}, + "sync_pick", + specific_deployment, + messages, + input, + ) + picked: Final = selection.run_sync() + return picked if isinstance(picked, Mapping) else {} + + def get_available_deployment_for_pass_through( + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the legacy signature + ) -> Mapping[str, object]: + """Deprecated: ``select(model, kind="pass_through")``.""" + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice + ) + + _note("get_available_deployment_for_pass_through", 'select(model, kind="pass_through")') + selection: Final = _Selection( + self, + model, + request_kwargs if request_kwargs is not None else {}, + "pass_through", + specific_deployment, + messages, + input, + ) + picked: Final = selection.run_sync() + return picked if isinstance(picked, Mapping) else {} + + # -- endpoints ------------------------------------------------------------------------------- + + async def call(self, op: str, request: Mapping[str, object]) -> object: + """Run one endpoint with the request's keyword payload, as ``getattr(core, op)(**request)``.""" + member: Final[object] = getattr(self, op) # pyright: ignore[reportAny] # litellm's untyped helpers and request values + if not callable(member): + raise TypeError(f"{op!r} is not an endpoint") + result: Final[object] = member(**request) + return await result if inspect.isawaitable(result) else result + + async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.acompletion(request_from_call("acompletion", args, kwargs)) + + async def atext_completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.atext_completion(request_from_call("atext_completion", args, kwargs)) + + async def aembedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._drive(plan_for("aembedding"), request_from_call("aembedding", args, kwargs)) + + async def aresponses(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.generic(plan_for("aresponses"), request_from_call("aresponses", args, kwargs)) + + async def aanthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.generic(plan_for("aanthropic_messages"), kwargs) + + async def anthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload + """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async, as in PythonCore.""" + return await self._endpoints.generic(plan_for("anthropic_messages"), kwargs) + + async def aimage_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._drive(plan_for("aimage_generation"), request_from_call("aimage_generation", args, kwargs)) + + async def atranscription(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._drive(plan_for("atranscription"), request_from_call("atranscription", args, kwargs)) + + async def aspeech(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + request: Final = dict(request_from_call("aspeech", args, kwargs)) + request.setdefault("voice", None) + return await self._drive(plan_for("aspeech"), request) + + async def arerank(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + # `arerank` sets kwargs["input"] = input: the Python builtin, which reaches the deployment call. + request: Final = {**request_from_call("arerank", args, kwargs), "input": input} + return await self._drive(plan_for("arerank"), request) + + async def amoderation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.amoderation(request_from_call("amoderation", args, kwargs)) + + def completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync(plan_for("completion"), request_from_call("completion", args, kwargs)) + + def text_completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._endpoints.text_completion(request_from_call("text_completion", args, kwargs)) + + def embedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync(plan_for("embedding"), request_from_call("embedding", args, kwargs)) + + def image_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync(plan_for("image_generation"), request_from_call("image_generation", args, kwargs)) + + @property + def total_calls(self) -> Mapping[str, int]: + return self._endpoints.counters.total_calls + + @property + def success_calls(self) -> Mapping[str, int]: + return self._endpoints.counters.success_calls + + @property + def fail_calls(self) -> Mapping[str, int]: + return self._endpoints.counters.fail_calls + + # -- driving a call -------------------------------------------------------------------------- + + def _start( + self, + plan: Plan, + request: Mapping[str, object], + overrides: Mapping[str, object] | None = None, + *, + prepared_already: bool = False, + ) -> "tuple[RouterCall, Mapping[str, object]]": + if self._retired: + raise RuntimeError("this router core was retired") + self._sync_routing_state() + unserved: Final = sorted(key for key in request if key.startswith("mock_testing_") and key not in MOCK_FLAGS) + if unserved: + raise RustCoreUnsupported(f"request keys not served by the Rust router core yet: {unserved}") + model: Final = request.get("model") + if not isinstance(model, str): + raise TypeError(f"{plan.endpoint}() needs a model name") + field: Final = plan.prepare_field or "metadata" + if request.get(field) is not None and not isinstance(request.get(field), Mapping): + raise TypeError(f"{field} must be a mapping") + prepared: Final = dict(request) if prepared_already else self._endpoints.prepare(plan, request, model) + mock: Final = mock_failure(request) + fields: Final[Mapping[str, object]] = { + **self._routing_fields(prepared), + **pinned_video_deployment(plan, request), + **unrouted_fields(plan), + # The sync path never runs pre-routing; a custom strategy picks on both. + **({} if plan.sync and self._custom_strategy is None else {"pre_routes": True}), + **({"mock_failure": mock} if mock is not None else {}), + # Every target of the walk is refused. + **({"fallbacks": []} if rejected_fallback_target(request) is not None else {}), + **self._retry_policy_raises(plan.endpoint, request), + **(overrides or {}), + } + return self._handle.start(plan.endpoint, _document({"model": model, **fields})), prepared + + def _process_renamed(self, model: str, group: str) -> str: + """``model`` as the error names it: a group the name resolved to by name, alias or routing + group is renamed through the process ``model_alias_map`` (``_common_checks_available_deployment``).""" + renames: Final = self._host.dynamic().model_alias_map + resolution: Final = self._handle.resolve(model) if group in renames else None + if resolution is None or resolution["kind"] not in ("Group", "Alias", "RoutingGroup"): + return model + return renames[group] + + def _max_parallel_requests_client(self, deployment: Mapping[str, object], kwargs: Mapping[str, object]) -> object: + """``_get_client(client_type="max_parallel_requests")``: the deployment's slot limit.""" + return max_parallel_requests_limit(self, deployment, kwargs) + + def _routable(self, model: str, request: Mapping[str, object]) -> Sequence[Mapping[str, object]]: + return self._handle.routable(model, _document(_start_fields(request))) + + def _routing_fields(self, request: Mapping[str, object]) -> Mapping[str, object]: + """The request fields the Rust core routes on, and whether its tag stage may apply: the + router's ``enable_tag_filtering``, the request's own, or a deployment's + ``model_info.enable_tag_filtering`` (``get_deployments_for_tag``'s effective setting).""" + filter_tags: Final = ( + self._resolved.enable_tag_filtering + or request.get("enable_tag_filtering") is True + or self._model_tag_filtering() + ) + return { + **_start_fields(request), + **self._group_policy_fields(request), + **({"filter_tags": True} if filter_tags else {}), + **({"has_image": True} if request_carries_image(None, request) else {}), + } + + def _group_policies(self, request: Mapping[str, object]) -> tuple[object, GroupPolicies]: + """The ``model_group_retry_policy`` the retry loop reads for this request (the request's own, + even ``None``, replaces the router's: ``kwargs.pop(..., self.model_group_retry_policy)``), + resolved for every group name the call may look up.""" + policies: Final = request.get("model_group_retry_policy", self._resolved.model_group_retry_policy) + return policies, group_policies(policies, self._policy_lookup_names(request)) + + def _policy_lookup_names(self, request: Mapping[str, object]) -> frozenset[str]: + """The model names a retry policy lookup may run on (each hop looks up the name it calls): + the requested one and every name in the fallback lists the walk may take.""" + resolved: Final = self._resolved + lists: Final = ( + resolved.fallbacks, + resolved.context_window_fallbacks, + resolved.content_policy_fallbacks, + *(request.get(key) for key in _REQUEST_FALLBACK_KEYS), + ) + return frozenset((*_strings(request.get("model")), *_strings(lists))) + + def _retry_policy_raises(self, endpoint: str, request: Mapping[str, object]) -> Mapping[str, object]: + """The groups whose ``model_group_retry_policy`` lookup raises on a hop's first failure, for + router-core to end that hop with the error the lookup raises (see ``_Driver.take_events``).""" + raising: Final = raising_groups(endpoint, request, *self._group_policies(request)) + return {"retry_policy_raises": list(raising)} if raising else {} + + def _group_policy_fields(self, request: Mapping[str, object]) -> Mapping[str, object]: + """The request's ``model_group_retry_policy`` for router-core; a router policy that is not a + mapping of policies rides on every request, resolved for it.""" + policies, resolved = self._group_policies(request) + if "model_group_retry_policy" not in request and is_policy_per_group(policies): + return {} + return {"model_group_retry_policy": _plain(resolved.document)} + + def _model_tag_filtering(self) -> bool: + """Whether any deployment sets ``model_info.enable_tag_filtering``, cached per registry version.""" + cached: Final = self._tag_filtering_models + if cached is not None and cached[0] == self._version: + return cached[1] + found: Final = any( + (as_str_mapping((as_str_mapping(record) or {}).get("model_info")) or {}).get("enable_tag_filtering") + is not None + for record in self._handle.deployments() + ) + self._tag_filtering_models = (self._version, found) + return found + + async def _drive( + self, + plan: Plan, + request: Mapping[str, object], + overrides: Mapping[str, object] | None = None, + *, + passthrough: "Callable[[Mapping[str, object]], Awaitable[object]] | None" = None, + prepared_already: bool = False, + ) -> object: + call, prepared = self._start(plan, request, overrides, prepared_already=prepared_already) + driver, step = await self._run(plan, call, prepared) + if step["kind"] == "failed": + if passthrough is not None and driver.attempted == 0: + return await passthrough(prepared) + raise driver.failed(step) + return await driver.finish_async(step) + + async def _run( + self, plan: Plan, call: "RouterCall", prepared: Mapping[str, object] + ) -> "tuple[_Driver, Mapping[str, object]]": + """Drive ``call`` to its ``done`` or ``failed`` step, answering every step on the host.""" + driver: Final = _Driver(self, plan, prepared) + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome against the step it answers + driver.take_events(step.get("events")) + await await_all(driver.outcomes.take_pending()) + match step["kind"]: + case "attempt": + outcome = await driver.attempt(step) + case "sleep": + await self._host.sleep(_sleep_seconds(step)) + outcome = None + case "allow_fallback": + outcome = await driver.allow_fallback(str(step["target"])) + case "cache": + outcome = await self._answer_cache(step) + case "select": + outcome = await driver.select(step) + case "pre_route": + outcome = await self._pre_route(step, prepared) + case "done" | "failed": + return driver, step + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + def _call_sync( + self, plan: Plan, request: Mapping[str, object], overrides: Mapping[str, object] | None = None + ) -> object: + call, prepared = self._start(plan, request, overrides) + driver, step = self._run_sync(plan, call, prepared) + if step["kind"] == "failed": + raise driver.failed(step) + return driver.finish(step) + + def _run_sync( + self, plan: Plan, call: "RouterCall", prepared: Mapping[str, object] + ) -> "tuple[_Driver, Mapping[str, object]]": + driver: Final = _Driver(self, plan, prepared) + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # see _run + driver.take_events(step.get("events")) + if pending := driver.outcomes.take_pending(): + from litellm.litellm_core_utils import asyncify + + asyncify.run_async_function(await_all, pending) # pyright: ignore[reportUnknownMemberType] # the untyped sync bridge + match step["kind"]: + case "attempt": + outcome = driver.attempt_sync(step) + case "sleep": + self._host.sleep_sync(_sleep_seconds(step)) + outcome = None + case "allow_fallback": + from litellm.litellm_core_utils import asyncify + + outcome = asyncify.run_async_function(driver.allow_fallback, str(step["target"])) # pyright: ignore[reportUnknownMemberType] # the untyped sync bridge + case "cache": + outcome = self._answer_cache_sync(step) + case "select": + outcome = driver.select_sync(step) + case "pre_route": + outcome = self._custom_pick_sync(step, prepared) + case "done" | "failed": + return driver, step + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + # -- stream recovery (the host's stream wrappers, router_rust.invoke.streaming) --------------------- + + def _recovery_call( + self, + recovery: str, + op: str, + request: Mapping[str, object], + origin: Mapping[str, object], + failure: BaseException, + ) -> "RouterCall": + start: Final = self._handle.retry_stream if recovery == "retry" else self._handle.fallback_stream + routing: Final = _document( + { + "model": request.get("model"), + **self._routing_fields(request), + **self._retry_policy_raises(op, request), + **continuation_routing(self, request), + } + ) + return start(op, routing, _document(origin), attempt_failure(failure)) + + def stream_policy( + self, request: Mapping[str, object], origin: Mapping[str, object], trigger: BaseException | None + ) -> Mapping[str, object]: + routing: Final = _document( + {"model": request.get("model"), **_start_fields(request), **self._group_policy_fields(request)} + ) + kind: Final = None if trigger is None else str(attempt_failure(trigger)["kind"]) + return self._handle.stream_policy(routing, _document(origin), kind) + + async def recover_stream( + self, + recovery: str, + op: str, + request: Mapping[str, object], + origin: Mapping[str, object], + failure: BaseException, + ) -> object: + """The fallback walk for a stream that failed before content: the fallback's response with the + router's headers, or the error the walk ended on.""" + call: Final = self._recovery_call(recovery, op, request, origin, failure) + driver, step = await self._run(plan_for(op), call, request) + if step["kind"] == "failed": + raise driver.failed(step) + return await driver.finish_async(step) + + def recover_stream_sync( + self, + recovery: str, + op: str, + request: Mapping[str, object], + origin: Mapping[str, object], + failure: BaseException, + ) -> object: + call: Final = self._recovery_call(recovery, op, request, origin, failure) + driver, step = self._run_sync(plan_for(op), call, request) + if step["kind"] == "failed": + raise driver.failed(step) + return driver.finish(step) + + async def retry_stream( + self, op: str, request: Mapping[str, object], origin: Mapping[str, object], failure: BaseException + ) -> "tuple[object, int, int] | BaseException": + """A same-group retry of a stream that failed before content: the retried response (without the + router's headers) with the retry count and budget, or the newest failure (``failure`` itself + when no retry ran).""" + call: Final = self._recovery_call("retry", op, request, origin, failure) + driver, step = await self._run(plan_for(op), call, request) + if step["kind"] == "failed": + return driver.last_error if driver.last_error is not None else failure + attempted: Final = step.get("attempted_retries") + budget: Final = step.get("max_retries") + return ( + driver.response, + attempted if isinstance(attempted, int) else 0, + budget if isinstance(budget, int) else 0, + ) + + def params_sent_to_deployment(self, request: Mapping[str, object], deployment_id: object) -> Mapping[str, object]: + """The request as a deployment received it: its params, the router defaults, its tool choice, then + the request (``_params_sent_to_deployment``).""" + record: Final = self._handle.deployment(deployment_id) if isinstance(deployment_id, str) else None + params: Final = as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {} + defaults: Final = {k: v for k, v in self._resolved.default_litellm_params.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}) + + def _drive_sync_steps(self, call: "RouterCall") -> Mapping[str, object]: + """Drive a call whose only steps are cache ops (``report_failure``) to its end.""" + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # see _drive + match step["kind"]: + case "cache": + outcome = self._answer_cache_sync(step) + case "done": + return step + case "failed": + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + case other: + raise RuntimeError(f"router-core asked for {other!r} while recording an outcome") + + async def _drive_cache_steps(self, call: "RouterCall") -> Mapping[str, object]: + """``_drive_sync_steps`` answered through the async cache (``report_success``).""" + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(call.next(outcome)) # pyright: ignore[reportArgumentType] # see _drive + match step["kind"]: + case "cache": + outcome = await self._answer_cache(step) + case "done": + return step + case "failed": + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + case other: + raise RuntimeError(f"router-core asked for {other!r} while recording an outcome") + + # -- pre-route steps ------------------------------------------------------------------------- + + async def _pre_route( + self, + step: Mapping[str, object], + request: Mapping[str, object], + messages: object = _NO_INPUT, + input: object = _NO_INPUT, + specific_deployment: object = _NO_INPUT, + ) -> object: + """The host's pre-routing of one pick (router_rust.steps.prerouting): the model group to resolve + instead of the step's ``model``, a custom strategy's deployment, or ``None``.""" + assigned: Final[object] = self.__dict__.get("_core_async_get_available_deployment") + if callable(assigned): + # The async pick assigned through the deprecated attribute answers, as Python's + # replaced `_core_async_get_available_deployment` does. + answer: Final[object] = assigned( + model=str(step.get("model")), + request_kwargs=request, + messages=request.get("messages") if messages is _NO_INPUT else messages, + input=request.get("input") if input is _NO_INPUT else input, + specific_deployment=request.get("specific_deployment") + if specific_deployment is _NO_INPUT + else specific_deployment, + ) + chosen: Final[object] = await answer if inspect.isawaitable(answer) else answer + return _custom_pick(chosen) + custom: Final = self._custom_strategy + if custom is not None: + picked: Final[object] = await custom.async_get_available_deployment( # pyright: ignore[reportUnknownMemberType] # the strategy's untyped pick + model=str(step.get("model")), + messages=request.get("messages") if messages is _NO_INPUT else messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + input=request.get("input") if input is _NO_INPUT else input, # pyright: ignore[reportArgumentType] # the request's input, as given + specific_deployment=request.get("specific_deployment") + if specific_deployment is _NO_INPUT + else specific_deployment, # pyright: ignore[reportArgumentType] # same + request_kwargs=request, # pyright: ignore[reportArgumentType] # the request kwargs, as given + ) + return _custom_pick(picked) + if not isinstance(request, dict): + return None + depth: Final = step.get("fallback_depth") + r: Final = self._resolved + owner: Final[object] = self.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return await self._prerouter.pre_route( + str(step.get("model")), + request, # pyright: ignore[reportUnknownArgumentType] # the request kwargs the hook stamps + messages=request.get("messages") if messages is _NO_INPUT else messages, + input=request.get("input") if input is _NO_INPUT else input, + fallback_depth=depth if isinstance(depth, int) else 0, + router_strategy=normalize_strategy(r.routing_strategy), + enable_tag_filtering=r.enable_tag_filtering, + match_any=r.tag_filtering_match_any, + public_router=public if public is not None else self, + run_plugins=partial( + self._run_plugins, request, request.get("messages") if messages is _NO_INPUT else messages + ) # pyright: ignore[reportUnknownArgumentType] # the request kwargs the plugins stamp + if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list + else None, + ) + + async def _run_plugins( + self, + request: dict[str, object], # mutable-ok: the plugins stamp their outcome on the caller's request + messages: object, + model: str, + ) -> None: + """``_run_routing_plugins`` over the pre-routing hook's session-resolved ``model``.""" + await run_plugins( + tuple(self._hosted._core_routing_plugins), # pyright: ignore[reportPrivateUsage] # the live plugin list + self._hosted, # pyright: ignore[reportArgumentType] # the host registry serves the listing + self._resolved.model_group_alias, + model, + request, + messages, + ) + + def _custom_pick_sync( + self, + step: Mapping[str, object], + request: Mapping[str, object], + messages: object = _NO_INPUT, + input: object = _NO_INPUT, + specific_deployment: object = _NO_INPUT, + ) -> object: + """A custom strategy's sync pick (``get_available_deployment``); without one, no pre-routing.""" + custom: Final = self._custom_strategy + if custom is None: + return None + picked: Final[object] = custom.get_available_deployment( # pyright: ignore[reportUnknownMemberType] # the strategy's untyped pick + model=str(step.get("model")), + messages=request.get("messages") if messages is _NO_INPUT else messages, # pyright: ignore[reportArgumentType] # see _pre_route + input=request.get("input") if input is _NO_INPUT else input, # pyright: ignore[reportArgumentType] # same + specific_deployment=request.get("specific_deployment") + if specific_deployment is _NO_INPUT + else specific_deployment, # pyright: ignore[reportArgumentType] # same + request_kwargs=request, # pyright: ignore[reportArgumentType] # same + ) + return _custom_pick(picked) + + # -- select steps ---------------------------------------------------------------------------- + + def _strategy_inputs(self, step: Mapping[str, object], request: Mapping[str, object]) -> "_StrategyInputs | None": + """The strategy, its selector, the name it keys state by and the candidate deployments of a + ``strategy`` select step; ``None`` for a stage the host passes through.""" + if step.get("stage") != "strategy": + return None + strategy: Final = str(step.get("strategy")) + group: Final = step.get("group") + selector: Final = self._selectors.selector_for( + strategy, str(step.get("scope")), str(group) if group is not None else None, request + ) + candidates: Final = [ + dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(step.get("candidates")) + ] + return strategy, selector, str(step.get("model")), candidates + + def _records_of(self, ids: object) -> list[dict[str, object]]: # mutable-ok: the deployment dicts filters take + return [dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(ids)] + + async def _answer_tags(self, step: Mapping[str, object], request: Mapping[str, object]) -> object: + """``get_deployments_for_tag`` over the step's candidates (bridged: not ported to Rust yet).""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + + kept: Final = await get_deployments_for_tag( + ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the filter reads registry() and settings() + model=str(step.get("model")), + healthy_deployments=self._records_of(step.get("candidates")), # pyright: ignore[reportArgumentType] # deployment dicts + request_kwargs=dict(request), # pyright: ignore[reportArgumentType] # the request kwargs + metadata_variable_name=get_metadata_variable_name_from_kwargs(dict(request)), + request_deployments=self._records_of(step.get("pool")), # pyright: ignore[reportArgumentType] # deployment dicts + ) + return [found for record in kept if (found := picked_id(record)) is not None] + + def _pre_call_facts( + self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object + ) -> PreCallFacts: + hosted: Final = self._hosted + + def record(id: str) -> Mapping[str, object]: + return hosted._core_get_model_info(id) or as_str_mapping(self._handle.deployment(id)) or {} # pyright: ignore[reportPrivateUsage] # see __init__ + + return PreCallFacts(self._host, hosted, step, record, request, messages, input) # pyright: ignore[reportArgumentType] # the host registry serves the reads + + def _text_only(self, step: Mapping[str, object]) -> list[str]: # mutable-ok: the step's answer + """The candidates that declare no vision support (bridged: needs the cost map).""" + group: Final = str(step.get("model_group")) + return [ + str(id) + for id in _as_list(step.get("candidates")) + if declares_no_vision(as_str_mapping(self._handle.deployment(str(id))) or {}, group) + ] + + async def _answer_select( + self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object + ) -> object: + if step.get("stage") == "pre_call": + return await self._pre_call_facts(step, request, messages, input).answer_async() + if step.get("stage") == "routing_plugins": + # The plugins already ran in this pick's pre_route step (see `_pre_route`), on the + # session-resolved name and before a strategy was selected, as Python's hook runs them. + return candidate_models(request) + if step.get("stage") == "tags": + return await self._answer_tags(step, request) + if step.get("stage") == "text_only": + return self._text_only(step) + inputs: Final = self._strategy_inputs(step, request) + if inputs is None: + return step["candidates"] + strategy, selector, model, candidates = inputs + if selector is None: + return [] + picked: Final = await pick_async( + strategy, + selector, # pyright: ignore[reportArgumentType] # a selector from selector_for + model, + candidates, + messages, + input, + dict(request), + ) + if picked is not None and self._selectors.is_override(strategy, selector): + await selector.async_pre_call_check(dict(as_str_mapping(picked) or {}), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + found: Final = picked_id(picked) + return [found] if found is not None else [] + + def _answer_select_sync( + self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object + ) -> object: + if step.get("stage") == "pre_call": + return self._pre_call_facts(step, request, messages, input).answer() + if step.get("stage") == "routing_plugins": + return candidate_models(request) + if step.get("stage") == "callback_filters": + # Only the async candidate lists run the callback filters. + return step["candidates"] + if step.get("stage") == "text_only": + return self._text_only(step) + inputs: Final = self._strategy_inputs(step, request) + if inputs is None: + return step["candidates"] + strategy, selector, model, candidates = inputs + if selector is None: + return [] + picked: Final = pick_sync( + strategy, + selector, # pyright: ignore[reportArgumentType] # a selector from selector_for + model, + candidates, + messages, + input, + dict(request), + ) + if picked is not None and self._selectors.is_override(strategy, selector): + selector.pre_call_check(dict(as_str_mapping(picked) or {})) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # see above + found: Final = picked_id(picked) + return [found] if found is not None else [] + + # -- cache steps ----------------------------------------------------------------------------- + + def _store_for(self, key: str) -> "DualCache": + """Cooldowns live in the cooldown store; every other routing key in the routing cache.""" + return self._cooldowns.cooldown_store if key.endswith(":cooldown") else self._cache + + async def _answer_cache(self, step: Mapping[str, object]) -> object: + with service_target(_cache_target(str(step.get("key")))): + return await self._answer_cache_in_target(step) + + async def _answer_cache_in_target(self, step: Mapping[str, object]) -> object: + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + local: Final = step.get("local") is True + match step.get("op"): + case "get": + return _plain(await self._store_for(key).async_get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + return _plain( + [await self._store_for(str(k)).async_get_cache(str(k)) for k in _as_list(step.get("keys"))] # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + ) + case "set": + store: Final = self._store_for(key) + if ttl is None: + await store.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + else: + await store.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + return None + case "increment": + by: Final = _amount(step.get("value")) + counter: Final = self._store_for(key) + total: Final = ( + await counter.async_increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType] # see above + if ttl is None + else await counter.async_increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + ) + return by if total is None else total + case "delete": + await self._store_for(key).async_delete_cache(key) + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + def _answer_cache_sync(self, step: Mapping[str, object]) -> object: + with service_target(_cache_target(str(step.get("key")))): + return self._answer_cache_sync_in_target(step) + + def _answer_cache_sync_in_target(self, step: Mapping[str, object]) -> object: + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + local: Final = step.get("local") is True + match step.get("op"): + case "get": + return _plain(self._store_for(key).get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + return _plain([self._store_for(str(k)).get_cache(str(k)) for k in _as_list(step.get("keys"))]) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + case "set": + store: Final = self._store_for(key) + if ttl is None: + store.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + else: + store.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + return None + case "increment": + by: Final = _amount(step.get("value")) + counter: Final = self._store_for(key) + # DualCache.increment_cache is annotated int but adds floats as given. + return ( + counter.increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above + if ttl is None + else counter.increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above + ) + case "delete": + self._store_for(key).delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + # -- errors the core raises when no attempt failed ------------------------------------------ + + def _failure_error( + self, + step: Mapping[str, object], + model: str, + last: BaseException | None, + *, + pass_through: str = "", + request: Mapping[str, object] | None = None, + ) -> BaseException: + """The error a call or pick that failed before any attempt raises; ``pass_through`` names + the selection kind (``pass_through``/``async_pass_through``) whose errors it words.""" + import litellm.exceptions as litellm_exceptions + + kind: Final = step.get("error_kind") + aliases: Final = self._resolved.model_group_alias + refused: Final = selection_refusal(step, model, request or {}) + if refused is not None: + return refused + if kind == "paused": + # The message names the deployment's litellm model, not the id that was asked for; the + # default deployment's copy is named by the requested name. + record: Final = as_str_mapping(self._handle.deployment(model)) + litellm_model: Final = ( + str((as_str_mapping(record.get("litellm_params")) or {}).get("model")) if record is not None else model + ) + return litellm_exceptions.ServiceUnavailableError( + message=f"Model '{litellm_model}' is currently paused and cannot accept requests.", + model=litellm_model, + llm_provider="", + ) + if kind == "only_markers": + named_model: Final = str(step.get("error_model") or model) + return litellm_exceptions.BadRequestError( + message=f"You passed in model={named_model}. {RouterErrors.only_strategy_marker_deployments.value}", + model=named_model, + llm_provider="", + ) + if kind == "ptu_reserved": + reserved: Final = str(step.get("error_model") or model) + return litellm_exceptions.BadRequestError( + message=f"Deployment {reserved} is reserved for the teams holding a PTU share of it", + model=reserved, + llm_provider="", + ) + if kind == "not_pass_through": + return litellm_exceptions.BadRequestError( + message=f"Deployment {model} does not support pass-through endpoint (use_in_pass_through=False)", + model=model, + llm_provider="", + ) + if kind == "no_pass_through_deployments": + message: Final = ( + f"Model {model} has no deployments configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true to the deployment configuration" + if pass_through == "async_pass_through" + else f"Model {model} has no deployment configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true in the deployment configuration" + ) + return litellm_exceptions.BadRequestError(message=message, model=model, llm_provider="") + if kind == "context_window_exceeded": + return 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{step.get('error_detail')}", + model=str(step.get("error_model") or model), + llm_provider="", + ) + if kind == "unsupported_params": + params: Final = sorted(str(param) for param in _as_list(step.get("error_params"))) + unsupported_model: Final = str(step.get("error_model") or model) + return litellm_exceptions.UnsupportedParamsError( + message=( + f"no deployment of model={unsupported_model} supports {params}. " + "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " + f"or send allowed_openai_params={params} to pass them through" + ), + model=unsupported_model, + ) + if kind == "routing_plugins_left_nothing": + left: Final = [str(name) for name in _as_list(step.get("error_params"))] + return ValueError(f"No deployments left after routing-plugin filtering. candidate_models={left}") + if kind == "routing_plugins_need_async": + return ValueError(SYNC_PATH_MESSAGE) + looked_up: Final = _lookup_error(kind, str(step.get("error_model") or model), self._resolved.default_deployment) + if looked_up is not None: + return looked_up + if kind == "unknown_model": + named: Final = _alias_target(aliases.get(model)) or model if model in aliases else model + return litellm_exceptions.BadRequestError( + message=f"You passed in model={named}. {RouterErrors.no_healthy_deployments.value}", + model=named, + llm_provider="", + ) + if kind == "no_healthy_deployments": + group: Final = self.registry().resolve(model) + group_ids: Final = self.registry().names(str(group), kind="ids") + # A listed pass-through pick names only the pass-through deployments. + ids: Final = ( + [id for id in group_ids if _uses_pass_through(self._handle.deployment(id))] + if pass_through == "pass_through" + else list(group_ids) + ) + active: Final = self._cooldowns.get_active_cooldowns(list(ids), None) # pyright: ignore[reportArgumentType] # a list of ids + cooling: Final = [id for id, _ in active] + min_cooldown: Final = min( + (_seconds(entry.get("cooldown_time"), self._cooldowns.default_cooldown_time) for _, entry in active), + default=_seconds(None, self._cooldowns.default_cooldown_time), + ) + return RouterRateLimitError( + model=self._process_renamed(model, str(group)), + cooldown_time=min_cooldown, + enable_pre_call_checks=self._resolved.enable_pre_call_checks, + cooldown_list=cooling, + model_ids=list(ids), # pyright: ignore[reportArgumentType] # a list of ids + ) + if last is not None: + return last + return RuntimeError(f"router-core ({kind}): {step.get('error')}") + + +def _lookup_error(kind: object, model: str, default_deployment: object) -> BaseException | None: + """The plain Python error a pick's lookup raised for ``model``, or None for any other failure.""" + if kind == "default_deployment_raises": + return default_deployment_copy_error(default_deployment, model) + if kind == "index_out_of_range": + # `self._core_model_list[index]` for a written index position past the list. + return IndexError("list index out of range") + if kind == "missing_routing_group": + # `_get_routing_context` reads the stale member map's group from the replaced groups. + return KeyError(model) + return None + + +def _sleep_seconds(step: Mapping[str, object]) -> float: + """A retry's sleep: ``_time_to_sleep_before_retry`` returns a plain ``0`` when another + deployment is healthy, else a backoff in float seconds.""" + seconds: Final = float(str(step["seconds"])) + return 0 if seconds == 0 else seconds + + +def _uses_pass_through(record: object) -> bool: + params: Final = as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {} + return bool(params.get("use_in_pass_through")) + + +def _seconds(value: object, default: float) -> float | int: + """A cooldown time as stored: an int stays an int, so error texts read ``60`` like Python's.""" + if isinstance(value, int | float) and not isinstance(value, bool): + return value + return int(default) if float(default).is_integer() else default + + +def _ttl(step: Mapping[str, object]) -> float | None: + ttl: Final = step.get("ttl") + return ttl if isinstance(ttl, int | float) else None + + +def _amount(value: object) -> float | int: + return value if isinstance(value, int | float) and not isinstance(value, bool) else 1 + + +class _NotSynchronous(TypeError): + """A sync attempt's host endpoint returned an awaitable: a host bug, never an attempt failure.""" + + +class _Selection: + """One ``select``: the Rust selection call, its select and cache steps answered on the host.""" + + def __init__( + self, + core: RustCore, + model: str, + request: dict[str, object], # mutable-ok: the caller's request, whose narrowing keys are popped + kind: str, + specific_deployment: bool | None, + messages: object, + input: object, + *, + probe: bool = False, + ) -> None: + self._core: Final = core + self._model: Final = model + self._request: Final = request + self._kind: Final = kind + self._messages: Final = messages + self._input: Final = input + fields: Final = { + **core._routing_fields(request), # pyright: ignore[reportPrivateUsage] # the selection serves its own core + **_narrowing_fields(request, kind), + **({"has_image": True} if request_carries_image(messages, request) else {}), + } # pyright: ignore[reportPrivateUsage] # the selection serves its own core + if specific_deployment: + fields["specific_deployment"] = True + if probe: + fields["health_check_probe"] = True + if kind == "pick" or (kind == "sync_pick" and core._custom_strategy is not None): # pyright: ignore[reportPrivateUsage] # the selection serves its own core + fields["pre_routes"] = True + self._specific_deployment: Final = specific_deployment + # The error a callback filter raised, which a failed selection raises. + self._filter_error: BaseException | None = None + if kind == "pass_through": + fields["pass_through"] = "listed" + elif kind == "async_pass_through": + fields["pass_through"] = "healthy" + core._sync_routing_state() # pyright: ignore[reportPrivateUsage] # see above + handle: Final = core._handle # pyright: ignore[reportPrivateUsage] # see above + self._call: Final = (handle.healthy if kind == "healthy" else handle.select)(model, _document(fields)) + + def _done(self, step: Mapping[str, object]) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + if self._kind == "healthy": + if step.get("single"): + # A deployment named by its id is listed alone, as a dict. + (record,) = _as_list(step.get("deployments")) + return for_requested_model(as_str_mapping(record) or {}, self._model) + # The deployments as served: a priority group's members carry their priority as order. + return [ + for_requested_model(as_str_mapping(record) or {}, self._model) + for record in _as_list(step.get("deployments")) + ] + if step.get("deployment_id") in (DEFAULT_DEPLOYMENT_FAILING_ID, DEFAULT_DEPLOYMENT_COPIED_ID): + # A default deployment every call to fails, or whose id is not a string, is picked as Python copies it. + picked: Final = default_deployment_copy(self._core._resolved.default_deployment, self._model) # pyright: ignore[reportPrivateUsage] # the selection serves its own core + return picked # pyright: ignore[reportReturnType] # the assigned value's copy, as the Python core returns it + return for_requested_model(as_str_mapping(step.get("deployment")) or {}, self._model) + + def _failed(self, step: Mapping[str, object]) -> BaseException: + core: Final = self._core + error: Final = core._failure_error( # pyright: ignore[reportPrivateUsage] # see above + step, self._model, self._filter_error, pass_through=self._kind, request=self._request + ) + logging_obj: Final = self._request.get("litellm_logging_obj") + dispatch: Final[object] = getattr(logging_obj, "dispatch_failure_handlers", None) + if callable(dispatch) and self._kind not in ("pass_through", "sync_pick"): + import asyncio + import traceback + + # A refused pick is logged to the request's logging object (`_async_get_available_deployment`). + trace: Final = "Traceback (most recent call last):\n" + "".join( + traceback.format_exception(type(error), error, error.__traceback__) + ) + pending: Final[object] = dispatch(exception=error, traceback_exception=trace, prefer_async_handlers=True) + if inspect.isawaitable(pending): + _ = asyncio.ensure_future(pending) + return error + + async def run(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + core: Final = self._core + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(self._call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome + match step["kind"]: + case "cache": + outcome = await core._answer_cache(step) # pyright: ignore[reportPrivateUsage] # see above + case "select" if step.get("stage") == "callback_filters": + outcome, self._filter_error = await core._answer_callback_filters( # pyright: ignore[reportPrivateUsage] # see above + step, self._request, self._messages + ) + case "select": + outcome = await core._answer_select(step, self._request, self._messages, self._input) # pyright: ignore[reportPrivateUsage] # see above + case "pre_route": + outcome = await core._pre_route( # pyright: ignore[reportPrivateUsage] # see above + step, self._request, self._messages, self._input, self._specific_deployment + ) + case "done": + done = self._done(step) + if self._kind != "pick" or not isinstance(done, Mapping): + return done + # `_core_async_get_available_deployment` resolves its pick to a pod. + return await self._resolved_pod(core, done) + case "failed": + raise self._failed(step) + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + async def _resolved_pod(self, core: "RustCore", done: Mapping[str, object]) -> Mapping[str, object]: + discovery: Final = core._pick_discovery(sync=False) # pyright: ignore[reportPrivateUsage] # the selection serves its own core + if discovery is None: + return done + resolved: Final[object] = await discovery.async_resolve_deployment(done, self._request) + return resolved # pyright: ignore[reportReturnType] # the deployment, resolved + + def _resolved_pod_sync(self, core: "RustCore", done: Mapping[str, object]) -> Mapping[str, object]: + discovery: Final = core._pick_discovery(sync=True) # pyright: ignore[reportPrivateUsage] # see _resolved_pod + return done if discovery is None else discovery.resolve_deployment(done, self._request) + + def run_sync(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + core: Final = self._core + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = _STEP.validate_python(self._call.next(outcome)) # pyright: ignore[reportArgumentType] # see run + match step["kind"]: + case "cache": + outcome = core._answer_cache_sync(step) # pyright: ignore[reportPrivateUsage] # see above + case "select": + outcome = core._answer_select_sync(step, self._request, self._messages, self._input) # pyright: ignore[reportPrivateUsage] # see above + case "pre_route": + outcome = core._custom_pick_sync( # pyright: ignore[reportPrivateUsage] # see above + step, self._request, self._messages, self._input, self._specific_deployment + ) + case "done": + done = self._done(step) + if self._kind != "sync_pick" or not isinstance(done, Mapping): + return done + # `get_available_deployment` resolves its pick to a pod. + return self._resolved_pod_sync(core, done) + case "failed": + raise self._failed(step) + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + +def _custom_pick(picked: object) -> Mapping[str, object]: + """A custom strategy's answer as the deployment document the Rust core serves.""" + document: Final = as_str_mapping(picked) + if document is None: + raise TypeError(f"the custom routing strategy returned {type(picked).__name__}, not a deployment") + return _document(document) + + +def _narrowing_fields(request: dict[str, object], kind: str) -> Mapping[str, object]: # mutable-ok: popped + """``_target_order``, ``_excluded_deployment_ids`` and ``_retry_skipped_deployment_ids``, popped + from the request as the healthy chain pops them (a listed pass-through pick only reads the order).""" + from litellm.litellm_core_utils.routing_common_utils import ( + as_excluded_deployment_ids, + as_retry_skipped_deployment_ids, + ) + + if kind == "pass_through": + order: Final = request.get("_target_order") + return {"target_order": order} if isinstance(order, int) and not isinstance(order, bool) else {} + target_order: Final = request.pop("_target_order", None) + excluded: Final = as_excluded_deployment_ids(request.pop("_excluded_deployment_ids", None)) + skipped: Final = as_retry_skipped_deployment_ids(request.pop("_retry_skipped_deployment_ids", None)) + return { + **( + {"target_order": target_order} + if isinstance(target_order, int) and not isinstance(target_order, bool) + else {} + ), + **({"excluded_ids": list(excluded)} if excluded else {}), + **({"retry_skipped_ids": list(skipped)} if skipped else {}), + } + + +_NO_ROUTING_PRE_CALL_CHECKS: Final = frozenset({"search", "rerank", "moderation"}) +# Attempts whose deployment is no pick (`_core_async_get_available_deployment` / +# `get_available_deployment`), so Kubernetes pod discovery never resolves it. +_UNPICKED_KINDS: Final = frozenset({"search", "create_file"}) + + +class _Driver: + """Answers the attempt and fallback steps of one call, and shapes its result. + + Each attempt is one host call against the deployment Rust picked. The driver keeps the error + each fallback level ended on: a level that fails raises its own error, whatever its fallbacks + raised, as the Python router does. + """ + + def __init__(self, core: RustCore, plan: Plan, request: Mapping[str, object]) -> None: + self._core: Final = core + self._plan: Final = plan + self._op: Final = plan.endpoint + self._request: Final = request + self._last_response: object = None + self._last_error: BaseException | None = None + self._level_errors: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call runs + # The retry counts an exhausted retry loop stamps on the pick error its level then raises. + self._pending_stamps: Final[dict[int, tuple[object, object]]] = {} # mutable-ok: filled as the call runs + # Each level's current hop's first failure, once it retries: the error its retry loop raises + # last is raised while that one is handled, so it chains to it. + self._hop_first: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call runs + self._served: tuple[str | None, Mapping[str, object]] = (None, MappingProxyType({})) + self._attempts: int = 0 + self._walk: Final = WalkMetadata(request, plan.metadata_field, self._params_base_model) + self.outcomes: Final = FallbackOutcomes(request) + # The fallback level the walk is at, from its fallback events. + self._level: int = 0 + + @property + def attempted(self) -> int: + """How many attempts reached the host.""" + return self._attempts + + def _blocked(self, response: object, endpoint_name: str) -> str | None: + """How the provider blocked a served response, the first check of + ``_should_raise_content_policy_error`` (a chat response with no choices, or whose first + choice was content-filtered) and of ``_should_raise_anthropic_refusal_error`` (a + /v1/messages safeguard refusal carrying ``stop_details``).""" + from litellm.llms.anthropic.messages.native.utils import get_safeguard_refusal_stop_details + from litellm.types.utils import ModelResponse + + if endpoint_name == "anthropic_messages" and get_safeguard_refusal_stop_details(response) is not None: + return "refusal" + if self._op not in ("acompletion", "completion") or not isinstance(response, ModelResponse): + return None + return ( + "content_filter" if not response.choices or response.choices[0].finish_reason == "content_filter" else None + ) + + def _blocked_error(self, event: Mapping[str, object]) -> BaseException: + """The error the router raises for a blocked response a fallback would serve.""" + import litellm.exceptions as litellm_exceptions + from litellm.llms.anthropic.messages.native.utils import ( + get_safeguard_refusal_stop_details, + safeguard_refusal_error, + ) + + model: Final = str(event.get("model")) + stop_details: Final = get_safeguard_refusal_stop_details(self._last_response) + if event.get("blocked") == "refusal" and stop_details is not None: + return safeguard_refusal_error(model=model, stop_details=stop_details) + return litellm_exceptions.ContentPolicyViolationError( + message="Response output was blocked.", model=model, llm_provider="" + ) + + def _params_base_model(self, deployment_id: str) -> str | None: + """``_deployment_params_base_model``: the ``base_model`` under the deployment's ``litellm_params``.""" + record: Final = as_str_mapping(self._core._handle.deployment(deployment_id)) or {} # pyright: ignore[reportPrivateUsage] # the driver serves its own core + base_model: Final = (as_str_mapping(record.get("litellm_params")) or {}).get("base_model") + return base_model if isinstance(base_model, str) else None + + def _logged_error(self, event: Mapping[str, object]) -> BaseException | None: + """The error a ``retry_logged`` event records.""" + return self._cause_error(event) + + def _policy_raised(self, event: Mapping[str, object]) -> None: + """A hop's first failure looked its group's retry policy up and the lookup raised: the hop + ends with that error, raised from the failure the event names.""" + model: Final = str(event.get("model")) + original: Final = self._cause_error(event) or RuntimeError(f"router-core: the hop serving {model} failed") + policies: Final = self._request.get( + "model_group_retry_policy", + self._core._resolved.model_group_retry_policy, # pyright: ignore[reportPrivateUsage] # the driver serves its own core + ) + raised: Final = policy_lookup_error(original, model, self._core._resolved.retry_policy, policies) # pyright: ignore[reportPrivateUsage] # see above + level: Final = event.get("level") + self._last_error = raised + self._level_errors[level if isinstance(level, int) else 0] = raised + + def _cause_error(self, event: Mapping[str, object]) -> BaseException | None: + """The error the ``cause`` of a ``retry_logged`` or ``retry_policy_raised`` event names.""" + match event.get("cause"): + case "level": + level: Final = event.get("level") + return self._level_errors.get(level if isinstance(level, int) else 0, self._last_error) + case "pick": + return self._picked_error(event) + case _: + return self._last_error + + def _picked_error(self, event: Mapping[str, object]) -> BaseException: + """The error a pick an event names raised, as a failed step with that error raises it.""" + model: Final = str(event.get("failed_model") or event.get("model")) + return self._core._failure_error(event, model, None, request=self._request) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + + def _level_failed(self, event: Mapping[str, object]) -> None: + """A level raises the error a pick raised (``level_failed``).""" + level: Final = event.get("level") + at_level: Final = level if isinstance(level, int) else 0 + picked: Final = self._picked_error(event) + if event.get("handling") is True: + # Raised by the level's fallback handling, while the level's error was handled. + picked.__context__ = self._level_errors.get(at_level, self._last_error) + self._level_errors[at_level] = picked + return + stamp: Final = self._pending_stamps.pop(at_level, None) + if stamp is not None and type(picked) in self._core._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # the driver serves its own core + # The exhausted retry loop raised this pick's error: stamped as `retries_exhausted` stamps. + picked.max_retries, picked.num_retries = stamp # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes + self._chain_to_hop_first(picked, at_level) + self._level_errors[at_level] = picked + + def take_events(self, events: object) -> None: + for raw_event in _as_list(events): + event = as_str_mapping(raw_event) + if event is None: + continue + if event.get("kind") == "mock_failed": + resolved = self._core._resolved # pyright: ignore[reportPrivateUsage] # the driver serves its own core + mocked = mock_error( + str(event.get("mock")), + str(event.get("model")), + event.get("num_retries"), + self._request, + { + "fallbacks": resolved.fallbacks, + "context_window_fallbacks": resolved.context_window_fallbacks, + "content_policy_fallbacks": resolved.content_policy_fallbacks, + }, + ) + level = event.get("level") + self._last_error = mocked + self._level_errors[level if isinstance(level, int) else 0] = mocked + continue + if event.get("kind") == "retry_policy_raised": + self._policy_raised(event) + continue + if event.get("kind") == "level_failed": + self._level_failed(event) + continue + if event.get("kind") == "fallback": + hop_to = event.get("level") + self._hop_first.pop(hop_to if isinstance(hop_to, int) else 0, None) + self._walk.hop(event) + hop_level = event.get("level") + self._level = hop_level if isinstance(hop_level, int) else self._level + continue + if event.get("kind") == "blocked": + level = event.get("level") + blocked = self._blocked_error(event) + self._last_error = blocked + self._level_errors[level if isinstance(level, int) else 0] = blocked + continue + if event.get("kind") == "fallback_done": + level = event.get("level") + at = level if isinstance(level, int) else 0 + self._level = at + if event.get("succeeded") is not True: + self._hint_nested_failure(self._level_errors.get(at + 1), str(event.get("to"))) + self.outcomes.done( + event, + self._level_errors.get(at, self._last_error), + self._level_errors.get(at + 1), + self._walk.hop_kwargs(event), + ) + continue + if event.get("kind") == "retry_logged": + logged = self._logged_error(event) + self._walk.log_retry(event, logged) + logged_level = event.get("level") + if event.get("attempted_retries") == 0 and logged is not None: + self._hop_first[logged_level if isinstance(logged_level, int) else 0] = logged + continue + if event.get("kind") != "retries_exhausted": + continue + if event.get("picked") is True: + # The level raises the last retry's pick error, which its `level_failed` names. + exhausted_level = event.get("level") + self._pending_stamps[exhausted_level if isinstance(exhausted_level, int) else 0] = ( + event.get("max"), + event.get("attempted"), + ) + continue + error: BaseException | None = self._last_error + # `async_function_with_retries` stamps the error its exhausted retry loop raises. + if error is not None and type(error) in self._core._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # the driver serves its own core + error.max_retries = event.get("max") # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes + error.num_retries = event.get("attempted") # pyright: ignore[reportAttributeAccessIssue] # same + + def _call(self, step: Mapping[str, object]) -> "tuple[Mapping[str, object], AttemptCall]": + name: Final = step.get("model") + requested: Final = str(name) if name is not None else str(self._request.get("model")) + deployment: Final = ( + # A default deployment whose id is not a string is called as Python copies it, id and all. + as_str_mapping(default_deployment_copy(self._core._resolved.default_deployment, requested)) or {} # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if step.get("deployment_id") == DEFAULT_DEPLOYMENT_COPIED_ID + else for_requested_model(as_str_mapping(step.get("deployment")) or {}, requested) + ) + self._attempts += 1 + level: Final = step.get("fallback_level") + # The request params the walk's dict targets merged in (`kwargs.update(target)`). + params: Final = as_str_mapping(step.get("params")) + call: Final = attempt_call( + self._core, + self._plan, + self._request if params is None else {**self._request, **params}, + deployment, + group=requested, + requested=requested, + fallback_level=level if isinstance(level, int) else 0, + ) + shaped: Final = shape_attempt_kwargs(call.kwargs, deployment) + # The walk's metadata for the attempt's level (retries, fallback hops) and its hop kwargs. + call.kwargs = self._walk.stamp(step, shaped) + field: Final = self._plan.metadata_field + if isinstance(call.served, dict) and field in call.kwargs: + call.served[field] = call.kwargs[field] + self._served = (requested, call.served) + self._core._endpoints.counters.started(self._plan, call) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + return deployment, call + + @property + def response(self) -> object: + return self._last_response + + @property + def last_error(self) -> BaseException | None: + return self._last_error + + def _stream_attempt( + self, step: Mapping[str, object], kwargs: Mapping[str, object], slot: "contextlib.ExitStack | None" = None + ) -> StreamAttempt: + def number(key: str) -> int: + value: Final = step.get(key) + return value if isinstance(value, int) else 0 + + order: Final = step.get("target_order") + deployment_id: Final = step.get("deployment_id") + return StreamAttempt( + core=self._core, + op=self._op, + request=self._request, + kwargs=kwargs, + group=str(step.get("model") or step.get("model_group")), + order=order if isinstance(order, int) else None, + depth=number("fallback_depth"), + level=number("fallback_level"), + deployment_id=deployment_id if isinstance(deployment_id, str) else None, + attempted_retries=number("retry"), + slot=slot, + ) + + def _failed( + self, + step: Mapping[str, object], + deployment: Mapping[str, object], + call: "AttemptCall | None", + error: BaseException, + ) -> Mapping[str, object]: + served: Final = call.served if call is not None else {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + stamp_error(self._plan, error, deployment, served) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + self._core._endpoints.attempt_failed(self._plan, error, deployment, call, str(step.get("model"))) # pyright: ignore[reportPrivateUsage] # see above + self._last_error = error + level: Final = step.get("fallback_level") + retry: Final = step.get("retry") + if isinstance(retry, int) and retry > 0: + self._chain_to_hop_first(error, level if isinstance(level, int) else 0) + self._level_errors[level if isinstance(level, int) else 0] = error + return attempt_failure(error) + + def _chain_to_hop_first(self, error: BaseException, level: int) -> None: + """A retry's failure is raised while the hop's first failure is handled (the retry loop + runs in its ``except`` block), so it chains to it.""" + first: Final = self._hop_first.get(level) + if first is not None and first is not error: + error.__context__ = first # rebind-ok: chained as the retry loop's raise chains it + + def _succeeded( + self, call: "AttemptCall", response: object, host_op: str, stream: "StreamAttempt | None" = None + ) -> Mapping[str, object]: + self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see above + self._last_response = response + tokens: Final = total_tokens(response) + # A wrapped stream reports its success or failure when it ends; the attempt counts the + # request and the tokens read so far, which the end's report leaves out. + deferred: Final = stream is not None and stream.reports_outcome + if stream is not None and deferred: + stream.counted_tokens = tokens # rebind-ok: what the end's report leaves out + return {"ok": True, "total_tokens": tokens, "blocked": self._blocked(response, host_op), "deferred": deferred} + + async def _refused( + self, + step: Mapping[str, object], + deployment: Mapping[str, object], + call: "AttemptCall", + error: BaseException, + ) -> Mapping[str, object]: + """A routing pre-call check refused the attempt (``async_routing_strategy_pre_call_checks``): + a rate limit cools the deployment down (``set_cooldown_deployments``); the attempt's failure + is recorded by nothing else.""" + import litellm.exceptions as litellm_exceptions + + core: Final = self._core + if isinstance(error, litellm_exceptions.RateLimitError): + deployment_id: Final = (as_str_mapping(deployment.get("model_info")) or {}).get("id") + if isinstance(deployment_id, str): + report: Final = core._handle.report_failure(deployment_id, dict(attempt_failure(error)), uncounted=True) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + await core._drive_cache_steps(report) # pyright: ignore[reportPrivateUsage] # see above + return {**self._failed(step, deployment, call, error), "recorded": True} + + def _pick_kwargs(self, step: Mapping[str, object]) -> Mapping[str, object]: + """The request kwargs the pick sees: the request with its level's retry count.""" + count: Final = self._walk.request_retry_count(step) + if count is None: + return self._request + field: Final = self._plan.metadata_field + metadata: Final = as_str_mapping(self._request.get(field)) or {} + return {**self._request, field: {**metadata, "request_retry_count": count}} + + def _with_deployment(self, step: Mapping[str, object], resolved: object) -> Mapping[str, object]: + return step if resolved is step.get("deployment") else {**step, "deployment": resolved} + + async def _resolve_pods(self, step: Mapping[str, object]) -> Mapping[str, object]: + """``async_resolve_pods_after``: a picked deployment with Kubernetes pod discovery on is + sent to one of its service's pods.""" + if self._plan.kind in _UNPICKED_KINDS or step.get("deployment") is None: + return step + discovery: Final = self._core._pick_discovery(sync=False) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if discovery is None: + return step + resolved: Final[object] = await discovery.async_resolve_deployment( + step.get("deployment"), self._pick_kwargs(step) + ) + return self._with_deployment(step, resolved) + + def _resolve_pods_sync(self, step: Mapping[str, object]) -> Mapping[str, object]: + """``resolve_pods_after``, on the sync pick.""" + if self._plan.kind in _UNPICKED_KINDS or step.get("deployment") is None: + return step + discovery: Final = self._core._pick_discovery(sync=True) # pyright: ignore[reportPrivateUsage] # see _resolve_pods + if discovery is None: + return step + return self._with_deployment( + step, discovery.resolve_deployment(step.get("deployment"), self._pick_kwargs(step)) + ) + + def _default_call_error(self, step: Mapping[str, object]) -> BaseException | None: + """The error a call to a default deployment every call fails to raises (its first + reads of the copy, before any endpoint runs): nothing records it against a deployment.""" + if step.get("deployment_id") != DEFAULT_DEPLOYMENT_FAILING_ID: + return None + return default_deployment_call_error(self._core._resolved.default_deployment) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + + async def attempt(self, step: Mapping[str, object]) -> Mapping[str, object]: + step = await self._resolve_pods(step) # rebind-ok: the picked deployment, resolved to a pod + deployment: Mapping[str, object] = as_str_mapping(step.get("deployment")) or {} # rebind-ok: the picked one + call: AttemptCall | None = None # rebind-ok: built inside the try + refused_call: Final = self._default_call_error(step) + if refused_call is not None: + return {**self._failed(step, deployment, None, refused_call), "recorded": True} + try: + deployment, call = self._call(step) # rebind-ok: rebound as the Python router rebinds it + host: Final = self._core._host # pyright: ignore[reportPrivateUsage] # the driver serves its own core + # The deployment's max_parallel_requests slot; an open chat stream takes it along. + with enter_deployment_slot(self._core, self._plan, deployment, call.kwargs) as slot: + # A search tool is no deployment: nothing to check; Python's `_arerank` and + # `amoderation` run no routing pre-call checks either. + if self._plan.kind not in _NO_ROUTING_PRE_CALL_CHECKS: + try: + await routing_pre_call_checks( + self._core._checks.routing_callbacks(), # pyright: ignore[reportPrivateUsage] # see above + as_str_mapping(step.get("deployment")) or {}, + ) + except Exception as refused: # noqa: BLE001 # handed to the Rust call loop as the attempt's failure + return await self._refused(step, deployment, call, refused) + pending: Final[object] = ( + self._core._endpoints.create_file_everywhere( # pyright: ignore[reportPrivateUsage] # see above + call.kwargs, call.model_name, self._request.get("specific_deployment") + ) + if self._plan.kind == "create_file" + else host.invoke(call.host_op)(**call.kwargs) + ) + # The generic helper awaits whatever the handler returns, as the Python router does. + awaited: Final = self._plan.kind == "generic" or inspect.isawaitable(pending) + answered: Final[object] = await pending if awaited else pending # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType] # see above + self._core._endpoints.check_response(self._plan, answered, call) # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType] # see above + stream: Final = self._stream_attempt(step, call.kwargs, slot) + response: Final = await open_attempt_stream(stream, answered) # pyright: ignore[reportUnknownArgumentType] # see above + except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop + return self._failed(step, deployment, call, error) + return self._succeeded(call, response, call.host_op, stream) + + def attempt_sync(self, step: Mapping[str, object]) -> Mapping[str, object]: + step = self._resolve_pods_sync(step) # rebind-ok: see attempt + deployment: Mapping[str, object] = as_str_mapping(step.get("deployment")) or {} # rebind-ok: see attempt + call: AttemptCall | None = None # rebind-ok: see attempt + refused_call: Final = self._default_call_error(step) + if refused_call is not None: + return {**self._failed(step, deployment, None, refused_call), "recorded": True} + try: + deployment, call = self._call(step) # rebind-ok: rebound as the Python router rebinds it + try: + self._core._endpoints.sync_pre_call_checks(self._plan, deployment, self._request) # pyright: ignore[reportPrivateUsage] # see above + except Exception as refused: # noqa: BLE001 # handed to the Rust call loop like any attempt failure + # The sync path records no outcome for a refused pre-call check: no cooldown. + return {**self._failed(step, deployment, call, refused), "recorded": True} + answered: Final = self._core._host.invoke(call.host_op)(**call.kwargs) # pyright: ignore[reportPrivateUsage] # a sync endpoint + if inspect.isawaitable(answered): + raise _NotSynchronous(f"the host's {call.host_op} endpoint is not synchronous") + self._core._endpoints.check_response(self._plan, answered, call) # pyright: ignore[reportPrivateUsage] # see above + stream: Final = self._stream_attempt(step, call.kwargs) + response: Final = open_attempt_stream_sync(stream, answered) + except _NotSynchronous: + raise + except Exception as error: # noqa: BLE001 # see attempt + return self._failed(step, deployment, call, error) + return self._succeeded(call, response, "", stream) + + async def select(self, step: Mapping[str, object]) -> object: + request: Final = self._request + if step.get("stage") == "callback_filters": + answer, error = await self._core._answer_callback_filters(step, request, request.get("messages")) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if error is not None: + # The try failed with the filter's error: the level raises it if nothing else serves. + self._last_error = error + self._level_errors[self._level] = error + return answer + return await self._core._answer_select(step, request, request.get("messages"), request.get("input")) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + + def select_sync(self, step: Mapping[str, object]) -> object: + request: Final = self._request + return self._core._answer_select_sync(step, request, request.get("messages"), request.get("input")) # pyright: ignore[reportPrivateUsage] # see above + + async def allow_fallback(self, target: str) -> bool: + """`_is_fallback_target_authorized`, then `_is_fallback_target_within_budget`.""" + if target == self._request.get("model"): + return True + from litellm.types.router import FallbackSkipReason + + ctx: Final = RoutingContext(self._core) # pyright: ignore[reportArgumentType] # the checks read registry() and settings() + request: Final = dict(self._request) + if not await ctx.fallback_access_allowed(model=target, request_kwargs=request): + self._walk.targets().record_skip(target, FallbackSkipReason.NOT_AUTHORIZED) + return False + if not await ctx.fallback_budget_allowed(model=target, request_kwargs=request): + self._walk.targets().record_skip(target, FallbackSkipReason.OVER_BUDGET) + return False + return True + + def finish(self, step: Mapping[str, object]) -> object: + """The response with the router's headers: the deployment id, the model group it was + served as, and the retry and fallback counts (``set_response_headers``).""" + if not self._plan.sets_headers: + return self._last_response + response: object = prepare_response_for_header_attachment(self._last_response) # rebind-ok: gains headers + if response is None: + return self._last_response + name, kwargs = self._served + additional: Final = ensure_response_additional_headers(response) + apply_response_model_id(response, find_deployment_metadata(kwargs)) + additional["x-litellm-model-group"] = name + retries: Final = step.get("attempted_retries") + if isinstance(retries, int): + max_retries: Final = step.get("max_retries") + response = add_retry_headers_to_response( # rebind-ok: see above + response=response, + attempted_retries=retries, + max_retries=max_retries if isinstance(max_retries, int) and retries > 0 else None, + ) + fallbacks: Final = step.get("attempted_fallbacks") + attempted: Final = fallbacks if isinstance(fallbacks, int) else 0 + return add_fallback_headers_to_response( + response=self.outcomes.add_headers(response, attempted), attempted_fallbacks=attempted + ) + + async def finish_async(self, step: Mapping[str, object]) -> object: + response: Final = self.finish(step) + name: Final = self._served[0] + if name is not None and response is not None and self._plan.sets_headers: + remaining: Final = await self._core.health().usage(name, kind="remaining") + if isinstance(remaining, Mapping): + apply_remaining_usage_headers(ensure_response_additional_headers(response), dict(remaining)) + return response + + def failed(self, step: Mapping[str, object]) -> BaseException: + model: Final = str(self._request.get("model")) + primary: Final = self._level_errors.get(0) + if primary is None: + return self._core._failure_error(step, model, self._last_error, request=self._request) # pyright: ignore[reportPrivateUsage] # see above + if self._core._host.dynamic().expose_router_debug_in_errors and hasattr(primary, "message"): # pyright: ignore[reportPrivateUsage] # see above + primary.message += self._explanation(primary, model) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + return primary + + def _hint_nested_failure(self, error: BaseException | None, model: str) -> None: + """A fallback level that fails on a context-window or content-policy error with no list + of its own appends the hint to set one, as every level of the Python walk does.""" + from litellm.router_utils.fallback_event_handlers import check_non_standard_fallback_format + + if ( + error is None + or not hasattr(error, "message") + or not self._core._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # the driver serves its own core + ): + return + resolved: Final = self._core._resolved # pyright: ignore[reportPrivateUsage] # see above + fallbacks: Final = self._request.get("fallbacks", resolved.fallbacks) + if check_non_standard_fallback_format(fallbacks=_as_list(fallbacks)): + return + hint: Final = self._kind_hint(error, model, fallbacks) + if hint: + error.message += hint # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the router appends to the error it raises + + def _kind_hint(self, error: BaseException, model: str, fallbacks: object) -> str: + """The hint a context-window or content-policy error gets when its own list is unset.""" + import litellm.exceptions as litellm_exceptions + + resolved: Final = self._core._resolved # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if isinstance(error, litellm_exceptions.ContextWindowExceededError): + if self._request.get("context_window_fallbacks", resolved.context_window_fallbacks) is not None: + return "" + return ( + f"\nmodel={model}. context_window_fallbacks=None. fallbacks={fallbacks}.\n\n" + "Set 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + ) + if isinstance(error, litellm_exceptions.ContentPolicyViolationError): + if self._request.get("content_policy_fallbacks", resolved.content_policy_fallbacks) is not None: + return "" + return "\n" + format_content_policy_fallback_hint(model, None, fallbacks) # pyright: ignore[reportArgumentType] # the list as configured + return "" + + def _explanation(self, primary: BaseException, model: str) -> str: + """What ``_async_function_with_fallbacks_common_utils`` appends to the error it re-raises + under ``expose_router_debug_in_errors``, from the resolved settings and what the call did.""" + import litellm.exceptions as litellm_exceptions + from litellm.router_utils.fallback_event_handlers import ( + check_non_standard_fallback_format, + get_fallback_model_group_for_lookup_groups, + ) + + resolved: Final = self._core._resolved # pyright: ignore[reportPrivateUsage] # the driver serves its own core + request: Final = self._request + # The request's own lists replace the router's (`kwargs.get("fallbacks", router)`). + fallbacks: Final = request.get("fallbacks", resolved.fallbacks) + window_list: Final = request.get("context_window_fallbacks", resolved.context_window_fallbacks) + policy_list: Final = request.get("content_policy_fallbacks", resolved.content_policy_fallbacks) + lookup: Final = (model,) + if check_non_standard_fallback_format(fallbacks=_as_list(fallbacks)): + chain_list: Final = _as_list(fallbacks) + rejected: Final = rejected_fallback_target(request) + detail: Final = ( + truncate_fallback_error_detail(redact_string(str(rejected))) + if rejected is not None + else self._fallback_detail(chain_list) + ) + return format_fallback_outcome_message(model, chain_list, detail, skipped=self._walk.targets().skipped) + context_window: Final = isinstance(primary, litellm_exceptions.ContextWindowExceededError) + content_policy: Final = isinstance(primary, litellm_exceptions.ContentPolicyViolationError) + # The error kind's own list when it is set, else the generic one with a hint to set it. + own: Final = window_list if context_window else policy_list if content_policy else None + hint: Final = ( + f"\nmodel={model}. context_window_fallbacks=None. fallbacks={fallbacks}.\n\n" + "Set 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + if context_window and own is None + else "\n" + format_content_policy_fallback_hint(model, None, fallbacks) # pyright: ignore[reportArgumentType] # the list as configured + if content_policy and own is None + else "" + ) + source: Final = own if own is not None else fallbacks + entries: Final = [dict(m) for entry in _as_list(source) if (m := as_str_mapping(entry))] + found, generic = ( + get_fallback_model_group_for_lookup_groups(entries, lookup) # pyright: ignore[reportArgumentType] # the fallbacks config shape + if entries + else (None, None) + ) + chain: Final[Sequence[object] | None] = ( + found if found is not None else _as_list(entries[generic].get("*")) if generic is not None else None + ) + if chain is None and entries and own is None: + return hint + format_no_fallback_group_message(lookup, entries) # pyright: ignore[reportArgumentType] # the fallbacks config shape + if chain is None: + return hint + format_fallback_outcome_message(model, None, "") + return hint + format_fallback_outcome_message( + model, chain, self._fallback_detail(chain), skipped=self._walk.targets().skipped + ) + + def _fallback_detail(self, chain: Sequence[object]) -> str: + """The error the fallback walk ended on: the first fallback level's last error, or the + unknown-model error of a target with no deployment, which fails before any attempt.""" + last: Final = str(chain[-1]) if chain else None + attempted: Final = self._level_errors.get(1) + unknown: Final = last is not None and self._core._handle.resolve(last) is None # pyright: ignore[reportPrivateUsage] # the driver serves its own core + error: Final = ( + attempted + if attempted is not None or not unknown or last is None + else self._core._failure_error({"error_kind": "unknown_model"}, last, None) # pyright: ignore[reportPrivateUsage] # see above + ) + return truncate_fallback_error_detail(redact_string(str(error))) if error is not None else "" + + +__all__ = [ + "RUST_SERVED_OPS", + "Resolved", + "RustCore", + "RustCoreUnsupported", + "RustHealth", + "RustRegistry", + "RustSettings", + "attempt_failure", + "failure_kind", + "total_tokens", + "unsupported_arguments", +] diff --git a/litellm/router_rust/invoke/README.md b/litellm/router_rust/invoke/README.md new file mode 100644 index 00000000000..d8562f99359 --- /dev/null +++ b/litellm/router_rust/invoke/README.md @@ -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. diff --git a/litellm/router_rust/invoke/__init__.py b/litellm/router_rust/invoke/__init__.py new file mode 100644 index 00000000000..faa2c13198d --- /dev/null +++ b/litellm/router_rust/invoke/__init__.py @@ -0,0 +1 @@ +"""The host side of ``Host::invoke``: one router attempt against one deployment.""" diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py new file mode 100644 index 00000000000..01c48b517fb --- /dev/null +++ b/litellm/router_rust/invoke/endpoints.py @@ -0,0 +1,2117 @@ +""" +The Rust router core's endpoints: what each public router method does before, around and after +the routed call, as the Python router's endpoint methods do it. + +An endpoint is either routed (the Rust core picks the deployment, retries and falls back, and the +host runs each attempt) or served directly (the Python router calls litellm's function without +routing: no model group, an id that names no deployment, a fan-out over every deployment). Every +call out goes through the host: ``host.invoke(op)``, where an op is a litellm function name or a +dotted path to one (``vector_stores.main.asearch``), looked up when it is called. +""" + +import asyncio +import copy +import dataclasses +import threading +import uuid +from collections import defaultdict +from collections.abc import Awaitable, Callable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._internal_context import service_target +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name +from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs +from litellm.responses.retrieve_stream_query import without_retrieve_stream_defaults +from litellm.router_rust.invoke.shaping import ( + ROUTER_CLIENTS_TARGET, + STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + Shaping, + as_dict, + deployment_params_with_request_reasoning_override, + drop_unsupported_classifier_reasoning_effort, + metadata_bucket, + stamp_deployment_num_retries, + stamp_failed_deployment_id, + stamp_retry_skip_deployment_id, + update_kwargs_with_deployment, +) + +if TYPE_CHECKING: + from litellm.router_rust.core import RustCore + +Kind = Literal[ # rebind-ok: rebound as the Python router rebinds it + "chat", + "text", + "embedding", + "image", + "transcription", + "speech", + "rerank", + "adapter", + "generic", + "generic_sync", + "batch", + "cancel_batch", + "vector_store_create", + "moderation", + "create_file", + "search", +] +Stamp = Literal["deployment", "retry_skip", "none"] # rebind-ok: rebound as the Python router rebinds it +CountKey = Literal["deployment", "group", "empty"] # rebind-ok: rebound as the Python router rebinds it + +# 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", + "original_function", + # The walk's own settings, which `async_function_with_fallbacks` pops before the attempt. + "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", + } +) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Plan: + """How one endpoint runs its routed attempts.""" + + endpoint: str + host_op: str + kind: Kind + sync: bool = False + # The bucket `_update_kwargs_before_fallbacks` writes the requested group into; None for the + # endpoints that write `metadata.model_group` alone (no trace id, no alias). + prepare_field: str | None = "metadata" + function_name: str | None = None + splits_stream_timeout: bool = False + stamp: Stamp = "retry_skip" + total_key: CountKey | None = "deployment" + success_key: CountKey | None = "deployment" + fail_key: CountKey | None = "deployment" + # One attempt, no retries and no fallbacks. + single_attempt: bool = False + # The named parameter the endpoint takes beside **kwargs, sent between the deployment params and + # the request (`{**data, name: value, **kwargs}`). + named: tuple[str, ...] = () + + @property + def sets_headers(self) -> bool: + """Whether the response gets the router's headers (the sync generic wrapper adds none).""" + return self.kind != "generic_sync" + + @property + def metadata_field(self) -> str: + """Where `_update_kwargs_with_deployment` writes the deployment it routed to.""" + return get_router_metadata_variable_name(function_name=self.function_name) + + +def _generic(endpoint: str, host_op: str | None = None, *, sync: bool = False) -> Plan: + if sync: + return Plan( + endpoint=endpoint, + host_op=host_op or endpoint, + kind="generic_sync", + sync=True, + prepare_field="litellm_metadata", + function_name="generic_api_call", + splits_stream_timeout=_handler_name(host_op or endpoint) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + stamp="none", + fail_key="group", + single_attempt=True, + ) + return Plan( + endpoint=endpoint, + host_op=host_op or endpoint, + kind="generic", + prepare_field="litellm_metadata", + function_name="_ageneric_api_call_with_fallbacks", + splits_stream_timeout=_handler_name(host_op or endpoint) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + stamp="deployment", + fail_key="group", + ) + + +def _handler_name(host_op: str) -> str: + return host_op.rsplit(".", 1)[-1] + + +def invoked_handler_name(core: "RustCore", plan: Plan) -> str | None: + """The generic handler's ``__name__`` (`original_generic_function.__name__`): the function the + host serves for the op, which decides the stream-timeout split and the responses socket.""" + function: Final[object] = core._host.invoke(plan.host_op) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the core's host + name: Final[object] = getattr(function, "__name__", None) + return name if isinstance(name, str) else None + + +PLANS: Final[Mapping[str, Plan]] = MappingProxyType( + { + "acompletion": Plan( + "acompletion", "acompletion", "chat", splits_stream_timeout=True, stamp="deployment", named=("messages",) + ), + "completion": Plan( + "completion", + "completion", + "chat", + sync=True, + splits_stream_timeout=True, + stamp="deployment", + total_key=None, + success_key=None, + fail_key=None, + named=("messages",), + ), + "atext_completion": Plan( + "atext_completion", + "atext_completion", + "text", + splits_stream_timeout=True, + fail_key="group", + named=("prompt",), + ), + # `_asearch_with_fallbacks`: the search tool name routes as the model, with no deployment + # picked; its helper picks among the router's search tools. + "asearch": Plan( + "asearch", + "search.asearch", + "search", + prepare_field="litellm_metadata", + stamp="none", + total_key=None, + success_key=None, + fail_key=None, + ), + "aembedding": Plan("aembedding", "aembedding", "embedding", named=("input",)), + "embedding": Plan("embedding", "embedding", "embedding", sync=True, named=("input",)), + "aimage_generation": Plan("aimage_generation", "aimage_generation", "image", named=("prompt",)), + "image_generation": Plan( + "image_generation", + "image_generation", + "image", + sync=True, + prepare_field=None, + total_key="empty", + success_key="empty", + fail_key="empty", + named=("prompt",), + ), + "atranscription": Plan( + "atranscription", + "atranscription", + "transcription", + total_key="group", + success_key="group", + fail_key="group", + named=("file",), + ), + "aspeech": Plan( + "aspeech", + "aspeech", + "speech", + total_key="group", + success_key="group", + fail_key="group", + named=("input",), + ), + "arerank": Plan("arerank", "arerank", "rerank"), + "aadapter_completion": Plan( + "aadapter_completion", + "aadapter_completion", + "adapter", + prepare_field=None, + splits_stream_timeout=True, + fail_key="group", + ), + "acreate_batch": Plan( + "acreate_batch", + "acreate_batch", + "batch", + prepare_field="litellm_metadata", + function_name="_acreate_batch", + fail_key="group", + ), + "acancel_batch": Plan( + "acancel_batch", + "acancel_batch", + "cancel_batch", + prepare_field="litellm_metadata", + function_name="_acancel_batch", + fail_key="group", + ), + "acreate_file": Plan( + "acreate_file", + "acreate_file", + "create_file", + function_name="acreate_file", + stamp="none", + total_key=None, + success_key=None, + fail_key="group", + ), + "avector_store_create": Plan( + "avector_store_create", + "vector_stores.acreate", + "vector_store_create", + prepare_field=None, + function_name="avector_store_create", + stamp="none", + fail_key="group", + single_attempt=True, + ), + "amoderation": Plan( + "amoderation", + "amoderation", + "moderation", + stamp="none", + total_key=None, + success_key=None, + fail_key=None, + single_attempt=True, + ), + # The Python router's anthropic messages host op is `anthropic_messages`. + "aanthropic_messages": _generic("aanthropic_messages", "anthropic_messages"), + "anthropic_messages": _generic("anthropic_messages", "anthropic_messages"), + "aresponses": _generic("aresponses"), + "responses": _generic("responses", sync=True), + } +) + +# --- the generic endpoints ------------------------------------------------------------------ + +# endpoint -> host op, for the async generic endpoints routed through +# `_ageneric_api_call_with_fallbacks` when they carry a model. +_HOST_GENERIC: Final[Mapping[str, str]] = MappingProxyType( + { + "afile_delete": "afile_delete", + "afile_content": "afile_content", + "afile_list": "afile_list", + "aget_responses": "aget_responses", + "acancel_responses": "acancel_responses", + "acompact_responses": "acompact_responses", + "adelete_responses": "adelete_responses", + "alist_input_items": "alist_input_items", + "_arealtime": "_arealtime", + "acreate_realtime_client_secret": "acreate_realtime_client_secret", + "arealtime_calls": "arealtime_calls", + "acreate_realtime_transcription_session": "acreate_realtime_transcription_session", + "_aresponses_websocket": "_aresponses_websocket", + "acreate_fine_tuning_job": "acreate_fine_tuning_job", + "acancel_fine_tuning_job": "acancel_fine_tuning_job", + "alist_fine_tuning_jobs": "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job": "aretrieve_fine_tuning_job", + "aimage_edit": "aimage_edit", + "allm_passthrough_route": "allm_passthrough_route", + "acreate_skill": "acreate_skill", + "alist_skills": "alist_skills", + "aget_skill": "aget_skill", + "adelete_skill": "adelete_skill", + "agenerate_content": "google_genai.agenerate_content", + "agenerate_content_stream": "google_genai.agenerate_content_stream", + "aocr": "ocr.aocr", + "adecisions": "decisions.adecisions", + "asystemone": "systemone.asystemone", + "aextraction": "extraction.aextraction", + **{f"avideo_{name}": f"videos.avideo_{name}" for name in ("generation", "list", "status", "content", "remix")}, + **{ + name: f"containers.{name}" + for name in ("acreate_container", "alist_containers", "aretrieve_container", "adelete_container") + }, + **{ + name: f"containers.endpoint_factory._generated_endpoints.{name}" + for name in ( + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", + ) + }, + **{ + f"avector_store_{name}": f"vector_stores.main.a{name}" + for name in ("search", "retrieve", "list", "update", "delete") + }, + **{ + f"avector_store_file_{name}": f"vector_store_files.main.a{target}" + for name, target in ( + ("create", "create"), + ("list", "list"), + ("retrieve", "retrieve"), + ("content", "retrieve_content"), + ("update", "update"), + ("delete", "delete"), + ) + }, + "acreate_interaction": "interactions.acreate", + "aget_interaction": "interactions.aget", + "adelete_interaction": "interactions.adelete", + "acancel_interaction": "interactions.acancel", + } +) + +_SYNC_GENERIC: Final[Mapping[str, str]] = MappingProxyType( + { + "generate_content": "google_genai.generate_content", + "generate_content_stream": "google_genai.generate_content_stream", + "ocr": "ocr.ocr", + # The sync wrapper routes a search tool name as a model group, unlike `asearch`. + "search": "search.search", + "decisions": "decisions.decisions", + "systemone": "systemone.systemone", + "extraction": "extraction.extraction", + **{f"video_{name}": f"videos.video_{name}" for name in ("generation", "list", "status", "content", "remix")}, + **{ + name: f"containers.{name}" + for name in ("create_container", "list_containers", "retrieve_container", "delete_container") + }, + **{ + f"vector_store_{name}": f"vector_stores.main.{name}" + for name in ("search", "create", "retrieve", "list", "update", "delete") + }, + } +) + +_CONTAINER_FILE_ROUTED: Final = frozenset( + { + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", + } +) +_VIDEO_DIRECT: Final = frozenset({"avideo_create_character", "avideo_get_character", "avideo_edit", "avideo_extension"}) +_SYNC_VIDEO_UNBRANCHED: Final = frozenset( + {"video_create_character", "video_get_character", "video_edit", "video_extension"} +) +_SYNC_CONTAINER_FILES_UNBRANCHED: Final = frozenset( + { + "upload_container_file", + "list_container_files", + "retrieve_container_file", + "delete_container_file", + "retrieve_container_file_content", + } +) +_UNBRANCHED: Final = frozenset( + { + "aadapter_generate_content", + *_SYNC_VIDEO_UNBRANCHED, + *_SYNC_CONTAINER_FILES_UNBRANCHED, + "get_interaction", + "delete_interaction", + "cancel_interaction", + } +) +_ASSISTANTS: Final = frozenset( + { + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", + } +) +_MANAGED_AGENTS: Final[Mapping[str, str]] = MappingProxyType( + { # comprehension-ok: a flat walk of nested rows + f"{prefix}{name}": f"interactions.agents.{prefix}{target}" + for prefix in ("a", "") + for name, target in ( + ("create_agent", "create"), + ("list_agents", "list"), + ("get_agent", "get"), + ("delete_agent", "delete"), + ("list_agent_versions", "list_versions"), + ) + } +) +_SYNC_VECTOR_STORE_FILES: Final[Mapping[str, str]] = MappingProxyType( + { + f"vector_store_file_{name}": f"vector_store_files.main.{target}" + for name, target in ( + ("create", "create"), + ("list", "list"), + ("retrieve", "retrieve"), + ("content", "retrieve_content"), + ("update", "update"), + ("delete", "delete"), + ) + } +) + + +def plan_for(endpoint: str) -> Plan: + plan: Final = generic_plan(endpoint) + if plan is None: + raise KeyError(endpoint) + return plan + + +def generic_plan(endpoint: str) -> Plan | None: + if endpoint in PLANS: + return PLANS[endpoint] + if endpoint in _HOST_GENERIC: + return _generic(endpoint, _HOST_GENERIC[endpoint]) + if endpoint in _VIDEO_DIRECT: + return _generic(endpoint, f"videos.{endpoint}") + if endpoint in _SYNC_GENERIC: + return _generic(endpoint, _SYNC_GENERIC[endpoint], sync=True) + return None + + +# --- preparing the request ------------------------------------------------------------------- + + +def prepare_request( + plan: Plan, request: Mapping[str, object], model: str, aliases: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the call's kwargs + """What the router adds once before the first attempt (`_update_kwargs_before_fallbacks`): a + trace id and the requested group, in the caller's own metadata dict.""" + prepared: Final = dict(request) + if plan.kind == "vector_store_create": + return prepared # routed straight to a deployment, with no group recorded + if plan.prepare_field is None: + metadata_bucket(prepared, "metadata").update({"model_group": model}) + return prepared + prepared.setdefault("litellm_trace_id", str(uuid.uuid4())) + metadata_bucket(prepared, plan.prepare_field).update( + {"model_group": model, "model_group_alias": model if model in aliases else None} + ) + return prepared + + +# --- one attempt ---------------------------------------------------------------------------- + + +@dataclasses.dataclass(slots=True) +class AttemptCall: + """One attempt's host call: the op, its kwargs, and the request kwargs the router kept (the + ones its response headers are read from).""" + + host_op: str + kwargs: Mapping[str, object] + served: Mapping[str, object] + model_name: str + group: str + + +class Counters: + """`total_calls`, `success_calls` and `fail_calls`, keyed as each endpoint keys them.""" + + __slots__ = ("fail_calls", "success_calls", "total_calls") + + def __init__(self) -> None: + self.total_calls: defaultdict[str, int] = defaultdict(int) # mutable-ok: the router's live counters + self.success_calls: defaultdict[str, int] = defaultdict(int) # mutable-ok: same + self.fail_calls: defaultdict[str, int] = defaultdict(int) # mutable-ok: same + + @staticmethod + def _key(which: CountKey, model_name: str, group: str) -> str: + return model_name if which == "deployment" else group if which == "group" else "" + + def started(self, plan: Plan, call: AttemptCall) -> None: + if plan.total_key is not None: + self.total_calls[self._key(plan.total_key, call.model_name, call.group)] += 1 + + def succeeded(self, plan: Plan, call: AttemptCall) -> None: + if plan.success_key is not None: + self.success_calls[self._key(plan.success_key, call.model_name, call.group)] += 1 + + def failed(self, plan: Plan, model_name: str | None, group: str) -> None: + if plan.fail_key is None: + return + if plan.fail_key == "deployment" and model_name is None: + return + self.fail_calls[self._key(plan.fail_key, model_name or "", group)] += 1 + + +def attempt_call( + core: "RustCore", + plan: Plan, + request: Mapping[str, object], + deployment: Mapping[str, object], + *, + group: str, + requested: str, + fallback_level: int = 0, +) -> AttemptCall: + """The host call for one attempt of ``plan`` against ``deployment``: the request (as the attempt + helper receives it), shaped for the deployment and merged with its litellm params.""" + shaping: Final = core._shaping # pyright: ignore[reportPrivateUsage] # the core's own settings + kwargs: Final = {key: value for key, value in request.items() if key not in ROUTING_KEYS} + named: Final = {name: kwargs.pop(name) for name in plan.named if name in kwargs} + data: Final = as_dict(deployment.get("litellm_params")) + model_name: Final = str(data.get("model")) + cache: Final = core._cache # pyright: ignore[reportPrivateUsage] # the routing cache the shaping refunds into + caching: Final = shaping.cache_responses + # A fallback hop routes the target group in a copy of the metadata (`run_async_fallback`). + if fallback_level > 0: + field: Final = plan.prepare_field or "metadata" + kwargs[field] = {**as_dict(kwargs.get(field)), "model_group": requested} + match plan.kind: + case "chat": + host = core._host # pyright: ignore[reportPrivateUsage] # the core's host # rebind-ok: rebound as the Python router rebinds it + drop_unsupported_classifier_reasoning_effort(deployment, requested, kwargs, host) + params = ( # rebind-ok: rebound as the Python router rebinds it + deployment_params_with_request_reasoning_override(data, kwargs) + ) + silent_model = params.pop("silent_model", None) # rebind-ok: rebound as the Python router rebinds it + start_silent_experiments(core, plan, silent_model, named.get("messages"), kwargs) + if "messages" in named: + kwargs.setdefault("messages", named["messages"]) + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, splits_stream_timeout=True) + kwargs.pop("silent_model", None) + final = without_router_only_kwargs( # rebind-ok: rebound as the Python router rebinds it + { + **params, + **named, + "caching": caching, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + ) + return AttemptCall(plan.host_op, final, kwargs, str(params.get("model")), group) + case "text" | "embedding" | "image" | "transcription" | "rerank" | "adapter": + update_kwargs_with_deployment( + shaping, deployment, kwargs, cache=cache, splits_stream_timeout=plan.splits_stream_timeout + ) + if plan.kind == "embedding": + raise_embedding_batch_error(deployment, named.get("input"), kwargs, requested) + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + **named, + "caching": caching, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + return AttemptCall(plan.host_op, without_router_only_kwargs(final), kwargs, model_name, group) + case "speech": + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + voice = kwargs.pop("voice", None) # rebind-ok: rebound as the Python router rebinds it + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + **named, + "voice": data.get("voice") if voice is None else voice, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + return AttemptCall(plan.host_op, final, kwargs, model_name, group) + case "batch" | "cancel_batch" | "vector_store_create": + if plan.kind == "cancel_batch": + data.update(_resolved_credentials(core, deployment, requested)) + data.pop("litellm_credential_name", None) + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, function_name=plan.function_name) + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + "custom_llm_provider": _provider(core, data), + "caching": caching, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + return AttemptCall(plan.host_op, final, kwargs, str(data.get("model")), group) + case "generic": + return _generic_attempt(core, plan, kwargs, deployment, data, requested, group) + case "generic_sync": + update_kwargs_with_deployment( + shaping, + deployment, + kwargs, + cache=cache, + function_name="generic_api_call", + splits_stream_timeout=invoked_handler_name(core, plan) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + ) + endpoint = kwargs.get("endpoint") # rebind-ok: rebound as the Python router rebinds it + if endpoint and isinstance(endpoint, str): + kwargs["endpoint"] = endpoint.replace(requested, model_name) + kwargs["model"] = model_name + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + "custom_llm_provider": _generic_provider(data), + "caching": caching, + **kwargs, + } + return AttemptCall(plan.host_op, final, kwargs, model_name, group) + case "moderation": + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + kwargs["model"] = model_name + kwargs.update(data) + return AttemptCall(plan.host_op, kwargs, kwargs, model_name, group) + case "search": + return _search_attempt(core, plan, kwargs, group) + case "create_file": + # `_acreate_file` creates the file on every healthy deployment of the hop's name, not + # only the picked one: the driver hands the request to `create_file_everywhere`. + return AttemptCall(plan.host_op, kwargs, kwargs, requested, group) + + +def _generic_attempt( + core: "RustCore", + plan: Plan, + kwargs: dict[str, object], # mutable-ok: the attempt's own kwargs + deployment: Mapping[str, object], + data: dict[str, object], # mutable-ok: the attempt's copy of the deployment params + requested: str, + group: str, +) -> AttemptCall: + """`_ageneric_api_call_with_fallbacks_helper` for one deployment.""" + from litellm.llms.base_llm.passthrough.transformation import replace_path_segment + from litellm.router_rust.invoke.websocket import turn_admission + + handler: Final = invoked_handler_name(core, plan) + socket: Final = handler == "_aresponses_websocket" + update_kwargs_with_deployment( + core._shaping, # pyright: ignore[reportPrivateUsage] # the core's own settings + deployment, + kwargs, + cache=core._cache, # pyright: ignore[reportPrivateUsage] # the core's routing cache + function_name="_ageneric_api_call_with_fallbacks", + splits_stream_timeout=handler in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + opens_responses_socket=socket, + ) + model_name: Final = str(data.get("model")) + endpoint: Final = kwargs.get("endpoint") + if endpoint: + try: + stripped, _, _, _ = core._host.invoke("get_llm_provider")( # pyright: ignore[reportPrivateUsage] # the core's host + model=model_name, + custom_llm_provider=kwargs.get("custom_llm_provider"), + api_base=kwargs.get("api_base"), + ) + replacement = str(stripped) # rebind-ok: rebound as the Python router rebinds it + except Exception: # noqa: BLE001 # as the Python router: an unresolvable model is used as is + replacement = model_name # rebind-ok: rebound as the Python router rebinds it + kwargs["endpoint"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + replace_path_segment(str(endpoint), requested, replacement) + ) + provider: Final = _generic_provider(data) + final: Final = { + **without_retrieve_stream_defaults(data, kwargs), + "caching": core._shaping.cache_responses, # pyright: ignore[reportPrivateUsage] # the core's own settings + **kwargs, + "model": model_name, + **_session_with_model(kwargs.get("session"), model_name), + **({"turn_admission": turn_admission(core, requested, deployment, kwargs)} if socket else {}), + } + if provider is not None: + final["custom_llm_provider"] = provider + return AttemptCall(plan.host_op, final, kwargs, model_name, group) + + +_SEARCH_PROTECTED_PARAMS: Final = frozenset({"search_provider", "api_key", "api_base"}) + + +def _search_attempt( + core: "RustCore", + plan: Plan, + kwargs: dict[str, object], # mutable-ok: the attempt's own kwargs + group: str, +) -> AttemptCall: + """`SearchAPIRouter.async_search_with_fallbacks_helper`: one of the search tools named + ``group`` at random, called with its provider, credentials and params under the request's.""" + import random + + tools: Final = [ + tool + for tool in core._endpoints.search_tools # pyright: ignore[reportPrivateUsage] # the core's live search tools + if as_dict(tool).get("search_tool_name") == group + ] + if not tools: + raise ValueError(f"Search tool '{group}' not found in router.search_tools") + selected: Final = random.choice(tools) # load balancing, as the Python router picks + params: Final = as_dict(as_dict(selected).get("litellm_params")) + search_provider: Final = params.get("search_provider") + if not search_provider: + raise ValueError(f"search_provider not found in litellm_params for search tool '{group}'") + search_params: Final = { + key: value + for source in (params, kwargs) + for key, value in source.items() + if key not in _SEARCH_PROTECTED_PARAMS and value is not None + } # comprehension-ok: the tool's params, then the request's + final: Final = { + "search_provider": search_provider, + "api_key": params.get("api_key"), + "api_base": params.get("api_base"), + **search_params, + } + return AttemptCall(plan.host_op, final, kwargs, group, group) + + +def unrouted_fields(plan: Plan) -> Mapping[str, object]: + """A search tool name is no model group: the Rust core hands every attempt to the host as is.""" + return {"unrouted": True} if plan.kind == "search" else {} + + +def _session_with_model(session: object, model_name: str) -> Mapping[str, object]: + if not isinstance(session, Mapping): + return {} + typed: Final = as_dict(session) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + if "model" not in typed: + return {} + return {"session": {**typed, "model": model_name}} + + +def _generic_provider(data: Mapping[str, object]) -> str | None: + from litellm.litellm_core_utils.routing_common_utils import provider_for_generic_call + + return provider_for_generic_call(data) # pyright: ignore[reportArgumentType] # deployment params + + +def model_client( + core: "RustCore", plan: Plan, deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> object: + """The SDK client cached for the deployment (`_get_async_openai_model_client`; `_get_client` + with the sync keys for sync completion and embedding), dropped when the request carries + another api key.""" + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + + if as_dict(deployment.get("litellm_params")).get("kubernetes_pod_discovery") is True: + return None + model_id: Final = as_dict(deployment.get("model_info")).get("id") + sync_keys: Final = plan.sync and plan.kind in ("chat", "embedding") + stream: Final = "stream_" if kwargs.get("stream") is True else "" + key: Final = f"{model_id}_{stream}client" if sync_keys else f"{model_id}_{stream}async_client" + span: Final = get_parent_otel_span_from_kwargs(dict(kwargs)) + client: Final = _cached_client(core, key, span, local_only=not sync_keys) + dynamic_api_key: Final = kwargs.get("api_key") + if dynamic_api_key is not None and client is not None and dynamic_api_key != getattr(client, "api_key", None): + return None + return client + + +def _cached_client(core: "RustCore", key: str, span: object, *, local_only: bool) -> object: + cache: Final = core._cache # pyright: ignore[reportPrivateUsage] # the core's routing cache + with service_target(ROUTER_CLIENTS_TARGET): + if local_only: + return cache.get_cache(key=key, local_only=True, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # the cache holds untyped values + return cache.get_cache(key=key, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # see above + + +def _provider(core: "RustCore", data: Mapping[str, object]) -> object: + """The deployment's declared provider, else the one its model names (`get_llm_provider`).""" + declared: Final = data.get("custom_llm_provider") + _, inferred, _, _ = core._host.invoke("get_llm_provider")( # pyright: ignore[reportPrivateUsage] # the core's host + model=data["model"], custom_llm_provider=declared + ) + return declared or inferred + + +def _resolved_credentials(core: "RustCore", deployment: Mapping[str, object], requested: str) -> Mapping[str, object]: + """`get_deployment_credentials_with_provider` for the picked deployment (`_acancel_batch`).""" + registry: Final = core.registry() + model_id: Final = as_dict(deployment.get("model_info")).get("id") + resolved: Final = registry.credentials(str(model_id or requested), form="with_provider") + return as_dict(resolved) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + + +def raise_embedding_batch_error( + deployment: Mapping[str, object], input: object, kwargs: Mapping[str, object], model: str +) -> None: + from litellm.router_utils.embedding_batch_size_check import embedding_batch_too_large_error + + error: Final = embedding_batch_too_large_error( + deployment=deployment, + input=input, + request_kwargs=kwargs, + model=model, # pyright: ignore[reportArgumentType] # opaque values + ) + if error is not None: + raise error + + +def stamp_error( + plan: Plan, error: BaseException, deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> None: + """What each attempt helper stamps on the error it raises.""" + match plan.stamp: + case "deployment": + stamp_deployment_num_retries(error, deployment) + stamp_failed_deployment_id(error, kwargs.get("model_info") or deployment.get("model_info")) + case "retry_skip": + stamp_retry_skip_deployment_id(error, kwargs) + case "none": + pass + + +# --- silent experiments --------------------------------------------------------------------- + + +_SilentModel: TypeAlias = str | list[str] +_SILENT_MODEL_ADAPTER: Final[TypeAdapter[_SilentModel]] = TypeAdapter(_SilentModel) + + +def _silent_targets(silent_model: object) -> tuple[str, ...]: + """`_silent_experiment_targets`: none inside a mirrored call or for a malformed value.""" + from litellm._internal_context import in_mirrored_call + + if silent_model is None or in_mirrored_call(): + return () + try: + targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model) + except ValidationError: + verbose_router_logger.warning( + "silent_model must be a model name or a list of model names, got %r; skipping shadow traffic", + silent_model, + ) + return () + return (targets,) if isinstance(targets, str) else tuple(targets) + + +def _silent_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]: # mutable-ok: the experiment's own kwargs + """`_get_silent_experiment_kwargs`.""" + from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + + silent: Final = safe_deep_copy(dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + original_metadata: Final = kwargs.get("metadata") + if original_metadata is not None and silent.get("metadata") is original_metadata: # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + silent["metadata"] = dict(original_metadata) # pyright: ignore[reportArgumentType, reportCallIssue] # the request's metadata dict + if "metadata" not in silent: + silent["metadata"] = {} + metadata: Final = silent["metadata"] # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + if isinstance(metadata, dict): + metadata.pop("litellm_parent_otel_span", None) # pyright: ignore[reportUnknownMemberType] # the request's metadata dict + metadata["is_silent_experiment"] = True + silent.pop("litellm_call_id", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + silent.pop("litellm_logging_obj", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + silent.pop("standard_logging_object", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + return silent # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + + +def _snapshot(kwargs: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a snapshot handed to a background task + snapshot: Final = dict(kwargs) + metadata: Final = snapshot.get("metadata") + if isinstance(metadata, Mapping): + snapshot["metadata"] = dict(as_dict(metadata)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + return snapshot + + +async def _run_silent(core: "RustCore", silent_model: str, messages: object, kwargs: Mapping[str, object]) -> None: + from litellm._internal_context import mirrored_call + + remaining: Final = {key: value for key, value in kwargs.items() if key != "stream"} + with mirrored_call(): + response: Final = await core.acompletion( + model=silent_model, messages=messages, stream=bool(kwargs.get("stream", False)), **remaining + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + if not isinstance(response, CustomStreamWrapper): + return + async for _ in response: + pass + + +async def _silent_async(core: "RustCore", silent_model: str, messages: object, kwargs: Mapping[str, object]) -> None: + try: + if as_dict(kwargs.get("metadata")).get("is_silent_experiment", False): + return + silent: Final = _silent_kwargs(kwargs) + as_dict_metadata: Final = silent["metadata"] + if isinstance(as_dict_metadata, dict): + as_dict_metadata["model_group"] = silent_model + await _run_silent(core, silent_model, copy.deepcopy(messages), silent) + except Exception: # noqa: BLE001, S110 # a silent experiment never fails the primary call + pass + + +def _silent_thread(core: "RustCore", silent_model: str, messages: object, kwargs: Mapping[str, object]) -> None: + try: + if as_dict(kwargs.get("metadata")).get("is_silent_experiment", False): + return + silent: Final = _silent_kwargs(kwargs) + metadata: Final = silent["metadata"] + if isinstance(metadata, dict): + metadata["model_group"] = silent_model + loop: Final = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + + async def run() -> None: + await _run_silent(core, silent_model, copy.deepcopy(messages), silent) + pending: Final = asyncio.all_tasks() + current: Final = asyncio.current_task() + if current is not None: + pending.discard(current) + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + loop.run_until_complete(run()) + finally: + loop.close() + except Exception: # noqa: BLE001, S110 # a silent experiment never fails the primary call + pass + + +def start_silent_experiments( + core: "RustCore", plan: Plan, silent_model: object, messages: object, kwargs: Mapping[str, object] +) -> None: + """Mirror the call to each silent model in the background: a task on the async path, a daemon + thread on the sync one.""" + for target in _silent_targets(silent_model): + if plan.sync: + threading.Thread( + target=_silent_thread, args=(core, target, messages, _snapshot(kwargs)), daemon=True + ).start() + else: + _ = asyncio.create_task( + _silent_async(core, target, messages, _snapshot(kwargs)) + ) # fire and forget, as the Python router + + +# --- the endpoints ---------------------------------------------------------------------------- + +_KEEPS_CLIENT: Final = frozenset( + {"acreate_realtime_client_secret", "arealtime_calls", "acreate_realtime_transcription_session"} +) +_KEEPS_PROVIDER_AND_CLIENT: Final = frozenset({"afile_delete", "afile_content"}) +_RESPONSE_ID_ENDPOINTS: Final = frozenset( + {"aget_responses", "acancel_responses", "acompact_responses", "adelete_responses", "alist_input_items"} +) +_CONTAINER_ENDPOINTS: Final = frozenset( + {"acreate_container", "alist_containers", "aretrieve_container", "adelete_container", *_CONTAINER_FILE_ROUTED} +) +_VECTOR_STORE_ENDPOINTS: Final = frozenset( + { + *(f"avector_store_{name}" for name in ("search", "create", "retrieve", "list", "update", "delete")), + *(f"avector_store_file_{name}" for name in ("create", "list", "retrieve", "content", "update", "delete")), + } +) +_INTERACTION_ENDPOINTS: Final = frozenset( + {"acreate_interaction", "aget_interaction", "adelete_interaction", "acancel_interaction"} +) +_MISSING: Final = object() + + +def pinned_video_deployment(plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: + """`_video_call_pin`: the deployment a video id names, for a call on that video only.""" + from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider + + video_id: Final = request.get("video_id") + if plan.endpoint not in VIDEO_ID_CALL_TYPES or not isinstance(video_id, str): + return {} + deployment_id: Final = decode_video_id_with_provider(video_id).get("deployment_id") + return {"pinned_deployment_id": deployment_id} if deployment_id else {} + + +async def _await(value: object) -> object: + """``await value``, raising the TypeError Python raises for a value that is not awaitable.""" + return await value # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType] # deliberately awaited as the Python router does + + +def _missing_model() -> TypeError: + return TypeError("_ageneric_api_call_with_fallbacks() missing 1 required positional argument: 'model'") + + +class RustEndpoints: + """The public endpoint methods of one Rust core, beyond the routed call loop the core drives.""" + + def __init__(self, core: "RustCore", config: "RouterConfig") -> None: + from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout + from litellm.router_utils.request_priority import InvalidPriority, parse_default_priority + + self._core: Final = core + self.counters: Final = Counters() + explicit: Final = config.get("timeout") + self._explicit_timeout: Final = explicit + self._request_timeout: Final = get_configured_request_timeout() if explicit is not None else None + self._stream_timeout: Final = config.get("stream_timeout") + # apply() flips it after construction, as on the Python core + self._cache_responses: bool = bool(config.get("cache_responses")) + self._assistants_config: Final = ( + as_dict(config.get("assistants_config")) if config.get("assistants_config") else None + ) + # These three are assigned by apply() after construction, as on the Python core. + tools: Final = list(_as_items(config.get("search_tools"))) + self._search_tools: list[object] = tools # mutable-ok: the router's live list + self._guardrails: list[object] = list(_as_items(config.get("guardrail_list"))) # mutable-ok: same + parsed: Final = parse_default_priority(config.get("default_priority")) + self._default_priority: Final = None if isinstance(parsed, InvalidPriority) else parsed + self._polling_interval: Final = config.get("polling_interval") + self._caching_groups: Final = config.get("caching_groups") + self._default_max_parallel_requests: object = config.get("default_max_parallel_requests") + self._scheduler: object = None + + @property + def cache_responses(self) -> bool: + return self._cache_responses + + def set_cache_responses(self, value: bool) -> None: + """``cache_responses`` assigned after construction (the proxy turns it on once its cache + exists): later calls pass ``caching=value``, as the Python core's flag does.""" + self._cache_responses = value + + def assign_setting(self, name: str, value: object) -> None: + """``search_tools``, ``guardrail_list`` or ``default_max_parallel_requests`` assigned after + construction, kept as given: later searches, guardrail picks and slot limits built from + then on read it, as the Python core's attributes are read.""" + if name == "search_tools": + self._search_tools = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it + elif name == "guardrail_list": + self._guardrails = value # pyright: ignore[reportAttributeAccessIssue] # same + elif name == "default_max_parallel_requests": + self._default_max_parallel_requests = value + else: + raise ValueError(f"not an endpoint setting: {name}") + + @property + def default_max_parallel_requests(self) -> object: + """The setting as given: ``settings()`` reports it and the slot limits read it.""" + return self._default_max_parallel_requests + + @property + def search_tools(self) -> list[object]: # mutable-ok: the router's live list, as the Python router reports it + return self._search_tools + + @property + def guardrail_list(self) -> list[object]: # mutable-ok: same + return self._guardrails + + @staticmethod + def prepare_host(config: "RouterConfig", host: "PythonHost") -> None: + """`cache_responses` sets the process cache, unless one is set already (the proxy's).""" + if not config.get("cache_responses") or host.global_cache() is not None: + return + from litellm.caching.caching import Cache + + redis_url: Final = config.get("redis_url") + redis_host: Final = config.get("redis_host") + redis_port: Final = config.get("redis_port") + cache_config: Final[dict[str, object]] = {} # mutable-ok: the request kwargs the Python router mutates in place + cache_type: str = "local" # rebind-ok: redis when the router is given one + if redis_url is not None or (redis_host is not None and redis_port is not None): + cache_type = "redis" # rebind-ok: rebound as the Python router rebinds it + if redis_url is not None: + cache_config["url"] = redis_url + if redis_host is not None: + cache_config["host"] = redis_host + if redis_port is not None: + cache_config["port"] = str(redis_port) + if config.get("redis_password") is not None: + cache_config["password"] = config.get("redis_password") + if config.get("redis_db") is not None: + cache_config["db"] = str(config.get("redis_db")) + cache_config.update(as_dict(config.get("cache_kwargs"))) + host.set_global_cache(Cache(type=cache_type, **cache_config)) # pyright: ignore[reportArgumentType] # the Python router's cache config + + def shaping(self, resolved: "Resolved") -> Shaping: + return Shaping( + default_litellm_params=self.default_litellm_params(resolved), + timeout=resolved.timeout, + explicit_timeout=self._explicit_timeout, + request_timeout=self._request_timeout, + stream_timeout=self._stream_timeout, + cache_responses=self._cache_responses, + ) + + def default_litellm_params(self, resolved: "Resolved") -> Mapping[str, object]: + """The defaults the Python constructor completes before any call: the router ``timeout`` + argument, no client retries, and the caching groups in the default metadata.""" + configured: Final = resolved.default_litellm_params + metadata: Final = configured.get("metadata") + return MappingProxyType( + { + "timeout": self._explicit_timeout, + "max_retries": 0, + **configured, + "metadata": { + **(as_dict(metadata) if isinstance(metadata, Mapping) else {}), # pyright: ignore[reportUnknownArgumentType] # the configured metadata + "caching_groups": self._caching_groups, + }, + } + ) + + def prepare( + self, plan: Plan, request: Mapping[str, object], model: str + ) -> dict[str, object]: # mutable-ok: the call's kwargs + return prepare_request(plan, request, model, self._core._resolved.model_group_alias) # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + + # -- host and registry --------------------------------------------------------------------- + + def _invoke(self, op: str) -> Callable[..., object]: + return self._core._host.invoke(op) # pyright: ignore[reportPrivateUsage] # see above + + async def _direct(self, op: str, kwargs: Mapping[str, object]) -> object: + return await _await(self._invoke(op)(**kwargs)) + + def _group(self, name: str) -> list[Mapping[str, object]]: # mutable-ok: a fresh list + """`_core_get_model_list(model_name=name)`: the deployments of the group ``name``.""" + records: Final = self._core._handle.deployments() # pyright: ignore[reportPrivateUsage] # see above + return [record for record in (as_dict(r) for r in records) if record.get("model_name") == name] + + def _is_model_group(self, name: str) -> bool: + return bool(self._group(name)) + + # -- members ------------------------------------------------------------------------------ + + def member(self, name: str) -> Callable[..., object] | None: + """The public endpoint ``name``, bound to this core, or None when the core has none.""" + if name in _UNBRANCHED: + return _answers_none + if name in _ASSISTANTS: + return self._assistants(name) + if name in _MANAGED_AGENTS: + return self._managed_agents(name) + if name in _SYNC_VECTOR_STORE_FILES: + return self._sync_vector_store_files(name) + if name == "create_interaction": + return self._factory("create_interaction", "interactions.create") + direct: Final = { + "avector_store_create": self.avector_store_create, + "acreate_file": self.acreate_file, + "acreate_batch": self.acreate_batch, + "acancel_batch": self.acancel_batch, + "aretrieve_batch": self.aretrieve_batch, + "alist_batches": self.alist_batches, + "aadapter_completion": self.aadapter_completion, + "abatch_completion": self.abatch_completion, + "abatch_completion_one_model_multiple_requests": self.abatch_completion_one_model_multiple_requests, + "abatch_completion_fastest_response": self.abatch_completion_fastest_response, + "schedule_acompletion": self.schedule_acompletion, + "priority_header_update": self.priority_header_update, + "aguardrail": self.aguardrail, + "asearch": self.asearch, + }.get(name) + if direct is not None: + return direct + if name == "responses": + return self._sync_factory(name, PLANS["responses"]) + plan: Final = generic_plan(name) + if plan is None or name in PLANS: + return None + if plan.sync: + return self._sync_factory(name, plan) + return self._factory(name, plan.host_op) + + # -- the factory wrappers ------------------------------------------------------------------ + + def _factory(self, name: str, host_op: str) -> Callable[..., Awaitable[object]]: + async def async_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the endpoint's request + ) -> object: + return await self._wrapped(name, host_op, custom_llm_provider, client, dict(kwargs)) + + async_wrapper.__name__ = name + return async_wrapper + + async def _wrapped( + self, + name: str, + host_op: str, + custom_llm_provider: str | None, + client: object, + kwargs: dict[str, object], # mutable-ok: the wrapper's own kwargs + ) -> object: + plan: Final = generic_plan(name) or _generic(name, host_op) + if name in _CONTAINER_ENDPOINTS: + return await self._containers(plan, custom_llm_provider, kwargs) + if name in _RESPONSE_ID_ENDPOINTS: + from litellm.responses.utils import ResponsesAPIRequestUtils + + model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) # pyright: ignore[reportArgumentType] # any id the caller sends + if model_id is not None: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "model" + ] = model_id + return await self._ageneric(plan, kwargs) + if name in _VECTOR_STORE_ENDPOINTS: + return await self._vector_stores(name, plan, custom_llm_provider, kwargs) + if name in _INTERACTION_ENDPOINTS or name == "create_interaction": + return await self._interactions(plan, custom_llm_provider, kwargs) + if name in _KEEPS_CLIENT: + return await self._ageneric(plan, {**kwargs, "client": client}) + if name in _KEEPS_PROVIDER_AND_CLIENT: + return await self._ageneric(plan, {**kwargs, "custom_llm_provider": custom_llm_provider, "client": client}) + if name == "allm_passthrough_route": + if client: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "client" + ] = client + return await self._ageneric(plan, kwargs, passthrough=True) + return await self._ageneric(plan, kwargs) + + async def generic(self, plan: Plan, request: Mapping[str, object]) -> object: + """A dispatched generic endpoint called with its whole request (the wrapper's + ``custom_llm_provider`` and ``client`` are consumed by the wrapper).""" + kwargs: Final = {key: value for key, value in request.items() if key not in ("custom_llm_provider", "client")} + return await self._ageneric(plan, kwargs) + + async def _ageneric( + self, + plan: Plan, + kwargs: Mapping[str, object], + *, + passthrough: bool = False, + ) -> object: + """`_ageneric_api_call_with_fallbacks`.""" + if "model" not in kwargs: + raise _missing_model() + model: Final = kwargs["model"] + + async def passthrough_call(prepared: Mapping[str, object]) -> object: + sent: Final = {key: value for key, value in prepared.items() if key not in ROUTING_KEYS} + return await self._direct(plan.host_op, {"model": model, **sent}) + + return await self._core._drive(plan, kwargs, passthrough=passthrough_call if passthrough else None) # pyright: ignore[reportPrivateUsage] # see above + + def _sync_factory(self, name: str, plan: Plan) -> Callable[..., object]: + is_vector_store: Final = name.startswith("vector_store_") + + def sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the endpoint's request + ) -> object: + if not is_vector_store: + if "model" not in kwargs: + raise TypeError( + "_generic_api_call_with_fallbacks() missing 1 required positional argument: 'model'" + ) + return self._core._call_sync(plan, kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above + provider_kwargs: dict[ # mutable-ok: the request kwargs the Python router mutates in place # rebind-ok: rebound as the Python router rebinds it + str, object + ] = ( + {**kwargs, "custom_llm_provider": custom_llm_provider} + if custom_llm_provider and "custom_llm_provider" not in kwargs + else dict(kwargs) + ) + if name == "vector_store_search": + provider_kwargs = { # rebind-ok: rebound as the Python router rebinds it + **provider_kwargs, + "_direct_vector_store_embedding_executor": self._embedding_executor(kwargs), + } + if provider_kwargs.get("model"): + return self._core._call_sync(plan, provider_kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above + if name == "vector_store_search": + return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core}) + return self._invoke(plan.host_op)(**provider_kwargs) + + sync_wrapper.__name__ = name + return sync_wrapper + + def _embedding_executor(self, kwargs: Mapping[str, object]) -> object: + from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, + vector_store_request_metadata, + ) + + return RouterVectorStoreEmbeddingExecutor(router=self._core, metadata=vector_store_request_metadata(kwargs)) # pyright: ignore[reportArgumentType] # the router serves aembedding + + async def _containers( + self, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> object: + """`_init_containers_api_endpoints`.""" + from litellm.responses.utils import ResponsesAPIRequestUtils + + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + container_id: Final = kwargs.get("container_id") + forwarded: Final = kwargs.get("model_id") + if isinstance(container_id, str): + decoded: Final = ResponsesAPIRequestUtils.decode_container_id(container_id) + original_id: Final = decoded.get("response_id", container_id) + if original_id != container_id: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "container_id" + ] = original_id + decoded_provider: Final = decoded.get("custom_llm_provider") + if decoded_provider and kwargs.get("custom_llm_provider") == "openai": + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = decoded_provider + model_id: Final = decoded.get("model_id") or ( + forwarded.strip() if isinstance(forwarded, str) and forwarded.strip() else None + ) + if model_id: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "model" + ] = model_id + return await self._ageneric(plan, kwargs) + requested: Final = kwargs.get("model") + if isinstance(requested, str) and requested.strip(): + return await self._ageneric(plan, kwargs, passthrough=True) + return await self._direct(plan.host_op, kwargs) + + async def _vector_stores( + self, + name: str, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: see _wrapped + ) -> object: + """`_init_vector_store_api_endpoints`.""" + if name == "avector_store_search": + kwargs = { # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + **kwargs, + "_direct_vector_store_embedding_executor": self._embedding_executor(kwargs), + } + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + if kwargs.get("model"): + return await self._ageneric(plan, kwargs) + if name == "avector_store_search": + return await self._direct(plan.host_op, {**kwargs, "router": self._core}) + return await self._direct(plan.host_op, kwargs) + + async def _interactions( + self, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> object: + """`_init_interactions_api_endpoints`.""" + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + if "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = "gemini" + if kwargs.get("agent") and kwargs.get("model") == kwargs.get("agent"): + kwargs["model"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + None + ) + if kwargs.get("model"): + return await self._ageneric(plan, kwargs) + return await self._direct(plan.host_op, kwargs) + + def _managed_agents(self, name: str) -> Callable[..., object]: + op: Final = _MANAGED_AGENTS[name] + + def prepared( + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> dict[str, object]: # mutable-ok: the call's kwargs + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + if "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = "gemini" + return kwargs + + if name.startswith("a"): + + async def managed_agents_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return await self._direct(op, prepared(custom_llm_provider, dict(kwargs))) + + return managed_agents_wrapper + + def managed_agents_sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return self._invoke(op)(**prepared(custom_llm_provider, dict(kwargs))) + + return managed_agents_sync_wrapper + + def _sync_vector_store_files(self, name: str) -> Callable[..., object]: + op: Final = _SYNC_VECTOR_STORE_FILES[name] + + def vector_store_file_sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return self._invoke(op)(custom_llm_provider=custom_llm_provider, client=client, **kwargs) + + return vector_store_file_sync_wrapper + + def _assistants(self, name: str) -> Callable[..., Awaitable[object]]: + async def assistants_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + """`_pass_through_assistants_endpoint_factory`.""" + provider: object = custom_llm_provider # rebind-ok: the assistants config's when none is given + sent: Final = dict(kwargs) + if provider is None: + if self._assistants_config is None: + raise Exception( # noqa: TRY002 # the Python router raises a bare Exception + "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" + ) + provider = self._assistants_config[ # rebind-ok: rebound as the Python router rebinds it + "custom_llm_provider" + ] + sent.update(as_dict(self._assistants_config.get("litellm_params"))) + return await self._direct(name, {"custom_llm_provider": provider, "client": client, **sent}) + + return assistants_wrapper + + # -- completions, priority and prompt management ---------------------------------------- + + def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool: + from litellm.litellm_core_utils.routing_common_utils import get_request_team_id + from litellm.router_utils.request_priority import request_drops_params + + def deployment_params() -> tuple[Mapping[str, object], ...]: + names: Final = {as_dict(r).get("model_name") for r in self._core._handle.deployments()} # pyright: ignore[reportPrivateUsage] # see above + pinned: Final = self._core._handle.deployment(model) if model not in names else None # pyright: ignore[reportPrivateUsage] # see above + if pinned is not None: + return (as_dict(as_dict(pinned).get("litellm_params")),) + if kwargs.get("specific_deployment") is True: + candidates = [ # rebind-ok: rebound as the Python router rebinds it + as_dict(r) + for r in self._core._handle.deployments() # pyright: ignore[reportPrivateUsage] # see above + if as_dict(as_dict(r).get("litellm_params")).get("model") == model + ] + else: + team_id = get_request_team_id(kwargs) # rebind-ok: rebound as the Python router rebinds it + candidates = [ # rebind-ok: rebound as the Python router rebinds it + record + for record in self._group(model) + if team_id is None or as_dict(record.get("model_info")).get("team_id") in (None, team_id) + ] + return tuple(as_dict(candidate.get("litellm_params")) for candidate in candidates) + + return request_drops_params( + kwargs, + self._core._resolved.default_litellm_params, # pyright: ignore[reportPrivateUsage] # see above + deployment_params(), + ) + + def _split_priority( + self, model: str, kwargs: Mapping[str, object] + ) -> tuple[int | None, dict[str, object]]: # mutable-ok: the call's kwargs + """`_core_split_request_priority`.""" + from litellm.router_utils.request_priority import InvalidPriority, priority_error, resolve_request_priority + + priority: Final = resolve_request_priority( + requested=kwargs.get("priority"), + default_priority=self._default_priority, + drops_params=lambda: self._drops_params(model=model, kwargs=kwargs), + ) + if isinstance(priority, InvalidPriority): + raise priority_error(model=model, message=priority.message) + return priority, {key: value for key, value in kwargs.items() if key != "priority"} + + def priority_header_update( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + from litellm.router_utils.request_priority import priority_header_fields + + return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs)) + + def _get_scheduler(self) -> object: + if self._scheduler is None: + from litellm.scheduler import Scheduler + + self._scheduler = ( + Scheduler(polling_interval=self._polling_interval) # pyright: ignore[reportArgumentType] # as configured + if self._polling_interval is not None + else Scheduler() + ) + return self._scheduler + + async def _schedule( + self, model: str, priority: int, run: Callable[[], Awaitable[object]], request: Mapping[str, object] + ) -> object: + """`_schedule_factory`: wait for the scheduler's turn, run, mark the response.""" + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.scheduler import FlowItem, Scheduler + + scheduler: Final = self._get_scheduler() + assert isinstance(scheduler, Scheduler) + + async def healthy_deployments() -> Sequence[object]: + healthy: Final = await self._core.select(model, dict(request), kind="healthy") + return list(healthy) if isinstance(healthy, Sequence) else [healthy] + + await scheduler.wait_for_turn( + request=FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model), + timeout=self._core._resolved.timeout, # pyright: ignore[reportPrivateUsage, reportArgumentType] # the router timeout + get_healthy_deployments=healthy_deployments, + ) + try: + response: Final = await run() + except Exception as error: + setattr(error, "priority", priority) # noqa: B010 # the Python router stamps the same attribute + raise + hidden: Final = get_hidden_params(response) + if hidden is not None: + headers: Final = hidden.setdefault("additional_headers", {}) + if isinstance(headers, dict): + headers.update({"x-litellm-request-prioritization-used": True}) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + return response + + async def acompletion(self, request: Mapping[str, object]) -> object: + model: Final = request.get("model") + if not isinstance(model, str): + raise TypeError("acompletion() needs a model name") + kwargs: Final = {**request, "stream": request.get("stream", False)} + plan: Final = PLANS["acompletion"] + priority, rest = self._split_priority(model, kwargs) + from litellm.router_rust.invoke.prompt_management import ( + is_prompt_management_model, + prompt_management_acompletion, + ) + + if is_prompt_management_model(self._core, model): + prepared: Final = self.prepare(plan, rest, model) + messages: Final = rest.get("messages") + + def factory() -> Awaitable[object]: + return prompt_management_acompletion(self._core, model, messages, prepared) + + if priority is None: + return await factory() + return await self._schedule(model, priority, factory, prepared) + if priority is None: + return await self._core._drive(plan, rest) # pyright: ignore[reportPrivateUsage] # see above + return await self._schedule(model, priority, lambda: self._core._drive(plan, rest), rest) # pyright: ignore[reportPrivateUsage] # see above + + async def atext_completion(self, request: Mapping[str, object]) -> object: + model: Final = request.get("model") + if not isinstance(model, str): + raise TypeError("atext_completion() needs a model name") + kwargs: Final = {k: v for k, v in request.items() if k not in ("is_retry", "is_fallback", "is_async")} + plan: Final = PLANS["atext_completion"] + priority, rest = self._split_priority(model, kwargs) + if priority is None: + return await self._core._drive(plan, rest) # pyright: ignore[reportPrivateUsage] # see above + return await self._schedule(model, priority, lambda: self._core._drive(plan, rest), rest) # pyright: ignore[reportPrivateUsage] # see above + + async def schedule_acompletion( + self, + model: str, + messages: object, + priority: int | None = None, + stream: bool = False, + **kwargs: object, # kwargs-ok: the request + ) -> object: + from litellm.router_utils.request_priority import priority_error, unqueueable_priority_message + + queued, _ = self._split_priority(model, {**kwargs, "priority": priority}) + if queued is None: + raise priority_error(model=model, message=unqueueable_priority_message(priority)) + return await self._core.acompletion(model=model, messages=messages, stream=stream, priority=queued, **kwargs) + + # -- batch completions ------------------------------------------------------------------- + + async def _no_exceptions(self, **kwargs: object) -> object: # kwargs-ok: the request + try: + return await self._core.acompletion(**kwargs) + except Exception as error: # noqa: BLE001 # returned in place, as the Python router does + return error + + async def abatch_completion( + self, + models: list[str], # mutable-ok: the request kwargs the Python router mutates in place + messages: object, + **kwargs: object, # kwargs-ok: the endpoint's open request + ) -> object: + if isinstance(messages, list) and all(isinstance(m, dict) for m in messages): # pyright: ignore[reportUnknownVariableType] # checked item by item + return await asyncio.gather( + *(self._no_exceptions(model=model, messages=messages, **kwargs) for model in models) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + ) + if isinstance(messages, list) and all(isinstance(m, list) for m in messages): # pyright: ignore[reportUnknownVariableType] # same + rows: Final[list[object]] = list(messages) # pyright: ignore[reportUnknownArgumentType] # same # mutable-ok: the request kwargs the Python router mutates in place + + async def indexed(model: str, message: object, idx: int) -> tuple[object, int]: + return await self._no_exceptions(model=model, messages=message, **kwargs), idx + + responses: Final = await asyncio.gather( + *( + indexed(model, message, idx) for idx, message in enumerate(rows) for model in models + ) # comprehension-ok: a flat walk of nested rows + ) + final: Final[list[list[object]]] = [ # mutable-ok: the request kwargs the Python router mutates in place + [] for _ in rows + ] + for response, idx in responses: + final[idx].append(response) + return final + return None + + async def abatch_completion_one_model_multiple_requests( + self, + model: str, + messages: list[object], # mutable-ok: the request kwargs the Python router mutates in place + **kwargs: object, # kwargs-ok: the endpoint's open request + ) -> object: + return await asyncio.gather(*(self._no_exceptions(model=model, messages=m, **kwargs) for m in messages)) + + async def abatch_completion_fastest_response( + self, + model: str, + messages: object, + stream: bool = False, + **kwargs: object, # kwargs-ok: the request + ) -> object: + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.types.utils import ModelResponse + + models: Final = [m.strip() for m in model.split(",")] + + async def no_exceptions(model_name: str) -> object: + try: + return await self._core.acompletion(model=model_name, messages=messages, stream=stream, **kwargs) + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 # see _no_exceptions + return error + + pending: set[ # mutable-ok: the request kwargs the Python router mutates in place # rebind-ok: rebound as the Python router rebinds it + asyncio.Task[object] + ] = {asyncio.create_task(no_exceptions(m)) for m in models} + while pending: + done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) + for task in done: + result = task.result() + if isinstance(result, (ModelResponse, CustomStreamWrapper)): + for other in pending: + other.cancel() + hidden = get_hidden_params(result) + if hidden is not None: + hidden["fastest_response_batch_completion"] = True + return result + raise Exception("All tasks failed") # noqa: TRY002 # the Python router raises a bare Exception + + # -- moderation, text completion and adapter -------------------------------------------- + + async def amoderation(self, request: Mapping[str, object]) -> object: + """`_pass_through_moderation_endpoint_factory`: one pick when the model is a group.""" + kwargs: Final = {k: v for k, v in request.items() if k not in ("custom_llm_provider", "client")} + plan: Final = PLANS["amoderation"] + model: Final = kwargs.get("model", "") + if model and isinstance(model, str) and self._is_model_group(model): + return await self._core._drive(plan, kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above + prepared: Final = prepare_request(plan, kwargs, str(model), self._core._resolved.model_group_alias) # pyright: ignore[reportPrivateUsage] # see above + return await self._direct("amoderation", prepared) + + def text_completion(self, request: Mapping[str, object]) -> object: + """The sync `text_completion`: one pick, no retries, the group name as the model.""" + from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin + + kwargs: Final = {k: v for k, v in request.items() if k not in ("is_retry", "is_fallback", "is_async")} + model: Final = kwargs.get("model") + prompt: Final = kwargs.get("prompt") + if not isinstance(model, str): + raise TypeError("text_completion() missing 1 required positional argument: 'model'") + kwargs["num_retries"] = kwargs.get("num_retries", self._core._resolved.num_retries) # pyright: ignore[reportPrivateUsage] # see above + metadata_bucket(kwargs, "metadata").update({"model_group": model}) + specific: Final = kwargs.pop("specific_deployment", None) + deployment: Final = self.select_sync(model, kwargs, specific=bool(specific)) + data: Final = as_dict(deployment.get("litellm_params")) + for key, value in self._core._resolved.default_litellm_params.items(): # pyright: ignore[reportPrivateUsage] # see above + if key not in kwargs: + kwargs[key] = value + elif key == "metadata": + bucket = kwargs[key] + if isinstance(bucket, dict): + bucket.update(as_dict(value)) # pyright: ignore[reportUnknownMemberType] # the request's metadata dict + pinned: Final = deployment_model_id_pin(data) + if pinned is not None: + kwargs["model_id"] = pinned + return self._invoke("text_completion")(**{**data, "prompt": prompt, "caching": self._cache_responses, **kwargs}) + + def select_sync(self, model: str, request: Mapping[str, object], *, specific: bool = False) -> Mapping[str, object]: + """Pick a deployment for ``model`` on the sync path (`get_available_deployment`).""" + from litellm.router_attempt import for_requested_model + + core: Final = self._core + fields: Final = {"specific_deployment": True} if specific else {} + call: Final = core._handle.select(model, _document_of({**_start_fields_of(request), **fields})) # pyright: ignore[reportPrivateUsage] # see above + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = as_dict(call.next(outcome)) # pyright: ignore[reportArgumentType] # Rust validates each outcome against the step it answers + match step.get("kind"): + case "cache": + outcome = core._answer_cache_sync(step) # pyright: ignore[reportPrivateUsage] # see above + case "select": + outcome = step["candidates"] + case "done": + record = core._handle.deployment(str(step["deployment_id"])) # pyright: ignore[reportPrivateUsage] # see above + return for_requested_model(as_dict(record), model) + case "failed": + raise core._failure_error(step, model, None) # pyright: ignore[reportPrivateUsage] # see above + case other: + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + async def aadapter_completion( + self, + adapter_id: str, + model: str, + is_retry: bool | None = False, + is_fallback: bool | None = False, + is_async: bool | None = False, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return await self._core._drive( # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + PLANS["aadapter_completion"], {**kwargs, "model": model, "adapter_id": adapter_id} + ) # pyright: ignore[reportPrivateUsage] # see above + + # -- files, batches and vector stores --------------------------------------------------- + + async def acreate_batch(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + return await self._core._drive(PLANS["acreate_batch"], {**kwargs, "model": model}) # pyright: ignore[reportPrivateUsage] # see above + + async def acancel_batch(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + return await self._core._drive(PLANS["acancel_batch"], {**kwargs, "model": model}) # pyright: ignore[reportPrivateUsage] # see above + + async def avector_store_create(self, model: str | None, **kwargs: object) -> object: # kwargs-ok: the request + if model is None: + return await self._factory("avector_store_create", "vector_stores.main.acreate")(**kwargs) + return await self._core._drive(PLANS["avector_store_create"], {**kwargs, "model": model}, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above + + async def aretrieve_batch(self, model: str | None = None, **kwargs: object) -> object: # kwargs-ok: the request + """Ask every deployment of the group (or of the router) for the batch.""" + from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + from litellm.types.utils import LiteLLMBatch + + shaping: Final = self._core._shaping # pyright: ignore[reportPrivateUsage] # see above + metadata_variable_name: Final = get_router_metadata_variable_name(function_name="aretrieve_batch") + if model is not None: + specific: Final = kwargs.pop("specific_deployment", None) + healthy = await self._core.select( # rebind-ok: rebound as the Python router rebinds it + model, dict(kwargs), kind="healthy", specific_deployment=bool(specific) or None + ) + else: + healthy = [as_dict(r) for r in self._core._handle.deployments()] # pyright: ignore[reportPrivateUsage] # see above # rebind-ok: rebound as the Python router rebinds it + errors: Final[list[BaseException]] = [] # mutable-ok: the request kwargs the Python router mutates in place + + async def try_retrieve(deployment: Mapping[str, object]) -> object: + try: + data: Final = as_dict(deployment.get("litellm_params")) + deployment_model: Final = data.get("model") + provider: object = data.get("custom_llm_provider") # rebind-ok: inferred when not declared + if deployment_model is None: + raise Exception(f"Model not found in litellm_params for deployment: {deployment}") # noqa: TRY002 # as the Python router + if not provider: + _, provider, _, _ = self._core._host.invoke("get_llm_provider")(model=deployment_model) # pyright: ignore[reportPrivateUsage] # see above # rebind-ok: rebound as the Python router rebinds it + new_kwargs: Final = safe_deep_copy(dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + update_kwargs_with_deployment( + shaping, + deployment, + new_kwargs, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + cache=self._core._cache, # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + function_name="aretrieve_batch", # pyright: ignore[reportPrivateUsage] # see above + ) + bucket: Final = metadata_bucket(new_kwargs, metadata_variable_name) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + if not bucket.get("model_group"): + bucket["model_group"] = model or deployment.get("model_name") + new_kwargs.pop("custom_llm_provider", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + data.pop("custom_llm_provider", None) + return await self._direct("aretrieve_batch", {**data, "custom_llm_provider": provider, **new_kwargs}) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + except Exception as error: # noqa: BLE001 # collected; the first one is raised when no deployment has the batch + errors.append(error) + return None + + deployments: Final = list(healthy) if isinstance(healthy, Sequence) else [healthy] + if not deployments: + raise Exception("No healthy deployments found.") # noqa: TRY002 # as the Python router + results: Final = await asyncio.gather(*(try_retrieve(as_dict(d)) for d in deployments), return_exceptions=True) + for result in results: + if isinstance(result, LiteLLMBatch): + return result + if errors: + raise errors[0] + raise Exception(f"Unable to find batch in any model. Received errors - {errors}") # noqa: TRY002 # as the Python router + + async def alist_batches(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + deployments: Final = self._group(model) + + async def try_list(deployment: Mapping[str, object]) -> object: + try: + return await self._direct("alist_batches", {**as_dict(deployment.get("litellm_params")), **kwargs}) + except Exception: # noqa: BLE001 # a failing deployment is left out, as the Python router does + return None + + results: Final = await asyncio.gather(*(try_list(d) for d in deployments)) + final: Final[dict[str, object]] = { # mutable-ok: the request kwargs the Python router mutates in place + "object": "list", + "data": [], + "first_id": None, + "last_id": None, + "has_more": False, + } + data: Final[list[object]] = [] # mutable-ok: the request kwargs the Python router mutates in place + final["data"] = data + for result in results: + if result is None: + continue + if final["first_id"] is None and hasattr(result, "first_id"): + final["first_id"] = getattr(result, "first_id") # noqa: B009 # as the Python router + final["last_id"] = getattr(result, "last_id") # noqa: B009 # same + data.extend(getattr(result, "data")) # noqa: B009 # same # pyright: ignore[reportAny] # litellm's untyped helpers and request values + if getattr(result, "has_more", False) is True: + final["has_more"] = True + return final + + async def acreate_file(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + return await self._core._drive(PLANS["acreate_file"], {**kwargs, "model": model}) # pyright: ignore[reportPrivateUsage] # see above + + async def create_file_everywhere(self, kwargs: Mapping[str, object], requested: str, specific: object) -> object: + """`_acreate_file`: create the file on every healthy deployment of the group.""" + from litellm.litellm_core_utils import batch_utils + from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.litellm_core_utils.routing_common_utils import ( + add_model_file_id_mappings, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + from litellm.types.utils import LlmProviders + + healthy: Final = await self._core.select( + requested, dict(kwargs), kind="healthy", specific_deployment=bool(specific) or None + ) + deployments: Final = list(healthy) if isinstance(healthy, Sequence) else [healthy] + + async def create_for(deployment: Mapping[str, object]) -> object: + kwargs_copy: Final = safe_deep_copy(dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + update_kwargs_with_deployment( + self._core._shaping, # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + deployment, + kwargs_copy, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + cache=self._core._cache, # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + function_name="acreate_file", # pyright: ignore[reportPrivateUsage] # see above + ) + data: Final = as_dict(deployment.get("litellm_params")) + model_name: Final = str(data.get("model")) + self.counters.total_calls[model_name] += 1 + declared: Final = data.get("custom_llm_provider") + stripped, inferred, _, _ = self._core._host.invoke("get_llm_provider")( # pyright: ignore[reportPrivateUsage] # see above + model=data["model"], custom_llm_provider=declared + ) + provider: Final = declared or inferred + purpose: Final = kwargs.get("purpose") + file: object = kwargs.get("file") # rebind-ok: rewritten for the deployment + if not file or not purpose: + raise Exception("file and file_purpose are required for create_file") # noqa: TRY002 # as the Python router + replaces: Final = batch_utils.should_replace_model_in_jsonl( + purpose=purpose, # pyright: ignore[reportArgumentType] # the caller's purpose + passthrough=kwargs.get("passthrough") is True, + ) + if replaces: + file = batch_utils.replace_model_in_jsonl( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # an untyped helper # rebind-ok: rewritten for the deployment + file_content=file, # pyright: ignore[reportArgumentType] # the caller's file + new_model_name=stripped, + keep_line_numbers=provider == LlmProviders.VERTEX_AI.value, + ) + kwargs_copy["file"] = file + if provider == LlmProviders.LITELLM_PROXY.value: + kwargs_copy["extra_body"] = MappingProxyType( + {**as_dict(kwargs_copy.get("extra_body")), "target_model_names": stripped} # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # litellm's untyped helpers and request values + ) + if "gcs_bucket_name" in data: + metadata_bucket(kwargs_copy, "litellm_metadata")["gcs_bucket_name"] = data["gcs_bucket_name"] # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + from litellm.router_rust.invoke.slots import deployment_slot + + with deployment_slot(self._core, deployment, kwargs_copy): # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + response: Final = await self._direct( + "acreate_file", + { + **data, + "custom_llm_provider": provider, + "caching": self._cache_responses, + "client": model_client(self._core, PLANS["acreate_file"], deployment, kwargs_copy), # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + **kwargs_copy, + }, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + ) + self.counters.success_calls[model_name] += 1 + return response + + responses: Final = await asyncio.gather(*(create_for(as_dict(d)) for d in deployments)) + if len(responses) == 0: + raise Exception("No healthy deployments found.") # noqa: TRY002 # as the Python router + mapping: Final = add_model_file_id_mappings(healthy_deployments=deployments, responses=responses) # pyright: ignore[reportArgumentType, reportUnknownVariableType] # deployment records + returned: Final = responses[0] + hidden: Final = get_hidden_params(returned) + if hidden is not None: + hidden["model_file_id_mapping"] = mapping + return returned + + # -- search tools ------------------------------------------------------------------------- + + async def asearch( + self, + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + """`_asearch_with_fallbacks`: the search tool name (``search_tool_name``, else ``model``) + routes as the model; the factory wrapper drops ``custom_llm_provider`` and ``client``.""" + request: Final = dict(kwargs) + search_tool_name: Final = request.get("search_tool_name", request.get("model")) + if not search_tool_name: + raise ValueError("search_tool_name or model parameter is required for search") + request["model"] = search_tool_name + return await self._core._drive(PLANS["asearch"], request) # pyright: ignore[reportPrivateUsage] # the endpoints drive their own core + + # -- guardrails --------------------------------------------------------------------------- + + async def aguardrail( + self, + guardrail_name: str, + original_function: Callable[..., Awaitable[object]], + **kwargs: object, # kwargs-ok: the request + ) -> object: + """Pick one of the guardrails named ``guardrail_name`` (weighted) and run it. A guardrail name + is no model group, so the router's retries never apply.""" + request: Final = dict(kwargs) + request.setdefault("litellm_trace_id", str(uuid.uuid4())) + metadata_bucket(request, "litellm_metadata").update( + { + "model_group": guardrail_name, + "model_group_alias": guardrail_name + if guardrail_name in self._core._resolved.model_group_alias # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + else None, # pyright: ignore[reportPrivateUsage] # see above + } + ) + return await original_function(**request, selected_guardrail=self.available_guardrail(guardrail_name)) + + def available_guardrail(self, guardrail_name: str) -> object: + """``_core_get_available_guardrail``: the one guardrail configured under the name, else one of + them picked by weight (``simple_shuffle``, the name resolved through the alias map).""" + from litellm.router_strategy.simple_shuffle import simple_shuffle + + candidates: Final = [g for g in self._guardrails if as_dict(g).get("guardrail_name") == guardrail_name] + if not candidates: + raise ValueError(f"No guardrail found with name: {guardrail_name}") + if len(candidates) == 1: + return candidates[0] + hosted: Final = self._core._hosted # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + return simple_shuffle( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helper + resolve_model_alias=hosted._core_get_model_from_alias, # pyright: ignore[reportPrivateUsage] # the host registry's alias lookup + healthy_deployments=candidates, # pyright: ignore[reportArgumentType] # guardrail configs + model=guardrail_name, + request_kwargs=None, + ) + + # -- attempts ----------------------------------------------------------------------------- + + def attempt_failed( + self, + plan: Plan, + error: BaseException, + deployment: Mapping[str, object], + call: AttemptCall | None, + requested: str, + ) -> None: + self.counters.failed(plan, call.model_name if call is not None else None, requested) + if plan.kind != "chat": + return + from litellm import exceptions as litellm_exceptions + from litellm.types.router import LiteLLM_Params + + if isinstance(error, litellm_exceptions.Timeout) and self._core._host.dynamic().expose_router_debug_in_errors: # pyright: ignore[reportPrivateUsage] # see above + params: Final = as_dict(deployment.get("litellm_params")) + try: + typed: Mapping[str, object] = ( # rebind-ok: rebound as the Python router rebinds it + LiteLLM_Params.model_validate(params).model_dump() + ) + except Exception: # noqa: BLE001 # the params as configured, when they do not validate + typed = params # rebind-ok: rebound as the Python router rebinds it + error.message += ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + f"\n\nDeployment Info: request_timeout: {typed.get('request_timeout')}" + f"\ntimeout: {typed.get('timeout')}" + f"\nstream_timeout: {typed.get('stream_timeout')}" + ) + + def check_response(self, plan: Plan, response: object, call: AttemptCall) -> None: + """What the attempt helper checks of a response before it counts it as served.""" + + def sync_pre_call_checks(self, plan: Plan, deployment: Mapping[str, object], request: Mapping[str, object]) -> None: + """`routing_strategy_pre_call_checks` on the sync path: the process callbacks' pre-call checks. + Sync completion runs them for a group but not for a deployment id.""" + from litellm.integrations.custom_logger import CustomLogger + + if plan.kind == "chat": + model: Final = str(request.get("model")) + names: Final = {as_dict(r).get("model_name") for r in self._core._handle.deployments()} # pyright: ignore[reportPrivateUsage] # see above + if model not in names and self._core._handle.deployment(model) is not None: # pyright: ignore[reportPrivateUsage] # see above + return + for callback in tuple(self._core._checks.routing_callbacks()): # pyright: ignore[reportPrivateUsage] # `_routing_callbacks` + if isinstance(callback, CustomLogger): + callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + + +_SINGLE_ATTEMPT: Final[Mapping[str, object]] = MappingProxyType({"num_retries": 0, "disable_fallbacks": True}) + + +async def _answers_none(*_args: object, **_kwargs: object) -> None: # kwargs-ok: any request + """An endpoint the Python router's async wrapper has no branch for.""" + return + + +def _as_items(value: object) -> list[object]: # mutable-ok: a fresh list + return list(value) if isinstance(value, list) else [] # pyright: ignore[reportUnknownArgumentType] # a config list + + +def _start_fields_of(request: Mapping[str, object]) -> Mapping[str, object]: + from litellm.router_rust.core import ( + _start_fields, # pyright: ignore[reportPrivateUsage] # the core's own request fields + ) + + return _start_fields(request) + + +def _document_of(value: Mapping[str, object]) -> Mapping[str, object]: + from litellm.router_rust.core import _document # pyright: ignore[reportPrivateUsage] # the core's plain-JSON rule + + return _document(value) + + +if TYPE_CHECKING: + from litellm.router_config import RouterConfig + from litellm.router_python_host import PythonHost + from litellm.router_rust.core import Resolved + + +__all__ = [ + "PLANS", + "ROUTING_KEYS", + "AttemptCall", + "Counters", + "Plan", + "RustEndpoints", + "Shaping", + "attempt_call", + "generic_plan", + "pinned_video_deployment", + "plan_for", + "prepare_request", + "stamp_error", + "unrouted_fields", +] diff --git a/litellm/router_rust/invoke/prompt_management.py b/litellm/router_rust/invoke/prompt_management.py new file mode 100644 index 00000000000..9af34869c5d --- /dev/null +++ b/litellm/router_rust/invoke/prompt_management.py @@ -0,0 +1,131 @@ +""" +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 TYPE_CHECKING, Final + +from litellm.router_rust.invoke.endpoints import PLANS +from litellm.router_rust.invoke.shaping import as_dict, update_kwargs_with_deployment + +if TYPE_CHECKING: + from litellm.router_rust.core import RustCore + +_PROMPT_MANAGEMENT_PARAMS: Final = frozenset( + {"bitbucket_config", "dotprompt_config", "prompt_id", "prompt_variables", "prompt_label", "prompt_version"} +) + + +def is_prompt_management_model(core: "RustCore", 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._hosted._core_get_model_list(model_name=model) # pyright: ignore[reportPrivateUsage] # the core's host registry + 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 # pyright: ignore[reportPrivateUsage] # the core's host + 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: "RustCore", + 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.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, # pyright: ignore[reportPrivateUsage] # the core's own settings + deployment, + kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs + cache=core._cache, # pyright: ignore[reportPrivateUsage] # the routing cache the shaping refunds into + 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 = {key: value for key, value in data.items() if key not in _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._hosted._core_get_model_list(model_name=str(rendered_model)) # pyright: ignore[reportPrivateUsage, 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[reportPrivateUsage, 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 = {key: value for key, value in final.items() if key not in ("timeout", "stream_timeout")} + return await core._drive(PLANS["acompletion"], {**routed, **timeouts}, prepared_already=True) # pyright: ignore[reportPrivateUsage] # the core drives its own call diff --git a/litellm/router_rust/invoke/shaping.py b/litellm/router_rust/invoke/shaping.py new file mode 100644 index 00000000000..08707dbe9f8 --- /dev/null +++ b/litellm/router_rust/invoke/shaping.py @@ -0,0 +1,634 @@ +""" +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.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 = {key: value for key, value in nested.items() if key != "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", +] diff --git a/litellm/router_rust/invoke/slots.py b/litellm/router_rust/invoke/slots.py new file mode 100644 index 00000000000..b0847cb2f7e --- /dev/null +++ b/litellm/router_rust/invoke/slots.py @@ -0,0 +1,108 @@ +""" +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.core import RustCore + from litellm.router_rust.invoke.endpoints import 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: "RustCore", 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._cache # pyright: ignore[reportPrivateUsage] # the core's 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[reportPrivateUsage, 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: "RustCore", 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: "RustCore", 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"] diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py new file mode 100644 index 00000000000..a736f497e1f --- /dev/null +++ b/litellm/router_rust/invoke/streaming.py @@ -0,0 +1,1794 @@ +""" +The stream wrappers of the Rust router core: host-side objects around a streamed response. + +Every attempt the Rust core runs that returns a stream is wrapped here before the core sees the +outcome, the way the Python router wraps each attempt's stream: + +- a chat stream (``CustomStreamWrapper``) has its first chunk read inside the attempt, so a stream + that fails before any token fails the attempt and the core retries or falls back as for any error +- a chat, ``/v1/messages`` or Responses stream that fails later, before the caller received content, + hands the failure back to the core: the core decides whether a same-group retry or a fallback takes + over (``RouterCoreHandle.stream_policy`` / ``retry_stream`` / ``fallback_stream``), and the wrapper + splices what that call returns into the stream the caller already holds + +Each wrapper also reports its stream's outcome to the core when the stream ends, as litellm's stream +logging reports it to the Python router's callbacks: the attempt that returned the stream counted +only its request (the attempt answer is ``deferred``), so the success with the stream's final tokens, +or a failure the stream raised after it opened, is recorded then (``StreamAttempt.report_success`` / +``report_failure``). A stream closed before it ends reports nothing. + +The wrappers decide nothing about routing: what may take over, the retry budget and every +deployment pick come from the core. They hold live Python iterators, frames and hidden params, +which is why they live in the host. Ported from the Python router's ``_acompletion_streaming_iterator``, +``_completion_streaming_iterator``, ``_aanthropic_messages_streaming_iterator`` and +``_aresponses_streaming_iterator``, without importing it. +""" + +import asyncio +import contextlib +import dataclasses +import json +import weakref +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Generator, + Iterable, + Iterator, + Mapping, + Sequence, +) +from typing import ( + TYPE_CHECKING, + Final, + Generic, + NoReturn, + TypeAlias, + TypeVar, + cast, # noqa: TID251 # untyped stream iterators and provider chunk lists +) + +import anyio + +import litellm.exceptions as litellm_exceptions +from litellm._logging import verbose_router_logger +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_retry_headers_to_response, + get_hidden_params_dict, + replace_complexity_router_headers, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.utils import ModelResponseStream, Usage + +if TYPE_CHECKING: + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + from litellm.router_python_host import PythonHost + from litellm.router_rust.core import RustCore + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ContinuationOutput, ContinuedPartial + from litellm.types.llms.openai import ResponseAPIUsage, ResponseIncompleteEvent + +# A stream that never emits content or an error could grow a held-frame buffer without bound, so +# hitting either cap commits the stream instead. +MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 +MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 + +# The hidden params and headers a fallback's response carries, as dicts the proxy writes into. +_Prepared: TypeAlias = "tuple[dict[str, object], dict[str, object]]" +_RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) + + +# --- the attempt a stream came from ------------------------------------------------------------- + + +@dataclasses.dataclass(slots=True) +class StreamAttempt: + """The attempt that returned a stream, and the core that can still recover it. + + ``request`` is the call's prepared request, the one every attempt of the call is built from; + ``kwargs`` are this attempt's own kwargs. The origin fields are what the core reads to decide + what may take over a failure of this stream. + """ + + core: "RustCore" + op: str + request: Mapping[str, object] + kwargs: Mapping[str, object] + group: str + order: int | None + depth: int + level: int + deployment_id: str | None + attempted_retries: int + committed_budget: int | None = None + # The deployment's max_parallel_requests slot the attempt holds; an open chat stream takes it. + slot: contextlib.ExitStack | None = None + # Set by the wrapper that reports the stream's outcome when it ends; the attempt's answer is then + # deferred and counts only the request and the tokens read so far (``counted_tokens``). + reports_outcome: bool = False + counted_tokens: int = 0 + outcome_reported: bool = False + + @property + def model_group(self) -> str: + return self.group + + def origin(self) -> Mapping[str, object]: + return { + "group": self.group, + "order": self.order, + "depth": self.depth, + "level": self.level, + "deployment_id": self.deployment_id, + "attempted_retries": self.attempted_retries, + "committed_budget": self.committed_budget, + } + + def policy(self, trigger: BaseException | None = None) -> Mapping[str, object]: + return self.core.stream_policy(self.request, self.origin(), trigger) + + def can_retry(self) -> bool: + return self.policy().get("can_retry") is True + + def can_fall_back(self) -> bool: + return self.policy().get("can_fall_back") is True + + def refusal_fallback_available(self) -> bool: + return self.policy().get("refusal_fallback") is True + + def retry_budget(self, trigger: BaseException) -> tuple[int, bool]: + policy: Final = self.policy(trigger) + budget: Final = policy.get("retry_budget") + return (budget if isinstance(budget, int) else 0), policy.get("policy_applies") is True + + def should_retry_status(self, status_code: int) -> bool: + return bool(self.core._host.invoke("should_retry")(status_code)) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the wrapper serves its own core's host + + async def fallback(self, trigger: BaseException, request: Mapping[str, object] | None = None) -> object: + """The fallback walk for this stream's failure; raises when no fallback answers.""" + return await self.core.recover_stream("fallback", self.op, request or self.request, self.origin(), trigger) + + def fallback_sync(self, trigger: BaseException) -> object: + return self.core.recover_stream_sync("fallback", self.op, self.request, self.origin(), trigger) + + async def retry(self, trigger: BaseException) -> "tuple[object, int, int] | BaseException": + """A same-group retry: the retried response with the retry count and budget, or the newest failure.""" + return await self.core.retry_stream(self.op, self.request, self.origin(), trigger) + + # -- the stream's outcome, reported once when it ends -------------------------------------- + + def _logged(self) -> dict[str, object]: # mutable-ok: litellm's logging kwargs + """The logging kwargs litellm's stream logging hands the outcome callbacks, for this attempt.""" + from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY + + kwargs: Final = self.kwargs + metadata: Final = as_str_mapping(kwargs.get("litellm_metadata")) or as_str_mapping(kwargs.get("metadata")) or {} + return { + "call_type": self.op, + "litellm_params": { + "metadata": {**metadata, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: self.counted_tokens}, + "model_info": kwargs.get("model_info") or {}, + "api_base": kwargs.get("api_base"), + **({"cooldown_time": kwargs["cooldown_time"]} if "cooldown_time" in kwargs else {}), + }, + } + + def _claim(self) -> bool: + if not self.reports_outcome or self.outcome_reported: + return False + self.outcome_reported = True + return True + + def report_success(self, total_tokens: int) -> None: + """The stream ended: count its success and the tokens the attempt did not count + (``deployment_callback_on_success``), in the background.""" + if self._claim(): + logged: Final = {**self._logged(), "standard_logging_object": {"total_tokens": total_tokens}} + _record_in_background(self.core.deployment_callback_on_success(logged, None, None, None)) + + def report_failure(self, error: BaseException) -> None: + """The stream failed after it opened: count the failure, cool down by the rules and count one + more request (``async_deployment_callback_on_failure``), in the background.""" + if self._claim(): + logged: Final = {**self._logged(), "exception": _provider_error(error)} + _record_in_background(self.core.report_stream_failure(logged)) + + def report_success_sync(self) -> None: + """A sync stream ended: a sync call's success callback counts the success alone.""" + if self._claim(): + _record_now(lambda: self.core.sync_deployment_callback_on_success(self._logged(), None, None, None)) + + def report_failure_sync(self, error: BaseException) -> None: + """A sync stream failed after it opened: counted and cooled down, with no request counted.""" + if self._claim(): + logged: Final = {**self._logged(), "exception": _provider_error(error)} + _record_now(lambda: self.core.report_stream_failure_sync(logged)) + + +# --- reporting a stream's outcome --------------------------------------------------------------- + +# The background recordings still running, held so none is collected before it finishes. +_RECORDINGS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: tasks leave it when they finish + + +def _provider_error(error: BaseException) -> BaseException: + """The provider's own error, which the failure rules classify, rather than the router's envelope.""" + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None: + return error.original_exception + return error + + +def _record_in_background(recording: Coroutine[object, object, object]) -> None: + """Run a recording through the core's async cache driver without holding up the stream, as litellm + schedules its stream logging.""" + + async def run() -> None: + try: + _ = await recording + except Exception as error: # noqa: BLE001 # a recording never fails the stream + verbose_router_logger.warning("a stream's outcome was not recorded: %s", error) + + task: Final = asyncio.get_running_loop().create_task(run()) + _RECORDINGS.add(task) + task.add_done_callback(_RECORDINGS.discard) + + +def _record_now(recording: Callable[[], object]) -> None: + try: + _ = recording() + except Exception as error: # noqa: BLE001 # see _record_in_background + verbose_router_logger.warning("a stream's outcome was not recorded: %s", error) + + +_ItemT: Final = TypeVar("_ItemT") + + +async def _areported( + attempt: StreamAttempt, + items: AsyncIterable[_ItemT], + total_tokens: Callable[[object], int], + observe: Callable[[object], None] | None = None, +) -> AsyncGenerator[_ItemT, None]: + """``items``, reporting the attempt's stream once they end: a success with the tokens + ``total_tokens`` reads off the last item, or the failure they raise. Closed early, nothing is + reported. The wrapper marks the attempt (``reports_outcome``) when it wraps the stream, before + the attempt answers.""" + last: object = None # rebind-ok: the newest item + try: + async for item in items: + last = item + if observe is not None: + observe(item) + yield item + except Exception as error: + attempt.report_failure(error) + raise + attempt.report_success(total_tokens(last)) + + +def _reported(attempt: StreamAttempt, items: Iterable[_ItemT]) -> Generator[_ItemT, None, None]: + """``_areported`` for a sync stream, whose success counts no tokens.""" + try: + yield from items + except Exception as error: + attempt.report_failure_sync(error) + raise + attempt.report_success_sync() + + +def _usage_total_tokens(usage: object) -> int: + mapping: Final = as_str_mapping(usage) + tokens: Final[object] = mapping.get("total_tokens") if mapping is not None else getattr(usage, "total_tokens", None) + return tokens if isinstance(tokens, int) and not isinstance(tokens, bool) and tokens > 0 else 0 + + +def _chat_total_tokens(last: object) -> int: + """The total tokens of a finished chat stream: its usage chunk's, else the final usage its last + chunk's hidden params carry once the stream assembled its response.""" + return _usage_total_tokens(getattr(last, "usage", None)) or _usage_total_tokens( + (get_hidden_params(last) or {}).get("usage") + ) + + +def _responses_total_tokens(last: object, source: object) -> int: + """The total tokens of a finished Responses stream: its terminal event's response usage.""" + for event in (last, getattr(source, "completed_response", None)): + tokens = _usage_total_tokens(getattr(getattr(event, "response", None), "usage", None)) + if tokens: + return tokens + return 0 + + +class _AnthropicStreamUsage: + """The usage a ``/v1/messages`` stream reports in its frames, as litellm logs it: the prompt's input, + cache-read and cache-write tokens on ``message_start`` (a ``message_delta`` may restate them), and the + cumulative output tokens on each ``message_delta``. A frame split across byte chunks is read once its + line completes.""" + + __slots__ = ("_open_line", "_output", "_prompt") + + def __init__(self) -> None: + self._open_line = b"" + self._prompt: dict[str, int] = {} # mutable-ok: the newest count per prompt field + self._output = 0 + + def observe(self, chunk: object) -> None: + if not isinstance(chunk, (bytes, bytearray)): + return + *lines, self._open_line = (self._open_line + bytes(chunk)).split(b"\n") + for line in lines: + if line.startswith(b"data:") and b'"usage"' in line: + self._read(line[len(b"data:") :]) + + def _read(self, payload: bytes) -> None: + try: + event: Final[object] = json.loads(payload) # pyright: ignore[reportAny] # plain JSON + except ValueError: + return + frame: Final = as_str_mapping(event) or {} + kind: Final = frame.get("type") + usage: Final = as_str_mapping( + (as_str_mapping(frame.get("message")) or {}).get("usage") if kind == "message_start" else frame.get("usage") + ) + if usage is None or kind not in ("message_start", "message_delta"): + return + for field in ("input_tokens", "cache_creation_input_tokens", "cache_read_input_tokens"): + count = usage.get(field) + if isinstance(count, int) and not isinstance(count, bool): + self._prompt[field] = count + output: Final = usage.get("output_tokens") + if isinstance(output, int) and not isinstance(output, bool): + self._output = output + + def total_tokens(self, _last: object = None) -> int: + return sum(self._prompt.values()) + self._output + + +# --- the first chunk ---------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True, slots=True) +class _FirstChunk(Generic[_RestT]): + """The chunk an attempt already read off a stream (None: the provider closed it empty) and the + iterator it was read from, which carries on after it.""" + + chunk: object + rest: _RestT + + +def _mid_stream_retry_trigger(error: MidStreamFallbackError) -> Exception: + """The provider's own error, which is what the retry policy and the retry rules classify.""" + return error.original_exception if error.original_exception is not None else error + + +async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[AsyncIterator[object]]: + """Read the first chunk while the attempt still holds the request, so a stream that fails before + any token fails the attempt. Iterates the way ``async for`` does: one ``__aiter__`` call, then + ``__anext__`` on what it returned.""" + rest: Final = stream.__aiter__() + try: + return _FirstChunk(chunk=await rest.__anext__(), rest=rest) + except StopAsyncIteration: + return _FirstChunk(chunk=None, rest=rest) + except MidStreamFallbackError as e: + raise _mid_stream_retry_trigger(e) from e + + +def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[Iterator[object]]: + rest: Final = iter(stream) + try: + return _FirstChunk(chunk=next(rest), rest=rest) + except StopIteration: + return _FirstChunk(chunk=None, rest=rest) + except MidStreamFallbackError as e: + raise _mid_stream_retry_trigger(e) from e + + +def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk[Iterator[object]]) -> Iterator[object]: + if taken.chunk is None: + return + yield taken.chunk + yield from taken.rest + + +async def _achunks_with_first( + model_response: AsyncIterable[object], taken: _FirstChunk[AsyncIterator[object]] +) -> AsyncIterator[object]: + if taken.chunk is None: + return + yield taken.chunk + async for item in taken.rest: + yield item + + +# --- shared helpers ----------------------------------------------------------------------------- + + +_CONTENT_DELTA_KEYS: Final = ( + "content", + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", +) + + +def _delta_has_content(delta: object) -> bool: + get: Final[object] = getattr(delta, "get", None) + if not callable(get): + return False + for key in _CONTENT_DELTA_KEYS: + value: object = get(key) + if value: + return True + return False + + +def _stream_chunks_have_generated_content(chunks: Sequence[object]) -> bool: + for chunk in chunks: + choices: object = getattr(chunk, "choices", None) + if not isinstance(choices, Sequence) or not choices: + continue + delta: object = getattr(choices[0], "delta", None) # pyright: ignore[reportUnknownArgumentType] # a chunk choice + if _delta_has_content(delta): + return True + return False + + +def _raise_original_mid_stream_error(e: MidStreamFallbackError) -> NoReturn: + """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" + if e.original_exception is not None: + raise e.original_exception from e + raise e + + +def _stream_pre_content_error(error: Exception, model: str, llm_provider: str = "anthropic") -> MidStreamFallbackError: + """The envelope the fallback chain judges a failure by when the client has received no content yet.""" + return MidStreamFallbackError( + message=str(error), + model=model, + llm_provider=llm_provider, + original_exception=error, + is_pre_first_chunk=True, + ) + + +def _context_window_stream_fallback_error( + error: "MidStreamFallbackError | litellm_exceptions.ContextWindowExceededError", + model: object, + llm_provider: str, + has_output: bool, +) -> MidStreamFallbackError: + if isinstance(error, litellm_exceptions.ContextWindowExceededError): + if has_output: + raise error + return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) + if isinstance(error.original_exception, litellm_exceptions.ContextWindowExceededError) and ( + has_output or error.generated_content or error.emitted_disqualifying_content + ): + raise error.original_exception + return error + + +def _is_context_window(error: "MidStreamFallbackError | litellm_exceptions.ContextWindowExceededError") -> bool: + return isinstance(error, litellm_exceptions.ContextWindowExceededError) or isinstance( + error.original_exception, litellm_exceptions.ContextWindowExceededError + ) + + +def _fallback_trigger(e: MidStreamFallbackError, unwrapped: tuple[type[Exception], ...]) -> Exception: + """The error the fallback walk dispatches on: the provider's own for the kinds with a list of their own.""" + return e.original_exception if isinstance(e.original_exception, unwrapped) else e + + +def _prepare_fallback_hidden_params( + fallback_response: object, +) -> _Prepared: + fallback_hidden_params: Final = get_hidden_params_dict(fallback_response) + fallback_headers: Final = fallback_hidden_params.get("additional_headers") + if not isinstance(fallback_headers, dict): + return fallback_hidden_params, {} + return fallback_hidden_params, cast("dict[str, object]", fallback_headers) # cast-ok: a dict of headers + + +def _apply_fallback_hidden_params_to_item( + fallback_item: object, + prepared_fallback_hidden_params: _Prepared, +) -> None: + if fallback_item is None or not hasattr(fallback_item, HIDDEN_PARAMS_ATTR): + return + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + item_hidden_params: Final = get_hidden_params_dict(fallback_item) + item_headers: object = item_hidden_params.get("additional_headers") # rebind-ok: replaced when not a dict + if not isinstance(item_headers, dict): + item_headers = {} # rebind-ok: rebound as the stream advances + set_hidden_params( + fallback_item, + { + **item_hidden_params, + **fallback_hidden_params, + "additional_headers": {**cast("dict[str, object]", item_headers), **fallback_headers}, # cast-ok: checked + }, + ) + + +def _combine_fallback_usage(fallback_item: ModelResponseStream, complete_response_object_usage: Usage | None) -> None: + """Merge partial-stream usage with fallback-stream usage on the chunk. A usage given as a plain + dict comes out all-zero: ``combine_usage_objects`` reads token counts as attributes.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + + usage: Final = cast("Usage | None", getattr(fallback_item, "usage", None)) # cast-ok: a chunk's usage + usage_objects: Final = [usage] if usage is not None else [] + if complete_response_object_usage is not None: + usage_objects.append(complete_response_object_usage) + fallback_item.usage = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) # pyright: ignore[reportAttributeAccessIssue] # usage is not a declared chunk field # rebind-ok: the Python router rewrites the chunk's usage in place + + +# --- mid-stream continuation -------------------------------------------------------------------- + +CONTINUATION_FLAG: Final = "enable_mid_stream_fallback_continuation" +NO_FLAG: Final = object() +# The routing-body key naming the deployments a continuation may not land on (``continuation_deployments``). +CONTINUATION_EXCLUDED_IDS: Final = "continuation_excluded_ids" + + +def continuation_flag(value: object) -> bool: + """``_is_flag_enabled``: True, or a string ``str_to_bool`` reads as True.""" + from litellm.secret_managers.main import str_to_bool + + return value is True or (isinstance(value, str) and str_to_bool(value) is True) + + +def register_continuation_prefill_check(host: "PythonHost") -> None: + """``_register_continuation_prefill_check``: process-global on purpose, as in the Python router.""" + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ContinuationPrefillDeploymentCheck + + host.callbacks.add_litellm_callback(ContinuationPrefillDeploymentCheck()) + + +def continuation_routing(core: "RustCore", request: Mapping[str, object]) -> Mapping[str, object]: + """The routing fields of a continuation request: the deployments that cannot continue it, which the + core leaves out of selection the way ``continuation_deployments`` filters the healthy list.""" + from litellm.litellm_core_utils.routing_common_utils import MID_STREAM_CONTINUATION_KWARG + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_MARKER, + deployment_can_continue, + ) + + if request.get(MID_STREAM_CONTINUATION_KWARG) is not MID_STREAM_CONTINUATION_MARKER: + return {} + registry: Final = core.registry() + excluded: Final = [ + deployment.model_info.id + for record in registry.all() + if (deployment := registry.deployment(str(_record_id(record)))) is not None + and not deployment_can_continue(deployment.model_dump(exclude_none=True)) + ] + return {"body": {CONTINUATION_EXCLUDED_IDS: excluded}} + + +def _record_id(record: Mapping[str, object]) -> object: + model_info: Final = record.get("model_info") + return model_info.get("id") if isinstance(model_info, Mapping) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a plain record + + +def _names_deployment_that_cannot_continue(core: "RustCore", target: object) -> bool: + from litellm.router_utils.fallback_event_handlers import fallback_target_name + from litellm.router_utils.pre_call_checks.continuation_prefill_check import deployment_can_continue + + name: Final = fallback_target_name(target) + registry: Final = core.registry() + model_names: Final = {record.get("model_name") for record in registry.all()} + if name is None or name in model_names or not registry.has_deployment(name): + return False + deployment: Final = registry.deployment(name) + return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) + + +def _mid_stream_continuation_eligible( + attempt: StreamAttempt, e: MidStreamFallbackError, cut_deployment_id: object +) -> bool: + """``_mid_stream_continuation_eligible``: whether a stream that broke after plain assistant text may + continue via prefill.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.router_utils.fallback_event_handlers import ( + continuation_fallback_targets, + fallback_lookup_groups, + fallback_target_overrides_request_params, + ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import constrains_continuation + from litellm.types.guardrails import GuardrailEventHooks + + core: Final = attempt.core + if not core.enable_mid_stream_fallback_continuation: + return False + if not e.generated_content or e.emitted_disqualifying_content: + return False + if isinstance(e.original_exception, litellm_exceptions.ContentPolicyViolationError): + return False + if isinstance(cut_deployment_id, str) and not core.registry().has_deployment(cut_deployment_id): + return False + if constrains_continuation(_params_sent_to_deployment(attempt, cut_deployment_id)): + return False + request: Final = attempt.request + requested_group: Final = request.get("model") + lookup_groups: Final = fallback_lookup_groups( + request, requested_group if isinstance(requested_group, str) else None + ) + targets: Final = continuation_fallback_targets(core.settings().fallbacks, lookup_groups) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # fallbacks as Python spells them + if any(fallback_target_overrides_request_params(target) for target in targets): + return False + if any(_names_deployment_that_cannot_continue(core, target) for target in targets): + return False + return not any( + isinstance(callback, CustomGuardrail) + and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks + GuardrailEventHooks.post_call + ) + for callback in core._host.callback_list("callbacks") # pyright: ignore[reportPrivateUsage, reportArgumentType] # the core's own host; read as getattr(litellm, name) + ) + + +def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: + """Continue the last text block in place so the prefill reads as the text the caller already received.""" + last: Final = blocks[-1] if blocks else None + if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): + return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] + return [*blocks, {"type": "text", "text": text}] + + +def _build_completion_continuation_input( + messages: Sequence[Mapping[str, object]], generated_content: str +) -> Sequence[Mapping[str, object]] | None: + """The partial output as an assistant prefill, or a trailing ``prefix: True`` assistant turn extended in + place. None when the request ends in a plain assistant turn, in a prefill turn whose content this cannot + extend, or when the cut came mid-echo of the prefill.""" + from pydantic import TypeAdapter, ValidationError + + from litellm.router_utils.pre_call_checks.continuation_prefill_check import continued_text + + last: Final = messages[-1] if messages else None + if last is None or last.get("role") != "assistant": + return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] + if last.get("prefix") is not True: + return None + new_text: Final = continued_text(last, generated_content) + if new_text is None: + return None + content: Final = last.get("content") + if content is None or isinstance(content, str): + return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}] + try: + blocks: Final = TypeAdapter(list[dict[str, object]]).validate_python(content) + except ValidationError: + return None + return [*messages[:-1], {**last, "content": _extend_trailing_text_block(blocks, new_text), "prefix": True}] + + +def _price_cut_partial(cut_stream: CustomStreamWrapper, partial_response: object) -> "ContinuedPartial | None": + """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ContinuedPartial + + cut_logging: Final = cut_stream.logging_obj + partial_usage: Final = getattr(partial_response, "usage", None) + if partial_response is None or not isinstance(partial_usage, Usage): + return None + details: Final = cut_logging.model_call_details + inherited_usage: Final = details.get("continued_partial_usage") + inherited_cost: Final = details.get("continued_partial_cost") + inherited_id: Final = details.get("continued_response_id") + priced_partial: Final = partial_response.model_copy( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # a ModelResponse + update={"model": cut_stream.model} if cut_stream.model else {} # pyright: ignore[reportUnknownMemberType] # untyped attribute + ) + priced_partial._hidden_params = { # pyright: ignore[reportUnknownMemberType] # no public accessor + key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" + } + if LiteLLMLogging.usage_capture_callback is not None: + from litellm._internal_context import internal_sub_call + + with internal_sub_call(cut_logging.litellm_call_id): + cut_logging.capture_usage_before_logging(priced_partial) # pyright: ignore[reportUnknownArgumentType] # a ModelResponse + return ContinuedPartial( + usage=BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[partial_usage, inherited_usage] if isinstance(inherited_usage, Usage) else [partial_usage] + ), + response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0) # pyright: ignore[reportUnknownArgumentType] # a ModelResponse + + (inherited_cost if isinstance(inherited_cost, float) else 0.0), + response_id=inherited_id if isinstance(inherited_id, str) else str(getattr(partial_response, "id", "")), + ) + + +def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None: + """Bill what the cut legs consumed on the continuation's success log.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + fallback_logging: Final = getattr(fallback_response, "logging_obj", None) + if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging): + return + fallback_logging.record_continued_partial( + usage=cut_partial.usage, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # a Usage + response_cost=cut_partial.response_cost, + response_id=cut_partial.response_id, + ) + + +@dataclasses.dataclass(frozen=True, slots=True) +class _Continuation: + """What a chat continuation sends the fallback and splices back.""" + + request: Mapping[str, object] + output: "ContinuationOutput" + partial: "ContinuedPartial | None" + usage: Usage | None + + +async def _chat_continuation( + attempt: StreamAttempt, e: MidStreamFallbackError, model_response: CustomStreamWrapper +) -> _Continuation: + """The request a fallback continues a chat stream from, once it broke after content. Raises the + provider's error when the continuation is declined or the output ceiling is spent.""" + from litellm.constants import MID_STREAM_FAILURE_LOG_WAIT_SECONDS + from litellm.litellm_core_utils.routing_common_utils import MID_STREAM_CONTINUATION_KWARG + from litellm.main import ( + stream_chunk_builder, # pyright: ignore[reportUnknownVariableType] # untyped chunk list argument + ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutput, + continuation_output_ceilings, + ) + + request: Final = attempt.request + raw_messages: Final = request.get("messages") + messages: Final = cast("list[Mapping[str, object]]", raw_messages or []) # cast-ok: chat messages + complete: Final = stream_chunk_builder(chunks=model_response.chunks, messages=messages) # pyright: ignore[reportUnknownArgumentType] # provider chunk list + usage: Final = cast("Usage | None", getattr(complete, "usage", None)) # cast-ok: a built response's usage + continuation_messages: Final = _build_completion_continuation_input(messages, e.generated_content) + output: Final = ( + ContinuationOutput( + e.generated_content, + continuation_messages[-1], + str(request.get("model") or ""), + complete, + checks_post_call_rules=request.get(MID_STREAM_CONTINUATION_KWARG) is not MID_STREAM_CONTINUATION_MARKER, + ) + if continuation_messages is not None + else None + ) + partial: Final = _price_cut_partial(model_response, complete) if output is not None else None + emitted_tokens: Final = int(getattr(usage, "completion_tokens", 0) or 0) + reduced: Final = continuation_output_ceilings(request, emitted_tokens) + if reduced is None or continuation_messages is None or output is None: + _raise_original_mid_stream_error(e) + await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS) + return _Continuation( + request={ + **request, + **reduced, + "messages": continuation_messages, + MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER, + }, + output=output, + partial=partial, + usage=usage, + ) + + +# --- chat completions --------------------------------------------------------------------------- + + +class FallbackAwareStreamWrapper(CustomStreamWrapper): + """A chat stream built around the attempt the core picked first, which repoints itself when a + fallback takes over.""" + + fallback_headers_adopted: bool = False + + def adopt_fallback_response_headers( + self, + fallback_response: object, + prepared_fallback_hidden_params: _Prepared, + ) -> None: + """Repoint this wrapper at the deployment that served the stream. Replaces rather than merges, + so the failed attempt's request id, counters, ``model_id`` and ``api_base`` cannot leak.""" + self._response_headers = getattr(fallback_response, "_response_headers", None) + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + if fallback_hidden_params: + self._hidden_params = {**fallback_hidden_params, "additional_headers": dict(fallback_headers)} + self._base_hidden_params = {**self._hidden_params, "response_cost": None} + self.fallback_headers_adopted = True + + +def _adopt_fallback_response_headers( + wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", fallback_response: object +) -> _Prepared: + prepared: Final = _prepare_fallback_hidden_params(fallback_response) + adopting_wrapper: Final = wrapper_ref() + if adopting_wrapper is not None: + adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) + return prepared + + +def _chat_wrapper_init( + wrapper: FallbackAwareStreamWrapper, stream: object, model_response: CustomStreamWrapper +) -> None: + CustomStreamWrapper.__init__( # pyright: ignore[reportUnknownMemberType] # the base constructor is untyped + wrapper, + completion_stream=stream, + model=model_response.model, # pyright: ignore[reportUnknownMemberType] # an untyped attribute + custom_llm_provider=model_response.custom_llm_provider, + logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), + ) + inner_chunks: Final[object] = getattr(model_response, "chunks", None) + if isinstance(inner_chunks, list): + chunk_list: Final = cast("list[ModelResponseStream]", inner_chunks) # cast-ok: the inner stream's chunks + wrapper.chunks = chunk_list # rebind-ok: initializes the wrapper being built + hidden_params: Final = get_hidden_params(model_response) + if hidden_params is not None: + wrapper._hidden_params = dict(hidden_params) # pyright: ignore[reportPrivateUsage] # the wrapper's own field # rebind-ok: initializes the wrapper being built + + +class _AsyncChatStream(FallbackAwareStreamWrapper): + def __init__(self, generator: AsyncGenerator[object, None], model_response: CustomStreamWrapper) -> None: # pyright: ignore[reportMissingSuperCall] # _chat_wrapper_init calls it + _chat_wrapper_init(self, generator, model_response) + self._async_generator: Final = generator + + def __aiter__(self) -> "_AsyncChatStream": # pyright: ignore[reportIncompatibleMethodOverride] # iterates the generator + return self + + async def __anext__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields what the stream yields + return await self._async_generator.__anext__() + + +class _SyncChatStream(FallbackAwareStreamWrapper): + def __init__(self, generator: Generator[object, None, None], model_response: CustomStreamWrapper) -> None: # pyright: ignore[reportMissingSuperCall] # _chat_wrapper_init calls it + _chat_wrapper_init(self, generator, model_response) + self.chunks = [] # the sync twin keeps no inner chunk list + hidden_params: Final = get_hidden_params(model_response) + if hidden_params is not None: + self._hidden_params = dict(hidden_params) + self._sync_generator: Final = generator + + def __iter__(self) -> "_SyncChatStream": # pyright: ignore[reportIncompatibleMethodOverride] # iterates the generator + return self + + def __next__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields what the stream yields + return next(self._sync_generator) + + +def _chat_has_output(model_response: CustomStreamWrapper) -> bool: + return bool(getattr(model_response, "response_uptil_now", "")) or _stream_chunks_have_generated_content( + model_response.chunks + ) + + +def wrap_chat_stream( # noqa: C901 # one function per Python wrapper it ports + attempt: StreamAttempt, + model_response: CustomStreamWrapper, + first_chunk: _FirstChunk[AsyncIterator[object]], + held_slot: contextlib.ExitStack | None = None, +) -> FallbackAwareStreamWrapper: + """``_acompletion_streaming_iterator``: a mid-stream error before content falls back through the core; + after content, the provider's error is raised. ``held_slot`` (the deployment's max_parallel_requests + slot) is released when the stream is exhausted, closed, or falls back.""" + slot: Final = held_slot if held_slot is not None else contextlib.ExitStack() + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + + async def close_model_response() -> None: + if not hasattr(model_response, "aclose"): + return + try: + await model_response.aclose() + except BaseException as e: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks: error closing model_response: %s", e) + + async def stream_with_fallbacks() -> AsyncGenerator[object, None]: + fallback_response: object = None # rebind-ok: tracked for cleanup + try: + async for item in _areported(attempt, _achunks_with_first(model_response, first_chunk), _chat_total_tokens): + yield item + except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + _chat_has_output(model_response), + ) + if _is_context_window(stream_error) + else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error + ) + with anyio.CancelScope(shield=True): + await close_model_response() + slot.close() + committed: Final = bool( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ) + cut_deployment_id: Final = (get_hidden_params(model_response) or {}).get("model_id") + if committed and not _mid_stream_continuation_eligible(attempt, e, cut_deployment_id): + _raise_original_mid_stream_error(e) + trigger: Final = _fallback_trigger(e, (litellm_exceptions.ContextWindowExceededError,)) + try: + continuation: Final = await _chat_continuation(attempt, e, model_response) if committed else None + fallback_response = await attempt.fallback( # rebind-ok: rebound as the stream advances + trigger, continuation.request if continuation is not None else None + ) + if hasattr(fallback_response, "__aiter__"): + if continuation is not None: + _carry_partial_spend_to_fallback(fallback_response, continuation.partial) + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) # rebind-ok: re-adopted + settled = False # rebind-ok: rebound as the stream advances + fallback_stream: Final = cast("AsyncIterable[object]", fallback_response) # cast-ok: checked + fallback_items: Final = ( + continuation.output.splice(fallback_stream) if continuation is not None else fallback_stream + ) + async for fallback_item in fallback_items: + if not settled: + settled = True + # a fallback that failed over again only repoints itself once it yields + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) + _apply_fallback_hidden_params_to_item(fallback_item, prepared) + if ( + fallback_item + and isinstance(fallback_item, ModelResponseStream) + and hasattr(fallback_item, "usage") + ): + _combine_fallback_usage( + fallback_item, + continuation.usage if continuation is not None else None, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # a Usage + ) + yield fallback_item + else: + # A fallback that answers without a stream yields None (an oddity kept as is). + yield None + except Exception as fallback_error: + verbose_router_logger.error("Fallback also failed: %s", fallback_error) + if isinstance(trigger, MidStreamFallbackError) and trigger.original_exception is not None: + raise trigger.original_exception from fallback_error + raise trigger from fallback_error + finally: + with anyio.CancelScope(shield=True): + await close_model_response() + slot.close() + if fallback_response is not None and hasattr(fallback_response, "aclose"): + try: + await fallback_response.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + except BaseException as close_error: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug( + "stream_with_fallbacks: error closing fallback_response: %s", close_error + ) + + wrapped_response: Final = _AsyncChatStream(stream_with_fallbacks(), model_response) + # weak, so the generator closing over it does not keep the wrapper alive + wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) + return wrapped_response + + +def wrap_chat_stream_sync( + attempt: StreamAttempt, model_response: CustomStreamWrapper, first_chunk: _FirstChunk[Iterator[object]] +) -> FallbackAwareStreamWrapper: + """``_completion_streaming_iterator``: the sync twin, which never continues after content.""" + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + + def stream_with_fallbacks() -> Generator[object, None, None]: + fallback_response: object = None # rebind-ok: tracked for cleanup + try: + yield from _reported(attempt, _chunks_with_first(model_response, first_chunk)) + except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + _chat_has_output(model_response), + ) + if _is_context_window(stream_error) + else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error + ) + if attempt.request.get("disable_fallbacks") is True or ( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ): + _raise_original_mid_stream_error(e) + trigger: Final = _fallback_trigger(e, (litellm_exceptions.ContextWindowExceededError,)) + try: + fallback_response = attempt.fallback_sync(trigger) # rebind-ok: rebound as the stream advances + if hasattr(fallback_response, "__iter__"): + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) # rebind-ok: re-adopted + settled = False # rebind-ok: rebound as the stream advances + for fallback_item in cast("Iterable[object]", fallback_response): # cast-ok: checked + if not settled: + settled = True + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) + _apply_fallback_hidden_params_to_item(fallback_item, prepared) + if ( + fallback_item + and isinstance(fallback_item, ModelResponseStream) + and hasattr(fallback_item, "usage") + ): + _combine_fallback_usage(fallback_item, None) + yield fallback_item + else: + yield None + except Exception as fallback_error: + verbose_router_logger.error("Fallback also failed: %s", fallback_error) + if isinstance(trigger, MidStreamFallbackError) and trigger.original_exception is not None: + raise trigger.original_exception from fallback_error + raise trigger from fallback_error + finally: + for closable in (model_response, fallback_response): + close = getattr(closable, "close", None) + if callable(close): + try: + close() + except BaseException as close_error: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks: error closing: %s", close_error) + + wrapped_response: Final = _SyncChatStream(stream_with_fallbacks(), model_response) + wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) + return wrapped_response + + +# --- /v1/messages ------------------------------------------------------------------------------- + + +class FallbackAwareAnthropicMessagesStream: + """A ``/v1/messages`` byte stream that carries the ``_hidden_params`` the proxy reads its response + headers off, through retries and fallbacks of the stream it wraps.""" + + def __init__( + self, async_generator: AsyncGenerator[bytes, None], source_iterator: object, attempt: StreamAttempt + ) -> None: + self._async_generator: Final = async_generator + self._source_iterator: object = source_iterator + self.attempt: Final = attempt + self.fallback_headers_adopted = False + self._hidden_params: dict[str, object] = dict( # mutable-ok: the proxy writes headers into it + getattr(source_iterator, "_hidden_params", None) or {} + ) + self._followed_source_params: object = None + + @property + def has_buffered_provider_output(self) -> bool: + return getattr(self._source_iterator, "has_buffered_provider_output", False) is True + + @property + def chunks(self) -> object: + return getattr(self._source_iterator, "chunks", None) + + @property + def messages(self) -> object: + return getattr(self._source_iterator, "messages", None) + + @property + def model(self) -> object: + return getattr(self._source_iterator, "model", None) + + def adopt_fallback_source(self, fallback_response: object) -> None: + self._source_iterator = fallback_response + self.fallback_headers_adopted = True + + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect + + return partial_responses_incomplete_event_on_disconnect(self._source_iterator) + + def follow_source_attribution(self) -> None: + """Mirror the source's current attribution onto this wrapper whenever the source has adopted a + new one: a retry's or a fallback's stream carries a wrapper of its own.""" + source: Final = self._source_iterator + if not getattr(source, "fallback_headers_adopted", False): + return + source_params: Final = getattr(source, "_hidden_params", None) + if source_params is None or source_params is self._followed_source_params: + return + self._followed_source_params = source_params + hidden_params, headers = _prepare_fallback_hidden_params(source) + self.merge_fallback_hidden_params(hidden_params, headers) + + def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": + return self + + async def __anext__(self) -> bytes: + return await self._async_generator.__anext__() + + async def aclose(self) -> None: + await self._async_generator.aclose() + + def merge_fallback_hidden_params( + self, fallback_hidden_params: Mapping[str, object], fallback_headers: Mapping[str, object] + ) -> None: + existing_headers: Final = cast( # cast-ok: additional_headers is a dict when present + "dict[str, object]", self._hidden_params.get("additional_headers") or {} + ) + self._hidden_params = { + **self._hidden_params, + **fallback_hidden_params, + "additional_headers": dict(replace_complexity_router_headers(existing_headers, fallback_headers)), + } + + +def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool) -> bool: + from litellm.llms.anthropic.messages.native.streaming_iterator import is_anthropic_ping_chunk + + return not has_generated_content and is_anthropic_ping_chunk(chunk) + + +def _is_retriable_anthropic_status(status_code: int) -> bool: + return status_code == 429 or status_code >= 500 + + +def _anthropic_stream_error_is_gateway_verdict(chunk: object) -> bool: + from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import is_server_fulfilled_tool_leak_error + + return is_server_fulfilled_tool_leak_error(chunk) + + +def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: MidStreamFallbackError) -> bool: + return has_generated_content or not error.is_pre_first_chunk + + +def _stream_raised_error_status(error: Exception) -> int | None: + raw_status: Final = getattr(error, "status_code", None) + if isinstance(raw_status, int): + return raw_status + if isinstance(raw_status, str) and raw_status.isdigit(): + return int(raw_status) + response_status: Final = getattr(getattr(error, "response", None), "status_code", None) + return response_status if isinstance(response_status, int) else None + + +def _stream_fallback_error_for_raised( + attempt: StreamAttempt, error: Exception, model: str, has_generated_content: bool, llm_provider: str = "anthropic" +) -> MidStreamFallbackError | None: + """The pre-stream retry rule (408, 409, 429, 5xx); None means the raise propagates unchanged.""" + if has_generated_content: + return None + status_code: Final = _stream_raised_error_status(error) + if status_code is None or attempt.should_retry_status(status_code): + return _stream_pre_content_error(error, model, llm_provider) + return None + + +def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: + from litellm.llms.anthropic.messages.native.streaming_iterator import is_anthropic_content_delta_chunk + + if has_generated_content: + return False + return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + + +def _anthropic_recoverable_frame_error( + attempt: StreamAttempt, + error_event: tuple[str, str, int] | None, + chunk: object, + has_generated_content: bool, + model_group: str, +) -> Exception | None: + """The exception a provider ``event: error`` frame before content recovers through when a retry of + its class or a fallback can still take over; None forwards the frame as the provider sent it.""" + from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception + + if has_generated_content or error_event is None: + return None + error_type, message, status_code = error_event + if not _is_retriable_anthropic_status(status_code) or _anthropic_stream_error_is_gateway_verdict(chunk): + return None + frame_error: Final = anthropic_error_frame_exception(error_type, message, status_code, model_group) + budget, _ = attempt.retry_budget(frame_error) + if budget > attempt.attempted_retries: + return frame_error + if attempt.can_fall_back(): + return frame_error + return None + + +def wrap_anthropic_stream( + attempt: StreamAttempt, response: AsyncIterator[bytes] +) -> FallbackAwareAnthropicMessagesStream: + """``_aanthropic_messages_streaming_iterator``: lifecycle frames are held while a retry or a fallback + can still take over; an error frame, a raised error or a safeguard refusal before content is + retried in the group, then falls back, as the core decides.""" + from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + aclose_if_supported, + aclose_superseded_attempt, + parse_anthropic_error_event, + parse_anthropic_refusal_stop_details, + ) + from litellm.llms.anthropic.messages.native.utils import safeguard_refusal_error + + source_iterator: Final = response + replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) + model: Final = attempt.group + usage: Final = _AnthropicStreamUsage() + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + + async def close_source() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(source_iterator) + + async def close_superseded() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_superseded_attempt(source_iterator) + + async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: + policy: Final = attempt.policy() + has_generated_content = not ( # rebind-ok: set once content is seen, the buffer cap is hit, or nothing can take over + policy.get("can_retry") is True or policy.get("can_fall_back") is True + ) + buffered: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + try: + async for chunk in _areported(attempt, source_iterator, usage.total_tokens, usage.observe): + if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): + yield chunk + continue + if _stream_commits_now(chunk, has_generated_content, len(buffered)): + has_generated_content = True + # A transport can split one SSE data line across byte chunks, so pre-content detection + # parses the held frames plus the current chunk. + parse_window = ( + b"".join(c for c in (*buffered, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes + if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # see above + else chunk + ) + error_event = parse_anthropic_error_event(parse_window) + recoverable_frame_error = _anthropic_recoverable_frame_error( + attempt, error_event, chunk, has_generated_content, model + ) + refusal_stop_details = ( + parse_anthropic_refusal_stop_details(parse_window) + if not has_generated_content and error_event is None + else None + ) + if ( + refusal_stop_details is not None + and not replays_a_finished_agentic_turn + and attempt.refusal_fallback_available() + ): + refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) + raise MidStreamFallbackError( + message=refusal_error.message, + model=model, + llm_provider="anthropic", + original_exception=refusal_error, + is_pre_first_chunk=True, + ) + if not has_generated_content and error_event is None: + buffered = (*buffered, chunk) + continue + if recoverable_frame_error is not None and error_event is not None: + raise MidStreamFallbackError( + message=error_event[1], + model=model, + llm_provider="anthropic", + original_exception=recoverable_frame_error, + is_pre_first_chunk=True, + ) + for buffered_chunk in buffered: + yield buffered_chunk + buffered = () + yield chunk + for buffered_chunk in buffered: + yield buffered_chunk + except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate + async for item in _anthropic_recover_stream_error( + attempt, stream_error, has_generated_content, buffered, model, wrapper, close_superseded + ): + yield item + finally: + await close_source() + + wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator, attempt) + return wrapper + + +async def _anthropic_recover_stream_error( + attempt: StreamAttempt, + stream_error: Exception, + has_generated_content: bool, + buffered: tuple[bytes, ...], + model: str, + wrapper: FallbackAwareAnthropicMessagesStream, + close_superseded: Callable[[], Awaitable[None]], +) -> AsyncGenerator[bytes, None]: + """Turns a source-iterator failure into a retry, a fallback, or the error reaching the caller.""" + if isinstance(stream_error, MidStreamFallbackError) and _anthropic_stream_should_decline_fallback( + has_generated_content, stream_error + ): + for buffered_chunk in buffered: + yield buffered_chunk + if stream_error.original_exception is not None: + raise stream_error.original_exception from stream_error + raise stream_error + fallback_error: Final = ( + stream_error + if isinstance(stream_error, MidStreamFallbackError) + else _stream_fallback_error_for_raised(attempt, stream_error, model, has_generated_content) + ) + if fallback_error is None: + raise stream_error + # An error frame or a refusal the wrapper raised fails the attempt too (a raised error already did). + attempt.report_failure(fallback_error) + await close_superseded() + retried: Final = await attempt.retry(_mid_stream_retry_trigger(fallback_error)) + if isinstance(retried, BaseException): + last_error: Final = ( + fallback_error + if retried is _mid_stream_retry_trigger(fallback_error) + else _stream_pre_content_error( + retried if isinstance(retried, Exception) else fallback_error, model, fallback_error.llm_provider + ) + ) + async for item in _anthropic_fallback_attempt(attempt, last_error, wrapper): + yield item + return + response, attempted, budget = retried + inner: Final = response + if isinstance(inner, FallbackAwareAnthropicMessagesStream): + inner.attempt.attempted_retries = attempted + inner.attempt.committed_budget = budget + async for item in _anthropic_yield_recovered(inner, wrapper, (attempted, budget)): + yield item + + +async def _anthropic_yield_recovered( + recovered: object, + wrapper: FallbackAwareAnthropicMessagesStream, + retry_counters: tuple[int, int] | None = None, +) -> AsyncGenerator[bytes, None]: + """Hands a retry's or a fallback's response to the client through the wrapper, closing it after.""" + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + aclose_if_supported, + anthropic_messages_response_as_sse_events, + ) + + hidden_params, headers = _prepare_fallback_hidden_params(recovered) + wrapper.merge_fallback_hidden_params(hidden_params, headers) + wrapper.adopt_fallback_source(recovered) + if retry_counters is not None: + add_retry_headers_to_response( + response=wrapper, attempted_retries=retry_counters[0], max_retries=retry_counters[1] + ) + try: + if hasattr(recovered, "__aiter__"): + async for item in cast("AsyncIterator[bytes]", recovered): # cast-ok: __aiter__ checked above + wrapper.follow_source_attribution() + yield item + return + # A complete response for a streaming request is replayed as the SSE lifecycle a stream sends. + for event in anthropic_messages_response_as_sse_events(recovered): # pyright: ignore[reportArgumentType] # the non-streaming shape by elimination + yield event + finally: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(recovered) + + +async def _anthropic_fallback_attempt( + attempt: StreamAttempt, e: MidStreamFallbackError, wrapper: FallbackAwareAnthropicMessagesStream +) -> AsyncGenerator[bytes, None]: + """Re-enters the fallback walk for a ``/v1/messages`` stream error and yields what the fallback produces.""" + trigger: Final = _fallback_trigger( + e, (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError) + ) + try: + fallback_response: Final = await attempt.fallback(trigger) + except Exception as fallback_error: + verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) + if isinstance(trigger, MidStreamFallbackError) and trigger.original_exception is not None: + raise trigger.original_exception from fallback_error + raise trigger from fallback_error + async for fallback_item in _anthropic_yield_recovered(fallback_response, wrapper): + yield fallback_item + + +# --- Responses API ------------------------------------------------------------------------------ + + +def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: + from litellm.responses.streaming_iterator import PRE_OUTPUT_LIFECYCLE_EVENT_TYPES + + if held_event_count >= MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: + return False + return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES + + +def _responses_event_announces_background_run(item: object) -> bool: + return getattr(getattr(item, "response", None), "background", None) is True + + +def _mid_stream_error_before_output(error: MidStreamFallbackError) -> bool: + return not error.generated_content or error.is_pre_first_chunk + + +def _extract_partial_responses_usage(source_iterator: "BaseResponsesAPIStreamingIterator") -> "ResponseAPIUsage | None": + """Partial token usage of a Responses stream that failed, normalized to ``ResponseAPIUsage``.""" + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ) + + if isinstance(source_iterator, LiteLLMCompletionStreamingIterator): + chunks: Final = source_iterator.collected_chat_completion_chunks + if chunks: + try: + from litellm.main import ( + stream_chunk_builder, # pyright: ignore[reportUnknownVariableType] # untyped chunk list + ) + + built: Final = stream_chunk_builder(chunks=chunks) + chat: Final[object | None] = getattr(built, "usage", None) if built is not None else None + if chat is not None: + prompt: Final = int(getattr(chat, "prompt_tokens", 0) or 0) + completion: Final = int(getattr(chat, "completion_tokens", 0) or 0) + total: Final = int(getattr(chat, "total_tokens", prompt + completion) or (prompt + completion)) + return ResponseAPIUsage(input_tokens=prompt, output_tokens=completion, total_tokens=total) + except Exception as error: # noqa: BLE001 # the builder is best-effort + verbose_router_logger.debug("partial usage of a failed Responses stream: %s", error) + completed: Final = source_iterator.completed_response + if isinstance(completed, (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent)): + return completed.response.usage + return None + + +def _combine_responses_fallback_usage(fallback_item: object, partial_usage: "ResponseAPIUsage") -> None: + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ) + + if not isinstance(fallback_item, (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent)): + return + response: Final = fallback_item.response + if response.usage is None: + return + fb: Final = response.usage + response.usage = ResponseAPIUsage( + input_tokens=(partial_usage.input_tokens or 0) + (fb.input_tokens or 0), + output_tokens=(partial_usage.output_tokens or 0) + (fb.output_tokens or 0), + total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), + ) + + +def _params_sent_to_deployment(attempt: StreamAttempt, deployment_id: object) -> Mapping[str, object]: + """The request as the deployment received it, merged in the order the attempt applies.""" + return attempt.core.params_sent_to_deployment(attempt.request, deployment_id) + + +def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it ports + attempt: StreamAttempt, response: "BaseResponsesAPIStreamingIterator" +) -> "BaseResponsesAPIStreamingIterator": + """``_aresponses_streaming_iterator``: pre-output lifecycle events are held while a fallback can take + over; a mid-stream error falls back through the core, continuing after text.""" + import time + from datetime import datetime, timezone + + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, get_openai_response_types + + source_iterator: Final = response + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + openai_types: Final = get_openai_response_types() + terminal_event_types: Final = ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, + ) + + class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): + """Subclasses the base iterator only for isinstance checks; iteration is the generator's.""" + + fallback_headers_adopted: bool = False + + def __init__(self, async_generator: AsyncGenerator[object, None]) -> None: # pyright: ignore[reportMissingSuperCall] # no httpx response of its own + self._async_generator: Final = async_generator + self._billing_source: object = source_iterator + self.response = getattr(source_iterator, "response", None) # pyright: ignore[reportAttributeAccessIssue] # copied from the source + self.model = getattr(source_iterator, "model", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.logging_obj = getattr( + source_iterator, "logging_obj", getattr(source_iterator, "litellm_logging_obj", None) + ) # pyright: ignore[reportAttributeAccessIssue] # see above + self.finished = False + self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.completed_response = None + self.start_time = getattr(source_iterator, "start_time", None) or datetime.now(timezone.utc).replace( + tzinfo=None + ) + self._failure_handled = False + self._yielded_first_chunk = False + self._generated_content = "" + self._completed_response_cached = False + self._completed_response_logged = False + self._completed_response_cache_hit = None + self._persist_completed_response_before_logging = True + self._stream_created_time = time.time() + self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.request_data = getattr(source_iterator, "request_data", {}) or {} + self.call_type = getattr(source_iterator, "call_type", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) + + def adopt_fallback_headers(self, fallback_response: object) -> _Prepared: + prepared: Final = _prepare_fallback_hidden_params(fallback_response) + self._billing_source = fallback_response + self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} + self.fallback_headers_adopted = True + return prepared + + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect + + current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source) + if self._billing_source is source_iterator: + return current + primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator) + if current is None or primary is None: + return current or primary + if primary.response.usage is not None: + _combine_responses_fallback_usage(current, primary.response.usage) + return current + + def __aiter__(self) -> "FallbackResponsesStreamWrapper": + return self + + @property + def yielded_first_chunk(self) -> bool: + return self._yielded_first_chunk + + async def __anext__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields what the stream yields + try: + chunk: Final = await self._async_generator.__anext__() + except StopAsyncIteration: + if self.completed_response is None: + self.completed_response = getattr(source_iterator, "completed_response", None) + raise + if self.completed_response is None and getattr(chunk, "type", None) in terminal_event_types: + self.completed_response = chunk # pyright: ignore[reportAttributeAccessIssue] # the terminal event + self._yielded_first_chunk = True + return chunk + + async def aclose(self) -> None: + await self._async_generator.aclose() + + def background_stream_can_release(item: object) -> bool: + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + if ( + type(source_iterator) is not ResponsesAPIStreamingIterator + or source_iterator.custom_llm_provider != "openai" + ): + return False + if not _responses_event_announces_background_run(item): + return False + return not attempt.can_fall_back() + + async def stream_with_fallbacks() -> AsyncGenerator[object, None]: + held: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback + response_id_exposed = False # rebind-ok: set once, before an early-released event is yielded + try: + async for item in _areported( + attempt, + cast("AsyncIterator[object]", source_iterator), # cast-ok: an async stream + lambda last: _responses_total_tokens(last, source_iterator), + ): + if not response_id_exposed and background_stream_can_release(item): + response_id_exposed = True + if not response_id_exposed and _responses_stream_holds_event(item, len(held)): + held = (*held, item) + continue + for held_event in held: + yield held_event + held = () + yield item + for held_event in held: + yield held_event + except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + source_iterator.model or "", + source_iterator.custom_llm_provider or "", + wrapper.yielded_first_chunk, + ) + if _is_context_window(stream_error) + else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error + ) + if response_id_exposed: + _raise_original_mid_stream_error(e) + async with contextlib.aclosing( + _responses_fallback_attempt(attempt, e, source_iterator, wrapper.adopt_fallback_headers, held) + ) as fallback_stream: + async for fallback_item in fallback_stream: + yield fallback_item + except Exception: + for held_event in held: + yield held_event + raise + finally: + with anyio.CancelScope(shield=True): + if hasattr(source_iterator, "aclose"): + try: + await source_iterator.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + except Exception as exc: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks(aresponses): error closing source: %s", exc) + + wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) + return wrapper + + +def _responses_continuation_request( + attempt: StreamAttempt, e: MidStreamFallbackError, source_iterator: "BaseResponsesAPIStreamingIterator" +) -> Mapping[str, object] | None: + """The request a fallback continues from once a Responses stream broke after text, or None before + output. Raises the primary's error when the output is constrained, a tool call began, a content + filter cut it, or the output ceiling is spent.""" + from litellm.litellm_core_utils.token_counter import ( + token_counter, # pyright: ignore[reportUnknownVariableType] # untyped tokenizer argument + ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + RESPONSES_CONTINUATION_KWARG, + ResponsesContinuation, + constrains_responses_continuation, + continuation_output_ceilings, + responses_continuation_input, + ) + + if _mid_stream_error_before_output(e): + return None + if source_iterator.started_tool_call or isinstance( + e.original_exception, litellm_exceptions.ContentPolicyViolationError + ): + _raise_original_mid_stream_error(e) + cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id") + if constrains_responses_continuation(_params_sent_to_deployment(attempt, cut_deployment_id)): + _raise_original_mid_stream_error(e) + request: Final = attempt.request + earlier: Final = request.get(RESPONSES_CONTINUATION_KWARG) + continuation: Final = ( + earlier.extended(e.generated_content) + if isinstance(earlier, ResponsesContinuation) + else ResponsesContinuation(request.get("input"), e.generated_content) + ) + emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content) + reduced: Final = continuation_output_ceilings(request, emitted_tokens) + if reduced is None: + _raise_original_mid_stream_error(e) + return { + **request, + **reduced, + RESPONSES_CONTINUATION_KWARG: continuation, + "input": responses_continuation_input(continuation), + } + + +async def _responses_fallback_attempt( + attempt: StreamAttempt, + e: MidStreamFallbackError, + source_iterator: "BaseResponsesAPIStreamingIterator", + adopt_headers: Callable[[object], _Prepared], + held: tuple[object, ...], +) -> AsyncGenerator[object, None]: + """Re-enters the fallback walk for a mid-stream Responses error. The primary's held lifecycle events + reach the client only when no fallback lands.""" + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ResponsesContinuationSeam + + continued: Final = _responses_continuation_request(attempt, e, source_iterator) + seam: Final = ResponsesContinuationSeam(e.generated_content if continued is not None else "") + partial_usage: Final = _extract_partial_responses_usage(source_iterator) + trigger: Final = _fallback_trigger( + e, (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError) + ) + fallback_response: object = None # rebind-ok: set once a fallback answered, closed at the end + fallback_yielded = False # rebind-ok: flipped on the first fallback item + try: + fallback_response = await attempt.fallback(trigger, continued) # rebind-ok: rebound as the stream advances + prepared: Final = adopt_headers(fallback_response) + if hasattr(fallback_response, "__aiter__"): + spliced: Final = cast( # cast-ok: the fallback's events + "AsyncIterator[object]", + seam.splice(fallback_response), # pyright: ignore[reportArgumentType] # an async stream + ) + async for fallback_item in spliced: + _apply_fallback_hidden_params_to_item(fallback_item, prepared) + if partial_usage is not None: + _combine_responses_fallback_usage(fallback_item, partial_usage) + fallback_yielded = True + yield fallback_item + else: + fallback_yielded = True # rebind-ok: rebound as the stream advances + yield fallback_response + except Exception as fallback_error: + verbose_router_logger.error("Responses streaming fallback also failed: %s", fallback_error) + if not fallback_yielded: + for held_event in held: + yield held_event + if fallback_response is None: + if isinstance(trigger, MidStreamFallbackError) and trigger.original_exception is not None: + raise trigger.original_exception from fallback_error + raise trigger from fallback_error + if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: + raise fallback_error.original_exception from fallback_error + raise + finally: + if fallback_response is not None and hasattr(fallback_response, "aclose"): + with anyio.CancelScope(shield=True): + try: + await fallback_response.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + except Exception as exc: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks(aresponses): error closing fallback: %s", exc) + + +# --- the attempt hook --------------------------------------------------------------------------- + + +async def open_attempt_stream(attempt: StreamAttempt, response: object) -> object: + """What an attempt returns once its stream is wrapped: a chat stream's first chunk is read here, so + a stream that fails before any token raises inside the attempt.""" + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + + if attempt.op == "acompletion" and isinstance(response, CustomStreamWrapper): + if response.completion_stream is None and response.make_call is not None: # pyright: ignore[reportUnknownMemberType] # untyped attributes + await response.fetch_stream() + first_chunk: Final = await _aread_first_chunk(response) + held_slot: Final = attempt.slot.pop_all() if attempt.slot is not None else contextlib.ExitStack() + return wrap_chat_stream(attempt, response, first_chunk, held_slot) + if not attempt.kwargs.get("stream"): + return response + if attempt.op == "aresponses" and isinstance(response, BaseResponsesAPIStreamingIterator): + return wrap_responses_stream(attempt, response) + if attempt.op in ("aanthropic_messages", "anthropic_messages") and hasattr(response, "__aiter__"): + return wrap_anthropic_stream(attempt, cast("AsyncIterator[bytes]", response)) # cast-ok: a byte stream + return response + + +def open_attempt_stream_sync(attempt: StreamAttempt, response: object) -> object: + if attempt.op == "completion" and isinstance(response, CustomStreamWrapper): + if response.completion_stream is None and response.make_call is not None: # pyright: ignore[reportUnknownMemberType] # untyped attributes + response.fetch_sync_stream() + return wrap_chat_stream_sync(attempt, response, _read_first_chunk(response)) + return response + + +def shape_attempt_kwargs(kwargs: Mapping[str, object], deployment: Mapping[str, object]) -> Mapping[str, object]: + """A Responses continuation is shaped for the deployment the fallback landed on.""" + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + RESPONSES_CONTINUATION_KWARG, + responses_continuation_input_for, + ) + + if RESPONSES_CONTINUATION_KWARG not in kwargs: + return kwargs + shaped: Final = {key: value for key, value in kwargs.items() if key != RESPONSES_CONTINUATION_KWARG} + shaped["input"] = responses_continuation_input_for(kwargs, deployment) + return shaped + + +__all__ = [ + "MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS", + "MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS", + "FallbackAwareAnthropicMessagesStream", + "FallbackAwareStreamWrapper", + "StreamAttempt", + "open_attempt_stream", + "open_attempt_stream_sync", + "shape_attempt_kwargs", +] diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py new file mode 100644 index 00000000000..0051b0faa79 --- /dev/null +++ b/litellm/router_rust/invoke/walk.py @@ -0,0 +1,292 @@ +""" +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. +""" + +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 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", "_params_base_model", "_request", "_scopes") + + def __init__( + self, request: Mapping[str, object], field: str, params_base_model: Callable[[str], str | None] + ) -> 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 + + 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 + metadata["model_group"] = 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")) + return {**self._request, **scope.kwargs, "model": event.get("to"), 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 + 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 + + 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") + 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} + + +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 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"] diff --git a/litellm/router_rust/invoke/websocket.py b/litellm/router_rust/invoke/websocket.py new file mode 100644 index 00000000000..7461e514c5e --- /dev/null +++ b/litellm/router_rust/invoke/websocket.py @@ -0,0 +1,131 @@ +""" +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.router_rust.invoke.shaping import as_dict + +if TYPE_CHECKING: + from litellm.router_rust.core import RustCore + + +def turn_admission( + core: "RustCore", 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: "RustCore", 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._hosted._core_get_deployment(model_id) if isinstance(model_id, str) else None # pyright: ignore[reportPrivateUsage] # the core's host registry + return deployment if current is None else current.model_dump(exclude_none=True) + + +def _budget_limiter(core: "RustCore") -> 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: "RustCore", 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_rust.steps.selection import routing_pre_call_checks + 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._checks.routing_callbacks() # pyright: ignore[reportPrivateUsage] # `_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 routing_pre_call_checks(core._checks.routing_callbacks(), dict(without_io_token_limits(deployment))) # pyright: ignore[reportPrivateUsage] # `_routing_callbacks` + return io_token_admission + + +async def _admit_turn_through_routing_strategy( + core: "RustCore", 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 = {key: value for key, value in request_kwargs.items() if key != "litellm_logging_obj"} + override: Final = routing_kwargs.get("routing_strategy") + strategy, scope, group = core._handle.routing_context( # pyright: ignore[reportPrivateUsage] # the core's Rust handle + model_group, override if isinstance(override, str) else None + ) + selectors: Final = core._selectors # pyright: ignore[reportPrivateUsage] # the core's strategy 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"] diff --git a/litellm/router_rust/steps/README.md b/litellm/router_rust/steps/README.md new file mode 100644 index 00000000000..9b5c68bdba9 --- /dev/null +++ b/litellm/router_rust/steps/README.md @@ -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. diff --git a/litellm/router_rust/steps/__init__.py b/litellm/router_rust/steps/__init__.py new file mode 100644 index 00000000000..c47505e7468 --- /dev/null +++ b/litellm/router_rust/steps/__init__.py @@ -0,0 +1 @@ +"""Answers to router-core host steps that need live Python objects.""" diff --git a/litellm/router_rust/steps/access.py b/litellm/router_rust/steps/access.py new file mode 100644 index 00000000000..087a87e2c6b --- /dev/null +++ b/litellm/router_rust/steps/access.py @@ -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 diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py new file mode 100644 index 00000000000..1c5e7d4d2d8 --- /dev/null +++ b/litellm/router_rust/steps/callbacks.py @@ -0,0 +1,558 @@ +""" +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 never reaches a +Python router core. +""" + +from collections.abc import Iterable, Mapping, Sequence +from itertools import chain +from typing import TYPE_CHECKING, Final + +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.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, + 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.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, ...] = () + + # -- 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" + + 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, as it registers no outcome callbacks on the host.""" + 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 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.""" + 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 + + +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 _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.""" + live_check_owners.discard(self) + selectors: Final = (*self._sel.callback_selectors(), *self._sel.override_selectors()) + self._opt.retire_state(selectors) + self._opt.release(selectors) + self._sel.clear() + + 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 + try: + remaining: Final = await self._opt.filter_deployments( + str(step.get("model_group")), + records, # pyright: ignore[reportUnknownArgumentType] # see above + messages, + request, + ) + except Exception as error: # noqa: BLE001 # handed to the Rust core as the try's failure + return attempt_failure(error), error + return filtered_ids(remaining, candidates), None # pyright: ignore[reportUnknownArgumentType] # see above + + +__all__ = [ + "OptionalChecks", + "RustOptionalChecks", + "filtered_ids", + "register_router_selector", +] diff --git a/litellm/router_rust/steps/plugins.py b/litellm/router_rust/steps/plugins.py new file mode 100644 index 00000000000..4af73313838 --- /dev/null +++ b/litellm/router_rust/steps/plugins.py @@ -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"] diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py new file mode 100644 index 00000000000..d37f5026360 --- /dev/null +++ b/litellm/router_rust/steps/precall.py @@ -0,0 +1,318 @@ +""" +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. +""" + +from collections.abc import Callable, Mapping, Sequence +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.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, +) +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 + +# 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, +) -> 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.""" + 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, + ) + + input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=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 _count(self) -> int | None: + try: + return count_tokens(self._host, self._messages, self._input, self._request, self._tokenizer_model()) + 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() + ) + 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 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", "request_carries_image"] diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py new file mode 100644 index 00000000000..592bd01b21f --- /dev/null +++ b/litellm/router_rust/steps/prerouting.py @@ -0,0 +1,1097 @@ +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportAny=false, reportExplicitAny=false +# (the ported core code reads the untyped strategy routers, LiteLLM_Params dict fields and request dicts) +""" +Host-side pre-routing for ``RustCore``. + +Before an async pick resolves its name, the Rust core asks the host to pre-route it (a +``pre_route`` step, ``router-core/src/prerouting.rs``): the complexity, auto, adaptive and quality +routers registered from ``auto_router/...`` marker deployments, and the Claude Code session +binding, may send the request to another model group. They are live Python objects, so they stay +host-side for now (bridged, to port): this module registers them from the host registry's +deployments the way the Python router's ``init_*_router_deployment``, +``_finalize_adaptive_router_if_configured`` and ``_unregister_pre_routing_strategy_for_deployment`` +do, and answers the step the way ``_core_async_pre_routing_hook`` and ``_pin_tier_params_onto_request`` +do. It never reaches a Python router core. +""" + +import json +import logging +import re +from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast # noqa: TID251 # the strategy routers are untyped + +from litellm._internal_context import with_service_target +from litellm._logging import verbose_router_logger +from litellm.caching.caching import DualCache, InMemoryCache +from litellm.caching.redis_cache import log_redis_failure +from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + CONSUMED_REQUEST_TAGS_METADATA_KEY, + DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, + OUTPUT_TOKEN_CEILING_PARAMS, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, +) +from litellm.integrations.otel.routing import routing_decision_attributes +from litellm.integrations.otel.runtime import phase_attributes +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference +from litellm.router_api import deployment_id +from litellm.router_host import router_host +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + AUTO_ROUTER_MODEL_PREFIX, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capabilities, + classify_strategy_router_model, + count_capability_routers, +) +from litellm.router_strategy.complexity_router.context_compaction import ( + arm_compaction, + is_native_compaction_call, + reject_recursive_compactor, +) +from litellm.router_strategy.tag_based_routing import get_routing_tags_from_request_kwargs, is_valid_deployment_tag +from litellm.router_utils.fallback_event_handlers import record_pre_routing_selection +from litellm.types.router import ConsumedRequestTagsStamp, Deployment, LiteLLM_Params, TaggedPreRoutingStrategy +from litellm.types.utils import ( + PROMPT_QUOTING_ROUTING_DECISION_FIELDS, + CustomPricingLiteLLMParams, + StandardLoggingRoutingDecision, + all_litellm_params, +) +from litellm.utils import provider_rejectable_params + +if TYPE_CHECKING: + from litellm.router_python_host import PythonHost + from litellm.router_rust.surface.registry import HostRegistry + from litellm.router_strategy.auto_router.auto_router import PreRoutingHookResponse + +# The live registry dicts: model_name -> the strategies registered under it, one per tag set. +StrategyRegistry: TypeAlias = "dict[str, list[TaggedPreRoutingStrategy[Any]]]" + +_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( + {"model", "api_base", "api_key", "api_version", "responses_api_version"} +) +_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 +CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" +# `_async_get_available_deployment` runs the pre-routing hook only for the strategies with an async +# pick; any other router strategy takes the sync path, which never pre-routes. +_ASYNC_PICK_STRATEGIES: Final = frozenset( + {"usage-based-routing-v2", "simple-shuffle", "cost-based-routing", "latency-based-routing", "least-busy"} +) +TIER_PARAMS_NEVER_DROPPED: Final = frozenset(all_litellm_params) | frozenset( + { + "additional_drop_params", + "drop_params", + "messages", + "model", + "extra_headers", + "max_tokens", + "max_completion_tokens", + } +) + + +def _json_default_stable_id(value: object) -> str: + return f"{type(value).__module__}.{type(value).__qualname__}" + + +def _deployment_tags(deployment: Deployment) -> tuple[str, ...]: + return tuple(deployment.litellm_params.tags or ()) + + +def _has_registered_strategy(registry: StrategyRegistry, model_name: str, tags: tuple[str, ...]) -> bool: + return any(existing.tags == tags for existing in registry.get(model_name, [])) + + +def _routing_definition_fingerprint(deployment: Deployment) -> str | None: + """Identify configured routing fields, excluding external files and live strategy state.""" + try: + definition: Final = json.dumps( + deployment.litellm_params.model_dump( + include=set(STRATEGY_ROUTER_PARAM_FIELDS | {"model", "tags"}), exclude_none=True, warnings=False + ), + sort_keys=True, + default=_json_default_stable_id, + ) + except (TypeError, ValueError): + return None + return deployment_id(deployment.model_name, {"routing_definition": definition}) + + +def _unregister(registry: StrategyRegistry, model_name: str, tags: tuple[str, ...]) -> bool: + existing: Final = registry.get(model_name, []) + remaining: Final = [entry for entry in existing if entry.tags != tags] + if len(remaining) == len(existing): + return False + if remaining: + registry[model_name] = remaining # rebind-ok: the live registry, edited in place as the Python router does + else: + registry.pop(model_name, None) + return True + + +def is_strategy_marker(deployment: Mapping[str, object]) -> bool: + params: Final = deployment.get("litellm_params") + if not isinstance(params, Mapping): + return False + model: Final = params.get("model") + return isinstance(model, str) and classify_strategy_router_model(model) is not None + + +def _participates_in_adaptive_routing(params: LiteLLM_Params) -> bool: + kind: Final = classify_strategy_router_model(params.model) + if kind == "adaptive": + return True + if kind != "complexity": + return False + config: Final = params.complexity_router_config + return bool(config.get("adaptive")) if config else False + + +class PreRoutingStrategies: + """The pre-routing strategies of one router, kept on its host registry (``registry().plugins`` + reads ``_core_auto_routers`` and its siblings there).""" + + def __init__( + self, hosted: "HostRegistry", host: "PythonHost", ctx: object, capability_limit: object, cache: DualCache + ) -> None: + self._hosted: Final = hosted + self._host: Final = host + self._ctx: Final = ctx + self._capability_limit: object = capability_limit + self._cache: Final = cache + self._affinity: object = None + # The router's own `_ensure_deployment_affinity_callback` (its optional checks), when set. + self.ensure_affinity: Callable[[], None] | None = None + + # -- the registries ------------------------------------------------------------------------ + + @property + def auto(self) -> StrategyRegistry: + return self._hosted._core_auto_routers # pyright: ignore[reportPrivateUsage] # the registries live on the host registry + + @property + def complexity(self) -> StrategyRegistry: + return self._hosted._core_complexity_routers # pyright: ignore[reportPrivateUsage] # see above + + @property + def adaptive(self) -> StrategyRegistry: + return self._hosted._core_adaptive_routers # pyright: ignore[reportPrivateUsage] # see above + + @property + def quality(self) -> StrategyRegistry: + return self._hosted._core_quality_routers # pyright: ignore[reportPrivateUsage] # see above + + def all(self) -> tuple[StrategyRegistry, ...]: + """In the order ``_select_pre_routing_strategy`` searches them.""" + return (self.auto, self.complexity, self.adaptive, self.quality) + + def any_registered(self) -> bool: + return any(self.all()) + + # -- registry hooks (HostRegistry calls these) ----------------------------------------------- + + def reset(self) -> None: + """``set_model_list`` drops the auto, complexity and quality routers (never the adaptive ones).""" + self._hosted._core_quality_routers = {} # pyright: ignore[reportPrivateUsage] # see above + self._hosted._core_complexity_routers = {} # pyright: ignore[reportPrivateUsage] # see above + self._hosted._core_auto_routers = {} # pyright: ignore[reportPrivateUsage] # see above + + def register(self, deployment: Deployment) -> None: + """The end of ``_add_deployment``: an auto, complexity or quality marker registers its router + (adaptive ones wait for ``finalize``).""" + match classify_strategy_router_model(deployment.litellm_params.model): + case "semantic": + self._init_auto_router(deployment) + case "complexity": + self._init_complexity_router(deployment) + case "quality": + self._init_quality_router(deployment) + case _: + pass + + def release(self, deployment: Deployment) -> None: + """``_unregister_pre_routing_strategy_for_deployment``.""" + if not deployment.litellm_params.model.startswith("auto_router/"): + return + model_name: Final = deployment.model_name + tags: Final = _deployment_tags(deployment) + for registry in (self.auto, self.complexity, self.quality): + _ = _unregister(registry, model_name, tags) + if _unregister(self.adaptive, model_name, tags): + self._sync_adaptive_router_hooks() + + def release_entry(self, entry: Mapping[str, object]) -> None: + """``release`` of a deleted model_list entry; a failure only logs, as the delete has happened.""" + try: + self.release(Deployment(**entry)) # pyright: ignore[reportArgumentType] # a model_list entry + except Exception: # noqa: BLE001 # the delete already happened; a failed release only logs + verbose_router_logger.exception( + "delete_deployment: could not release pre-routing strategies for model_id=%s; " + "the deployment is out of the model_list and its indices are repaired", + as_dict(entry.get("model_info")).get("id"), + ) + + def upserted(self, deployment: Deployment, previous: Deployment | None) -> None: + """After an upsert's re-add: rebuild the adaptive routers when it touched one.""" + if _participates_in_adaptive_routing(deployment.litellm_params) or ( + previous is not None and _participates_in_adaptive_routing(previous.litellm_params) + ): + self.finalize() + + def without_capability_limit(self, action: Callable[[], object]) -> None: + """Run ``action`` without judging the capability ceiling: a rollback re-admits a serving router.""" + resolver: Final = self._capability_limit + self._capability_limit = None + try: + _ = action() + finally: + self._capability_limit = resolver + + def finalize(self) -> None: + """``_finalize_adaptive_router_if_configured``: an adaptive router for every adaptive marker, + and the adaptive companion of every complexity router configured with ``adaptive``.""" + for entry in self._hosted._core_model_list or []: # pyright: ignore[reportPrivateUsage] # the host registry's live list + params = entry.get("litellm_params") + lp_model = params.get("model") if isinstance(params, dict) else None + if not (isinstance(lp_model, str) and lp_model.startswith("auto_router/adaptive_router")): + continue + model_name = entry.get("model_name") + if not model_name or not params: + continue + deployment = Deployment( + model_name=str(model_name), + litellm_params=LiteLLM_Params.model_validate(params), + model_info=entry.get("model_info"), # pyright: ignore[reportArgumentType] # the entry's model_info + ) + if _has_registered_strategy(self.adaptive, str(model_name), _deployment_tags(deployment)): + continue + self._init_adaptive_router(deployment) + for model_name, tagged_complexity_routers in self.complexity.items(): + for tagged in tagged_complexity_routers: + complexity_router = tagged.strategy + if not complexity_router.config.adaptive: + continue + if _has_registered_strategy(self.adaptive, model_name, tagged.tags): + continue + adaptive_router = complexity_router._ensure_adaptive_router() + if adaptive_router is not None: + self.adaptive[model_name] = [ + *self.adaptive.get(model_name, []), + TaggedPreRoutingStrategy( + tags=tagged.tags, + strategy=adaptive_router, + definition_fingerprint=tagged.definition_fingerprint, + deployment=tagged.deployment, + ), + ] + self._sync_adaptive_router_hooks() + + def retire(self) -> None: + """Release the adaptive routers' post-call hooks this router registered.""" + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + own: Final = frozenset(id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive.values())) + hooks: Final[list[object]] = [ # mutable-ok: built here + hook + for hook in self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own + ] + if self._affinity is not None: + hooks.append(self._affinity) + self._affinity = None + if hooks: + self._host.callbacks.drop_callbacks(hooks) + + # -- building -------------------------------------------------------------------------------- + + def _register( + self, registry: StrategyRegistry, deployment: Deployment, strategy: object, strategy_label: str + ) -> None: + tags: Final = _deployment_tags(deployment) + if _has_registered_strategy(registry, deployment.model_name, tags): + raise ValueError( + f"{strategy_label} deployment {deployment.model_name} with tags {list(tags)} already exists. " + "Please use a different model name or set different tags." + ) + registry[deployment.model_name] = [ # rebind-ok: the live registry, edited in place as the Python router does + *registry.get(deployment.model_name, []), + TaggedPreRoutingStrategy( + tags=tags, + strategy=strategy, + definition_fingerprint=_routing_definition_fingerprint(deployment), + deployment=deployment, + ), + ] + + def _capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: + resolver: Final = self._capability_limit + limit: Final = resolver() if callable(resolver) else None + others: Final = count_capability_routers( + (entry for entry in self._hosted._core_model_list if isinstance(entry, Mapping)), # pyright: ignore[reportPrivateUsage, reportUnnecessaryIsInstance] # the live list + capability=capability, + ) + return capability_limit_violation(capability=capability, held=others + 1, limit=limit) # pyright: ignore[reportArgumentType] # the resolver's limit + + def _init_auto_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.auto_router.auto_router import AutoRouter + + params: Final = deployment.litellm_params + if params.auto_router_config_path is None and params.auto_router_config is None: + raise ValueError( + "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" + ) + if params.auto_router_default_model is None: + raise ValueError( + "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" + ) + if params.auto_router_embedding_model is None: + raise ValueError( + "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" + ) + max_input_chars: Final = params.auto_router_max_input_chars + router: Final = AutoRouter( + model_name=deployment.model_name, + auto_router_config_path=params.auto_router_config_path, + auto_router_config=params.auto_router_config, + default_model=params.auto_router_default_model, + embedding_model=params.auto_router_embedding_model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + max_input_chars=max_input_chars if max_input_chars is not None else DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, + ) + self._register(self.auto, deployment, router, "Auto-router") + + def _init_complexity_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + + config: Final = deployment.litellm_params.complexity_router_config + for capability in claimed_capabilities(config): + if (violation := self._capability_violation(capability)) is not None: + raise ValueError(violation) + default_model: Final = ( + ComplexityRouterConfig.model_validate(config).resolve_default_model( + deployment.litellm_params.complexity_router_default_model + ) + if config + else deployment.litellm_params.complexity_router_default_model + ) + if default_model is None: + raise ValueError( + "complexity_router_default_model is required for complexity-router deployments, " + "or configure tiers in complexity_router_config. Please set it in the litellm_params" + ) + router: Final = ComplexityRouter( + model_name=deployment.model_name, + default_model=default_model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + complexity_router_config=config, + ) + self._register(self.complexity, deployment, router, "Complexity-router") + if router._uses_deployment_pin: # pyright: ignore[reportPrivateUsage] # the Python router reads the same flag + self._ensure_deployment_affinity_callback() + + def _ensure_deployment_affinity_callback(self) -> None: + """``_ensure_deployment_affinity_callback``: the affinity check (every global flag off) a + complexity router's session pin needs, registered once.""" + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + if self.ensure_affinity is not None: + self.ensure_affinity() + return + if self._affinity is not None: + return + hosted: Final = self._hosted + + def is_priority_group(model: str) -> bool: + 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" + + affinity: Final = DeploymentAffinityCheck( + cache=self._cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=False, + model_group_affinity_config=None, + is_priority_group=is_priority_group, + session_affinity_repin_on_failover=False, + ) + self._affinity = affinity + self._host.callbacks.add_litellm_callback(affinity) + + def _init_quality_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.quality_router.quality_router import QualityRouter + + config: Final = deployment.litellm_params.quality_router_config + configured: Final = deployment.litellm_params.quality_router_default_model + default_model: Final = configured if configured is not None or not config else config.get("default_model") + if default_model is None: + raise ValueError( + "quality_router_default_model is required for quality-router deployments, " + "or set default_model in quality_router_config. Please configure it in the litellm_params" + ) + router: Final = QualityRouter( + model_name=deployment.model_name, + default_model=default_model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + quality_router_config=config, + ) + self._register(self.quality, deployment, router, "Quality-router") + + def _init_adaptive_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + from litellm.types.router import AdaptiveRouterConfig, AdaptiveRouterPreferences + + raw_config: Final = deployment.litellm_params.adaptive_router_config + if raw_config is None: + raise ValueError("adaptive_router_config is required for adaptive-router deployments.") + config: Final = AdaptiveRouterConfig(**raw_config) + model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} # mutable-ok: built here + model_to_cost: Final[dict[str, float]] = {} # mutable-ok: built here + for name in config.available_models: + deployments = self._hosted._core_deployments_named(name) # pyright: ignore[reportPrivateUsage] # the host registry's name index + if not deployments: + continue + first = deployments[0] + model_info = first.get("model_info") + info: Mapping[str, object] = model_info if isinstance(model_info, dict) else {} + prefs_raw = info.get("adaptive_router_preferences") + if isinstance(prefs_raw, dict): + model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) + cost = _adaptive_cost(as_dict(first.get("litellm_params")), info) + if isinstance(cost, int | float | str): + model_to_cost[name] = float(cost) + router: Final = AdaptiveRouter( + router_name=deployment.model_name, + config=config, + model_to_prefs=model_to_prefs, + model_to_cost=model_to_cost, + ) + self._register(self.adaptive, deployment, router, "Adaptive-router") + self._host.callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=router)) + verbose_router_logger.info( + "AdaptiveRouter[%s] initialized with %d models", deployment.model_name, len(config.available_models) + ) + + def _sync_adaptive_router_hooks(self) -> None: + """One post-call hook per registered adaptive router, plus the header hook while any exists.""" + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, + ) + + callbacks: Final = self._host.callbacks + for callback in callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): + callbacks.remove_callback_from_all_lists(callback) + for tagged in chain.from_iterable(self.adaptive.values()): + callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)) + if self.adaptive: + callbacks.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) + + +# -- the request side -------------------------------------------------------------------------- + + +def _stamp_or_clear_metadata_key( + request_kwargs: dict[str, object], # mutable-ok: the request kwargs, stamped in place + key: str, + value: object | None, +) -> None: + """Write a proxy-internal metadata key for this routing attempt, or clear it from both buckets.""" + if value is None: + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): + if isinstance(bucket, dict): + bucket.pop(key, None) + return + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket[key] = value + + +def _redact_prompt_text_if_needed( + request_kwargs: Mapping[str, object], routing_decision: StandardLoggingRoutingDecision +) -> StandardLoggingRoutingDecision: + from litellm.litellm_core_utils.redact_messages import should_redact_message_logging + + if not should_redact_message_logging( + { + "litellm_params": request_kwargs, + "standard_callback_dynamic_params": request_kwargs.get("standard_callback_dynamic_params"), + } + ): + return routing_decision + kept: Final = { + field: value for field, value in routing_decision.items() if field not in PROMPT_QUOTING_ROUTING_DECISION_FIELDS + } + return cast(StandardLoggingRoutingDecision, kept) # cast-ok: dropping optional keys preserves the type + + +def _record_routing_decision( + request_kwargs: dict[str, object], # mutable-ok: the request kwargs + routing_decision: StandardLoggingRoutingDecision | None, +) -> None: + """Make the request's metadata describe this routing attempt, and only this one.""" + from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters + from litellm.types.router import BaselineRouteStamp + + phase_attributes(routing_decision_attributes(routing_decision)) + baseline_model: Final = routing_decision.get("savings_baseline_model") if routing_decision else None + baseline_id: Final = routing_decision.get("savings_baseline_deployment_id") if routing_decision else None + router_name: Final = routing_decision.get("router_model_name") if routing_decision else None + caller_parameters: Final = capture_baseline_parameters(request_kwargs) if router_name and baseline_model else None + _stamp_or_clear_metadata_key( + request_kwargs, + "_autorouter_baseline_route", + BaselineRouteStamp(router_name, baseline_model, baseline_id, caller_parameters) + if router_name and baseline_model and baseline_id + else None, + ) + _stamp_or_clear_metadata_key( + request_kwargs, + "routing_decision", + None if routing_decision is None else _redact_prompt_text_if_needed(request_kwargs, routing_decision), + ) + + +def _as_output_cap(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int | float | str): + return None + try: + cap: Final = int(float(value)) + except (ValueError, OverflowError): + return None + return cap if cap >= 0 else None + + +def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return None + return MappingProxyType( + { + str(carrier): cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := _as_output_cap(value)) is not None + } + ) + + +def _restore_client_output_ceiling( + request_kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs, restored in place +) -> None: + callers_ceiling: Final = _stamped_client_ceiling(request_kwargs) + if callers_ceiling is None: + return + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) + request_kwargs.update(callers_ceiling) + + +def _stamp_client_output_ceiling(request_kwargs: dict[str, object]) -> None: # mutable-ok: the request kwargs + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, + ) + + +def _pop_effort_from_nested_carrier(request_kwargs: dict[str, object], carrier: str) -> None: # mutable-ok: same + nested: Final = request_kwargs.get(carrier) + if not isinstance(nested, Mapping): + return + sanitized: Final = {key: value for key, value in nested.items() if key != "effort"} + if sanitized: + request_kwargs[carrier] = sanitized # rebind-ok: copy-on-write of the request's nested carrier + else: + request_kwargs.pop(carrier, None) + + +def _tier_ceiling_under_the_surface_name( + tier_litellm_params: Mapping[str, object], responses_call: bool +) -> Mapping[str, object]: + surface_key: Final = "max_output_tokens" if responses_call else "max_tokens" + carried: Final = tuple( + key + for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens") + if key in tier_litellm_params + ) + if not carried: + return tier_litellm_params + return MappingProxyType( + { + **{k: v for k, v in tier_litellm_params.items() if k not in OUTPUT_TOKEN_CEILING_PARAMS}, + surface_key: tier_litellm_params[carried[0]], + } + ) + + +def _drop_client_carriers_a_tier_pin_supersedes( + request_kwargs: dict[str, object], # mutable-ok: the request kwargs + tier_litellm_params: Mapping[str, object], +) -> None: + if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): + _stamp_client_output_ceiling(request_kwargs) + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) + if "reasoning_effort" not in tier_litellm_params: + return + request_kwargs.pop("thinking", None) + _pop_effort_from_nested_carrier(request_kwargs, "output_config") + _pop_effort_from_nested_carrier(request_kwargs, "reasoning") + + +def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: + proxy_server_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return None + headers: Final = proxy_server_request.get("headers") + if not isinstance(headers, Mapping): + return None + return next( + ( + value + for key, value in headers.items() + if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) + ), + None, + ) + + +def _session_router_cache_key(request_kwargs: Mapping[str, object]) -> str | None: + session_id: Final = _request_header(request_kwargs, "x-claude-code-session-id") + if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: + return None + metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata: Final = request_kwargs.get(metadata_name) + if not isinstance(metadata, Mapping): + return None + caller_scope: Final = metadata.get("user_api_key_hash") + if not isinstance(caller_scope, str) or not caller_scope: + return None + return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + + +class PreRouter: + """Answers the Rust core's ``pre_route`` steps (``_core_async_pre_routing_hook``).""" + + def __init__( + self, + strategies: PreRoutingStrategies, + hosted: "HostRegistry", + host: "PythonHost", + ctx: object, + routing_cache: DualCache, + ) -> None: + self._strategies: Final = strategies + self._hosted: Final = hosted + self._host: Final = host + self._ctx: Final = ctx + self._session_cache: Final = DualCache(redis_cache=routing_cache.redis_cache, in_memory_cache=InMemoryCache()) + + # -- the step -------------------------------------------------------------------------------- + + async def pre_route( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: the request kwargs the hook stamps + *, + messages: object, + input: object, + fallback_depth: int, + router_strategy: str | None, + enable_tag_filtering: bool, + match_any: bool, + public_router: object, + run_plugins: Callable[[str], Awaitable[None]] | None = None, + ) -> str | None: + """The start of ``_async_get_available_deployment``: run the hook, record its pick and pin the + tier's litellm_params. The model group to resolve, or ``None`` to resolve ``model`` itself.""" + if router_strategy not in _ASYNC_PICK_STRATEGIES: + return None + responses_call: Final = input is not None and messages is None + response: Final = await self._hook( + model, + request_kwargs, + messages=messages, + input=input, + fallback_depth=fallback_depth, + enable_tag_filtering=enable_tag_filtering, + match_any=match_any, + public_router=public_router, + run_plugins=run_plugins, + ) + routed: Final = response.model if response is not None else model + if response is not None: + record_pre_routing_selection(request_kwargs, routed) + self._pin_tier_params( + routed, response.litellm_params if response is not None else None, request_kwargs, responses_call + ) + return routed if response is not None else None + + # -- the hook -------------------------------------------------------------------------------- + + async def _hook( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: see pre_route + *, + messages: object, + input: object, + fallback_depth: int, + enable_tag_filtering: bool, + match_any: bool, + public_router: object, + run_plugins: Callable[[str], Awaitable[None]] | None, + ) -> "PreRoutingHookResponse | None": + _restore_client_output_ceiling(request_kwargs) + _record_routing_decision(request_kwargs, None) + requested: Final = self._hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # the host registry's alias lookup + registered: Final = await self._resolve_session_router( + model, requested, request_kwargs, fallback_depth, enable_tag_filtering, match_any + ) + # The routing plugins run on the session-resolved name, before a strategy is selected, so a + # strategy reads their signals; router-core narrows the pick to what they left. + if run_plugins is not None: + await run_plugins(registered) + selected: Final = self._select_strategy(registered, request_kwargs, enable_tag_filtering, match_any) + if selected is None: + await arm_compaction(request_kwargs, None) + _stamp_or_clear_metadata_key(request_kwargs, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, None) + _stamp_or_clear_metadata_key(request_kwargs, CONSUMED_REQUEST_TAGS_METADATA_KEY, None) + return None + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + + marker: Final = self._selected_marker(registered, selected.tags, request_kwargs) + definition: Final = selected.deployment + config_id: Final = definition.model_info.id if definition else None + updated_at: Final = definition.model_info.updated_at if definition else None + router_kind: Final = classify_strategy_router_model(definition.litellm_params.model) if definition else None + provenance: Final[StandardLoggingRoutingDecision] = { + "router_model_name": definition.model_name if definition else registered, + **({"router_type": router_kind} if router_kind is not None else {}), + **({"router_config_id": config_id} if config_id is not None else {}), + **({"router_config_updated_at": updated_at.isoformat()} if updated_at is not None else {}), + **( + {"router_config_fingerprint": selected.definition_fingerprint} + if selected.definition_fingerprint is not None + else {} + ), + } # pyright: ignore[reportAssignmentType] # the optional keys, as the Python router builds them + phase_attributes(routing_decision_attributes(provenance)) + reject_recursive_compactor(registered) + strategy: Final = selected.strategy + complexity: Final = strategy if isinstance(strategy, ComplexityRouter) else None + await arm_compaction( + request_kwargs, + complexity.config.context_compaction if complexity is not None else None, + _tier_members(complexity.config.tiers) if complexity is not None else (), # pyright: ignore[reportArgumentType] # the tiers' model names + parent_model=model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + allow_escalation=complexity is not None and complexity.config.enable_context_window_escalation, + messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + ) + await router_host().authorize_strategy_inference( + deployment=marker, + request_kwargs=request_kwargs, + llm_router=public_router, # pyright: ignore[reportArgumentType] # the litellm.Router this core serves + ) + compression_policy: Final = router_host().compression_policy( + llm_router=public_router, # pyright: ignore[reportArgumentType] # same + model_alias=registered, + request_kwargs=request_kwargs, + request_tags=get_routing_tags_from_request_kwargs(request_kwargs), + ) + independent_compression: Final = compression_policy is not None and not ( + compression_policy.is_same + and compression_policy.model is not None + and router_host().model_hop_compression_armed() + ) + routing_messages: Final = ( + await router_host().messages_for_routing( + policy=compression_policy, + messages=messages, # pyright: ignore[reportArgumentType] # same + request_kwargs=request_kwargs, + ) + if independent_compression + else None + ) + routed: Final = await strategy.async_pre_routing_hook( + model=registered, + request_kwargs=request_kwargs, + messages=routing_messages if routing_messages is not None else messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + input=input, # pyright: ignore[reportArgumentType] # the request's input, as given + specific_deployment=False, + ) + hooked: Final[PreRoutingHookResponse | None] = routed + traced: Final = ( + hooked.model_copy(update={"routing_decision": {**provenance, **hooked.routing_decision}}) + if hooked is not None and hooked.routing_decision is not None + else hooked + ) + response: Final = ( + traced.model_copy(update={"messages": messages}) + if traced is not None and routing_messages is not None and traced.messages == routing_messages + else traced + ) + _record_routing_decision(request_kwargs, response.routing_decision if response else None) + _stamp_or_clear_metadata_key( + request_kwargs, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + response.session_affinity_ttl_seconds if response else None, + ) + request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) + consumed: Final = ( + ConsumedRequestTagsStamp(model_group=response.model, tags=selected.tags) + if response is not None + and selected.tags + and request_tags + and is_valid_deployment_tag(selected.tags, request_tags, match_any) # pyright: ignore[reportArgumentType] # the tags as the Python router passes them + else None + ) + _stamp_or_clear_metadata_key(request_kwargs, CONSUMED_REQUEST_TAGS_METADATA_KEY, consumed) + marker_params: Final = ( + self._forwardable_marker_params(registered, selected.tags, request_kwargs) if response is not None else () + ) + tier_param_keys: Final = tuple(response.litellm_params or ()) if response is not None else () + overridden: Final = frozenset(tier_param_keys) | ( + OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() + ) + newly_forwarded: Final = tuple( + (key, value) for key, value in marker_params if key not in request_kwargs and key not in overridden + ) + if response is not None: + _stamp_client_output_ceiling(request_kwargs) + request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) + request_kwargs.update(newly_forwarded) + if newly_forwarded: + request_kwargs.update({_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: tuple(key for key, _ in newly_forwarded)}) + return response + + # -- strategy selection ---------------------------------------------------------------------- + + def _deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[Mapping[str, object]]: + return self._hosted._core_deployments_for_request(model, request_kwargs) # pyright: ignore[reportPrivateUsage] # the host registry's lookup + + def _select_strategy( + self, model: str, request_kwargs: Mapping[str, object], enable_tag_filtering: bool, match_any: bool + ) -> "TaggedPreRoutingStrategy[Any] | None": + """``_select_pre_routing_strategy``.""" + registries: Final = self._strategies.all() + if not any(registries): + return None + deployments: Final = self._deployments_for_request(model, request_kwargs) + registered_names: Final = tuple(dict.fromkeys(str(d.get("model_name")) for d in deployments)) or (model,) + candidates: Final = ( + tuple( # comprehension-ok: registry by registry, name by name, as the Python router searches + tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) + ) + ) + if not candidates: + return None + request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) + if request_tags: + for tagged in candidates: + if tagged.tags and is_valid_deployment_tag(list(tagged.tags), request_tags, match_any): + return tagged + for tagged in candidates: + if "default" in tagged.tags: + return tagged + request_scoped: Final = request_kwargs.get("enable_tag_filtering") is True + if ( + (enable_tag_filtering or request_scoped) + and all(tagged.tags for tagged in candidates) + and any(not is_strategy_marker(d) for d in deployments) + ): + return None + return candidates[0] + + def _selected_marker( + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] + ) -> Mapping[str, object] | None: + markers: Final = tuple( + deployment + for deployment in self._deployments_for_request(model, request_kwargs) + if str(as_dict(deployment.get("litellm_params")).get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + and "model" in as_dict(deployment.get("litellm_params")) + ) + tag_matched: Final = tuple( + deployment + for deployment in markers + if tuple(_tags_of(as_dict(deployment.get("litellm_params")))) == strategy_tags + ) + return tag_matched[0] if tag_matched else (markers[0] if markers else None) + + def _forwardable_marker_params( + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] + ) -> tuple[tuple[str, object], ...]: + marker: Final = self._selected_marker(model, strategy_tags, request_kwargs) + if marker is None: + return () + return tuple( + (key, value) + for key, value in as_dict(marker.get("litellm_params")).items() + if key not in _ALIAS_PARAMS_NEVER_FORWARDED + and key not in CustomPricingLiteLLMParams.model_fields + and value is not None + ) + + # -- the Claude Code session binding --------------------------------------------------------- + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _delete_binding(self, cache_key: str) -> None: + try: + await self._session_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", + e, + ) + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _get_binding(self, cache_key: str) -> object: + try: + if self._session_cache.redis_cache is None: + return await self._session_cache.async_get_cache(key=cache_key) + return await self._session_cache.redis_cache.async_get_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis + log_redis_failure( + verbose_router_logger, + logging.WARNING, + "Failed to read Claude Code session router binding; using the requested model", + e, + ) + return None + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _resolve_session_router( + self, + model: str, + registered: str, + request_kwargs: dict[str, object], # mutable-ok: see pre_route + fallback_depth: int, + enable_tag_filtering: bool, + match_any: bool, + ) -> str: + """``_resolve_claude_code_session_router``.""" + if is_native_compaction_call() or not self._strategies.any_registered(): + return registered + cache_key: Final = _session_router_cache_key(request_kwargs) + if cache_key is None: + return registered + depth: Final = fallback_depth or request_kwargs.get("fallback_depth") + if depth not in (None, 0): + return registered + agent_id: Final = _request_header(request_kwargs, "x-claude-code-agent-id") + if agent_id is not None: + bound_model: Final = await self._get_binding(cache_key) + if not isinstance(bound_model, str): + return registered + bound_registered: Final = self._hosted._core_get_model_from_alias(model=bound_model) or bound_model # pyright: ignore[reportPrivateUsage] # see _hook + if self._select_strategy(bound_registered, request_kwargs, enable_tag_filtering, match_any) is None: + await self._delete_binding(cache_key) + return registered + await self._session_cache.async_set_cache( + key=cache_key, value=bound_model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS + ) + _stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) + return bound_registered + if self._select_strategy(registered, request_kwargs, enable_tag_filtering, match_any) is None: + return registered + await self._session_cache.async_set_cache( + key=cache_key, value=model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS + ) + return registered + + # -- tier litellm_params --------------------------------------------------------------------- + + def _tier_params_the_target_accepts( + self, model: str, tier_params: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: + from litellm.router_rust.invoke.shaping import deployment_accepts_param + + deployments: Final = self._hosted._core_get_model_list(model_name=model) or () # pyright: ignore[reportPrivateUsage] # the host registry's lookup + if not deployments: + return tier_params + allowlisted: Final = _allowlist(tier_params) | _allowlist(request_kwargs) + candidates: Final = provider_rejectable_params(tier_params) - TIER_PARAMS_NEVER_DROPPED - allowlisted + unsupported: Final = frozenset( + param + for param in candidates + if not any(deployment_accepts_param(deployment, model, param, self._host) for deployment in deployments) + ) + if not unsupported: + return tier_params + verbose_router_logger.warning( + "litellm.router.py: dropping tier params %s for model=%s, no deployment behind it declares them", + ", ".join(sorted(unsupported)), + model, + ) + return MappingProxyType({key: value for key, value in tier_params.items() if key not in unsupported}) + + def _pin_tier_params( + self, + model: str, + tier_litellm_params: Mapping[str, object] | None, + request_kwargs: dict[str, object], # mutable-ok: see pre_route + responses_call: bool, + ) -> None: + """``_pin_tier_params_onto_request``.""" + if not tier_litellm_params: + return + reject_api_key_secret_reference(tier_litellm_params) + accepted: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs) + surface: Final = _tier_ceiling_under_the_surface_name(accepted, responses_call=responses_call) + _drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface) + request_kwargs.update(surface) + + +def as_dict(value: object) -> Mapping[str, object]: + if not isinstance(value, Mapping): + return {} + return cast(Mapping[str, object], value) # cast-ok: a str-keyed config or request dict + + +def _adaptive_cost(params: Mapping[str, object], info: Mapping[str, object]) -> object: + """A model's input price for the adaptive router: litellm_params win over model_info.""" + cost: Final = params.get("input_cost_per_token") + return cost if cost is not None else info.get("input_cost_per_token") + + +def _tier_members(tiers: Mapping[str, object]) -> tuple[object, ...]: + """Every model the tier pools name, once, in order.""" + pools: Final = ( + (pool,) if isinstance(pool, str) or not isinstance(pool, Iterable) else pool for pool in tiers.values() + ) + return tuple(dict.fromkeys(chain.from_iterable(pools))) + + +def _tags_of(params: Mapping[str, object]) -> Sequence[object]: + tags: Final = params.get("tags") + return tags if isinstance(tags, list | tuple) else () + + +def _allowlist(params: Mapping[str, object]) -> frozenset[str]: + declared: Final = params.get("allowed_openai_params") + if not isinstance(declared, list | tuple | set | frozenset): + return frozenset() + return frozenset(entry for entry in declared if isinstance(entry, str)) + + +__all__ = ["PreRouter", "PreRoutingStrategies", "is_strategy_marker"] diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py new file mode 100644 index 00000000000..1e2e0386fc8 --- /dev/null +++ b/litellm/router_rust/steps/selection.py @@ -0,0 +1,405 @@ +""" +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 LeastBusyLoggingHandler +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.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") -> Selector | None: + """``_build_strategy_selector`` without registering: ``None`` for simple-shuffle and every + strategy without a selector.""" + routing_args: Final = dict(args) + match strategy: + case "least-busy": + return LeastBusyLoggingHandler(router_cache=cache) + 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: + self._host: Final = host + self._cache: Final = cache + self._strategy: 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]) -> None: + """``routing_strategy_init``: drop the old default and override selectors, build the new one.""" + self._unregister([self._default, *self._overrides.values()]) + self._overrides = {} + self._strategy = normalize_strategy(strategy) + self._default = build_selector(self._strategy, args, self._cache) + 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: + return + try: + rebuilt: Final = build_selector(self._strategy, args, self._cache) + 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)) 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: + return self._default + + 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._strategy: + return self._default + with self._lock: + if strategy not in self._overrides: + self._overrides[strategy] = build_selector(strategy, {}, self._cache) + 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._strategy,) if self._strategy 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 +) -> object: + """``_select_deployment_async``.""" + match selector: + case LeastBusyLoggingHandler() if strategy == "least-busy": + return await selector.async_get_available_deployments(model_group=model, healthy_deployments=deployments) # 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 +) -> 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) # 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)) diff --git a/litellm/router_rust/surface/README.md b/litellm/router_rust/surface/README.md new file mode 100644 index 00000000000..71ef500c252 --- /dev/null +++ b/litellm/router_rust/surface/README.md @@ -0,0 +1,13 @@ +# 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 +- `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. diff --git a/litellm/router_rust/surface/__init__.py b/litellm/router_rust/surface/__init__.py new file mode 100644 index 00000000000..74704783ecb --- /dev/null +++ b/litellm/router_rust/surface/__init__.py @@ -0,0 +1 @@ +"""The Python ``Router`` API kept over router-core's state, for Python callers.""" diff --git a/litellm/router_rust/surface/alerting.py b/litellm/router_rust/surface/alerting.py new file mode 100644 index 00000000000..760eecb35f4 --- /dev/null +++ b/litellm/router_rust/surface/alerting.py @@ -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"] diff --git a/litellm/router_rust/surface/fallbacks.py b/litellm/router_rust/surface/fallbacks.py new file mode 100644 index 00000000000..140f57eff52 --- /dev/null +++ b/litellm/router_rust/surface/fallbacks.py @@ -0,0 +1,132 @@ +""" +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(request: Mapping[str, object]) -> Exception | None: + """The error ``run_async_fallback`` raises before its walk when a client-side request list + carries a dict target the 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 + + fallbacks: Final = request.get("fallbacks") + if not isinstance(fallbacks, list) or not check_non_standard_fallback_format(fallbacks=fallbacks): # pyright: ignore[reportUnknownArgumentType] # a request list as given + return None + for target in fallbacks: # pyright: ignore[reportUnknownVariableType] # see above + if isinstance(target, dict): + try: + reject_server_owned_wif_params(target) # pyright: ignore[reportUnknownArgumentType] # see above + 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", +] diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py new file mode 100644 index 00000000000..d4975e2a445 --- /dev/null +++ b/litellm/router_rust/surface/health.py @@ -0,0 +1,220 @@ +""" +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 typing import TYPE_CHECKING, Final + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +if TYPE_CHECKING: + from litellm.rust_bridge._native import ( # pyright: ignore[reportPrivateUsage] # the stub's handle + RouterCall, + RouterCoreHandle, + ) + + +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 cooled_in_running_loop(cooled: bool) -> bool: + """``set_cooldown_deployments`` schedules the cooldown event with ``asyncio.create_task`` + after writing the cooldown, so outside a running loop it raises ``RuntimeError`` then.""" + import asyncio + + if cooled: + asyncio.get_running_loop() + return cooled + + +def _configured_cooldown(model_info: object, litellm_params: Mapping[str, object]) -> float | None: + """``_first_present(model_info, litellm_params, key="cooldown_time")``.""" + for source in (as_str_mapping(model_info), litellm_params): + value = source.get("cooldown_time") if source is not None else None + if isinstance(value, int | float) and not isinstance(value, bool): + return float(value) + return None + + +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" + + def _drive_sync_steps(self, call: "RouterCall") -> Mapping[str, object]: + raise NotImplementedError + + async def _drive_cache_steps(self, call: "RouterCall") -> Mapping[str, object]: + 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_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) + return call is not None and cooled_in_running_loop(bool(self._drive_sync_steps(call).get("cooled"))) + + async def report_stream_failure(self, kwargs: Mapping[str, object]) -> 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``).""" + call: Final = self._failure_call(kwargs, counts_request=True) + if call is not None: + _ = await self._drive_cache_steps(call) + + def report_stream_failure_sync(self, kwargs: Mapping[str, object]) -> None: + """``report_stream_failure`` for a sync stream: a sync call's failure callback counts no request.""" + call: Final = self._failure_call(kwargs) + if call is not None: + _ = self._drive_sync_steps(call) + + def _failure_call(self, kwargs: Mapping[str, object], counts_request: bool = False) -> "RouterCall | None": + """The Rust call recording a logged failure, or None for one the Python router does not hold + against the deployment.""" + from litellm.litellm_core_utils import core_helpers + from litellm.router_rust.core import attempt_failure + from litellm.router_utils.cooldown_handlers import ( + 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, + ) + from litellm.types.router import Deployment + + exception: Final = kwargs.get("exception") + if is_failure_outside_selected_deployment(exception if isinstance(exception, BaseException) else None): + return None + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + model_info: Final = litellm_params.get("model_info", {}) + if isinstance(exception, Exception) and is_background_response_poll_not_found(exception, litellm_params): + return None + if is_video_lookup_not_found(exception, kwargs): + return None + exception_status: Final[object] = getattr(exception, "status_code", "") + if is_caller_timeout_408(kwargs, exception_status): # pyright: ignore[reportArgumentType] # the status as raised + return None + info: Final = as_str_mapping(model_info) + if info is None: + return None + deployment_id: Final = info.get("id") + if isinstance(deployment_id, str): + record: Final = self._handle.deployment(deployment_id) + deployment: Final = Deployment(**dict(record)) if record is not None else None # pyright: ignore[reportArgumentType] # a model_list entry + if is_caller_scoped_auth_failure(deployment, exception_status): # pyright: ignore[reportArgumentType] # see above + return None + if deployment_id is None: + return None + described: Final[Mapping[str, object]] = ( + attempt_failure(exception) if isinstance(exception, BaseException) else {"ok": False, "kind": "APIError"} + ) + failure: Final = {**described, "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 + 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, + ) + + 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.""" + 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])).get("key") + return key if isinstance(key, str) else None + + 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.""" + 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 + 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 {}) + 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, + ) + ) + + +__all__ = [ + "RustHealthCallbacks", + "cooled_in_running_loop", + "failure_status", +] diff --git a/litellm/router_rust/surface/health_routing.py b/litellm/router_rust/surface/health_routing.py new file mode 100644 index 00000000000..bcb863ff41b --- /dev/null +++ b/litellm/router_rust/surface/health_routing.py @@ -0,0 +1,157 @@ +""" +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 isinstance(groups, Sequence | set) else None, # pyright: ignore[reportUnknownArgumentType] # the group names as passed + ) + + 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, + "background_health_check_model_groups": sorted(groups) 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", +] diff --git a/litellm/router_rust/surface/indexes.py b/litellm/router_rust/surface/indexes.py new file mode 100644 index 00000000000..b711abf0c06 --- /dev/null +++ b/litellm/router_rust/surface/indexes.py @@ -0,0 +1,499 @@ +# 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 + +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) -> object: + """A view's contents as router-core takes them: plain JSON, a set as a list.""" + mapping: Final = _as_mapping(value) + if mapping or isinstance(value, Mapping): + return {str(key): _json(item) for key, item in mapping.items()} + if isinstance(value, list | tuple | set | frozenset): + return [_json(item) for item in cast(Iterable[object], value)] # cast-ok: any collection a caller wrote + 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 request is None: + return None + 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 + + 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()] diff --git a/litellm/router_rust/surface/members.py b/litellm/router_rust/surface/members.py new file mode 100644 index 00000000000..a71ed818f95 --- /dev/null +++ b/litellm/router_rust/surface/members.py @@ -0,0 +1,95 @@ +""" +``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 and ``print_deployment``. + +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.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 + + +class RustRegistryMembers: + """Mixed into ``RustCore``, which supplies its host registry and host.""" + + @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 print_deployment(self, deployment: Mapping[str, object]) -> Mapping[str, object]: + """A copy of the deployment with its api key masked (``print_deployment``).""" + try: + printed: Final = dict(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] # the deployment's params dict + elif isinstance(litellm_params, dict) and "api_key" in litellm_params: + litellm_params["api_key"] = litellm_params["api_key"][:2] + "*" * 10 + except Exception as e: + verbose_router_logger.debug("Error occurred while printing deployment - %s", e) + raise + else: + return printed diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py new file mode 100644 index 00000000000..f5b0d8d18f7 --- /dev/null +++ b/litellm/router_rust/surface/model_info.py @@ -0,0 +1,832 @@ +# 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, 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 +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 ModelInfo +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 + +_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", 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: Mapping[str, object] = call.next(outcome) # pyright: ignore[reportArgumentType, reportAssignmentType] # plain JSON answers; the next step + 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 {}) + + +# 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) + + 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 diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py new file mode 100644 index 00000000000..7b9dd57bb47 --- /dev/null +++ b/litellm/router_rust/surface/model_list.py @@ -0,0 +1,319 @@ +""" +``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 +from collections.abc import Callable, Iterable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, SupportsIndex, TypeAlias, cast, overload + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +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]] + + +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): + held = self._stable.get(key) if key is not None else None + kept = held is not None and held[0] == record + entry = held[1] if held is not None and kept else _HostedEntry(_record(record), self._entry_written) + if key is not None: + stable[key] = (held[0] if held is not None and kept else copy.deepcopy(record), 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),)) diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py new file mode 100644 index 00000000000..09d5170f820 --- /dev/null +++ b/litellm/router_rust/surface/registry.py @@ -0,0 +1,2037 @@ +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportAny=false +# (the ported core code reads untyped litellm helpers: PatternMatchRouter, Deployment.to_json, host.invoke) +""" +``HostRegistry``: the Rust core's live ``model_list`` and the indexes ``registry()`` reads. + +router-core builds the deployments (``litellm-rust/crates/router-core/src/build.rs``): it walks +the list, decides which values are ``os.environ/`` secret references, assigns the id of a +deployment that declares none and adds one deployment per ``organization``. What needs live Python +objects is a host step this registry answers, one deployment at a time (``_answer_build_step``): +``secrets`` (``get_secret``), ``validate`` (the PTU, access-window and strict param checks, and the +pricing each deployment writes into ``litellm.model_cost``), ``provider`` (``get_llm_provider``) +and ``added`` (pass-through credentials, pre-routing strategies, the patterns and indexes below). + +router-core holds the deployments. ``_core_model_list`` is ``HostedModelList`` +(litellm/router_rust/surface/model_list.py): router-core's records read back through the handle, every +write a changeset (a build serves what it added once it ends, a removal is a delete, any other +write a replace), each deployment one Python object while its record is unchanged. + +The indexes are router-core's too (``host_index.rs``), kept by the Python router's history rules: +the id, group and team position maps, ``model_names``, the only-growing ``deployment_names``, +the team public names, the provider default ids and the wildcard pattern routers. The registry +hands them out as the live views of litellm/router_rust/surface/indexes.py, whose writes are changesets. + +The lookups run in router-core (``lookup.rs``): the deployments by name (team public name first), +the model list with its alias, routing group and wildcard rows, the ids, names and access groups +it lists, the candidate ids for a route, the routable deployments, the credential deployment picks +and a team model's alternatives. Each answers with rows, positions in the list: this registry +reads its stored entries at them (``_rows``), so identity and in-place reads stay its own, and a +position router-core could not read (an index a caller wrote junk into) is read here and raises +as the Python router's read raised. What stays here needs live Python objects: the alias map and +routing groups a lookup is scoped by (``_list_scope``), the provider a wildcard request is tried +under, pydantic ``Deployment`` values, the credentials, the PTU ceilings, search tool support, a +team model's provider pick (``provider_for_generic_call``) and the SDK clients. + +Model info and group info are router-core's rules too (``model_info.rs``, ``group_info.rs``), +answered from this process through ``HostModelInfo`` (litellm/router_rust/surface/model_info.py). + +The attribute and method names follow the router core's (``_core_model_list``, +``_core_get_model_list`` ...), so ``litellm.router_api.RouterRegistry`` reads this object as its +router and every ``registry()`` query dispatches the way it does on the Python core. +""" + +import copy +import dataclasses +import os +import re +import threading +from collections.abc import Callable, Iterable, Mapping, Sequence +from functools import lru_cache, partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast + +from pydantic import TypeAdapter + +from litellm import exceptions as litellm_exceptions +from litellm import get_secret_str +from litellm._internal_context import service_target +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 DEFAULT_MAX_LRU_CACHE_SIZE +from litellm.litellm_core_utils.access_windows import access_windows_config_error +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import ( + declared_authenticating_provider, + is_registered_custom_provider, + resolve_model_and_provider_for_metadata, +) +from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter +from litellm.litellm_core_utils.ptu_pricing import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + declares_ptu, + declares_ptu_shares, + is_ptu_cost_attribution_enabled, + ptu_config_error, + ptu_identity_error, + ptu_terms, + zeroed_ptu_pricing, +) +from litellm.litellm_core_utils.ptu_shares import ptu_capacity_warning +from litellm.litellm_core_utils.routing_common_utils import ( + provider_for_generic_call, + resolve_model_group_alias, + team_may_use_deployment, + warn_on_ignored_model_capability_flags, + warn_on_provider_credential_mismatch, + with_admin_api_key_secret_reference, + without_model_capability_flags, +) +from litellm.litellm_core_utils.routing_groups import apply_routing_group_priority +from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS +from litellm.responses.bridge_check import keep_registered_mode +from litellm.router_api import RouterRegistry +from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET +from litellm.router_rust.surface.indexes import HostedIndexes +from litellm.router_rust.surface.model_info import HostModelInfo +from litellm.router_rust.surface.model_list import HostedModelList +from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model +from litellm.router_utils.endpoint_eligibility import warn_on_unenforced_endpoint_flag +from litellm.types.router import ( + VALID_LITELLM_ENVIRONMENTS, + CredentialLiteLLMParams, + Deployment, + DeploymentTypedDict, + LiteLLM_Params, + ModelGroupInfo, + RouterModelGroupAliasItem, + RoutingGroup, + RoutingPlugin, + TaggedPreRoutingStrategy, + validate_litellm_params_strictly, +) +from litellm.types.utils import ( + CustomPricingLiteLLMParams, + LlmProviders, + is_server_derived_pricing_key, + shared_backend_model_info, +) +from litellm.utils import ( + ModelInfoLookupMiss, + get_secret, + lookup_model_info_for_registration, + model_info_with_pricing_overrides, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.router_python_host import PythonHost + from litellm.router_rust.steps.callbacks import OptionalChecks + from litellm.router_rust.steps.prerouting import PreRoutingStrategies + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + from litellm.rust_bridge._native import RouterCoreHandle + +# The PTU share teams and the team ceilings read so far. +_PTUMemo: TypeAlias = tuple[ + frozenset[str], + dict[tuple[str, str], "PTUTeamCeiling | None"], +] +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +_CACHE_PRICING_FIELDS: Final = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", +) +_OFF_PEAK_BASE_RATE_FIELDS: Final = ( + "input_cost_per_token", + "input_cost_per_second", + "cost_per_second", + "tiered_pricing", +) + +# A stored model_list entry: the live dict the router core keeps. +UntypedEntry: TypeAlias = dict[str, object] +_TeamIndex: TypeAlias = dict[tuple[str, str], list[int]] + +# How a build's deployments join router-core's registry (`build::Install`). +BuildInstall: TypeAlias = Literal["replace", "append", "upsert"] +# Runs one router-core build (`RouterCoreHandle.build`): the entries, the path, whether it resolves +# secrets, how it installs, and the registry's answer to each step. ``RustCore`` supplies it. +BuildRunner: TypeAlias = Callable[ + [Sequence[object], str, bool, BuildInstall, Callable[[Mapping[str, object]], object]], None +] + + +@dataclasses.dataclass(slots=True) +class _BuildEntry: + """One ``set_model_list`` entry as its loop holds it (the router core pops ``model_name``, + ``litellm_params`` and ``model_info`` off the entry and keeps the rest), which router-core's + build steps name by index.""" + + extras: UntypedEntry + model_name: str + litellm_params: UntypedEntry + model_info: UntypedEntry + declared_id: str | None + # `organization` is a list, which router-core adds one deployment for each of. + per_organization: bool = False + # A retried entry keeps whether its id was declared twice in the list it came from. + duplicate: bool | None = None + + +@dataclasses.dataclass(slots=True) +class _Build: + """The state of one router-core build, which ``_answer_build_step`` reads and writes.""" + + path: Literal["create", "add"] + # `create`: the entries (a deep copy, which the build consumes), and each one once router-core + # reaches it. + raw: Sequence[UntypedEntry] = () + # Filled as the build reaches each entry. + entries: dict[int, _BuildEntry] = dataclasses.field(default_factory=dict) # mutable-ok: filled in step order + # `add`: the deployment to add, and the entry it is stored as. + deployment: Deployment | None = None + stored: UntypedEntry | None = None + # The deployment the build is on, its provider and its data-source values. + pending: Deployment | None = None + provider: str | None = None + data_source_values: tuple[tuple[int, str, str], ...] = () + # `add`: the deployment it added. + result: Deployment | None = None + + +def model_info_is_active_for_environment(model_info: Mapping[str, object] | None) -> bool: + """The environment-gating rule (``model_info_is_active_for_environment``): a deployment whose + model_info names ``supported_environments`` loads only where LITELLM_ENVIRONMENT is in it.""" + if model_info is None: + return True + supported_environments: Final = model_info.get("supported_environments") + if supported_environments is None: + return True + if not isinstance(supported_environments, list | tuple): + raise ValueError( # noqa: TRY004 # the core raises ValueError here + f"supported_environments must be a list of {VALID_LITELLM_ENVIRONMENTS}. " + f"but set as: {supported_environments} for model_info: {model_info}" + ) + litellm_environment: Final = get_secret_str(secret_name="LITELLM_ENVIRONMENT") + if litellm_environment is None: + raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env") + if litellm_environment not in VALID_LITELLM_ENVIRONMENTS: + raise ValueError( + f"LITELLM_ENVIRONMENT must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {litellm_environment}" + ) + for environment in supported_environments: # pyright: ignore[reportUnknownVariableType] # a configured list + if environment not in VALID_LITELLM_ENVIRONMENTS: + raise ValueError( + f"supported_environments must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {environment} " + f"for model_info: {model_info}" + ) + return litellm_environment in supported_environments + + +# -- pricing written into litellm.model_cost ----------------------------------------------------- + + +def _backend_model_info( + host: "PythonHost", backend_model: str, custom_llm_provider: str | None +) -> Mapping[str, object]: + """``get_model_info`` for the backend model; raises like it for an unmapped model.""" + info: Final = host.invoke("get_model_info")(model=backend_model, custom_llm_provider=custom_llm_provider) + return _MODEL_INFO_ADAPTER.validate_python(info) + + +def _backend_cost_entry( + host: "PythonHost", backend_model: str, custom_llm_provider: str | None +) -> Mapping[str, object] | None: + """The backend model's raw ``litellm.model_cost`` entry, or None when it has none.""" + try: + info: Final = _backend_model_info(host, backend_model, custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model + return None + entry: Final[object] = host.dynamic().model_cost.get(str(info.get("key") or "")) + return _MODEL_INFO_ADAPTER.validate_python(entry) if isinstance(entry, dict) else None + + +def _inherited_cache_pricing( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The cache prices a custom-priced entry inherits (``_inherit_builtin_cache_pricing``).""" + inherited_fields: Final = ( + _CACHE_PRICING_FIELDS + if model_info.get("cache_read_input_token_cost") is not None + else (*_CACHE_PRICING_FIELDS, "explicit_cache_read_input_token_cost") + ) + missing_fields: Final = tuple(field for field in inherited_fields if model_info.get(field) is None) + if not missing_fields: + return {} + try: + backend_info: Final = _backend_model_info(host, backend_model, provider) + except Exception: # noqa: BLE001 # an unmapped backend model contributes nothing + return {} + return {field: backend_info[field] for field in missing_fields if backend_info.get(field) is not None} + + +def _inherited_service_tier_pricing( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The service-tier rates a custom-priced entry inherits (``_inherit_builtin_service_tier_pricing``).""" + if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled(): + return {} + if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")): + return {} + backend_entry: Final = _backend_cost_entry(host, backend_model, provider) + if backend_entry is None: + return {} + return { + field: copy.deepcopy(value) + for field, value in backend_entry.items() + if field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES) and model_info.get(field) is None and value is not None + } + + +def _inherited_off_peak_base_rates( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The whole backend entry an off-peak-only entry copies (``_inherit_builtin_base_rates_for_off_peak``).""" + if not model_info.get("off_peak_pricing"): + return {} + if any(model_info.get(field) is not None for field in _OFF_PEAK_BASE_RATE_FIELDS): + return {} + backend_entry: Final = _backend_cost_entry(host, backend_model, provider) + if backend_entry is None: + return {} + return { + field: copy.deepcopy(value) + for field, value in backend_entry.items() + if model_info.get(field) is None and value is not None + } + + +def _inherited_tiered_output_rate( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The flat output rate a tier table without one inherits (``_inherit_builtin_tiered_output_rate``).""" + tiers: Final = model_info.get("tiered_pricing") + if not isinstance(tiers, list) or not tiers: + return {} + if model_info.get("output_cost_per_token") is not None: + return {} + if all(isinstance(tier, dict) and "output_cost_per_token" in tier for tier in tiers): # pyright: ignore[reportUnknownVariableType] # a configured tier table + return {} + try: + backend_info: Final = _backend_model_info(host, backend_model, provider) + except Exception: # noqa: BLE001 # an unmapped backend model contributes nothing + return {} + backend_rate: Final = backend_info.get("output_cost_per_token") + return {"output_cost_per_token": backend_rate} if backend_rate else {} + + +def _bedrock_profile(params: Mapping[str, object], model: str) -> bool: + return params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/") + + +def _search_tool_support(deployment: Mapping[str, object]) -> bool | None: + """A deployment's answer to Codex's client ``tool_search``.""" + from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request + from litellm.utils import ProviderConfigManager + + configured: Final = _model_info_of(deployment).get("supports_search_tool") + if isinstance(configured, bool): + return configured + params: Final = _params_of(deployment) + model: Final = params.get("model") + provider: Final = params.get("custom_llm_provider") + if not isinstance(model, str): + return None + try: + resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( + model=model, custom_llm_provider=provider if isinstance(provider, str) else None + ) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=resolved_model, provider=LlmProviders(resolved_provider) + ) + except (litellm_exceptions.BadRequestError, ValueError): + return None + return True if chat_config_builds_anthropic_request(chat_config) else None + + +def cost_map_backend_model(deployment: Deployment) -> str: + """The model a deployment's prices are inherited from (``_cost_map_backend_model``).""" + model_info_base_model: Final = deployment.model_info.base_model + if isinstance(model_info_base_model, str) and model_info_base_model: + return model_info_base_model + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + params_base_model: Final = params.get("base_model") + if isinstance(params_base_model, str) and params_base_model: + return params_base_model + model: Final = deployment.litellm_params.model + invoked_model: Final = params.get("model_id") + if isinstance(invoked_model, str) and invoked_model and _bedrock_profile(params, model): + profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" + profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") + if not isinstance(profile_info, ModelInfoLookupMiss): + return profile_model + return model + + +def deployment_model_cost_payload(host: "PythonHost", deployment: Deployment) -> Mapping[str, object]: + """The ``model_info`` a deployment contributes to ``litellm.model_cost`` + (``_deployment_model_cost_payload``): its model_info, its custom pricing, then the inherited rates.""" + declared: Final = dict(_MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True))) + custom_pricing: Final = { + field: value + for field in CustomPricingLiteLLMParams.model_fields + if (value := deployment.litellm_params.get(field)) is not None + } + backend_model: Final = cost_map_backend_model(deployment) + provider: Final = deployment.model_info.pricing_provider or deployment.litellm_params.custom_llm_provider + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + priced: Final = {**declared, **custom_pricing} + profile_pricing: Final[Mapping[str, object]] = ( + _bedrock_profile_pricing(backend_model, priced) + if params.get("model_id") + and backend_model != deployment.litellm_params.model + and _bedrock_profile(params, deployment.litellm_params.model) + else {} + ) + with_profile: Final = {**priced, **profile_pricing} + off_peak: Final = {**with_profile, **_inherited_off_peak_base_rates(host, with_profile, backend_model, provider)} + cached: Final = ( + {**off_peak, **_inherited_cache_pricing(host, off_peak, backend_model, provider)} + if off_peak.get("input_cost_per_token") is not None + else off_peak + ) + tiers: Final = {**cached, **_inherited_service_tier_pricing(host, cached, backend_model, provider)} + return {**tiers, **_inherited_tiered_output_rate(host, tiers, backend_model, provider)} + + +def _bedrock_profile_pricing(backend_model: str, priced: Mapping[str, object]) -> Mapping[str, object]: + profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock") + if isinstance(profile_info, ModelInfoLookupMiss): + return {} + explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(priced) + priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info) # pyright: ignore[reportArgumentType] # the explicit model_info, as the core passes it + return {key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)} + + +def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: + """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under.""" + backend_key: Final = model if custom_llm_provider is None else f"{custom_llm_provider}/{model}" + if "responses/" in backend_key: + return (backend_key, backend_key.replace("responses/", "")) + return (backend_key,) + + +def register_deployment_in_model_cost( + host: "PythonHost", + *, + model_id: str | None, + model_info: Mapping[str, object], + model: str, + custom_llm_provider: str | None, + pricing_provider: str | None = None, +) -> None: + """Write a deployment's metadata into ``litellm.model_cost`` (``_register_deployment_in_model_cost``): + the whole entry under its id, the schema fields under its backend model key(s).""" + registered: Final = ( + {key: value for key, value in model_info.items() if key not in CustomPricingLiteLLMParams.model_fields} + if classify_strategy_router_model(model) is not None + else dict(model_info) + ) + register_model: Final = host.invoke("register_model") + if model_id is not None: + register_model( + model_cost={model_id: registered}, + persist_across_reloads=False, + warning_display_name=model, + custom_llm_provider=pricing_provider or custom_llm_provider, + add_to_provider_model_sets=False, + ) + shared_model_info: Final = shared_backend_model_info(registered) + register_model( + model_cost={ + key: keep_registered_mode(key, shared_model_info) + for key in _backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) + }, + persist_across_reloads=False, + ) + + +def register_deployment_pricing(host: "PythonHost", deployment: Deployment) -> None: + """Register a deployment's custom and inherited pricing (``_register_deployment_pricing``).""" + register_deployment_in_model_cost( + host, + model_id=deployment.model_info.id, + model_info=deployment_model_cost_payload(host, deployment), + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + pricing_provider=deployment.model_info.pricing_provider, + ) + + +# -- the registry -------------------------------------------------------------------------------- + + +def _model_info_of(entry: Mapping[str, object]) -> Mapping[str, object]: + model_info: Final[object] = entry.get("model_info") + return _MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else {} + + +def _params_of(entry: Mapping[str, object]) -> Mapping[str, object]: + params: Final[object] = entry.get("litellm_params") + return _LITELLM_PARAMS_ADAPTER.validate_python(params) if isinstance(params, dict) else {} + + +def _renamed(entry: UntypedEntry, model_name: str) -> UntypedEntry: + return {**entry, "model_name": model_name} + + +def _as_routing_group_row(deployment: Mapping[str, object]) -> UntypedEntry: + """A member deployment under its group's name, without the member's ``access_groups``.""" + model_info: Final = {k: v for k, v in _model_info_of(deployment).items() if k != "access_groups"} + return {**deployment, "model_info": model_info} + + +def _ref_env_only(ref: object) -> bool: + return bool((as_str_mapping(ref) or {}).get("env_only")) + + +def _duplicate_of(step: Mapping[str, object], entry: _BuildEntry) -> bool: + """Whether the entry's declared id is declared twice in its list: a retried entry keeps the + answer its first build had.""" + return entry.duplicate if entry.duplicate is not None else step.get("duplicate_id") is True + + +class HostRegistry(HostModelInfo): + """The live model_list and its indexes, kept the way the Python router core keeps them.""" + + def __init__( + self, + host: "PythonHost", + *, + model_group_alias: "dict[str, str | RouterModelGroupAliasItem]", # mutable-ok: the router's live alias map + ignore_invalid_deployments: bool, + on_change: Callable[[], None], + ) -> None: + self._host: Final = host + self._on_change: Final = on_change + self.ignore_invalid_deployments: bool = ignore_invalid_deployments + self._registry_version: int = 0 + self._deployment_index_lock: Final = threading.RLock() + self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip + # The live model_list: router-core's deployments read back, every write a changeset. + self._core_model_list: Final[HostedModelList] = HostedModelList(self._live_handle, self._hand_changes) + # The indexes the registry hands out: router-core's, read back, every write a changeset. + self._indexes: Final = HostedIndexes(self._live_handle, self._hand_changes) + self._core_deployment_latency_map: dict[str, float] = {} # mutable-ok: the live latency map + self._access_groups_cache: Mapping[str, Sequence[str]] | None = None + # The routing groups by name; RustCore sets them from its resolved settings. + self._core_routing_groups: dict[str, RoutingGroup] = {} # mutable-ok: the routing_groups attribute + self._routing_group_rows: tuple[UntypedEntry, ...] | None = None + self._provider_unresolved_deployments: tuple[Callable[[], None], ...] = () + # Runs router-core's deployment build; RustCore sets it before the first build. + self.build_runner: BuildRunner | None = None + # Hands router-core a write its build did not make (a removal, a swapped list) as a + # changeset, and reads its current handle; RustCore sets both before the first build. + self.changes_runner: Callable[[Mapping[str, object]], None] | None = None + self.handle_lookup: Callable[[], RouterCoreHandle | None] | None = None + # The pre-routing strategy registries (``registry().plugins``) and the hooks that keep them; + # RustCore installs its ``PreRoutingStrategies`` (litellm/router_rust/steps/prerouting.py). + self._core_auto_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: the live registry + self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self._core_quality_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self.strategies: PreRoutingStrategies | None = None + # The router's optional pre-call checks (deployment budgets, prompt caching pins); + # RustCore installs its ``OptionalChecks`` (litellm/router_rust/steps/callbacks.py). + self.checks: OptionalChecks | None = None + # The Rust core's resolution for a request (`_core_get_routable_deployments`), set by the core. + self.routable_lookup: Callable[[str, Mapping[str, object]], Sequence[Mapping[str, object]]] | None = None + # A deployment's cached `max_parallel_requests` limit (litellm/router_rust/invoke/slots.py), set by the core. + self.max_parallel_requests_lookup: Callable[[Mapping[str, object], Mapping[str, object]], object] | None = None + # `plugins("routing")`: the router's routing plugins, set by the core. + self._core_routing_plugins: list[RoutingPlugin] = [] # mutable-ok: the Python core hands out its live list + # `zero_cost()`: the per-deployment zero-cost lookup the auth checks fill, emptied with the + # group info caches. + self._zero_cost_cache: dict[str, bool] = {} # mutable-ok: the auth checks write into it + # The routing cache usage reads come from; RustCore sets it to the one its Rust core writes. + self._core_cache: DualCache = DualCache() + self._discovered_model_info_cache: InMemoryCache = InMemoryCache( + max_size_in_memory=1, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS + ) + # The PTU share teams and team ceilings, until the model list next changes. + self._ptu_cache: _PTUMemo | None = None + self._cached_get_model_group_info: Callable[[str], ModelGroupInfo | None] = lru_cache( # pyright: ignore[reportIncompatibleMethodOverride] # the per-router cache shadows the uncached method, as the core keeps it + maxsize=DEFAULT_MAX_LRU_CACHE_SIZE + )(self._core_get_model_group_info) + self._core_cached_deployment_model_info: Callable[[str, str], object] = lru_cache( + maxsize=DEFAULT_MAX_LRU_CACHE_SIZE + )(self._core_get_deployment_model_info) + + def __getattr__(self, name: str) -> object: + from litellm.router_rust.core import RustCoreUnsupported + + raise RustCoreUnsupported(f"registry().{name.removeprefix('_core_')} is not served by the Rust router core") + + # -- building --------------------------------------------------------------------------------- + + def build(self, model_list: Sequence[object] | None) -> None: + """The constructor's registry: ``set_model_list`` and the latency map, as the core builds them.""" + if model_list is None: + return + self._discovered_model_info_cache = InMemoryCache( + max_size_in_memory=max(len(model_list), 1), default_ttl=2 * MODEL_INFO_REFRESH_SECONDS + ) + self.set_model_list(model_list) + for entry in model_list: + params = _params_of(_LITELLM_PARAMS_ADAPTER.validate_python(entry)) + if "model" in params: + self._core_deployment_latency_map[str(params["model"])] = 0 + + async def arefresh_model_info(self, client: "AsyncHTTPHandler | None") -> None: + """``_core_arefresh_model_info``: token limits the served OpenAI-compatible deployments + advertise on their models endpoint, read with a few workers over a snapshot of the list.""" + import asyncio + + from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_CONCURRENCY + + deployments: Final = iter(tuple(self._core_model_list)) + + async def refresh_worker() -> None: + for raw_deployment in deployments: + try: + await self._arefresh_deployment_model_info(raw_deployment, client=client) + except Exception: # noqa: BLE001 # one invalid deployment must not prevent refreshing the others + verbose_router_logger.debug("Could not refresh deployment model info") + + _ = await asyncio.gather(*(refresh_worker() for _ in range(MODEL_INFO_REFRESH_CONCURRENCY))) + self._invalidate_model_group_info_cache() + + async def _arefresh_deployment_model_info( + self, raw_deployment: Mapping[str, object], *, client: "AsyncHTTPHandler | None" + ) -> None: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.llms.openai_like.model_info import ( + MODEL_INFO_DISCOVERY_PROVIDERS, + get_openai_compatible_model_info, + ) + from litellm.types.router import DiscoveredDeploymentModelInfo + from litellm.utils import aresolve_api_key_secret_reference + + deployment: Final = Deployment.model_validate(raw_deployment) + params: Final = LiteLLM_Params.model_validate( + MappingProxyType( + { + **deployment.litellm_params.model_dump(exclude_none=True), + **( + self._core_get_deployment_credentials_with_provider(deployment.model_info.id or "") + or MappingProxyType({}) + ), + } + ) + ) + model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")( # pyright: ignore[reportAny] # litellm's untyped tuple + model=params.model, litellm_params=params + ) + if provider not in MODEL_INFO_DISCOVERY_PROVIDERS: + return + if api_base is None or "*" in model or params.get("use_clientside_credentials"): # pyright: ignore[reportUnknownMemberType] # the params' untyped get + return + api_key: Final = await aresolve_api_key_secret_reference( + params.api_key or dynamic_api_key, # pyright: ignore[reportAny] # see above + model=model, # pyright: ignore[reportAny] # see above + llm_provider=provider, # pyright: ignore[reportAny] # see above + upstream=MappingProxyType({"api_base": params.api_base}), + ) + headers: Final = TypeAdapter(Mapping[str, str]).validate_python( + params.get("extra_headers") or params.get("headers") or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # see above + ) + auth_headers: Final = ( + MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({}) + ) + limits: Final = await get_openai_compatible_model_info( + model=model, # pyright: ignore[reportAny] # see above + api_base=api_base, # pyright: ignore[reportAny] # see above + headers=MappingProxyType( + {**auth_headers, **MappingProxyType({key.lower(): value for key, value in headers.items()})} + ), + client=client or get_async_httpx_client(llm_provider=LlmProviders.OPENAI), + cache=self._core_cache.in_memory_cache, + ) + model_id: Final = deployment.model_info.id + if not limits or model_id is None or self._core_get_model_info(model_id) is not raw_deployment: + return + self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1) + self._discovered_model_info_cache.delete_cache(model_id) + self._discovered_model_info_cache.set_cache( + model_id, + DiscoveredDeploymentModelInfo(deployment=raw_deployment, limits=limits), # pyright: ignore[reportArgumentType] # the model_list entry, as the core keeps it + ) + self._invalidate_model_group_info_cache() + + def _invalidate_model_group_info_cache(self) -> None: + self._ptu_cache = None + self._cached_get_model_group_info.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache + self._core_cached_deployment_model_info.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache + self._zero_cost_cache.clear() + self._routing_group_rows = None + + def set_routing_groups(self, groups: Sequence[RoutingGroup]) -> None: + """``_init_routing_groups``' registry side: the groups by name, and the caches they feed.""" + self._core_routing_groups = {group.group_name: group for group in groups} + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + + def _invalidate_access_groups_cache(self) -> None: + self._access_groups_cache = None + self._registry_version += 1 + + def set_model_list(self, model_list: Sequence[object]) -> None: + """Replace every deployment (``set_model_list``).""" + self._discovered_model_info_cache.flush_cache() + original_model_list: Final[list[UntypedEntry]] = copy.deepcopy(list(model_list)) # pyright: ignore[reportAssignmentType] # model_list entries # mutable-ok: a deep copy the build consumes + if self.strategies is not None: + self.strategies.reset() + self._provider_unresolved_deployments = () + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + ptu_declared: Final = tuple( + str(entry.get("model_name")) + for entry in original_model_list + if (ptu_info := _model_info_of(entry)).get("db_model") is not True and declares_ptu(ptu_info) + ) + if ptu_declared and not is_ptu_cost_attribution_enabled(): + verbose_router_logger.warning( + "PTU fields are set on config.yaml deployment(s) %s, but PTU cost attribution is disabled, so no " + "flat cost accrues and this traffic is billed per token. Set %s=True to enable it", + ", ".join(ptu_declared), + PTU_COST_ATTRIBUTION_ENV_VAR, + ) + # router-core walks the list (one deployment per organization, secrets resolved first, + # ids assigned), asking this registry about each deployment in turn; it starts the list + # and its id, name and team indexes over, and takes the given names as `model_names`. + self._run_build( + _Build(path="create", raw=original_model_list), + original_model_list, + resolve_secrets=True, + install="replace", + ) + if self.strategies is not None: + self.strategies.finalize() + self._on_change() + + # -- the build's steps (router-core's `build::build`) ------------------------------------------ + + def _run_build( + self, build: _Build, entries: Sequence[object], *, resolve_secrets: bool, install: BuildInstall + ) -> None: + """Have router-core build ``entries``, answering its steps from this registry.""" + if self.build_runner is None: + raise RuntimeError("the host registry builds deployments through its Rust core") + self._core_model_list.begin_build(replacing=install == "replace") + try: + self.build_runner(entries, build.path, resolve_secrets, install, partial(self._answer_shown, build)) + finally: + self._core_model_list.end_build() + + def _answer_shown(self, build: _Build, step: Mapping[str, object]) -> object: + """A build step's answer; a deployment the ``added`` step hands router-core shows in the + live list from then on (router-core serves it when the build ends).""" + answer: Final = self._answer_build_step(build, step) + stored: Final = as_str_mapping(answer).get("entry") if step["kind"] == "added" else None # pyright: ignore[reportOptionalMemberAccess] # an added step answers a mapping + if isinstance(stored, Mapping): + self._core_model_list.built(cast(Mapping[str, object], stored)) # cast-ok: the entry it stores + return answer + + def _answer_build_step(self, build: _Build, step: Mapping[str, object]) -> object: + """The answer to one of router-core's build steps. A ``set_model_list`` deployment that + fails is left out when the router ignores invalid deployments (``_create_deployment``), + and retried on the next price data reload when its provider did not resolve.""" + kind: Final = step["kind"] + index: Final = int(str(step["index"])) + refs: Final = step.get("refs") + if kind == "secrets" and isinstance(refs, list) and not any(_ref_env_only(ref) for ref in refs): + return self._build_secrets(build, index, refs) + if kind == "validate": + # A new deployment: the entry is popped where the loop pops it, outside the error + # handling, and nothing of the last deployment carries over. + build.pending = None # rebind-ok: the build's state, which each step advances + build.provider = None # rebind-ok: the build's state, which each step advances + build.data_source_values = () # rebind-ok: the build's state, which each step advances + if build.path == "create": + _ = self._build_entry(build, index) + try: + if kind == "secrets": + return self._build_data_sources(build, refs if isinstance(refs, list) else []) + if kind == "validate": + return self._build_validate(build, index, step) + if kind == "provider": + build.provider = self._build_provider(build) # rebind-ok: the build's state, which each step advances + return {"provider": build.provider} + return self._build_added(build) + except Exception as e: + if build.path != "create" or not self.ignore_invalid_deployments: + raise + if isinstance(e, litellm_exceptions.BadRequestError): + entry: Final = build.entries[index] + retried: Final = dataclasses.replace(entry, duplicate=_duplicate_of(step, entry)) + self._provider_unresolved_deployments = ( + *self._provider_unresolved_deployments, + partial(self._retry_create, retried), + ) + verbose_router_logger.exception( + "Error creating deployment: %s, ignoring and continuing with other deployments.", e + ) + return {"invalid": True} + + def _build_entry(self, build: _Build, index: int) -> _BuildEntry: + """Entry ``index`` as ``set_model_list``'s loop holds it, popped off the first time a step + names it (so a malformed entry raises where the loop raises).""" + held: Final = build.entries.get(index) + if held is not None: + return held + model: Final = build.raw[index] + model_name: Final = model.pop("model_name") + raw_params: Final = model.pop("litellm_params") + litellm_params: Final[UntypedEntry] = raw_params if isinstance(raw_params, dict) else {} # pyright: ignore[reportAssignmentType] # configured params + model_info: Final = model.pop("model_info", {}) + info_dict: Final[UntypedEntry] = model_info if isinstance(model_info, dict) else {} # pyright: ignore[reportAssignmentType] # a configured model_info + entry: Final = _BuildEntry( + extras=model, + model_name=str(model_name), + litellm_params=litellm_params, + model_info=info_dict, + declared_id=None if info_dict.get("id") is None else str(info_dict["id"]), + per_organization=isinstance(litellm_params.get("organization"), list), + ) + build.entries[index] = entry # rebind-ok: the build's state, which each step advances + return entry + + def _build_secrets(self, build: _Build, index: int, refs: Sequence[object]) -> Sequence[object]: + """``set_model_list``'s ``get_secret`` for each top-level ``os.environ/`` value router-core + found, written into the entry's params.""" + entry: Final = self._build_entry(build, index) + references: Final = tuple(as_str_mapping(ref) or {} for ref in refs) + values: Final = tuple(get_secret(str(reference["name"])) for reference in references) + entry.litellm_params.update( + {str(reference["key"]): value for reference, value in zip(references, values, strict=True)} + ) + return values + + def _build_data_sources(self, build: _Build, refs: Sequence[object]) -> Sequence[object]: + """``_add_deployment``'s environment reads for ``dataSources[].parameters``, applied to the + deployment where it applied them (``_build_added``).""" + references: Final = tuple(as_str_mapping(ref) or {} for ref in refs) + build.data_source_values = tuple( # rebind-ok: the build's state, which each step advances + ( + int(str(reference["source"])), + str(reference["key"]), + os.environ.get(str(reference["name"]).replace("os.environ/", ""), ""), + ) + for reference in references + ) + return tuple(value for _, _, value in build.data_source_values) + + def _build_validate(self, build: _Build, index: int, step: Mapping[str, object]) -> Mapping[str, object]: + if build.path == "add": + return self._validate_added_deployment(build) + entry: Final = self._build_entry(build, index) + routed: Final = as_str_mapping(step["entry"]) or {} + routed_info: Final = as_str_mapping(routed.get("model_info")) or {} + if "id" not in entry.model_info: + entry.model_info["id"] = routed_info.get("id") + if entry.per_organization: + entry.litellm_params["organization"] = (as_str_mapping(routed.get("litellm_params")) or {}).get( + "organization" + ) + duplicate: Final = _duplicate_of(step, entry) + deployment: Final = self._validated_deployment(entry, taken=duplicate) + if deployment is None: + return {"skip": True} + served: Final = with_admin_api_key_secret_reference(deployment) + build.pending = served # rebind-ok: the build's state, which each step advances + return {"entry": build.pending.to_json(exclude_none=True)} + + def _validated_deployment(self, entry: _BuildEntry, *, taken: bool) -> Deployment | None: + """``_create_deployment`` up to ``_add_deployment``: the config checks, the deployment, its + pricing, and None when the environment gate leaves it out.""" + model_name: Final = entry.model_name + litellm_params: Final = entry.litellm_params + model_info: Final = entry.model_info + config_sourced: Final = model_info.get("db_model") is not True + identity_error: Final = ( + ptu_identity_error( + declared_id=entry.declared_id, + taken=taken, + current_id=model_info.get("id"), # pyright: ignore[reportArgumentType] # the configured id + model_name=model_name, + ) + if config_sourced and ptu_terms(model_info) is not None + else None + ) + ptu_error: Final = ( + (ptu_config_error(model_info, model_name=model_name) or identity_error) if config_sourced else None + ) + if ptu_error is not None and (declares_ptu_shares(model_info) or is_ptu_cost_attribution_enabled()): + raise ValueError(ptu_error) + access_windows_error: Final = access_windows_config_error(model_info, model_name=model_name) + if access_windows_error is not None: + raise ValueError(access_windows_error) + capacity_warning: Final = ( + ptu_capacity_warning( + model_name, MappingProxyType({"model_info": model_info, "litellm_params": litellm_params}) + ) + if is_ptu_cost_attribution_enabled() + else None + ) + if capacity_warning is not None: + verbose_router_logger.warning(capacity_warning) + zeroed_pricing: Final = zeroed_ptu_pricing(model_info, litellm_params) if config_sourced else None + merged_params: Final[Mapping[str, object]] = ( + litellm_params if zeroed_pricing is None else MappingProxyType({**litellm_params, **zeroed_pricing}) + ) + params: Final[LiteLLM_Params] = validate_litellm_params_strictly(LiteLLM_Params, dict(**merged_params)) + declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params) + warn_on_provider_credential_mismatch(model_name=model_name, litellm_params=declared_params) + warn_on_ignored_model_capability_flags(model_name=model_name, litellm_params=declared_params) + warn_on_unenforced_endpoint_flag(model_name=model_name, model_info=model_info) + deployment: Final = without_model_capability_flags( + Deployment(**entry.extras, model_name=model_name, litellm_params=params, model_info=model_info) # pyright: ignore[reportArgumentType] # the entry's extra fields, as the core passes them + ) + register_deployment_pricing(self._host, deployment) + if model_info_is_active_for_environment(model_info=deployment.model_info) is not True: # pyright: ignore[reportArgumentType] # ModelInfo reads like a mapping + verbose_router_logger.warning( + "Ignoring deployment %s as it is not active for environment %s", + deployment.model_name, + deployment.model_info["supported_environments"], + ) + return None + for pattern in deployment.litellm_params.get("tag_regex") or []: + try: + re.compile(pattern) + except re.error as exc: + raise ValueError( + f"Invalid regex in tag_regex for model '{deployment.model_name}': {pattern!r} — {exc}" + ) from exc + return deployment + + def _validate_added_deployment(self, build: _Build) -> Mapping[str, object]: + """``add_deployment`` up to ``_add_deployment``: an id the registry holds is skipped.""" + deployment: Final = build.deployment + if deployment is None: + raise RuntimeError("an add build holds its deployment") + if deployment.model_info.id and self._core_has_model_id(deployment.model_info.id): + return {"skip": True} + declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True) + warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params) + warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params) + warn_on_unenforced_endpoint_flag( + model_name=deployment.model_name, model_info=deployment.model_info.model_dump() + ) + served: Final = with_admin_api_key_secret_reference(without_model_capability_flags(deployment)) + build.pending = served # rebind-ok: the build's state, which each step advances + build.stored = served.to_json(exclude_none=True) # rebind-ok: the build's state, which each step advances + return {"entry": build.stored} + + def _build_provider(self, build: _Build) -> str | None: + """``_add_deployment``'s provider: None for a prompt-management model, else the provider + ``get_llm_provider`` resolves, refused when litellm does not serve it.""" + deployment: Final = build.pending + if deployment is None: + raise RuntimeError("a provider step follows a validated deployment") + litellm_model: Final = deployment.litellm_params.model + if isinstance(deployment.litellm_params.drop_params, str): + verbose_router_logger.warning( + "model=%s drop_params=%r is not a flag value, treating it as unset", + deployment.model_name, + deployment.litellm_params.drop_params, + ) + dynamic: Final = self._host.dynamic() + if "/" in litellm_model and litellm_model.split("/")[0] in dynamic.known_custom_logger_compatible_callbacks: + return None + if declared_authenticating_provider(litellm_model, deployment.litellm_params.custom_llm_provider) is not None: + _, provider = self._model_group_llm_provider(deployment.litellm_params) + return provider + _, custom_llm_provider, _, _ = self._host.invoke("get_llm_provider")( + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), + api_base=deployment.litellm_params.api_base, + litellm_params=deployment.litellm_params, + ) + if ( + custom_llm_provider not in dynamic.provider_list + and not JSONProviderRegistry.exists(custom_llm_provider) + and not is_registered_custom_provider(custom_llm_provider) + ): + raise Exception(f"Unsupported provider - {custom_llm_provider}") # noqa: TRY002 # the core raises a plain Exception + return custom_llm_provider + + def _build_added(self, build: _Build) -> Mapping[str, object]: + """The rest of ``_add_deployment``: the entry router-core stores, which joins the list and + the indexes (``host_index.rs``) when the step returns.""" + deployment: Final = build.pending + if deployment is None: + raise RuntimeError("an added step follows a validated deployment") + self._register_deployment(deployment, build.provider, build.data_source_values) + build.pending = None # rebind-ok: the build's state, which each step advances + build.data_source_values = () # rebind-ok: the build's state, which each step advances + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + if build.path == "create": + return {"entry": deployment.to_json(exclude_none=True)} + entry: Final = build.stored if build.stored is not None else deployment.to_json(exclude_none=True) + register_deployment_pricing(self._host, deployment) + build.result = deployment # rebind-ok: the build's state, which each step advances + return {"entry": entry} + + def _retry_create(self, entry: _BuildEntry) -> None: + """A ``set_model_list`` deployment left out when its provider did not resolve, built again + and added at the end.""" + build: Final = _Build(path="create") + build.entries[0] = entry + self._run_build( + build, + [ + { + **entry.extras, + "model_name": entry.model_name, + "litellm_params": entry.litellm_params, + "model_info": entry.model_info, + } + ], + resolve_secrets=False, + install="append", + ) + + @staticmethod + def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]: + try: + return resolve_model_and_provider_for_metadata( + model=litellm_params.model, custom_llm_provider=litellm_params.custom_llm_provider + ) + except litellm_exceptions.BadRequestError as e: + verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) + return "", "" + + def _register_deployment( + self, + deployment: Deployment, + custom_llm_provider: str | None, + data_source_values: Sequence[tuple[int, str, str]], + ) -> None: + """``_add_deployment`` after the provider: its data-source values, pass-through + credentials and pre-routing strategies (router-core registers its names and patterns).""" + if deployment.litellm_params.rpm is None and getattr(deployment, "rpm", None) is not None: + deployment.litellm_params.rpm = getattr(deployment, "rpm") # noqa: B009 # an extra field of the entry # rebind-ok: the core sets it on the deployment it adds + if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None: + deployment.litellm_params.tpm = getattr(deployment, "tpm") # noqa: B009 # an extra field of the entry # rebind-ok: see above + data_sources: Final = deployment.litellm_params.get("dataSources", []) or [] + for source, key, value in data_source_values: + data_sources[source]["parameters"][key] = value + if custom_llm_provider is not None: + self._initialize_deployment_for_pass_through(deployment, custom_llm_provider) + if self.strategies is not None: + self.strategies.register(deployment) + + def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str) -> None: + """Register vertex credentials for a pass-through deployment (``_initialize_deployment_for_pass_through``).""" + if deployment.litellm_params.use_in_pass_through is not True or custom_llm_provider != "vertex_ai": + return + credential_name: Final = deployment.litellm_params.litellm_credential_name + credential_values: Final = ( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} + ) + project: Final = credential_values.get("vertex_project") or deployment.litellm_params.vertex_project + location: Final = credential_values.get("vertex_location") or deployment.litellm_params.vertex_location + credentials: Final = credential_values.get("vertex_credentials") or deployment.litellm_params.vertex_credentials + if project is None or location is None: + raise ValueError( + "vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints." + ) + from litellm.router_host import router_host + + router_host().add_vertex_pass_through_credentials( + project_id=str(project), + location=str(location), + vertex_credentials=credentials, # pyright: ignore[reportArgumentType] # the configured credentials + ) + + # -- writes ----------------------------------------------------------------------------------- + + def add_deployment(self, deployment: Deployment) -> Deployment | None: + """Add one deployment unless its id exists (``add_deployment``), built by router-core.""" + build: Final = _Build(path="add", deployment=deployment) + self._run_build(build, [deployment.to_json(exclude_none=True)], resolve_secrets=False, install="upsert") + if build.result is None: + return None + if self.checks is not None: + self.checks.sync_deployment_budget(build.result) + self._on_change() + return build.result + + def _core_upsert_deployment(self, deployment: Deployment) -> Deployment | None: + """Add or replace one deployment (``_upsert_deployment``).""" + model_id: Final = deployment.model_info.id or "" + on_router: Final = self._core_get_deployment(model_id=model_id) + served: Final = without_model_capability_flags(deployment) + try: + if on_router is not None: + if ( + deployment.model_name == on_router.model_name + and served.litellm_params == on_router.litellm_params + and deployment.model_info == on_router.model_info + ): + return None + with self._deployment_index_lock: + removal_idx: Final = self.model_id_to_deployment_index_map.get(model_id) + if removal_idx is not None: + _ = self._core_model_list[removal_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does + self._remove(model_id) + if self.strategies is not None: + self.strategies.release(on_router) + self._on_change() + self.add_deployment(deployment=deployment) + if self.strategies is not None: + self.strategies.upserted(deployment, on_router) + except Exception as e: + if not self.ignore_invalid_deployments: + raise + verbose_router_logger.warning( + "Error upserting deployment %s (id=%s): %s. Dropping it and continuing with other deployments.", + deployment.model_name, + deployment.model_info.id, + e, + ) + self._restore_deployment_after_failed_upsert(on_router, model_id) + return None + else: + return deployment + + def _restore_deployment_after_failed_upsert(self, previous: Deployment | None, model_id: str) -> None: + if previous is None or self._core_has_model_id(model_id): + return + try: + if self.strategies is None: + self.add_deployment(deployment=previous) + else: + self.strategies.without_capability_limit(lambda: self.add_deployment(deployment=previous)) + verbose_router_logger.info( + "Restored deployment %s (id=%s); it keeps serving its previous configuration.", + previous.model_name, + model_id, + ) + except Exception as restore_error: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload + verbose_router_logger.warning( + "Could not restore previously served deployment %s (id=%s) after the failed upsert: %s", + previous.model_name, + model_id, + restore_error, + ) + + def _live_handle(self) -> "RouterCoreHandle | None": + return None if self.handle_lookup is None else self.handle_lookup() + + def _hand_changes(self, changes: Mapping[str, object]) -> None: + if self.changes_runner is None: + raise RuntimeError("the host registry writes deployments through its Rust core") + self.changes_runner(changes) + self._core_model_list.refresh() + + def refresh(self) -> None: + """Read router-core's list again after a write (only what was appended, when that is all + that changed). The index views are read again when an entry point hands them out, which + keeps one ``add_deployment`` from reading every index back.""" + self._core_model_list.refresh() + + def _remove(self, model_id: str) -> None: + """Remove the deployment at the id map's position for ``model_id``, every index shifted + (``_update_deployment_indices_after_removal``, in router-core).""" + self._hand_changes({"remove": model_id}) + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + + # -- the indexes (router-core's, read back as live views) ------------------------------------- + + @property + def model_id_to_deployment_index_map(self) -> dict[str, int]: # mutable-ok: the live id index + self._indexes.refresh() + return cast(dict[str, int], self._indexes.id_map) # cast-ok: id to position, as router-core keeps it + + @property + def _core_model_name_to_deployment_indices(self) -> dict[str, list[int]]: # mutable-ok: the live name index + self._indexes.refresh() + return cast(dict[str, list[int]], self._indexes.model_name) # cast-ok: name to positions + + @_core_model_name_to_deployment_indices.setter + def _core_model_name_to_deployment_indices(self, value: object) -> None: + self._indexes.assign("model_name", value) + + @property + def _core_team_model_to_deployment_indices( + self, + ) -> dict[tuple[str, str], list[int]]: # mutable-ok: the live team index + self._indexes.refresh() + return cast(_TeamIndex, self._indexes.team_model) # cast-ok: (team, public name) to positions + + @_core_team_model_to_deployment_indices.setter + def _core_team_model_to_deployment_indices(self, value: object) -> None: + self._indexes.assign("team_model", value) + + @property + def _core_model_names(self) -> set[str]: # mutable-ok: the live group set + self._indexes.refresh() + return cast(set[str], self._indexes.model_names) # cast-ok: group names + + @_core_model_names.setter + def _core_model_names(self, value: object) -> None: + self._indexes.assign("model_names", value) + + @property + def _core_deployment_names(self) -> list[str]: # mutable-ok: the live name list + self._indexes.refresh() + return cast(list[str], self._indexes.deployment_names) # cast-ok: litellm model names + + @_core_deployment_names.setter + def _core_deployment_names(self, value: object) -> None: + self._indexes.assign("deployment_names", value) + + @property + def _core_team_public_model_names(self) -> frozenset[str]: + self._indexes.refresh() + return self._indexes.team_public_model_names + + @_core_team_public_model_names.setter + def _core_team_public_model_names(self, value: object) -> None: + self._indexes.assign("team_public_model_names", value) + + @property + def _core_provider_default_deployment_ids(self) -> list[str]: # mutable-ok: the live id list + self._indexes.refresh() + return cast(list[str], self._indexes.provider_default_ids) # cast-ok: deployment ids + + @_core_provider_default_deployment_ids.setter + def _core_provider_default_deployment_ids(self, value: object) -> None: + self._indexes.assign("provider_default_ids", value) + + @property + def _core_pattern_router(self) -> PatternMatchRouter: + self._indexes.refresh() + return self._indexes.pattern_router + + @_core_pattern_router.setter + def _core_pattern_router(self, value: object) -> None: + self._indexes.pattern_router.patterns = getattr(value, "patterns", {}) # pyright: ignore[reportAttributeAccessIssue] # a router a caller assigned + + @property + def _core_team_pattern_routers(self) -> dict[str, PatternMatchRouter]: # mutable-ok: the live per-team routers + self._indexes.refresh() + return cast(dict[str, PatternMatchRouter], self._indexes.team_pattern_routers) # cast-ok: team id to its router + + @_core_team_pattern_routers.setter + def _core_team_pattern_routers(self, value: object) -> None: + routers: Final = as_str_mapping(value) or {} + self._indexes.write( + { + "team_patterns": [ + [ + str(team), + [ + [regex, deployments] + for regex, deployments in (as_str_mapping(getattr(router, "patterns", None)) or {}).items() + ], + ] + for team, router in routers.items() + ] + } + ) + + def replace_model_list(self, model_list: "list[UntypedEntry]") -> None: # mutable-ok: the list as assigned + """``router.model_list = [...]``: the list is swapped as given, without reindexing, and + router-core takes it whole.""" + self._core_model_list.replace(model_list) + self._registry_version += 1 + self._on_change() + + def _core_delete_deployment(self, id: str) -> UntypedEntry | None: + """Remove one deployment and return its stored entry (``_delete_deployment``).""" + try: + with self._deployment_index_lock: + deployment_idx: Final = self.model_id_to_deployment_index_map.get(id) + if deployment_idx is None: + return None + item: Final = self._core_model_list[deployment_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does + self._remove(id) + if self.checks is not None: + self.checks.unregister_deployment_budget(id) + if self.strategies is not None: + self.strategies.release_entry(item) + except Exception: # noqa: BLE001 # the core answers None for any failed delete + return None + self._on_change() + return item + + def _core_get_router_deployment_budget_limiter(self) -> "RouterBudgetLimiting | None": + """The router's deployment-budget callback (``registry().plugins("budget_limiter")``).""" + return self.checks.budget_limiter() if self.checks is not None else None + + @property + def _core_prompt_caching_pin_cache(self) -> DualCache: + """The prompt caching pins (``registry().plugins("prompt_caching_pins")``).""" + if self.checks is None: + return self.__getattr__("_core_prompt_caching_pin_cache") # pyright: ignore[reportReturnType] # raises RustCoreUnsupported + return self.checks.prompt_caching_pin_cache + + def replay_model_cost_registrations(self) -> None: + """Re-register every served deployment's cost-map entries (``replay_model_cost_registrations``).""" + provider_unresolved: Final = self._provider_unresolved_deployments + self._provider_unresolved_deployments = () + for create_deployment in provider_unresolved: + create_deployment() + for entry in tuple(self._core_model_list): + try: + deployment = Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry + except Exception: # noqa: BLE001 # a malformed entry must not abort the rest of the rebuild + verbose_router_logger.exception( + "Router: could not rebuild cost-map entry for a deployment during a price data reload" + ) + continue + register_deployment_pricing(self._host, deployment) + self._invalidate_model_group_info_cache() + if provider_unresolved: + self._on_change() + + # -- one deployment --------------------------------------------------------------------------- + + def _core_has_model_id(self, candidate_id: str) -> bool: + handle: Final = self._live_handle() + return handle is not None and handle.has_model_id(candidate_id) + + def _core_get_model_info(self, id: str) -> UntypedEntry | None: + handle: Final = self._live_handle() + idx: Final[object] = None if handle is None else handle.id_position(id) + return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it + + def _core_get_deployment(self, model_id: str) -> Deployment | None: + entry: Final = self._core_get_model_info(model_id) + return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry + + def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: + handle: Final = self._live_handle() + indices: Final[object] = None if handle is None else handle.named_positions(model_group_name) + model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them + return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry + + def _at(self, indices: object) -> tuple[UntypedEntry, ...]: + """The entries at ``indices`` (what an index holds), read as Python's list reads them.""" + return tuple(self._core_model_list[idx] for idx in indices) # pyright: ignore[reportGeneralTypeIssues, reportArgumentType, reportCallIssue] # positions, or what a caller wrote + + def _core_deployments_named(self, model_name: str) -> tuple[UntypedEntry, ...] | None: + handle: Final = self._live_handle() + indices: Final[object] = None if handle is None else handle.named_positions(model_name) + return None if indices is None else self._at(indices) + + def team_public_model_names_of(self, team_id: str) -> frozenset[str]: + handle: Final = self._live_handle() + return frozenset(() if handle is None else handle.team_public_names_of(team_id)) + + def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[UntypedEntry, ...] | None: + handle: Final = self._live_handle() + indices: Final[object] = None if handle is None else handle.team_positions(team_id, public_model_name) + return None if indices is None else self._at(indices) + + # -- lookups (router-core `lookup.rs`): rows read at the stored entries -------------------- + + def _rows(self, rows: object) -> list[UntypedEntry]: # mutable-ok: a fresh list + """The stored entries router-core's lookup ``rows`` name, read as Python's list reads them: + a row it could not read is indexed here, and raises as the Python router's read raised.""" + entries: Final[list[UntypedEntry]] = [] # mutable-ok: filled per row + for row in cast(Sequence[Mapping[str, object]], rows): # cast-ok: router-core's lookup rows + if "raw" in row: + entries.append(self._core_model_list[row["raw"]]) # pyright: ignore[reportArgumentType, reportCallIssue] # what a caller wrote, read as Python reads it + elif "raw_positions" in row: + entries.extend(self._at(row["raw_positions"])) + elif "pattern" in row: + entries.append({**DeploymentTypedDict(**row["pattern"]), "model_name": row["name"]}) # pyright: ignore[reportCallIssue] # a pattern router's stored entry + else: + entries.append(self._row_entry(row)) + return entries + + def _row_entry(self, row: Mapping[str, object]) -> UntypedEntry: + held: Final[UntypedEntry] = self._core_model_list[int(str(row["at"]))] + name: Final = row.get("name") + named: Final = held if name is None else _renamed(held, str(name)) + group: Final = row.get("group") + if group is None: + return named + return _as_routing_group_row( + apply_routing_group_priority(self._core_routing_groups[str(group)], str(row["member"]), named) # pyright: ignore[reportArgumentType] # a stored entry + ) + + def _raise_at(self, raw: object) -> None: + """Read the list at the row router-core could not read it at, which raises.""" + if raw is not None: + _ = self._rows([raw]) + + def _list_scope(self, model_name: str | None, *, only: str | None = None) -> Mapping[str, object]: + """What ``get_model_list`` lists beyond the deployments: the aliases (read here, where an + alias item that is no mapping or lacks ``hidden`` raises as the Python router's read + does), the callable routing groups and the wildcard requests for ``model_name``.""" + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + aliases: Final[list[Mapping[str, str]]] = [] # mutable-ok: filled per alias + alias_map: Final = self._core_model_group_alias or {} + if model_name is None or model_name in alias_map: + items: Final[Sequence[tuple[str, object]]] = ( + [(model_name, alias_map[model_name])] if model_name is not None else list(alias_map.items()) + ) + for model_alias, model_value in items: + if isinstance(model_value, str): + aliases.append({"alias": model_alias, "target": model_value}) + elif isinstance(model_value, dict): + item = RouterModelGroupAliasItem(**model_value) # pyright: ignore[reportCallIssue, reportArgumentType] # an alias item dict + if item["hidden"] is True and model_name is None: + continue + aliases.append({"alias": model_alias, "target": item["model"]}) + groups: Final = [ + {"name": name, "members": list(group.models)} + for name in (self._core_routing_groups or {}) + if (group := self.get_routing_group(name)) is not None + ] + wildcard: Final = ( + [model_name, f"{inferred_provider(model_name)}/{model_name}"] if model_name is not None else [] + ) + return {"aliases": aliases, "groups": groups, "wildcard": wildcard, "only": only} + + def _lookup(self) -> "RouterCoreHandle": + handle: Final = self._live_handle() + if handle is None: + raise RuntimeError("the host registry's lookups run in its Rust core") + return handle + + # -- deployments by name ---------------------------------------------------------------------- + + @staticmethod + def _core_should_include_deployment( + model_name: str, model: Mapping[str, object], team_id: str | None = None + ) -> bool: + info: Final = _model_info_of(model) + if team_id is not None and info.get("team_id") == team_id and model_name == info.get("team_public_model_name"): + return True + if model_name is not None and model.get("model_name") == model_name: # pyright: ignore[reportUnnecessaryComparison] # the core checks it + model_team_id: Final = info.get("team_id") + return team_id is None or model_team_id is None or model_team_id == team_id + return False + + @staticmethod + def _core_deployment_usable_by_team(model: Mapping[str, object] | Deployment, team_id: str | None) -> bool: + info: Final[object] = ( + model.model_info.model_dump() if isinstance(model, Deployment) else model.get("model_info") + ) + owner: Final = _MODEL_INFO_ADAPTER.validate_python(info).get("team_id") if isinstance(info, Mapping) else None + return team_may_use_deployment(owner, team_id) + + def _core_deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[Mapping[str, object]]: + """The deployments ``model`` names for this caller: alias, then team-first, then global, then + admin-across-teams (``_core_deployments_for_request``).""" + from litellm.litellm_core_utils.routing_common_utils import get_request_team_id, is_proxy_admin_request + + registered: Final = self._core_get_model_from_alias(model=model) or model + team_id: Final = get_request_team_id(request_kwargs) + deployments: Final = self._core_get_all_deployments(model_name=registered, team_id=team_id) + if deployments or team_id is not None or not is_proxy_admin_request(request_kwargs): + return deployments + return self._team_deployments_across_teams(registered) + + def _team_deployments_across_teams(self, model: str) -> list[UntypedEntry]: # mutable-ok: a fresh list + """Every team's deployments under public name ``model``, for a proxy admin calling without a team.""" + rows, teams = self._lookup().lookup_across_teams(model) + team_deployments: Final = self._rows(rows) + if teams > 1: + raise litellm_exceptions.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments + + def _core_get_all_deployments( + self, model_name: str, model_alias: str | None = None, team_id: str | None = None + ) -> list[UntypedEntry]: # mutable-ok: a fresh list, as the core returns it + return self._rows(self._lookup().lookup_all(model_name, team_id, model_alias)) + + def get_routing_group(self, model_name: str) -> RoutingGroup | None: + """The routing group callable as ``model_name``: a deployment model_name or an alias of the + same name shadows it.""" + if not self._core_routing_groups: + return None + group: Final = self._core_routing_groups.get(model_name) + handle: Final = self._live_handle() + if ( + group is None + or (handle is not None and handle.named_positions(model_name) is not None) + or model_name in (self._core_model_group_alias or {}) + ): + return None + return group + + def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: + """A callable routing group whose member union spans more than one deployment.""" + if model_group is None: + return False + resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group + group: Final = self.get_routing_group(resolved) + if group is None: + return False + handle: Final = self._lookup() + + def held(member: str) -> Sequence[object]: + return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk + + return sum(len(held(member)) for member in group.models) > 1 + + def _core_get_model_list_from_model_alias( + self, model_name: str | None = None + ) -> list[UntypedEntry]: # mutable-ok: a fresh list + return self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) + ) + + def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[UntypedEntry]: + """Callable routing groups as model-list rows: each member deployment under the group's name.""" + if model_name is None and self._routing_group_rows is not None: + return self._routing_group_rows + rows: Final = tuple( + self._rows(self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups"))) + ) + if model_name is None: + self._routing_group_rows = rows + return rows + + def _core_get_model_list( + self, model_name: str | None = None, team_id: str | None = None + ) -> list[UntypedEntry] | None: # mutable-ok: a fresh list, as get_model_list returns + return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) + + def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + rows: Final = self._lookup().lookup_wildcard_route( + model, team_id, [model, f"{inferred_provider(model)}/{model}"] + ) + routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows + return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: each carries its routed copy + + def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_owned: Final = ( + team_router.get_models_by_pattern(model=model) if team_router is not None else frozenset[str]() + ) + return team_owned or self._core_pattern_router.get_models_by_pattern(model=model) + + def _core_config_deployments(self) -> Iterable[Mapping[str, object]]: + return iter(self._rows(self._lookup().lookup_config())) + + # -- names ------------------------------------------------------------------------------------ + + def _core_get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: a fresh list + names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) + self._raise_at(raw) + return names + + def _core_get_model_ids( + self, model_name: str | None = None, exclude_team_models: bool = False + ) -> list[str]: # mutable-ok: a fresh list + ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) + self._raise_at(raw) + return ids + + # -- resolving a name ------------------------------------------------------------------------- + + def _core_get_model_from_alias(self, model: str) -> str | None: + return resolve_model_group_alias(self._core_model_group_alias, model) + + def _core_routable_model_group(self, model_name: str) -> str: + target: Final = self._core_get_model_from_alias(model_name) + return target if target is not None else model_name + + def _core_is_recognized_model(self, model: str) -> bool: + handle: Final = self._live_handle() + return ( + (handle is not None and handle.has_model_name(model)) + or self._core_has_model_id(model) + or model in self._core_model_group_alias + or self.get_routing_group(model) is not None + ) + + def _core_get_model_group(self, id: str) -> list[UntypedEntry] | None: # mutable-ok: a fresh list + model_info: Final = self._core_get_model_info(id=id) + if model_info is None: + return None + return self._core_get_model_list(model_name=str(model_info["model_name"])) + + def _core_resolve_model_name_from_model_id( + self, model_id: str | None, custom_llm_provider: str | None = None + ) -> str | None: + if not model_id: + return None + handle: Final = self._lookup() + if handle.has_model_name(model_id): + return model_id + if self._core_has_model_id(model_id): + deployment: Final = self._core_get_deployment(model_id=model_id) + if deployment is not None and deployment.model_name: + return deployment.model_name + return model_id + listed, serving, raw = handle.lookup_name_serving(model_id, self._list_scope(None)) + self._raise_at(raw) + if not listed: + return None + if serving is not None: + return serving + if self._core_is_recognized_model(model_id): + return model_id + if custom_llm_provider is None: + return None + if self._core_pattern_router.routes_to_provider(model_id, custom_llm_provider): + return model_id + provider_prefixed: Final = f"{custom_llm_provider}/{model_id}" + if not model_id.startswith(f"{custom_llm_provider}/") and self._core_pattern_router.routes_to_provider( + provider_prefixed, custom_llm_provider + ): + return provider_prefixed + return None + + def _get_deployment_by_litellm_model(self, model: str) -> list[UntypedEntry]: # mutable-ok: a fresh list + handle: Final = self._live_handle() + return list(self._at(() if handle is None else handle.litellm_model_positions(model))) + + # -- credentials ------------------------------------------------------------------------------ + + @staticmethod + def _core_is_deployment_blocked(deployment: Deployment) -> bool: + model_info: Final[object] = getattr(deployment, "model_info", None) + return model_info is not None and getattr(model_info, "blocked", None) is True + + def _core_get_deployment_credentials(self, model_id: str) -> UntypedEntry | None: + deployment: Final = self._core_get_deployment(model_id=model_id) + if deployment is None or self._core_is_deployment_blocked(deployment): + return None + return CredentialLiteLLMParams.model_validate( + deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) + + def _get_model_group_deployment_usable_by_team( + self, model_group_name: str, team_id: str | None + ) -> Deployment | None: + position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + + def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: + if team_id is None: + return None + position, raw = self._lookup().lookup_team_public_name(model_id, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + + def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_models: Final = team_router.route(model_id) if team_router else None + global_models: Final = tuple( + model + for model in (self._core_pattern_router.route(model_id) or ()) + if self._core_deployment_usable_by_team(model, team_id) + ) + candidates: Final = team_models or global_models + if not candidates: + return None + return Deployment(**candidates[0]) # a pattern router holds stored entry dicts + + def _core_get_credential_deployment(self, model_id: str, team_id: str | None = None) -> Deployment | None: + deployment: Final = ( + self._core_get_deployment(model_id=model_id) + or self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id) + or self._get_team_public_name_deployment(model_id=model_id, team_id=team_id) + or self._get_wildcard_deployment_usable_by_team(model_id=model_id, team_id=team_id) + ) + if deployment is None or self._core_is_deployment_blocked(deployment): + return None + return deployment + + def _core_get_deployment_credentials_with_provider( + self, model_id: str, team_id: str | None = None + ) -> UntypedEntry | None: + deployment: Final = self._core_get_credential_deployment(model_id=model_id, team_id=team_id) + if deployment is None: + return None + base: Final = CredentialLiteLLMParams.model_validate( + deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) + credential_name: Final = deployment.litellm_params.litellm_credential_name + resolved: Final[UntypedEntry] = dict(base) + if credential_name is not None: + credential_values: Final = CredentialAccessor.get_credential_values(credential_name) + if not credential_values: + verbose_router_logger.warning("Credential '%s' not found in credential_list", credential_name) + resolved.update(credential_values) + _ = resolved.pop("litellm_credential_name", None) + model: Final = deployment.litellm_params.model + resolved["model"] = model + resolved["custom_llm_provider"] = ( + deployment.litellm_params.custom_llm_provider or (model.split("/")[0] if "/" in model else None) or "openai" + ) + return resolved + + # -- routable names --------------------------------------------------------------------------- + + def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[UntypedEntry, ...]: + return tuple(self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id))) + + def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": + """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, + else None (``team_ptu_ceiling``), kept until the model list next changes.""" + from litellm.litellm_core_utils.ptu_shares import ptu_share_teams, team_ptu_ceiling + + listed: Final = tuple(self._core_get_model_list() or ()) + if self._ptu_cache is None: + self._ptu_cache = (ptu_share_teams(listed), {}) + share_teams, ceilings = self._ptu_cache + if team_id not in share_teams: + return None + key: Final = (team_id, model_group) + if key not in ceilings: + ceilings[key] = team_ptu_ceiling( + listed, + self._core_model_list, + team_id, + model_group, + lambda name: self._core_wildcard_route_deployments(name, team_id), + ) + return ceilings[key] + + def _get_client( + self, deployment: Mapping[str, object], kwargs: Mapping[str, object], client_type: str | None = None + ) -> object: + """The SDK client cached in the routing cache for a deployment (``_get_client``); a pod + discovery deployment has none. ``max_parallel_requests`` is the deployment's slot limit, + built on first use.""" + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + from litellm.router_rust.core import RustCoreUnsupported + + # `deployment["model_info"]["id"]`: the Python router reads the id first, raising for a + # deployment without one whatever the client type. + _ = deployment["model_info"]["id"] # pyright: ignore[reportIndexIssue] # read as the Python router reads it + if client_type == "max_parallel_requests": + limit_lookup: Final = self.max_parallel_requests_lookup + if limit_lookup is None: + raise RustCoreUnsupported( + "registry().client(client_type='max_parallel_requests') is not served by the Rust router core" + ) + return limit_lookup(deployment, kwargs) + if _params_of(deployment).get("kubernetes_pod_discovery") is True: + return None + model_id: Final = _model_info_of(deployment).get("id") + span: Final = get_parent_otel_span_from_kwargs(dict(kwargs)) + stream: Final = "_stream" if kwargs.get("stream") is True else "" + with service_target(ROUTER_CLIENTS_TARGET): + if client_type == "async": + return self._core_cache.get_cache( + key=f"{model_id}{stream}_async_client", local_only=True, parent_otel_span=span + ) # pyright: ignore[reportUnknownMemberType] # DualCache's reads are untyped + return self._core_cache.get_cache(key=f"{model_id}{stream}_client", parent_otel_span=span) # pyright: ignore[reportUnknownMemberType] # see above + + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: + """Per routable deployment: its ``supports_search_tool`` when a bool, else True when its route + builds an Anthropic Messages request, else None (``get_search_tool_support``).""" + return tuple(_search_tool_support(deployment) for deployment in self._routable_deployments(model_name, team_id)) + + def _core_get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: + """The deployment ids that could serve ``model`` for ``team_id``: the alias target, then a + routing group, then a team public name or wildcard route for a name that is no model + group, then the model group (``get_candidate_model_ids_for_route``). Read-only.""" + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + resolved: Final = self._core_get_model_from_alias(model=model) or model + group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None + ids, raw = self._lookup().lookup_candidate_ids( + resolved, + team_id, + None if group is None else list(group.models), + [resolved, f"{inferred_provider(resolved)}/{resolved}"], + ) + self._raise_at(raw) + return frozenset(ids) + + def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: + """Whether a team deployment's public name has another unpaused deployment of that team.""" + found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) + self._raise_at(raw) + return found + + def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: + """``served_model`` when ``model_name`` routes to it (the alias target, a routing group + member, or the name itself), else ``model_name``.""" + target: Final = self._core_get_model_from_alias(model_name) or model_name + group: Final = self.get_routing_group(target) + members: Final = group.models if group is not None else () + return served_model if served_model == target or served_model in members else model_name + + def _core_resolved_litellm_models( + self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False + ) -> tuple[str, ...]: + """The provider model strings a call to ``model_name`` can be served by (``resolved_litellm_models``).""" + listed: Final = self._core_get_model_list(model_name=model_name, team_id=team_id) or () + if use_base_model: + return tuple( + backend for entry in listed if (backend := cost_map_backend_model(Deployment.model_validate(entry))) + ) + return tuple(model for entry in listed if isinstance(model := _params_of(entry).get("model"), str) and model) + + def _core_get_routable_deployments( + self, model: str, request_kwargs: Mapping[str, object] = MappingProxyType({}), call_type: str | None = None + ) -> tuple[Mapping[str, object], ...]: + """The deployments routing would consider for ``request_kwargs``, as the Rust core resolves them + (endpoint eligibility for ``call_type`` is not applied: the Rust resolution has none).""" + _ = call_type + lookup: Final = self.routable_lookup + if lookup is None: + from litellm.router_rust.core import RustCoreUnsupported + + raise RustCoreUnsupported("registry().get_routable_deployments is not served by the Rust router core") + return tuple(lookup(model, request_kwargs)) + + def _core_get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: + deployment: Final = next(iter(self._routable_deployments(model_name, team_id)), None) + if deployment is None: + return None + model: Final = _params_of(deployment).get("model") + return None if model is None else str(model) + + def _core_map_team_model( + self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None + ) -> str | None: + """A team's public name to its group, or a deployment for a model-less call (``map_team_model``).""" + models: Final = self._core_get_model_list(model_name=team_model_name, team_id=team_id) + if not models: + return None + if team_model_name is None and custom_llm_provider is not None: + return self._team_model_for_provider(models, team_id, custom_llm_provider) + for model in models: + info = _model_info_of(model) + if info.get("team_id") == team_id: + if team_model_name is None: + public = info.get("team_public_model_name") or model.get("model_name") + return None if public is None else str(public) + return team_model_name + return None + + @staticmethod + def _team_model_for_provider(models: Sequence[UntypedEntry], team_id: str, custom_llm_provider: str) -> str | None: + def public_name(deployment: UntypedEntry) -> str | None: + info: Final = _model_info_of(deployment) + if info.get("team_id") != team_id: + return None + name: Final = info.get("team_public_model_name") + return name if isinstance(name, str) and name else str(deployment["model_name"]) + + def provider(deployment: UntypedEntry) -> object: + return provider_for_generic_call(_params_of(deployment)) # pyright: ignore[reportArgumentType] # the deployment's params + + matching: Final = next( + (m for m in models if public_name(m) is not None and provider(m) == custom_llm_provider), None + ) + if matching is None: + return None + name: Final = public_name(matching) + if name is None: + return None + if all(provider(m) == custom_llm_provider for m in models if public_name(m) == name): + return name + matched_id: Final = _model_info_of(matching).get("id") + return matched_id if isinstance(matched_id, str) else None + + # -- access groups ---------------------------------------------------------------------------- + + def _core_get_model_access_groups( + self, model_name: str | None = None, model_access_group: str | None = None, team_id: str | None = None + ) -> Mapping[str, Sequence[str]]: + use_cache: Final = model_name is None and model_access_group is None and team_id is None + if use_cache and self._access_groups_cache is not None: + return self._access_groups_cache + pairs, raw = self._lookup().lookup_access_groups( + model_name, model_access_group, team_id, self._list_scope(model_name) + ) + self._raise_at(raw) + groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict + if use_cache: + self._access_groups_cache = groups + return groups + + def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: + groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) + return any( + self._core_pattern_router.route(request=model) is not None for model in groups.get(model_access_group, []) + ) + + +# -- the registry() view ------------------------------------------------------------------------- + + +class RustRegistry(RouterRegistry): + """``registry()`` of the Rust core: ``RouterRegistry``'s queries over the host registry.""" + + __slots__ = () + + def includes( + self, + deployment: "DeploymentTypedDict | dict[str, object]", # mutable-ok: the RouterRegistry signature + *, + model_name: str | None = None, + team_id: str | None = None, + ) -> bool: + hosted: Final = self._router + if not isinstance(hosted, HostRegistry): + return super().includes(deployment, model_name=model_name, team_id=team_id) + if model_name is None: + return HostRegistry._core_deployment_usable_by_team(deployment, team_id) # pyright: ignore[reportPrivateUsage, reportArgumentType] # the view reads its own registry + return HostRegistry._core_should_include_deployment(model_name=model_name, model=deployment, team_id=team_id) # pyright: ignore[reportPrivateUsage, reportArgumentType] # see above + + +def hosted_registry_view(hosted: HostRegistry) -> RouterRegistry: + """The ``registry()`` view over ``hosted``.""" + return RustRegistry(hosted) # pyright: ignore[reportArgumentType] # the host registry answers every _core_ query the view makes + + +# The registry indexes a deprecated attribute write replaces (``router.model_names = ...`` reaches +# the host registry as ``_core_model_names``), by the field router-core reads each one as. +INDEX_ATTRIBUTES: Final[Mapping[str, str]] = MappingProxyType( + { + "_core_model_names": "model_names", + "_core_model_name_to_deployment_indices": "model_name", + "_core_team_model_to_deployment_indices": "team_model", + } +) + + +def _positions(indices: object, length: int) -> list[int]: # mutable-ok: a JSON list + """The model-list positions ``indices`` read, as Python's ``model_list[index]`` reads them: a + negative one counts from the end, and one before the start becomes ``length``, a position past + the end, so the lookup raises ``IndexError`` as Python's does.""" + if not isinstance(indices, Iterable) or isinstance(indices, str | bytes): + return [] + return [ + index if index >= 0 else (index + length if index + length >= 0 else length) + for index in indices + if isinstance(index, int) and not isinstance(index, bool) + ] + + +def index_document(hosted: HostRegistry, written: Iterable[str]) -> Mapping[str, object]: + """The written registry indexes, as the host registry holds them now, in the shape router-core + reads (``IndexOverrides``): names a lookup finds, and the deployment positions each model name + (or ``(team_id, public name)`` pair) indexes. Entries Python could not read as such are left out.""" + document: Final[dict[str, object]] = {} # mutable-ok: built here + length: Final = len(hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # positions read against the live list + for attribute in written: + value: object = getattr(hosted, attribute, None) + field: str = INDEX_ATTRIBUTES[attribute] + if field == "model_names": + names = value if isinstance(value, Iterable) and not isinstance(value, str | bytes | Mapping) else () + document[field] = sorted({name for name in names if isinstance(name, str)}) + elif field == "model_name": + entries = value if isinstance(value, Mapping) else {} + document[field] = { + name: _positions(indices, length) for name, indices in entries.items() if isinstance(name, str) + } + else: + pairs = value if isinstance(value, Mapping) else {} + document[field] = [ + [key[0], key[1], _positions(indices, length)] + for key, indices in pairs.items() + if isinstance(key, tuple) and len(key) == 2 and all(isinstance(part, str) for part in key) + ] + return document + + +DEFAULT_DEPLOYMENT_FAILING_ID: Final = "__litellm_default_deployment_failing_call__" +"""The id router-core serves a ``default_deployment`` under when every call to it fails.""" + +DEFAULT_DEPLOYMENT_COPIED_ID: Final = "__litellm_default_deployment_copied__" +"""The id router-core serves a callable ``default_deployment`` under when its ``model_info.id`` is +not a string: the host picks and calls Python's copy of it, so its id stays what was assigned.""" + + +def default_deployment_copy(value: object, model: str) -> object: + """``_try_early_resolve_deployments_for_model_not_in_names``' copy of the default deployment + for ``model``, raising what it raises for a value it cannot copy.""" + updated: Final = value.copy() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # any assigned value, read as Python reads it + updated["litellm_params"] = value["litellm_params"].copy() # pyright: ignore[reportIndexIssue, reportUnknownMemberType] # same + updated["litellm_params"]["model"] = model # pyright: ignore[reportUnknownMemberType] # same + return updated # pyright: ignore[reportUnknownVariableType] # same + + +def default_deployment_copy_error(value: object, model: str) -> BaseException: + """The error the default deployment's copy raises for ``model``.""" + try: + _ = default_deployment_copy(value, model) + except Exception as raised: # noqa: BLE001 # the error the pick raises + return raised + return RuntimeError(f"router-core: the default deployment copies for {model!r}") + + +def default_deployment_call_error(value: object) -> BaseException | None: + """The error every call to the default deployment raises, from its first reads + (``_update_kwargs_with_deployment``'s ``deployment["model_name"]``, ``_get_client``'s + ``deployment["model_info"]["id"]``); None when they read.""" + try: + served: Final = default_deployment_copy(value, "") + _ = served["model_name"] # pyright: ignore[reportIndexIssue] # read as Python reads it + _ = served["model_info"]["id"] # pyright: ignore[reportIndexIssue] # same + except Exception as raised: # noqa: BLE001 # the error the call raises + return raised + return None + + +def served_default_deployment(value: object) -> tuple[Mapping[str, object] | None, bool]: + """The ``default_deployment`` as router-core serves it, and whether resolving to it raises + (a value Python cannot copy). A deployment every call to fails is served under + ``DEFAULT_DEPLOYMENT_FAILING_ID``, one whose id is not a string under + ``DEFAULT_DEPLOYMENT_COPIED_ID``; names that are not strings are served as text.""" + if value is None: + return None, False + try: + copied: Final = default_deployment_copy(value, "") + except Exception: # noqa: BLE001 # any copy error: the pick raises it + return None, True + entry: Final = as_str_mapping(copied) or {} + model_info: Final = dict(as_str_mapping(entry.get("model_info")) or {}) + failing: Final = default_deployment_call_error(value) is not None + assigned_id: Final = model_info.get("id") + model_info["id"] = ( + DEFAULT_DEPLOYMENT_FAILING_ID + if failing + else assigned_id + if isinstance(assigned_id, str) + else DEFAULT_DEPLOYMENT_COPIED_ID + ) + name: Final = entry.get("model_name") + return { + **entry, + "model_name": name if isinstance(name, str) else str(name or ""), + "litellm_params": dict(as_str_mapping(entry.get("litellm_params")) or {}), + "model_info": model_info, + }, False diff --git a/litellm/router_rust/surface/replay.py b/litellm/router_rust/surface/replay.py new file mode 100644 index 00000000000..6c770e07883 --- /dev/null +++ b/litellm/router_rust/surface/replay.py @@ -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) diff --git a/litellm/router_rust/surface/retries.py b/litellm/router_rust/surface/retries.py new file mode 100644 index 00000000000..a763987b9b7 --- /dev/null +++ b/litellm/router_rust/surface/retries.py @@ -0,0 +1,243 @@ +""" +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 +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.types.router import RetryPolicy + +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()) + + +@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 + + +_STREAM_POLICY_ENDPOINTS: Final = frozenset({"aanthropic_messages", "anthropic_messages"}) + + +def raising_groups( + endpoint: str, request: Mapping[str, object], policies: 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``, or it is a mapping + on a streamed Anthropic messages call, which resolves such an entry to no policy at all + (``_anthropic_messages_resolved_retry_policy``).""" + if not resolved.raising or request.get("num_retries") == 0: + return () + if endpoint in _STREAM_POLICY_ENDPOINTS and request.get("stream") is True and isinstance(policies, Mapping): + 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 _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" + + 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 + if not is_policy_per_group(resolved.model_group_retry_policy): + return router_policy_retries( + exception, resolved.retry_policy, resolved.model_group_retry_policy, model_group + ) + return self._handle.retry_policy_retries(dict(attempt_failure(exception)), model_group) + + +__all__ = [ + "GroupPolicies", + "RustRetries", + "group_policies", + "is_policy_per_group", + "policy_lookup_error", + "raising_groups", + "resolved_retry_policy", + "router_policy_retries", + "static_group_policies", +] diff --git a/litellm/router_rust/surface/settings.py b/litellm/router_rust/surface/settings.py new file mode 100644 index 00000000000..e3b421737bd --- /dev/null +++ b/litellm/router_rust/surface/settings.py @@ -0,0 +1,218 @@ +""" +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.types.router import RetryPolicy + +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 + + +def read_settings_update(settings: Mapping[str, object]) -> SettingsUpdate: + """``_core_update_settings``' reading of each setting: 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).""" + 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 + 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": + 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: + values[name] = value + return SettingsUpdate( + values=MappingProxyType(values), + plain=MappingProxyType(plain), + continuation=continuation, + routing_groups=routing_groups, + optional_pre_call_checks=optional_checks, + ) + + +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", +] diff --git a/litellm/router_utils/global_callback_holders.py b/litellm/router_utils/global_callback_holders.py new file mode 100644 index 00000000000..53a7c20cff3 --- /dev/null +++ b/litellm/router_utils/global_callback_holders.py @@ -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", +] diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..345c877c10a 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import Mapping, Sequence -from typing import Never, final +from typing import Literal, Never, NotRequired, TypedDict, final import httpx from pydantic import JsonValue @@ -218,6 +218,8 @@ __all__ = [ "NativeDiagnosticProcessor", "ProcessReservedForForking", "ResponsesWebSocketConnection", + "RouterCall", + "RouterCoreHandle", "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", @@ -325,3 +327,312 @@ class NativeCacheHandle: def ping(self) -> Future[bool]: ... def disconnect(self) -> Future[None]: ... def delete(self, keys: Sequence[str]) -> Future[None]: ... + +# The FailureKind names an attempt outcome's `kind` takes (router-core `FailureKind`). +_FailureKind = Literal[ + "RateLimitError", + "ContextWindowExceededError", + "ContentPolicyViolationError", + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "BadRequestError", + "Timeout", + "APIConnectionError", + "InternalServerError", + "ServiceUnavailableError", + "GuardrailIntervention", + "Declined", + "CallerCredentialError", + "APIError", +] + +class _AttemptStep(TypedDict): + kind: Literal["attempt"] + deployment_id: str + # The deployment as a `model_list` entry, with its id in `model_info`. + deployment: dict[str, JsonValue] + # The name this attempt serves: the requested model or the fallback target. + model: str + model_group: str + target_order: int | None + candidate_ids: list[str] + retry: int + fallback_depth: int + fallback_level: int + timeout: float | None + +class _SleepStep(TypedDict): + kind: Literal["sleep"] + seconds: float + +class _AllowFallbackStep(TypedDict): + kind: Literal["allow_fallback"] + target: str + +class _SelectStep(TypedDict): + kind: Literal["select"] + stage: Literal["plugins", "strategy"] + strategy: str | None + model_group: str + candidates: list[str] + +# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`, +# which sets `keys`; `value` is the JSON to store (`set`) or the amount (`increment`); `ttl` is +# in seconds. Outcomes: get -> the JSON value or None; batch_get -> one JSON value or None per +# key, in order; increment -> the new int or float; set and delete -> None. +class _CacheStep(TypedDict): + kind: Literal["cache"] + op: Literal["get", "set", "increment", "batch_get", "delete"] + key: str | None + keys: list[str] | None + value: JsonValue + ttl: float | None + +class _DoneStep(TypedDict): + kind: Literal["done"] + deployment_id: str + model_group: str + attempted_retries: int + attempted_fallbacks: int + total_tokens: int + +class _FailedStep(TypedDict): + kind: Literal["failed"] + error: str + error_kind: Literal["config", "unknown_model", "no_healthy_deployments", "attempt", "cache", "retired"] + attempt_failed: bool + status: int | None + failure_kind: _FailureKind | None + +# The end of a `RouterCoreHandle.select`. +class _SelectedStep(TypedDict): + kind: Literal["done"] + deployment_id: str + model_group: str + +# The end of a `RouterCoreHandle.healthy`. +class _HealthyStep(TypedDict): + kind: Literal["done"] + deployment_ids: list[str] + +# The end of a `RouterCoreHandle.report_failure`. +class _ReportedStep(TypedDict): + kind: Literal["done"] + cooled: bool + +# 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(TypedDict): + kind: Literal["secrets"] + index: int + refs: list[_SecretRef] + +# The deployment build's steps for one deployment. `validate` and `added` are answered with +# `{"entry": }`, `{"skip": True}` or `{"invalid": True}`; `provider` with +# `{"provider": str | None}` or `{"invalid": True}`. +class _BuildStep(TypedDict): + 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(TypedDict): + kind: Literal["done"] + built: list[_BuiltDeployment] + +# `RouterCoreHandle.resolve`: what a name resolves to. +class _Resolution(TypedDict): + model_group: str + kind: str + ids: list[str] + +_RouterStep = _AttemptStep | _SleepStep | _AllowFallbackStep | _SelectStep | _CacheStep | _DoneStep | _FailedStep + +class _AttemptSucceeded(TypedDict): + ok: Literal[True] + total_tokens: NotRequired[int] + 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] + +class _AttemptFailed(TypedDict): + ok: Literal[False] + kind: _FailureKind + status: NotRequired[int | None] + retry_after: NotRequired[float | None] + num_retries: NotRequired[int | None] + detail: NotRequired[str | None] + +# Each document is a mapping or its JSON text. +_Document = Mapping[str, object] | str + +@final +class RouterCoreHandle: + def __new__(cls, config: _Document, dynamic: _Document) -> RouterCoreHandle: ... + def start(self, op: str, request: _Document) -> RouterCall: ... + # Driven like `start`; its `done` step is a `_SelectedStep`. + def select(self, model: str, request: _Document) -> RouterCall: ... + # Driven like `select`; its `done` step is a `_HealthyStep`. + def healthy(self, model: str, request: _Document) -> RouterCall: ... + # `{"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: ... + # 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: ... + # `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: ... + # `_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: ... + # `get_configured_token_limits`, driven like `router_model_info`; `value` is `[input, output]`. + def configured_token_limits(self, id: object, configured: _Document) -> RouterCall: ... + # `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: ... + # 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, + ) -> RouterCall: ... + # A success a logging callback saw, driven like `report_failure`; `done` carries `key`. + def report_success(self, id: str, total_tokens: float, requests: float) -> RouterCall: ... + # Count a success against any id, driven like `report_failure`; `done` carries `key`. + def count_success(self, id: str) -> RouterCall: ... + # [(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: ... + def fallback_stream( + self, op: str, request: _Document, origin: _Document, failure: Mapping[str, object] + ) -> RouterCall: ... + 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: ... + +@final +class RouterCall: + # attempt: _AttemptSucceeded | _AttemptFailed; sleep: None; allow_fallback: bool; + # select: the remaining candidate ids, pick first; cache: see _CacheStep; first call: None. + # A call from `select` steps through select and cache only, and ends in _SelectedStep or failed. + def next( + self, outcome: _AttemptSucceeded | _AttemptFailed | bool | Sequence[str] | JsonValue | None + ) -> _RouterStep | _SelectedStep | _HealthyStep | _ReportedStep | _SecretsStep | _BuildStep | _BuiltStep: ... diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 629d9098970..a080eb1c51d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -23,6 +23,7 @@ class Route(str, Enum): OCR = "ocr" TOKEN_COUNTER = "token_counter" TOKENIZER = "tokenizer" + ROUTER = "router" class Target(str, Enum): @@ -221,6 +222,7 @@ RULES: Final[Rules] = ( RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), + RouteRule(Route.ROUTER, Rollout.RUST_OPT_IN), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/utils.py b/litellm/utils.py index b0826f84006..8533f416225 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4165,6 +4165,13 @@ class _LiveDeploymentReplay: """ callback: Callable[[], None] | None = None + # Replays other router cores add beside the slot, run after it. + extra: tuple[Callable[[], None], ...] = () + + +def add_live_deployment_replay(replay: Callable[[], None]) -> None: + """Add a callback that re-asserts another router core's live deployments after a refresh.""" + _LiveDeploymentReplay.extra = (*_LiveDeploymentReplay.extra, replay) def set_live_deployment_replay(replay: Callable[[], None]) -> None: @@ -4198,6 +4205,8 @@ def reapply_runtime_model_cost_registrations() -> None: """ if _LiveDeploymentReplay.callback is not None: _LiveDeploymentReplay.callback() + for extra_replay in _LiveDeploymentReplay.extra: + extra_replay() if _runtime_registered_model_cost: register_model(model_cost=dict(_runtime_registered_model_cost)) diff --git a/tests/code_coverage_tests/router_litellm_globals_check.py b/tests/code_coverage_tests/router_litellm_globals_check.py index ff9fac0e3dc..172a3754000 100644 --- a/tests/code_coverage_tests/router_litellm_globals_check.py +++ b/tests/code_coverage_tests/router_litellm_globals_check.py @@ -21,6 +21,7 @@ CHECKED = ( "litellm/router_core.py", "litellm/router_deprecated.py", "litellm/router_python_core.py", + "litellm/router_rust/core.py", ) diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py index a8e0d6fa0d7..1c41526101a 100644 --- a/tests/code_coverage_tests/router_surface_check.py +++ b/tests/code_coverage_tests/router_surface_check.py @@ -11,7 +11,8 @@ modules listed there. The lists hold what is used today and may only shrink: also fails, so the entry is removed in the same change. "Outside" means every non-test file under litellm/ and enterprise/ except -litellm/router.py and the router's own modules. Uses from the router family +litellm/router.py and the router's own modules (including every +module under litellm/router_rust/, the Rust core's host-side glue). Uses from the router family (litellm/router_utils, litellm/router_strategy) are listed separately, because those modules become private parts of the router core in a later step. @@ -42,9 +43,12 @@ ROUTER_FILES = { "litellm/router_core.py", "litellm/router_python_core.py", "litellm/router_python_host.py", + "litellm/router_rust/core.py", "litellm/router_deprecated.py", "litellm/router_host.py", } +# The Rust core's host-side glue modules are part of the router. +ROUTER_FILE_PREFIXES = ("litellm/router_rust/",) ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_python_core.py", "litellm/router_deprecated.py"] ROUTER_CLASS_NAMES = {"Router", "PythonCore", "DeprecatedRouterMembers"} @@ -128,7 +132,7 @@ def scan() -> Dict[str, Dict[str, Set[str]]]: uses: Dict[str, Dict[str, Set[str]]] = {"outside": {}, "router_family": {}} imports: Dict[str, Set[str]] = {} for rel, path in _python_files(): - if rel in ROUTER_FILES or _is_test(rel): + if rel in ROUTER_FILES or rel.startswith(ROUTER_FILE_PREFIXES) or _is_test(rel): continue with open(path, encoding="utf-8") as f: src = f.read() diff --git a/tests/unit/test_router/core/conftest.py b/tests/unit/test_router/core/conftest.py index dffa938c441..3477cfc0c3d 100644 --- a/tests/unit/test_router/core/conftest.py +++ b/tests/unit/test_router/core/conftest.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +from collections.abc import Callable, Generator +from pathlib import Path +from typing import Final import pytest @@ -15,12 +17,53 @@ def host() -> ScriptedHost: def make_core(request: pytest.FixtureRequest, host: ScriptedHost) -> Callable[..., RouterCore]: """``make_core(**router_arguments)``: the core under test, on the ``host`` fixture. Every test that uses it runs once per core in ``CORE_BUILDERS`` (``-k python`` or ``-k rust`` picks one).""" - builder = CORE_BUILDERS[request.param] + name: Final[str] = str(request.param) # pyright: ignore[reportAny] # the fixture parameter + builder: Final = CORE_BUILDERS[name] def build(**arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments try: return builder(host, **arguments) except CoreDeclined as declined: - pytest.skip(f"{request.param} core declines this configuration: {declined}") + pytest.skip(f"{name} core declines this configuration: {declined}") return build + + +_RUST_PARITY_PENDING: Final = frozenset( + line.strip() + for line in (Path(__file__).parent / "rust_parity_pending.txt").read_text().splitlines() + if line.strip() and not line.startswith("#") +) + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + """Characterization tests the Rust core does not match yet are listed in + ``rust_parity_pending.txt`` and expected to fail (strictly: one that starts passing must leave the list).""" + here: Final = Path(__file__).parent + for item in items: + path: Final = Path(str(item.path)) + if not path.is_relative_to(here): + continue # the hook sees every collected item, not only this folder's + name: Final = "::".join((path.relative_to(here).as_posix(), *item.nodeid.split("::")[1:])) + if name in _RUST_PARITY_PENDING: + item.add_marker(pytest.mark.xfail(reason="Rust core parity pending", strict=True)) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, None]: + """A test that reaches a member or request the Rust core does not serve yet is skipped with + that member named, so the run separates "not served yet" from a wrong answer.""" + outcome: Final = yield + excinfo: Final[object] = getattr(outcome, "excinfo", None) + if not isinstance(excinfo, tuple): + return + info: Final[tuple[object, ...]] = excinfo # pyright: ignore[reportUnknownVariableType] # pluggy's (type, value, traceback) + if len(info) < 2: + return + from litellm.router_rust.core import RustCoreUnsupported + + error: Final[object] = info[1] + cause: Final[object] = getattr(error, "__cause__", None) + force_exception: Final[object] = getattr(outcome, "force_exception", None) + if (isinstance(error, RustCoreUnsupported) or isinstance(cause, RustCoreUnsupported)) and callable(force_exception): + force_exception(pytest.skip.Exception(f"not served by the Rust core yet: {error}", _use_item_location=True)) diff --git a/tests/unit/test_router/core/construction_precall_helpers.py b/tests/unit/test_router/core/construction_precall_helpers.py new file mode 100644 index 00000000000..c83bd37d692 --- /dev/null +++ b/tests/unit/test_router/core/construction_precall_helpers.py @@ -0,0 +1,98 @@ +""" +Helpers for the construction / pre-call characterization tests. + +``CallbackHost`` is a ``ScriptedHost`` whose callback manager also has +``remove_callback_from_list_by_object`` (litellm's ``LoggingCallbackManager`` has it; ``HostCallbacks`` +does not), so the path that switches optional pre-call checks off at runtime can run. +""" + +import copy +import gc +import uuid +from collections.abc import Generator, Mapping +from typing import Final + +import litellm +from tests.unit.test_router.core.scripted_host import HostCallbacks, ScriptedHost, deployment + + +class RemovingHostCallbacks(HostCallbacks): + def remove_callback_from_list_by_object( + self, + callback_list: object, + obj: object, + require_self: bool = True, + ) -> None: + """``LoggingCallbackManager.remove_callback_from_list_by_object``, on the host's lists.""" + if not isinstance(callback_list, list): + return + if require_self: + doomed = [c for c in callback_list if getattr(c, "__self__", None) is obj] + else: + doomed = [c for c in callback_list if c == obj] + for callback in doomed: + callback_list.remove(callback) + + +class CallbackHost(ScriptedHost): + __slots__ = () + + @property + def callbacks(self) -> RemovingHostCallbacks: # pyright: ignore[reportIncompatibleVariableOverride] + return RemovingHostCallbacks(self) + + +def callback_names(callbacks: list[object]) -> list[str]: # mutable-ok: a fresh list + """Each callback by name: a bound method by its function name, anything else by its class name.""" + names: Final[list[str]] = [] + for callback in callbacks: + function = getattr(callback, "__func__", None) + names.append(function.__name__ if function is not None else type(callback).__name__) + return names + + +def host_callback_names(host: ScriptedHost) -> dict[str, list[str]]: # mutable-ok: a fresh dict + return {name: callback_names(callbacks) for name, callbacks in host.lists.items()} + + +def windowed( + group: str, deployment_id: str, max_input_tokens: int, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry + """A deployment declaring ``max_input_tokens`` under a model name no other test uses: a deployment's + model_info is written into the process-wide cost map under its model name and id.""" + model_info: Final = {"max_input_tokens": max_input_tokens, **dict(params.pop("model_info", None) or {})} # pyright: ignore[reportArgumentType] + return deployment( + group, + deployment_id, + model=f"openai/cw-{deployment_id}-{uuid.uuid4().hex[:8]}", + model_info=model_info, + **params, + ) + + +def collect_other_routers() -> None: + """Drop unreachable routers from earlier tests: retire() and runtime pre-call-check changes look + at every live router in the process, not only the ones on this host.""" + gc.collect() + + +def restoring_model_cost(*names: str) -> Generator[None, None, None]: + """Undo what routers built meanwhile wrote into the process cost map: keys they added are removed + and the named entries get their old value back.""" + added_before: Final = frozenset(litellm.model_cost) + saved: Final = {name: copy.deepcopy(litellm.model_cost[name]) for name in names if name in litellm.model_cost} + yield + for key in frozenset(litellm.model_cost) - added_before: + litellm.model_cost.pop(key, None) + litellm.model_cost.update(saved) + + +__all__ = [ + "CallbackHost", + "RemovingHostCallbacks", + "callback_names", + "collect_other_routers", + "host_callback_names", + "restoring_model_cost", + "windowed", +] diff --git a/tests/unit/test_router/core/endpoints_helpers.py b/tests/unit/test_router/core/endpoints_helpers.py new file mode 100644 index 00000000000..ac629352e95 --- /dev/null +++ b/tests/unit/test_router/core/endpoints_helpers.py @@ -0,0 +1,210 @@ +""" +Helpers for the endpoint characterization tests (``test_endpoints_characterization*.py``). + +``EndpointHost`` is a ``ScriptedHost`` that also scripts the host ops ``ScriptedHost`` leaves to +litellm (assistants, skills, realtime, fine-tuning, batches, ...), so every ``host.invoke`` the core +makes is answered and recorded. + +Some endpoints do not go through ``host.invoke``: the Python core binds litellm's own functions for +them when it is built (vector stores, Google GenAI, OCR, search, videos, containers, interactions, +managed agents). ``patch_direct_endpoints`` replaces those litellm functions, before the core is +built, with recording stand-ins that answer from the same host script, under the same ``__name__`` +(the core routes on the handler's name). A call to one of them is recorded with the function's name +as its ``op``. +""" + +import importlib +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +from tests.unit.test_router.core.scripted_host import INFERENCE_OPS, ScriptedHost + +EXTRA_HOST_OPS: Final = frozenset( + { + "aadapter_completion", + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", + "acompact_responses", + "_arealtime", + "acreate_realtime_client_secret", + "arealtime_calls", + "acreate_realtime_transcription_session", + "_aresponses_websocket", + "acancel_fine_tuning_job", + "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job", + "afile_list", + "aimage_edit", + "allm_passthrough_route", + "acreate_skill", + "alist_skills", + "aget_skill", + "adelete_skill", + "aretrieve_batch", + "acancel_batch", + "alist_batches", + } +) + + +class EndpointHost(ScriptedHost): + """``ScriptedHost`` that scripts every host op the core's endpoints invoke.""" + + __slots__ = () + + def invoke(self, op: str) -> Callable[..., object]: + if op in INFERENCE_OPS or op not in EXTRA_HOST_OPS: + return super().invoke(op) + return recording_endpoint(self, op, is_async=True) + + +def recording_endpoint(host: ScriptedHost, name: str, *, is_async: bool) -> Callable[..., object]: + """A stand-in for litellm's ``name`` that records the call on ``host`` and answers from its script.""" + if is_async: + + async def endpoint(*args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request + return await host._answer_async(name, args, kwargs) # pyright: ignore[reportPrivateUsage] # the test host + + endpoint.__name__ = name + endpoint.__qualname__ = name + return endpoint + + def sync_endpoint(*args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request + return host._answer_sync(name, args, kwargs) # pyright: ignore[reportPrivateUsage] # the test host + + sync_endpoint.__name__ = name + sync_endpoint.__qualname__ = name + return sync_endpoint + + +# module -> {attribute on the module: the function's own name} +DIRECT_ENDPOINTS: Final[Mapping[str, Mapping[str, str]]] = { + "litellm.vector_stores.main": { + name: name + for name in ( + "asearch", + "search", + "acreate", + "create", + "aretrieve", + "retrieve", + "alist", + "list", + "aupdate", + "update", + "adelete", + "delete", + ) + }, + "litellm.vector_stores": {"acreate": "acreate"}, + "litellm.vector_store_files.main": { + name: name + for name in ( + "acreate", + "create", + "alist", + "list", + "aretrieve", + "retrieve", + "aretrieve_content", + "retrieve_content", + "aupdate", + "update", + "adelete", + "delete", + ) + }, + "litellm.google_genai": { + name: name + for name in ("agenerate_content", "generate_content", "agenerate_content_stream", "generate_content_stream") + }, + "litellm.ocr": {"aocr": "aocr", "ocr": "ocr"}, + "litellm.search": {"asearch": "asearch", "search": "search"}, + "litellm.decisions": {"adecisions": "adecisions", "decisions": "decisions"}, + "litellm.systemone": {"asystemone": "asystemone", "systemone": "systemone"}, + "litellm.extraction": {"aextraction": "aextraction", "extraction": "extraction"}, + "litellm.videos": { + f"{prefix}video_{name}": f"{prefix}video_{name}" + for prefix in ("a", "") + for name in ( + "generation", + "list", + "status", + "content", + "remix", + "create_character", + "get_character", + "edit", + "extension", + ) + }, + "litellm.containers": { + f"{prefix}{name}": f"{prefix}{name}" + for prefix in ("a", "") + for name in ("create_container", "list_containers", "retrieve_container", "delete_container") + }, + "litellm.interactions": { + f"{prefix}{name}": f"{prefix}{name}" for prefix in ("a", "") for name in ("create", "get", "delete", "cancel") + }, + "litellm.interactions.agents": { + f"{prefix}{name}": f"{prefix}{name}" + for prefix in ("a", "") + for name in ("create", "get", "delete", "list", "list_versions") + }, +} + +# Recorded op names for the direct endpoints whose bare names collide across modules. +RECORDED_NAMES: Final[Mapping[tuple[str, str], str]] = { + **{("litellm.vector_stores.main", name): f"vs.{name}" for name in DIRECT_ENDPOINTS["litellm.vector_stores.main"]}, + ("litellm.vector_stores", "acreate"): "vs.acreate", + **{ + ("litellm.vector_store_files.main", name): f"vsf.{name}" + for name in DIRECT_ENDPOINTS["litellm.vector_store_files.main"] + }, + **{("litellm.interactions", name): f"interactions.{name}" for name in DIRECT_ENDPOINTS["litellm.interactions"]}, + **{ + ("litellm.interactions.agents", name): f"agents.{name}" + for name in DIRECT_ENDPOINTS["litellm.interactions.agents"] + }, +} + + +def patch_direct_endpoints(monkeypatch: pytest.MonkeyPatch, host: ScriptedHost) -> None: + """Replace litellm's directly bound endpoint functions with recording stand-ins on ``host``. + + The stand-in keeps the function's ``__name__`` (the core routes on it); the call is recorded + under ``RECORDED_NAMES`` when the bare name is ambiguous, else under the function's name.""" + for module_name, attributes in DIRECT_ENDPOINTS.items(): + module: Final = importlib.import_module(module_name) + for attribute, function_name in attributes.items(): + original: Final[object] = getattr(module, attribute) + is_async: Final = function_name.startswith("a") + recorded: Final = RECORDED_NAMES.get((module_name, attribute), function_name) + stand_in: Final = recording_endpoint(host, recorded, is_async=is_async) + stand_in.__name__ = getattr(original, "__name__", function_name) + monkeypatch.setattr(module, attribute, stand_in) + from litellm.containers import endpoint_factory + + for name in list(endpoint_factory._generated_endpoints): # pyright: ignore[reportPrivateUsage] # litellm's registry + monkeypatch.setitem( + endpoint_factory._generated_endpoints, # pyright: ignore[reportPrivateUsage] # litellm's registry + name, + recording_endpoint(host, name, is_async=name.startswith("a")), + ) + + +__all__ = [ + "DIRECT_ENDPOINTS", + "EXTRA_HOST_OPS", + "EndpointHost", + "patch_direct_endpoints", + "recording_endpoint", +] diff --git a/tests/unit/test_router/core/fallbacks_helpers.py b/tests/unit/test_router/core/fallbacks_helpers.py new file mode 100644 index 00000000000..c88c31a9371 --- /dev/null +++ b/tests/unit/test_router/core/fallbacks_helpers.py @@ -0,0 +1,105 @@ +"""Helpers for the fallbacks characterization tests (``test_fallbacks_characterization*.py``).""" + +from collections.abc import Callable, Mapping +from typing import Final + +from litellm.router_core import RouterCore +from litellm.router_python_host import Endpoint +from tests.unit.test_router.core.scripted_host import Outcome, ScriptedHost, deployment + +MakeCore = Callable[..., RouterCore] + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +GROUPS_ABC: Final = [deployment("a", "a"), deployment("b", "b"), deployment("c", "c")] +GROUP_A_OF_THREE: Final = [deployment("a", "a1"), deployment("a", "a2"), deployment("a", "a3"), deployment("b", "b")] +TIERED_A: Final = [ + deployment("a", "a1", order=1), + deployment("a", "a2", order=2), + deployment("b", "b"), + deployment("c", "c"), +] + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_ABC, "num_retries": 0, **arguments}) + + +async def complete(core: RouterCore, model: str = "a", **request: object) -> object: # kwargs-ok: request kwargs + return await core.acompletion(model=model, messages=MESSAGES, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def complete_sync(core: RouterCore, model: str = "a", **request: object) -> object: # kwargs-ok: request kwargs + return core.completion(model=model, messages=MESSAGES, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def hidden(response: object) -> Mapping[str, object]: + return getattr(response, "_hidden_params") + + +def headers(response: object) -> Mapping[str, object]: + found: Final = hidden(response)["additional_headers"] + assert isinstance(found, Mapping) + return found + + +def metadata_of(kwargs: Mapping[str, object]) -> Mapping[str, object]: + found: Final = kwargs.get("metadata") or kwargs.get("litellm_metadata") or {} + assert isinstance(found, Mapping) + return found + + +def show_router_debug(host: ScriptedHost) -> None: + host.set_dynamic(expose_router_debug_in_errors=True) + + +def fails_once_then_answers(error: BaseException) -> Callable[..., Outcome]: + """One scripted outcome shared by several models: the first call to any of them raises + ``error``, every later call answers the default response.""" + calls: Final = [0] + + def outcome(**_request: object) -> Outcome: # kwargs-ok: the scripted request + calls[0] += 1 + return error if calls[0] == 1 else None + + return outcome + + +class NamedEndpointHost(ScriptedHost): + """A ``ScriptedHost`` whose ``anthropic_messages`` endpoint carries that name, as litellm's + own function does. The core recognises a /v1/messages call (and its safeguard-refusal gate) + by the endpoint's ``__name__``; ``ScriptedHost``'s endpoints are all named ``endpoint``.""" + + __slots__ = () + + def invoke(self, op: str) -> Endpoint: + endpoint: Final = super().invoke(op) + if op == "anthropic_messages": + endpoint.__name__ = op # pyright: ignore[reportFunctionMemberAccess] # a plain function + return endpoint + + +def anthropic_refusal(*, with_stop_details: bool = True) -> dict[str, object]: # mutable-ok: a provider response + response: Final[dict[str, object]] = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [], + "stop_reason": "refusal", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + if with_stop_details: + response["stop_details"] = {"type": "refusal", "category": None, "explanation": "no"} + return response + + +def anthropic_message(text: str = "fine") -> dict[str, object]: # mutable-ok: a provider response + return { + "id": "msg_2", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } diff --git a/tests/unit/test_router/core/gaps_l11100_helpers.py b/tests/unit/test_router/core/gaps_l11100_helpers.py new file mode 100644 index 00000000000..38fb2006503 --- /dev/null +++ b/tests/unit/test_router/core/gaps_l11100_helpers.py @@ -0,0 +1,112 @@ +"""Shared helpers for ``test_gaps_l11100*.py``: strategy-router marker deployments and a host whose +callback registration can also remove a callback from every list (what the adaptive router's hook +sync asks of litellm's ``LoggingCallbackManager``).""" + +from collections.abc import Callable, Mapping +from typing import Final + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import HostCallbacks, ScriptedHost + +MakeCore = Callable[..., RouterCore] + +TIERS: Final = {"SIMPLE": "small", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} + + +class RemovingHostCallbacks(HostCallbacks): + """``HostCallbacks`` plus ``remove_callback_from_all_lists``, as litellm's callback manager has it.""" + + def __init__(self, host: "RemovingHost") -> None: + super().__init__(host) + self._owner: Final = host + + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: + del require_self + for name, callbacks in self._owner.lists.items(): + self._owner.lists[name] = [callback for callback in callbacks if callback is not obj] + + +class RemovingHost(ScriptedHost): + """A ``ScriptedHost`` whose callbacks can be removed from every list.""" + + __slots__ = () + + @property + def callbacks(self) -> RemovingHostCallbacks: # pyright: ignore[reportIncompatibleVariableOverride] # same methods plus one + return RemovingHostCallbacks(self) + + +def complexity_marker( + name: str = "smart", + deployment_id: str = "marker", + *, + config: Mapping[str, object] | None = None, + **params: object, # kwargs-ok: extra marker litellm_params +) -> dict[str, object]: # mutable-ok: a model_list entry + """An ``auto_router/complexity_router`` marker with the shipped tiers and default model ``big``.""" + return { + "model_name": name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": dict(config) if config is not None else {"tiers": dict(TIERS)}, + "complexity_router_default_model": "big", + **params, + }, + "model_info": {"id": deployment_id}, + } + + +def adaptive_marker( + name: str, deployment_id: str, available_models: list[str], **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: extra marker litellm_params + """An ``auto_router/adaptive_router`` marker over ``available_models``.""" + return { + "model_name": name, + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": available_models}, + **params, + }, + "model_info": {"id": deployment_id}, + } + + +def marker( + name: str, deployment_id: str, model: str, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: marker litellm_params + """Any strategy-router marker: ``model`` is its ``auto_router/...`` litellm model.""" + return {"model_name": name, "litellm_params": {"model": model, **params}, "model_info": {"id": deployment_id}} + + +def registered_tags(core: RouterCore, kind: str) -> dict[str, list[tuple[str, ...]]]: # mutable-ok: a fresh dict + """``{router name: [tags of each registered strategy]}`` for a pre-routing strategy registry.""" + registry: Final = core.registry().plugins(kind) # pyright: ignore[reportCallIssue, reportArgumentType] # one of the strategy kinds + assert isinstance(registry, dict) + return {name: [entry.tags for entry in entries] for name, entries in registry.items()} # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # TaggedPreRoutingStrategy lists + + +def model_ids(core: RouterCore) -> list[str]: # mutable-ok: a fresh list + """The ``model_info.id`` of every deployment on the model list, in order.""" + return [str(entry["model_info"]["id"]) for entry in core.registry().all()] # pyright: ignore[reportAny] # model_list rows + + +class ModelsEndpoint: + """A stand-in for the ``AsyncHTTPHandler`` that model-info discovery reads ``/v1/models`` with: + answers ``body`` (or raises ``error``) and records each request's url and headers.""" + + def __init__(self, body: Mapping[str, object] | None = None, error: BaseException | None = None) -> None: + self.body: Final = body or {"data": []} + self.error: Final = error + self.requests: Final[list[tuple[str, dict[str, str]]]] = [] + + async def get( + self, url: str, headers: dict[str, str] | None = None, **_: object + ) -> object: # kwargs-ok: httpx options + import json + + import httpx + + self.requests.append((url, dict(headers or {}))) + if self.error is not None: + raise self.error + return httpx.Response(200, content=json.dumps(self.body).encode(), request=httpx.Request("GET", url)) diff --git a/tests/unit/test_router/core/registry_helpers.py b/tests/unit/test_router/core/registry_helpers.py new file mode 100644 index 00000000000..aa154454b8d --- /dev/null +++ b/tests/unit/test_router/core/registry_helpers.py @@ -0,0 +1,47 @@ +"""Shared helpers for the registry characterization tests (``test_registry_characterization*.py``).""" + +import uuid +from collections.abc import Callable, Iterable, Mapping + +from litellm.router_core import RouterCore +from litellm.types.router import Deployment +from tests.unit.test_router.core.scripted_host import deployment + +MakeCore = Callable[..., RouterCore] + +MESSAGES = [{"role": "user", "content": "hi"}] + + +def ids(deployments: Iterable[Mapping[str, object]] | None) -> list[str]: # mutable-ok: a fresh list + """The ``model_info.id`` of each deployment dict, in order.""" + return [str(entry["model_info"]["id"]) for entry in deployments or ()] # pyright: ignore[reportIndexIssue] # deployment dicts + + +def names_and_ids( + deployments: Iterable[Mapping[str, object]] | None, +) -> list[tuple[str, str]]: # mutable-ok: a fresh list + """``(model_name, model_info.id)`` of each deployment dict, in order.""" + return [(str(entry["model_name"]), str(entry["model_info"]["id"])) for entry in deployments or ()] # pyright: ignore[reportIndexIssue] # deployment dicts + + +def typed(model_name: str, deployment_id: str, **params: object) -> Deployment: # kwargs-ok: litellm_params + """A ``Deployment`` built from the same entry ``scripted_host.deployment`` returns.""" + return Deployment(**deployment(model_name, deployment_id, **params)) # pyright: ignore[reportArgumentType] # a model_list entry + + +def team_deployment( + model_name: str, deployment_id: str, team_id: str, public_name: str, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: litellm_params + """A team-owned deployment: ``model_info.team_id`` and ``team_public_model_name`` set.""" + extra_info: dict[str, object] = dict(params.pop("model_info", None) or {}) # pyright: ignore[reportArgumentType] # a model_info dict + return deployment( + model_name, + deployment_id, + model_info={"team_id": team_id, "team_public_model_name": public_name, **extra_info}, + **params, + ) + + +def unique_id(prefix: str = "reg") -> str: + """A deployment id no other test uses, for tests that write into the process-wide ``litellm.model_cost``.""" + return f"{prefix}-{uuid.uuid4().hex}" diff --git a/tests/unit/test_router/core/retries_health_helpers.py b/tests/unit/test_router/core/retries_health_helpers.py new file mode 100644 index 00000000000..eae0b836530 --- /dev/null +++ b/tests/unit/test_router/core/retries_health_helpers.py @@ -0,0 +1,198 @@ +""" +Helpers for the retries / deployment-health characterization tests (``test_retries_health_characterization*.py``). + +Backoff values carry litellm's random jitter (up to ``JITTER`` seconds), so sleeps are asserted inside +``[base, base + JITTER)``. +""" + +import asyncio +import copy +from collections.abc import Callable, Mapping +from typing import Final + +import httpx + +import litellm +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import Call, ScriptedHost + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + +MakeCore = Callable[..., RouterCore] +ErrorBuilder = Callable[[], Exception] + +_REQUEST: Final = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + + +def fresh(builder: ErrorBuilder) -> Callable[..., Exception]: + """A scripted outcome that raises a new error per call.""" + return lambda **_: builder() + + +def http_response(status: int, headers: Mapping[str, str] | None = None) -> httpx.Response: + return httpx.Response(status, headers=dict(headers or {}), request=_REQUEST) + + +def not_found() -> litellm.NotFoundError: + return litellm.NotFoundError(message="no such model", llm_provider="openai", model="gpt-4o") + + +def permission_denied() -> litellm.PermissionDeniedError: + return litellm.PermissionDeniedError( + message="forbidden", llm_provider="openai", model="gpt-4o", response=http_response(403) + ) + + +def unprocessable() -> litellm.UnprocessableEntityError: + return litellm.UnprocessableEntityError( + message="unprocessable", llm_provider="openai", model="gpt-4o", response=http_response(422) + ) + + +def service_unavailable() -> litellm.ServiceUnavailableError: + return litellm.ServiceUnavailableError(message="unavailable", llm_provider="openai", model="gpt-4o") + + +def bad_gateway() -> litellm.BadGatewayError: + return litellm.BadGatewayError(message="bad gateway", llm_provider="openai", model="gpt-4o") + + +def api_connection() -> litellm.APIConnectionError: + return litellm.APIConnectionError(message="connection reset", llm_provider="openai", model="gpt-4o") + + +def api_status(status: int) -> ErrorBuilder: + return lambda: litellm.APIError( + status_code=status, message=f"status {status}", llm_provider="openai", model="gpt-4o" + ) + + +def plain_error() -> ValueError: + return ValueError("not a provider error") + + +def rate_limited(retry_after: str | None = None) -> litellm.RateLimitError: + headers: Final = {"retry-after": retry_after} if retry_after is not None else {} + return litellm.RateLimitError( + message="slow down", llm_provider="openai", model="gpt-4o", response=httpx.Response(429, headers=headers) + ) + + +async def attempt(core: RouterCore, model: str = "g", **kwargs: object) -> object: # kwargs-ok: request kwargs + """The response, or the error the core raised. A request the Rust core does not serve yet propagates, so the + test is skipped rather than failed.""" + from litellm.router_rust.core import RustCoreUnsupported + + try: + return await core.acompletion(model=model, messages=MESSAGES, **kwargs) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + except RustCoreUnsupported: + raise + except Exception as error: + return error + + +def attempt_sync(core: RouterCore, model: str = "g", **kwargs: object) -> object: # kwargs-ok: request kwargs + from litellm.router_rust.core import RustCoreUnsupported + + try: + return core.completion(model=model, messages=MESSAGES, **kwargs) # pyright: ignore[reportAttributeAccessIssue] # every core serves completion + except RustCoreUnsupported: + raise + except Exception as error: + return error + + +def cooled_down(core: RouterCore, ids: list[str]) -> dict[str, Mapping[str, object]]: # mutable-ok: a fresh dict + return dict(core.health().cooldowns.get_active_cooldowns(ids, None)) + + +def headers_of(response: object) -> Mapping[str, object]: + hidden: Final = getattr(response, "_hidden_params", None) or {} + return hidden.get("additional_headers") or {} + + +class MetadataRecorder: + """A scripted outcome factory that snapshots the request metadata each call carried, then answers.""" + + KEYS: Final = ( + "attempted_retries", + "max_retries", + "model_group", + "model_group_size", + "request_retry_count", + "previous_models", + ) + + def __init__(self) -> None: + self.snapshots: Final[list[dict[str, object]]] = [] + + def then(self, builder: ErrorBuilder | None) -> Callable[..., object]: + def outcome(**kwargs: object) -> object: # kwargs-ok: the request + metadata: Final = kwargs.get("metadata") or kwargs.get("litellm_metadata") or {} + assert isinstance(metadata, Mapping) + self.snapshots.append({key: copy.deepcopy(metadata.get(key)) for key in self.KEYS}) + return builder() if builder is not None else None + + return outcome + + +class BackgroundSuccessLoggingHost(ScriptedHost): + """``ScriptedHost`` whose async success callbacks run in a background task after the response is returned, + as litellm's logging worker runs them (``ScriptedHost`` runs them before returning). Await ``drain()`` + before reading what they wrote.""" + + __slots__ = ("_pending",) + + def __init__(self) -> None: + super().__init__() + self._pending: list[asyncio.Task[None]] = [] + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + call, outcome = self._outcome(op, args, kwargs) + error, response = self._finish(op, call, outcome) + if error is not None: + logged_failure: Final = self._logging_kwargs(call, error, response) + for callback in list(self.lists.get("failure_callback", ())): + if callable(callback): + callback(logged_failure, None, 0.0, 0.0) + for callback in list(self.lists.get("_async_failure_callback", ())): + if callable(callback): + await callback(logged_failure, None, 0.0, 0.0) + raise error + + async def log_success(logged_call: Call) -> None: + logged: Final = self._logging_kwargs(logged_call, None, response) + for callback in list(self.lists.get("_async_success_callback", ())): + if callable(callback): + await callback(logged, response, 0.0, 0.0) + + self._pending.append(asyncio.ensure_future(log_success(call))) + return response + + async def drain(self) -> None: + while self._pending: + await self._pending.pop() + + +__all__ = [ + "MESSAGES", + "BackgroundSuccessLoggingHost", + "ErrorBuilder", + "MakeCore", + "MetadataRecorder", + "api_connection", + "api_status", + "attempt", + "attempt_sync", + "bad_gateway", + "cooled_down", + "fresh", + "headers_of", + "http_response", + "not_found", + "permission_denied", + "plain_error", + "rate_limited", + "service_unavailable", + "unprocessable", +] diff --git a/tests/unit/test_router/core/rust_parity_pending.txt b/tests/unit/test_router/core/rust_parity_pending.txt new file mode 100644 index 00000000000..ee2aa4e976c --- /dev/null +++ b/tests/unit/test_router/core/rust_parity_pending.txt @@ -0,0 +1,4 @@ +# unskipped by p3-selection +test_construction_precall_characterization.py::test_retiring_one_router_keeps_the_callbacks_of_another_on_the_same_host[rust] +# unskipped by p6-final +test_construction_precall_characterization.py::test_alerting_config_registers_slack_alerting_and_its_slow_response_callback[rust] diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index af28160b620..8c7f8df0648 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -324,6 +324,9 @@ class ScriptedHost(PythonHost): return done() + def sleep_sync(self, seconds: float) -> None: + self.sleeps.append(seconds) + def dynamic(self) -> RouterDynamicConfig: return self._dynamic @@ -398,9 +401,9 @@ class ScriptedHost(PythonHost): return outcome, None return None, _default_response(op, call.model) if outcome is None else outcome - async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: - call, outcome = self._outcome(op, args, kwargs) - error, response = self._finish(op, call, outcome) + async def log_async(self, call: Call, error: BaseException | None, response: object) -> None: + """Run the callbacks an async call's logging runs: the failure callbacks for ``error``, else + the async success callbacks for ``response``.""" logged: Final = self._logging_kwargs(call, error, response) if error is not None: for callback in list(self.lists.get("failure_callback", ())): @@ -409,24 +412,50 @@ class ScriptedHost(PythonHost): for callback in list(self.lists.get("_async_failure_callback", ())): if callable(callback): await callback(logged, None, 0.0, 0.0) - raise error + return for callback in list(self.lists.get("_async_success_callback", ())): if callable(callback): await callback(logged, response, 0.0, 0.0) - return response - def _answer_sync(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: - call, outcome = self._outcome(op, args, kwargs) - error, response = self._finish(op, call, outcome) + def log_sync(self, call: Call, error: BaseException | None, response: object) -> None: + """``log_async`` for a sync call: the failure callbacks, else the sync success callbacks.""" logged: Final = self._logging_kwargs(call, error, response) if error is not None: for callback in list(self.lists.get("failure_callback", ())): if callable(callback): callback(logged, None, 0.0, 0.0) - raise error + return for callback in list(self.lists.get("success_callback", ())): if callable(callback): callback(logged, response, 0.0, 0.0) + + def _logs_at_stream_end(self, call: Call, response: object) -> bool: + """A stream whose logging object logs through this host (``streaming_helpers.HostLogging``) + logs when it ends, as litellm's stream logging does, so the call logs nothing when it opens.""" + bind: Final[object] = getattr(getattr(response, "logging_obj", None), "bind_host", None) + if not callable(bind): + return False + bind(self, call) + return True + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + call, outcome = self._outcome(op, args, kwargs) + error, response = self._finish(op, call, outcome) + if error is not None: + await self.log_async(call, error, response) + raise error + if not self._logs_at_stream_end(call, response): + await self.log_async(call, None, response) + return response + + def _answer_sync(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + call, outcome = self._outcome(op, args, kwargs) + error, response = self._finish(op, call, outcome) + if error is not None: + self.log_sync(call, error, response) + raise error + if not self._logs_at_stream_end(call, response): + self.log_sync(call, None, response) return response @@ -452,25 +481,46 @@ def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwarg return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds -CoreBuilder = Callable[..., RouterCore] -"""``builder(host, **router_arguments)``: a router core on ``host``.""" - -CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core} -"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a -configuration raises ``CoreDeclined``, and the test is skipped for that core.""" - - class CoreDeclined(Exception): """The core does not serve this configuration (the router would stay on the Python core).""" +CoreBuilder = Callable[..., RouterCore] +"""``builder(host, **router_arguments)``: a router core on ``host``.""" + + +def rust_core(host: ScriptedHost, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + """A ``RustCore`` on ``host``, built the way ``make_core`` builds one when the Rust router is on.""" + from litellm.router_rust.core import RustCore + from litellm.rust_bridge._native import RustBridgeDeclined + + config: Final = build_router_config(arguments) # an invalid config raises here, as for the Python core + try: + return RustCore.build(config, host) + except RustBridgeDeclined as declined: # make_core keeps such a router on the Python core + raise CoreDeclined(str(declined)) from declined + + +CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core, "rust": rust_core} + + +def records_outcomes_through_callbacks(core: RouterCore) -> bool: + """Whether ``core`` learns attempt outcomes from callbacks it registers on its host. The Python + core does; the Rust core records every outcome from the attempt itself and registers none.""" + return isinstance(core, PythonCore) + + +"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a +configuration raises ``CoreDeclined``, and the test is skipped for that core.""" + + __all__ = [ - "Call", - "HostCallbacks", - "INFERENCE_OPS", "CORE_BUILDERS", + "INFERENCE_OPS", + "Call", "CoreBuilder", "CoreDeclined", + "HostCallbacks", "ScriptedHost", "auth_error", "bad_request", @@ -482,6 +532,7 @@ __all__ = [ "embedding_response", "python_core", "rate_limit", + "rust_core", "server_error", "timeout", ] diff --git a/tests/unit/test_router/core/selection_helpers.py b/tests/unit/test_router/core/selection_helpers.py new file mode 100644 index 00000000000..b38cfec1a19 --- /dev/null +++ b/tests/unit/test_router/core/selection_helpers.py @@ -0,0 +1,88 @@ +""" +Helpers for the selection characterization tests. + +``TrafficHost`` is a ``ScriptedHost`` whose async inference calls also feed the routing strategy +loggers the core registered on the host, the way litellm's logging does around a real call: input +callbacks see the call start (``log_pre_api_call``), and the ``CustomLogger`` instances in the +``callbacks`` list see it succeed or fail with real start and end times. A call can be held open +on a gate (an ``asyncio.Event``) or delayed for real, so in-flight counts and latencies are real +for any core. +""" + +import asyncio +import datetime +from collections.abc import Callable, Mapping +from typing import Final + +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import Call, ScriptedHost + +MakeCore = Callable[..., RouterCore] + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +class TrafficHost(ScriptedHost): + """``ScriptedHost`` plus the strategy-logger traffic litellm's logging produces around a call.""" + + def __init__(self) -> None: + super().__init__() + self.gates: dict[str, asyncio.Event] = {} + self.delays: dict[str, float] = {} + self.started: list[str] = [] # mutable-ok: the models whose calls have started, in order + + def _loggers(self, name: str) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in self.lists.get(name, ()) if isinstance(callback, CustomLogger)) + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + model: Final = str(kwargs.get("model") or (args[0] if args else "")) + logged: Final = self._logging_kwargs(Call(op=op, model=model, kwargs=dict(kwargs)), None, None) + for logger in self._loggers("input_callback"): + logger.log_pre_api_call(model=model, messages=kwargs.get("messages"), kwargs=logged) + self.started.append(model) + start: Final = datetime.datetime.now() + gate: Final = self.gates.get(model) + if gate is not None: + await gate.wait() + delay: Final = self.delays.get(model) + if delay: + await asyncio.sleep(delay) + try: + response = await super()._answer_async(op, args, kwargs) + except Exception as error: + failed = {**logged, "exception": error} + for logger in self._loggers("callbacks"): + await logger.async_log_failure_event(failed, None, start, datetime.datetime.now()) + raise + succeeded: Final = self._logging_kwargs(Call(op=op, model=model, kwargs=dict(kwargs)), None, response) + for logger in self._loggers("callbacks"): + await logger.async_log_success_event(succeeded, response, start, datetime.datetime.now()) + return response + + +async def picked_ids( + core: RouterCore, + model: str, + times: int = 20, + request_kwargs: Mapping[str, object] | None = None, + **select_arguments: object, # kwargs-ok: select's keyword options +) -> set[str]: + """The deployment ids ``select`` picks over ``times`` picks, each with a fresh copy of ``request_kwargs``.""" + picked: Final[set[str]] = set() + for _ in range(times): + chosen = await core.select(model, dict(request_kwargs or {}), **select_arguments) # pyright: ignore[reportCallIssue] # PythonCore's select takes keyword options + assert isinstance(chosen, Mapping) + picked.add(str(chosen["model_info"]["id"])) + return picked + + +def ids_of(deployments: object) -> list[str]: # mutable-ok: a fresh list + """The ids of a ``select(kind="healthy")`` result, in order (a single dict gives one id).""" + if isinstance(deployments, Mapping): + return [str(deployments["model_info"]["id"])] + assert isinstance(deployments, list) + return [str(deployment["model_info"]["id"]) for deployment in deployments] + + +__all__ = ["MESSAGES", "MakeCore", "TrafficHost", "ids_of", "picked_ids"] diff --git a/tests/unit/test_router/core/streaming_helpers.py b/tests/unit/test_router/core/streaming_helpers.py new file mode 100644 index 00000000000..c4764d859f9 --- /dev/null +++ b/tests/unit/test_router/core/streaming_helpers.py @@ -0,0 +1,469 @@ +""" +Scripted streams for the streaming characterization tests. + +``ScriptedHost`` outcomes may be callables, so each builder here returns a callable that the host +calls with the request kwargs and that answers with a fresh stream: + +- ``chat_stream``: a ``CustomStreamWrapper`` over chat chunks, optionally raising ``error`` after + ``fail_after`` chunks, and ending with a ``finish_reason`` chunk unless ``finish`` is False. +- ``anthropic_stream``: an async iterator of raw ``/v1/messages`` SSE byte frames, optionally raising. +- ``responses_stream``: a ``BaseResponsesAPIStreamingIterator`` over Responses API events, + optionally raising. + +A stream built with ``logged=True`` logs through the ``ScriptedHost`` that served it when it ends +(``HostLogging``), as litellm's stream logging does, rather than when the call returns it. +""" + +import datetime +import json +import sys +import types +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from litellm.types.utils import ModelResponseStream + +if TYPE_CHECKING: + from tests.unit.test_router.core.scripted_host import Call, ScriptedHost + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +def install_orjson_stub_if_missing() -> None: + """ + The /v1/messages stream wrapper imports proxy modules that import ``orjson`` and ``soundfile`` at load + time. Where either is not installed (some dev containers), a stand-in is registered so the wrapper can + load (json-backed for ``orjson``, empty for ``soundfile``, which only audio passthrough logging uses); + where they are installed (CI), nothing changes. + """ + try: + import orjson # noqa: F401 # pyright: ignore[reportUnusedImport] # only probing for the module + except ImportError: + stub: Final = types.ModuleType("orjson") + stub.loads = json.loads # pyright: ignore[reportAttributeAccessIssue] # module attribute + stub.dumps = lambda value, **_options: json.dumps(value).encode() # pyright: ignore[reportAttributeAccessIssue] # module attribute + stub.JSONDecodeError = json.JSONDecodeError # pyright: ignore[reportAttributeAccessIssue] # module attribute + sys.modules["orjson"] = stub + try: + import soundfile # noqa: F401 # pyright: ignore[reportUnusedImport, reportMissingImports] # only probing for the module + except (ImportError, OSError): + sys.modules["soundfile"] = types.ModuleType("soundfile") + + +def _logging_obj(model: str, call_type: str = "acompletion", logged: bool = False) -> LiteLLMLogging: + logging_obj: Final = (HostLogging if logged else LiteLLMLogging)( + model=model, + messages=MESSAGES, + stream=True, + call_type=call_type, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id=f"stream-{model}", + function_id=f"stream-{model}", + ) + logging_obj.update_environment_variables(model=model, user=None, optional_params={}, litellm_params={}) + return logging_obj + + +class HostLogging(LiteLLMLogging): + """A stream's logging object that runs the serving ``ScriptedHost``'s callbacks where litellm's + stream logging runs litellm's: once with the assembled response when the stream ends, and once + per failure the stream raises (``dispatch_attempt_failure_handlers``). The host binds itself and + the call when it returns the stream (``ScriptedHost._logs_at_stream_end``).""" + + def bind_host(self, host: "ScriptedHost", call: "Call") -> None: + self._scripted: tuple[ScriptedHost, Call] | None = (host, call) # pyright: ignore[reportUninitializedInstanceVariable] # set when the host binds + self._logged_success = False + + def _bound(self) -> "tuple[ScriptedHost, Call] | None": + return getattr(self, "_scripted", None) + + async def dispatch_success_handlers( # pyright: ignore[reportIncompatibleMethodOverride] # the arguments the stream passes + self, result: object = None, *_args: object, **_kwargs: object + ) -> None: + bound: Final = self._bound() + if bound is None or self._logged_success: + return + self._logged_success = True + await bound[0].log_async(bound[1], None, result) + + def dispatch_attempt_failure_handlers( + self, exception: Exception, traceback_exception: str + ) -> Coroutine[None, None, None]: + bound: Final = self._bound() + + async def dispatch() -> None: + if bound is not None: + await bound[0].log_async(bound[1], exception, None) + + return dispatch() + + def success_handler(self, result: object = None, *_args: object, **_kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # see dispatch_success_handlers + bound: Final = self._bound() + if bound is None or self._logged_success or not getattr(result, "usage", None): + return + self._logged_success = True + bound[0].log_sync(bound[1], None, result) + + def failure_handler(self, exception: Exception, *_args: object, **_kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # see dispatch_success_handlers + bound: Final = self._bound() + if bound is not None: + bound[0].log_sync(bound[1], exception, None) + + +class _EndLogged: + """What a hand-written scripted stream logs once when it ends: its failure, or a response carrying + ``total_tokens`` (the usage the provider's own stream logging would assemble).""" + + def __init__(self, model: str, call_type: str, logged: bool, total_tokens: int) -> None: + self.logging_obj: Final[LiteLLMLogging | None] = _logging_obj(model, call_type, True) if logged else None + self._total_tokens: Final = total_tokens + self._logged = False + + async def log_end(self, error: BaseException | None) -> None: + if not isinstance(self.logging_obj, HostLogging) or self._logged: + return + self._logged = True + if error is not None: + # The provider's own error, as a stream logs it before wrapping it for the router. + original: Final = getattr(error, "original_exception", None) + logged: Final = original if isinstance(original, Exception) else error + await self.logging_obj.dispatch_attempt_failure_handlers( + logged if isinstance(logged, Exception) else Exception(logged), "" + ) + return + usage: Final = types.SimpleNamespace(total_tokens=self._total_tokens) + await self.logging_obj.dispatch_success_handlers(types.SimpleNamespace(usage=usage)) + + +def text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": text}}]) + + +def role_chunk() -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}]) + + +def finish_chunk(reason: str = "stop", usage: Mapping[str, int] | None = None) -> ModelResponseStream: + chunk: Final = ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": reason}]) + if usage is not None: + setattr(chunk, "usage", dict(usage)) + return chunk + + +def chat_stream( + *texts: str, + error: BaseException | None = None, + fail_after: int | None = None, + finish: bool = True, + leading_role_chunk: bool = False, + model: str = "d1", + hidden_params: Mapping[str, object] | None = None, + sync: bool = False, + usage: Mapping[str, int] | None = None, + logged: bool = False, +) -> Callable[..., CustomStreamWrapper]: + """ + A stream of ``texts`` (one chunk each). With ``error``, raises it after ``fail_after`` chunks + (default: after every text chunk, before the finish chunk). ``leading_role_chunk`` sends a + role-only, empty-content chunk first, which counts as the first chunk but carries no content. + ``usage`` rides on the finish chunk; ``logged`` logs through the host when the stream ends. + """ + + def answer(**request: object) -> CustomStreamWrapper: # kwargs-ok: the host passes the request kwargs + produced: list[ModelResponseStream] = [] + if leading_role_chunk: + produced.append(role_chunk()) + produced.extend(text_chunk(text) for text in texts) + cut_at: Final = len(produced) if fail_after is None else fail_after + + def items() -> Iterator[ModelResponseStream]: + for index, chunk in enumerate(produced): + if error is not None and index == cut_at: + raise error + yield chunk + if error is not None: + raise error + if finish: + yield finish_chunk(usage=usage) + + async def achunks() -> AsyncIterator[ModelResponseStream]: + for chunk in items(): + yield chunk + + wrapper: Final = CustomStreamWrapper( + completion_stream=items() if sync else achunks(), + model=model, + custom_llm_provider="openai", + logging_obj=_logging_obj(model, "completion" if sync else "acompletion", logged), + ) + if hidden_params is not None: + wrapper._hidden_params.update(hidden_params) # pyright: ignore[reportPrivateUsage] # test seam + return wrapper + + return answer + + +async def collect_texts(stream: object) -> str: + """The text content of every chunk of an async chat stream, joined.""" + texts: Final[list[str]] = [] + async for chunk in stream: # pyright: ignore[reportGeneralTypeIssues] # an async stream + choices = getattr(chunk, "choices", None) + if choices: + content = choices[0].delta.content + if content: + texts.append(content) + return "".join(texts) + + +def collect_texts_sync(stream: object) -> str: + texts: Final[list[str]] = [] + for chunk in stream: # pyright: ignore[reportGeneralTypeIssues] # a sync stream + choices = getattr(chunk, "choices", None) + if choices: + content = choices[0].delta.content + if content: + texts.append(content) + return "".join(texts) + + +# --- /v1/messages ------------------------------------------------------------------------------ + + +def sse(event: str, data: Mapping[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +MESSAGE_START: Final = sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "content": []}}, +) +BLOCK_START: Final = sse( + "content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} +) +PING: Final = sse("ping", {"type": "ping"}) +MESSAGE_STOP: Final = sse("message_stop", {"type": "message_stop"}) + + +def text_delta(text: str) -> bytes: + return sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + +def error_frame(error_type: str = "overloaded_error", message: str = "Overloaded") -> bytes: + return sse("error", {"type": "error", "error": {"type": error_type, "message": message}}) + + +def refusal_delta(stop_details: Mapping[str, object] | None = None) -> bytes: + delta: Final[dict[str, object]] = {"stop_reason": "refusal"} + if stop_details is not None: + delta["stop_details"] = dict(stop_details) + return sse("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 1}}) + + +class ByteStream: + """An async iterator of SSE frames that records whether it was closed. ``logged`` logs through the host + when it ends, with ``total_tokens`` (``_EndLogged``).""" + + def __init__( + self, + frames: Sequence[bytes], + error: BaseException | None, + hidden_params: Mapping[str, object] | None, + logged: bool = False, + total_tokens: int = 0, + ) -> None: + self._frames: Final = list(frames) + self._error: Final = error + self._index = 0 + self.closed = False + self._end: Final = _EndLogged("d1", "anthropic_messages", logged, total_tokens) + if logged: + self.logging_obj = self._end.logging_obj + if hidden_params is not None: + self._hidden_params = dict(hidden_params) + + def __aiter__(self) -> "ByteStream": + return self + + async def __anext__(self) -> bytes: + if self._index < len(self._frames): + self._index += 1 + return self._frames[self._index - 1] + await self._end.log_end(self._error) + if self._error is not None: + raise self._error + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +def anthropic_stream( + *frames: bytes, + error: BaseException | None = None, + hidden_params: Mapping[str, object] | None = None, + opened: list[ByteStream] | None = None, + logged: bool = False, + total_tokens: int = 0, +) -> Callable[..., ByteStream]: + def answer(**_request: object) -> ByteStream: # kwargs-ok: the host passes the request kwargs + stream: Final = ByteStream(frames, error, hidden_params, logged, total_tokens) + if opened is not None: + opened.append(stream) + return stream + + return answer + + +async def collect_bytes(stream: object) -> list[bytes]: + return [frame async for frame in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + + +# --- Responses API ----------------------------------------------------------------------------- + + +class Event: + """A Responses API stream event: only ``type`` (and optionally ``response``) is read.""" + + def __init__(self, event_type: str, response: object = None, delta: str | None = None) -> None: + self.type = event_type + self.response = response + self.delta = delta + self._hidden_params: dict[str, object] = {} + + def __repr__(self) -> str: + return f"Event({self.type!r})" + + +class BackgroundResponse: + background = True + id = "resp_bg" + + +class ScriptedResponsesStream(BaseResponsesAPIStreamingIterator): + """A Responses stream over ``events`` that bypasses the provider plumbing of the base class.""" + + def __init__( # pyright: ignore[reportMissingSuperCall] # the base needs an httpx response + self, + events: Sequence[object], + error: BaseException | None, + model: str, + custom_llm_provider: str, + hidden_params: Mapping[str, object] | None, + started_tool_call: bool = False, + completed_response: object = None, + logged: bool = False, + total_tokens: int = 0, + ) -> None: + self._end: Final = _EndLogged(model, "aresponses", logged, total_tokens) + self._events: Final = list(events) + self._started_tool_call = started_tool_call + self._error: Final = error + self._index = 0 + self._hidden_params = dict(hidden_params or {}) + self.model = model + self.custom_llm_provider = custom_llm_provider + self.logging_obj = self._end.logging_obj or _logging_obj(model, "aresponses") + self.litellm_metadata = None + self.responses_api_provider_config = None + self.finished = False + self.completed_response = completed_response + self.response = None + self.start_time = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + self.request_data = {} + self.call_type = None + self.closed = False + + def __aiter__(self) -> "ScriptedResponsesStream": + return self + + async def __anext__(self) -> object: + if self._index < len(self._events): + self._index += 1 + return self._events[self._index - 1] + await self._end.log_end(self._error) + if self._error is not None: + raise self._error + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +def responses_stream( + *events: object, + error: BaseException | None = None, + model: str = "d1", + custom_llm_provider: str = "anthropic", + hidden_params: Mapping[str, object] | None = None, + opened: list[ScriptedResponsesStream] | None = None, + started_tool_call: bool = False, + completed_response: object = None, + logged: bool = False, + total_tokens: int = 0, +) -> Callable[..., ScriptedResponsesStream]: + """``completed_response`` is the terminal event the stream latched before failing, if any. ``logged`` + logs through the host when the stream ends, with ``total_tokens``.""" + + def answer(**_request: object) -> ScriptedResponsesStream: # kwargs-ok: the host passes the request kwargs + stream: Final = ScriptedResponsesStream( + events, + error, + model, + custom_llm_provider, + hidden_params, + started_tool_call, + completed_response, + logged, + total_tokens, + ) + if opened is not None: + opened.append(stream) + return stream + + return answer + + +async def drain_into(stream: object, seen: list[object]) -> None: + """Append every item of an async stream to ``seen``, letting an error the stream raises propagate.""" + async for item in stream: # pyright: ignore[reportGeneralTypeIssues] # an async stream + seen.append(item) + + +async def collect_event_types(stream: object) -> list[str]: + return [getattr(event, "type", repr(event)) async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + + +__all__ = [ + "BLOCK_START", + "MESSAGES", + "MESSAGE_START", + "MESSAGE_STOP", + "PING", + "BackgroundResponse", + "ByteStream", + "Event", + "HostLogging", + "ScriptedResponsesStream", + "anthropic_stream", + "chat_stream", + "collect_bytes", + "collect_event_types", + "collect_texts", + "collect_texts_sync", + "drain_into", + "error_frame", + "finish_chunk", + "install_orjson_stub_if_missing", + "refusal_delta", + "responses_stream", + "role_chunk", + "sse", + "text_chunk", + "text_delta", +] diff --git a/tests/unit/test_router/core/test_construction_precall_characterization.py b/tests/unit/test_router/core/test_construction_precall_characterization.py new file mode 100644 index 00000000000..a86e7ae5022 --- /dev/null +++ b/tests/unit/test_router/core/test_construction_precall_characterization.py @@ -0,0 +1,1271 @@ +""" +Characterization of a router core's construction and lifecycle: what each constructor argument sets +(and the process settings it falls back to), the validation errors, the caches and callbacks it +creates, alerting, optional pre-call check registration, ``observe``, ``register_plugin``, +``retire``, ``flush_cache``/``reset``, attribute writes, ``get_valid_args`` and ``print_deployment``. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import logging +from collections.abc import Callable, Generator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.router_api import ChangeSet, router_init_args +from litellm.router_core import RouterCore +from litellm.router_python_core import PythonCore +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.types.router import ( + AlertingConfig, + AllowedFailsPolicy, + CustomRoutingStrategyBase, + RetryPolicy, + RoutingStrategy, +) +from tests.unit.test_router.core.construction_precall_helpers import ( + CallbackHost, + collect_other_routers, + host_callback_names, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, records_outcomes_through_callbacks + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +@pytest.fixture +def host() -> ScriptedHost: + return CallbackHost() + + +def _one(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return make_core(model_list=[deployment("g", "a")], **arguments) + + +# --- defaults -------------------------------------------------------------------------------- + +SETTING_DEFAULTS: Final = { + "routing_strategy": "simple-shuffle", + "fallbacks": None, + "context_window_fallbacks": None, + "content_policy_fallbacks": None, + "default_deployment": None, + "default_max_parallel_requests": None, + "cooldown_time": 5, + "allowed_fails": 3, + "allowed_fails_policy": None, + "disable_cooldowns": None, + "enable_pre_call_checks": False, + "enable_health_check_routing": False, + "health_check_ignore_transient_errors": False, + "background_health_check_model_groups": None, + "enable_tag_filtering": False, + "tag_filtering_match_any": True, + "tag_routing_prefix": "", + "provider_budget_config": None, + "provider_default_deployment_ids": [], + "fallback_access_check": None, + "fallback_budget_check": None, + "search_tools": [], + "guardrail_list": [], + "cache_responses": False, +} + + +@pytest.mark.parametrize(("field", "expected"), sorted(SETTING_DEFAULTS.items())) +def test_settings_default_of_a_router_built_with_only_a_model_list( + make_core: MakeCore, field: str, expected: object +) -> None: + """Note: ``fallbacks`` defaults to ``[]`` in the signature, yet reads ``None``: an empty list falls + through ``fallbacks or ``.""" + assert getattr(_one(make_core).settings(), field) == expected + + +def test_default_litellm_params_default_to_no_timeout_no_client_retries_and_no_caching_groups( + make_core: MakeCore, +) -> None: + assert _one(make_core).settings().default_litellm_params == { + "timeout": None, + "max_retries": 0, + "metadata": {"caching_groups": None}, + } + + +def test_router_general_settings_default_to_async_only_and_pass_through_off(make_core: MakeCore) -> None: + general: Final = _one(make_core).settings().router_general_settings + assert (general.async_only_mode, general.pass_through_all_models) == (False, False) + + +def test_settings_as_dict_of_a_default_router(make_core: MakeCore) -> None: + assert dict(_one(make_core).settings().as_dict()) == { + "routing_strategy_args": {}, + "routing_strategy": "simple-shuffle", + "allowed_fails": 3, + "cooldown_time": 5, + "num_retries": 2, + "timeout": 6000.0, + "retry_after": 0, + "fallbacks": None, + "context_window_fallbacks": None, + "model_group_retry_policy": {}, + "retry_policy": None, + "model_group_alias": {}, + "enable_weighted_failover": False, + "enable_mid_stream_fallback_continuation": False, + "enable_tag_filtering": False, + "tag_routing_prefix": "", + "fallback_on_unsupported_params": True, + "routing_groups": [], + } + + +ATTRIBUTE_DEFAULTS: Final = { + "num_retries": 2, + "max_fallbacks": 5, + "timeout": 6000.0, + "stream_timeout": None, + "request_timeout": None, + "retry_after": 0, + "client_ttl": 3600, + "set_verbose": False, + "debug_level": "INFO", + "retry_policy": None, + "model_group_retry_policy": {}, + "alerting_config": None, + "routing_strategy_args": {}, + "deployment_affinity_ttl_seconds": 3600, + "default_priority": None, + "optional_callbacks": None, + "router_budget_logger": None, + "fallback_on_unsupported_params": True, + "enable_weighted_failover": False, + "enable_mid_stream_fallback_continuation": False, + "explicit_allowed_fails": None, + "ignore_invalid_deployments": False, + "auto_router_capability_limit": None, + "assistants_config": None, + "session_affinity_repin_on_failover": False, + "model_group_affinity_config": None, + "access_groups": None, + "leastbusy_logger": None, + "lowesttpm_logger": None, +} + + +@pytest.mark.parametrize(("name", "expected"), sorted(ATTRIBUTE_DEFAULTS.items())) +def test_attribute_default_of_a_router_built_with_only_a_model_list( + make_core: MakeCore, name: str, expected: object +) -> None: + assert getattr(_one(make_core), name) == expected + + +def test_a_router_without_a_model_list_has_no_deployments(make_core: MakeCore) -> None: + core: Final = make_core() + assert list(core.registry().all()) == [] + + +def test_construction_turns_off_the_host_debug_info(make_core: MakeCore, host: ScriptedHost) -> None: + _one(make_core) + assert host.suppressed is True + + +def test_model_group_alias_none_reads_as_an_empty_alias_map(make_core: MakeCore) -> None: + core: Final = _one(make_core, model_group_alias=None) + assert dict(core.settings().as_dict()["model_group_alias"]) == {} # pyright: ignore[reportArgumentType] + + +# --- process settings each argument falls back to ----------------------------------------- + + +@pytest.mark.parametrize(("argument", "host_value", "expected"), [(5, 4, 5), (0, 4, 0), (None, 4, 4), (None, None, 2)]) +def test_num_retries_comes_from_the_argument_then_the_process_then_the_openai_default( + make_core: MakeCore, host: ScriptedHost, argument: int | None, host_value: int | None, expected: int +) -> None: + host.set_dynamic(num_retries=host_value) + core: Final = _one(make_core, **({"num_retries": argument} if argument is not None else {})) + assert core.settings().as_dict()["num_retries"] == expected + + +@pytest.mark.asyncio +async def test_num_retries_from_the_process_settings_governs_the_attempts( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(num_retries=1) + host.answer("openai/a", litellm.InternalServerError(message="boom", llm_provider="openai", model="a")) + core: Final = _one(make_core) + with pytest.raises(litellm.InternalServerError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 2 + + +@pytest.mark.parametrize( + ("argument", "host_max", "host_router_max", "expected"), + [(3, 9, 11, 3), (0, 9, 11, 0), (None, 9, 11, 9), (None, None, 11, 11)], +) +def test_max_fallbacks_comes_from_the_argument_then_litellm_max_fallbacks_then_router_max_fallbacks( + make_core: MakeCore, + host: ScriptedHost, + argument: int | None, + host_max: int | None, + host_router_max: int, + expected: int, +) -> None: + host.set_dynamic(max_fallbacks=host_max, router_max_fallbacks=host_router_max) + core: Final = _one(make_core, **({"max_fallbacks": argument} if argument is not None else {})) + assert core.max_fallbacks == expected # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize(("argument", "expected", "explicit"), [(7, 7, 7), (0, 0, 0), (None, 8, None)]) +def test_allowed_fails_comes_from_the_argument_then_the_process_and_only_an_argument_counts_as_explicit( + make_core: MakeCore, host: ScriptedHost, argument: int | None, expected: int, explicit: int | None +) -> None: + host.set_dynamic(allowed_fails=8) + core: Final = _one(make_core, **({"allowed_fails": argument} if argument is not None else {})) + assert core.settings().allowed_fails == expected + assert core.explicit_allowed_fails == explicit # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize(("argument", "expected"), [(12.5, 12.5), (None, 33.0), (0, 33.0)]) +def test_timeout_comes_from_the_argument_else_the_process_request_timeout( + make_core: MakeCore, host: ScriptedHost, argument: float | None, expected: float +) -> None: + """Oddity: ``timeout=0`` is falsy and reads as the process ``request_timeout``.""" + host.set_dynamic(request_timeout=33.0) + core: Final = _one(make_core, **({"timeout": argument} if argument is not None else {})) + assert core.settings().as_dict()["timeout"] == expected + + +def test_timeout_zero_still_reaches_default_litellm_params_as_zero(make_core: MakeCore) -> None: + core: Final = _one(make_core, timeout=0) + assert core.settings().default_litellm_params["timeout"] == 0 + + +@pytest.mark.parametrize(("argument", "expected"), [(30, 30), (0.5, 0.5), (0, 5), (None, 5)]) +def test_cooldown_time_defaults_to_five_seconds_and_zero_reads_as_the_default( + make_core: MakeCore, argument: float | None, expected: float +) -> None: + """Oddity: ``cooldown_time=0`` is replaced by the 5 second default (``cooldown_time or DEFAULT``).""" + core: Final = _one(make_core, **({"cooldown_time": argument} if argument is not None else {})) + assert core.settings().cooldown_time == expected + + +def test_request_timeout_is_unset_without_a_router_timeout( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "request_timeout", 77.0) + assert _one(make_core).request_timeout is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_request_timeout_reads_litellm_request_timeout_directly_when_a_router_timeout_is_set( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + """Oddity: this one setting is read from ``litellm.request_timeout`` itself, not from the host's + process settings (which say 33 here).""" + host.set_dynamic(request_timeout=33.0) + monkeypatch.setattr(litellm, "request_timeout", 77.0) + assert _one(make_core, timeout=10).request_timeout == 77.0 # pyright: ignore[reportAttributeAccessIssue] + + +def test_request_timeout_is_unset_when_litellm_request_timeout_is_its_untouched_default( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "request_timeout", 6000) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False, raising=False) + assert _one(make_core, timeout=10).request_timeout is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_stream_timeout_is_kept_as_given(make_core: MakeCore) -> None: + assert _one(make_core, stream_timeout=3).stream_timeout == 3 # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize("field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) +def test_fallback_lists_come_from_the_process_when_the_argument_is_empty( + make_core: MakeCore, host: ScriptedHost, field: str +) -> None: + host.set_dynamic(**{field: [{"g": ["process"]}]}) + assert getattr(_one(make_core).settings(), field) == [{"g": ["process"]}] + + +@pytest.mark.parametrize("field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) +def test_fallback_list_arguments_win_over_the_process_lists( + make_core: MakeCore, host: ScriptedHost, field: str +) -> None: + host.set_dynamic(**{field: [{"g": ["process"]}]}) + assert getattr(_one(make_core, **{field: [{"g": ["argument"]}]}).settings(), field) == [{"g": ["argument"]}] + + +def test_default_fallbacks_are_appended_to_the_fallbacks_as_a_star_entry(make_core: MakeCore) -> None: + core: Final = _one(make_core, fallbacks=[{"g": ["b"]}], default_fallbacks=["c"]) + assert core.settings().fallbacks == [{"g": ["b"]}, {"*": ["c"]}] + + +def test_default_fallbacks_alone_become_the_only_star_entry(make_core: MakeCore) -> None: + assert _one(make_core, default_fallbacks=["c"]).settings().fallbacks == [{"*": ["c"]}] + + +def test_default_fallbacks_come_from_the_process_when_not_given(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(default_fallbacks=["d"]) + assert _one(make_core).settings().fallbacks == [{"*": ["d"]}] + + +def test_default_fallbacks_given_as_a_string_raise_type_error(make_core: MakeCore) -> None: + with pytest.raises(TypeError, match=r"^default_fallbacks must be a list of model names, got 'b'\. Use \['b'\]$"): + _one(make_core, default_fallbacks="b") + + +def test_default_fallbacks_are_appended_to_the_process_fallbacks_list_itself( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: with no ``fallbacks`` argument the router adopts the process list object and appends the + star entry to it, so the process setting is changed too.""" + process_fallbacks: Final[list[object]] = [{"g": ["e"]}] + host.set_dynamic(fallbacks=process_fallbacks) + core: Final = _one(make_core, default_fallbacks=["c"]) + assert core.settings().fallbacks == [{"g": ["e"]}, {"*": ["c"]}] + assert process_fallbacks == [{"g": ["e"]}, {"*": ["c"]}] + + +# --- validation -------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("field", ["fallbacks", "content_policy_fallbacks"]) +def test_a_fallback_entry_that_is_not_a_dict_raises_value_error(make_core: MakeCore, field: str) -> None: + with pytest.raises(ValueError, match=r"^Item 'x' is not a dictionary\.$"): + _one(make_core, **{field: ["x"]}) + + +@pytest.mark.parametrize("field", ["fallbacks", "content_policy_fallbacks"]) +def test_a_fallback_entry_with_two_keys_raises_value_error(make_core: MakeCore, field: str) -> None: + with pytest.raises(ValueError, match=r"must have exactly one key, but has 2 keys\.$"): + _one(make_core, **{field: [{"a": ["b"], "c": ["d"]}]}) + + +@pytest.mark.parametrize("field", ["fallbacks", "content_policy_fallbacks"]) +def test_a_fallback_target_given_as_a_string_raises_type_error(make_core: MakeCore, field: str) -> None: + with pytest.raises(TypeError) as error: + _one(make_core, **{field: [{"a": "b"}]}) + assert str(error.value) == "Fallbacks for 'a' must be a list of model names, got 'b'. Use {'a': ['b']}" + + +@pytest.mark.parametrize("targets", [[], [{"model": "b"}], ["b", "c"]]) +def test_fallback_targets_as_a_list_of_names_or_dicts_are_accepted(make_core: MakeCore, targets: list[object]) -> None: + assert _one(make_core, fallbacks=[{"g": targets}]).settings().fallbacks == [{"g": targets}] + + +@pytest.mark.parametrize("value", [["x"], [{"a": "b"}]]) +def test_context_window_fallbacks_are_not_validated(make_core: MakeCore, value: list[object]) -> None: + """Oddity: unlike ``fallbacks`` and ``content_policy_fallbacks``, malformed context window + fallbacks are accepted as given.""" + assert _one(make_core, context_window_fallbacks=value).settings().context_window_fallbacks == value + + +def test_malformed_process_fallbacks_raise_at_construction(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(fallbacks=[{"a": "e"}]) + with pytest.raises(TypeError, match=r"^Fallbacks for 'a' must be a list of model names"): + _one(make_core) + + +def test_validate_fallbacks_accepts_none(make_core: MakeCore) -> None: + assert _one(make_core).validate_fallbacks(fallback_param=None) is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_an_unknown_routing_strategy_raises_value_error(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=r"^Invalid routing_strategy: 'bogus'\.") as error: + _one(make_core, routing_strategy="bogus") + assert str(error.value).startswith("Invalid routing_strategy: 'bogus'. Valid options: [") + assert str(error.value).endswith( + "Check 'router_settings.routing_strategy' in your config.yaml " + "or the 'routing_strategy' parameter if using the Router SDK directly." + ) + + +def test_a_routing_strategy_enum_is_stored_as_its_string_value(make_core: MakeCore) -> None: + assert _one(make_core, routing_strategy=RoutingStrategy.LEAST_BUSY).settings().routing_strategy == "least-busy" + + +@pytest.mark.parametrize( + "strategy", + [ + "simple-shuffle", + "least-busy", + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", + "lar1", + ], +) +def test_each_routing_strategy_is_reported_in_settings(make_core: MakeCore, strategy: str) -> None: + assert _one(make_core, routing_strategy=strategy).settings().routing_strategy == strategy + + +# --- default_litellm_params ---------------------------------------------------------------- + + +def test_default_litellm_params_keep_given_values_and_gain_max_retries_and_caching_groups(make_core: MakeCore) -> None: + core: Final = _one( + make_core, + timeout=12.5, + default_litellm_params={"timeout": 99, "temperature": 0.1, "metadata": {"x": 1}}, + caching_groups=[("a", "b")], + ) + assert core.settings().default_litellm_params == { + "timeout": 99, + "temperature": 0.1, + "metadata": {"x": 1, "caching_groups": [("a", "b")]}, + "max_retries": 0, + } + + +def test_default_litellm_params_timeout_defaults_to_the_router_timeout_argument(make_core: MakeCore) -> None: + assert _one(make_core, timeout=12.5).settings().default_litellm_params["timeout"] == 12.5 + + +def test_default_litellm_params_argument_is_not_changed_by_construction(make_core: MakeCore) -> None: + given: Final = {"temperature": 0.1} + _one(make_core, default_litellm_params=given) + assert given == {"temperature": 0.1} + + +# --- policies, priority, scheduler ------------------------------------------------------- + + +def test_retry_policy_given_as_a_dict_becomes_a_retry_policy(make_core: MakeCore) -> None: + policy: Final = _one(make_core, retry_policy={"RateLimitErrorRetries": 3}).retry_policy # pyright: ignore[reportAttributeAccessIssue] + assert isinstance(policy, RetryPolicy) + assert policy.RateLimitErrorRetries == 3 + + +def test_retry_policy_given_as_a_retry_policy_is_kept(make_core: MakeCore) -> None: + policy: Final = RetryPolicy(TimeoutErrorRetries=1) + assert _one(make_core, retry_policy=policy).retry_policy == policy # pyright: ignore[reportAttributeAccessIssue] + + +def test_retry_policy_of_another_type_is_ignored(make_core: MakeCore) -> None: + assert _one(make_core, retry_policy="nope").retry_policy is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_allowed_fails_policy_given_as_a_dict_becomes_an_allowed_fails_policy(make_core: MakeCore) -> None: + policy: Final = ( + _one(make_core, allowed_fails_policy={"RateLimitErrorAllowedFails": 7}).settings().allowed_fails_policy + ) + assert isinstance(policy, AllowedFailsPolicy) + assert policy.RateLimitErrorAllowedFails == 7 + + +def test_allowed_fails_policy_of_another_type_is_ignored(make_core: MakeCore) -> None: + assert _one(make_core, allowed_fails_policy="nope").settings().allowed_fails_policy is None + + +def test_model_group_retry_policy_is_kept_as_given(make_core: MakeCore) -> None: + policy: Final = {"g": RetryPolicy(TimeoutErrorRetries=4)} + assert _one(make_core, model_group_retry_policy=policy).model_group_retry_policy == policy # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize( + ("given", "expected"), [(1, 1), ("2", 2), (-1, -1), (None, None), ("x", None), (1.5, None), (True, None)] +) +def test_default_priority_keeps_an_integer_and_ignores_anything_else( + make_core: MakeCore, given: object, expected: int | None +) -> None: + assert _one(make_core, default_priority=given).default_priority == expected # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize(("given", "expected"), [(None, 0.03), (0.5, 0.5)]) +def test_scheduler_polling_interval_defaults_to_three_hundredths_of_a_second( + make_core: MakeCore, given: float | None, expected: float +) -> None: + core: Final = _one(make_core, **({"polling_interval": given} if given is not None else {})) + assert core.scheduler.polling_interval == expected # pyright: ignore[reportAttributeAccessIssue] + + +def test_client_ttl_is_kept_as_given(make_core: MakeCore) -> None: + assert _one(make_core, client_ttl=10).client_ttl == 10 # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("routing_strategy_args", {"ttl": 5}), + ("deployment_affinity_ttl_seconds", 60), + ("session_affinity_repin_on_failover", True), + ("enable_weighted_failover", True), + ("fallback_on_unsupported_params", False), + ("ignore_invalid_deployments", True), + ("retry_after", 4), + ], +) +def test_plain_arguments_are_kept_as_attributes(make_core: MakeCore, name: str, value: object) -> None: + assert getattr(_one(make_core, **{name: value}), name) == value + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("enable_tag_filtering", True), + ("tag_filtering_match_any", False), + ("tag_routing_prefix", "team-"), + ("enable_pre_call_checks", True), + ("disable_cooldowns", True), + ("default_max_parallel_requests", 4), + ("enable_health_check_routing", True), + ("health_check_ignore_transient_errors", True), + ], +) +def test_setting_arguments_are_reported_in_settings(make_core: MakeCore, name: str, value: object) -> None: + assert getattr(_one(make_core, **{name: value}).settings(), name) == value + + +def test_background_health_check_model_groups_become_a_frozenset(make_core: MakeCore) -> None: + core: Final = _one(make_core, background_health_check_model_groups=["g", "g", "h"]) + assert core.settings().background_health_check_model_groups == frozenset({"g", "h"}) + + +@pytest.mark.parametrize(("given", "expected"), [(True, True), ("true", True), ("false", False), (False, False)]) +def test_mid_stream_fallback_continuation_accepts_a_bool_or_a_bool_string( + make_core: MakeCore, given: object, expected: bool +) -> None: + assert ( + _one(make_core, enable_mid_stream_fallback_continuation=given).enable_mid_stream_fallback_continuation + is expected + ) # pyright: ignore[reportAttributeAccessIssue] + + +# --- logging level ------------------------------------------------------------------------- + + +@pytest.fixture +def router_log_level() -> Generator[int, None, None]: + level: Final = verbose_router_logger.level + yield level + verbose_router_logger.setLevel(level) + + +@pytest.mark.parametrize( + ("set_verbose", "debug_level", "expected"), + [(True, "DEBUG", logging.DEBUG), (True, "INFO", logging.INFO)], +) +def test_set_verbose_sets_the_router_logger_to_the_debug_level( + make_core: MakeCore, router_log_level: object, set_verbose: bool, debug_level: str, expected: int +) -> None: + verbose_router_logger.setLevel(logging.WARNING) + _one(make_core, set_verbose=set_verbose, debug_level=debug_level) + assert verbose_router_logger.level == expected + + +def test_debug_level_without_set_verbose_leaves_the_router_logger_alone( + make_core: MakeCore, router_log_level: object +) -> None: + verbose_router_logger.setLevel(logging.WARNING) + _one(make_core, debug_level="DEBUG") + assert verbose_router_logger.level == logging.WARNING + + +# --- redis and response caching ---------------------------------------------------------- + + +def test_redis_host_and_port_give_the_routing_state_a_redis_cache(make_core: MakeCore) -> None: + """The port is passed on as a string.""" + core: Final = _one( + make_core, redis_host="localhost", redis_port=6399, redis_password="pw", cache_kwargs={"ssl": False} + ) + redis_cache: Final = core.health().store.redis_cache + assert isinstance(redis_cache, RedisCache) + assert {k: v for k, v in redis_cache.redis_kwargs.items() if k != "socket_timeout"} == { + "host": "localhost", + "port": "6399", + "password": "pw", + "ssl": False, + } + + +def test_redis_url_gives_the_routing_state_a_redis_cache(make_core: MakeCore) -> None: + redis_cache: Final = _one(make_core, redis_url="redis://localhost:6399").health().store.redis_cache + assert isinstance(redis_cache, RedisCache) + assert redis_cache.redis_kwargs["url"] == "redis://localhost:6399" + + +def test_redis_db_is_passed_on_as_a_string(make_core: MakeCore) -> None: + redis_cache: Final = _one(make_core, redis_host="localhost", redis_port=6399, redis_db=3).health().store.redis_cache + assert isinstance(redis_cache, RedisCache) + assert redis_cache.redis_kwargs["db"] == "3" + + +def test_redis_host_without_a_port_creates_no_redis_cache(make_core: MakeCore) -> None: + assert _one(make_core, redis_host="localhost").health().store.redis_cache is None + + +def test_without_redis_arguments_the_routing_state_is_in_memory_only(make_core: MakeCore) -> None: + assert _one(make_core).health().store.redis_cache is None + + +def test_the_routing_state_and_prompt_caching_pins_share_one_redis_cache(make_core: MakeCore) -> None: + core: Final = _one(make_core, redis_host="localhost", redis_port=6399) + assert core.registry().plugins("prompt_caching_pins").redis_cache is core.health().store.redis_cache + + +def test_applying_a_redis_cache_attaches_it_to_the_routing_state_and_the_pins(make_core: MakeCore) -> None: + core: Final = _one(make_core) + redis_cache: Final = RedisCache(host="localhost", port=6399) + core.apply(ChangeSet(redis_cache=redis_cache)) + assert core.health().store.redis_cache is redis_cache + assert core.registry().plugins("prompt_caching_pins").redis_cache is redis_cache + + +def test_cache_responses_installs_a_local_response_cache_on_the_host(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, cache_responses=True) + assert isinstance(host.cache, litellm.Cache) + assert host.cache.type == "local" + assert core.settings().cache_responses is True + + +def test_cache_responses_with_redis_installs_a_redis_response_cache(make_core: MakeCore, host: ScriptedHost) -> None: + _one(make_core, cache_responses=True, redis_host="localhost", redis_port=6399) + assert isinstance(host.cache, litellm.Cache) + assert host.cache.type == "redis" + + +def test_cache_responses_keeps_a_response_cache_the_host_already_has(make_core: MakeCore, host: ScriptedHost) -> None: + existing: Final = object() + host.cache = existing + _one(make_core, cache_responses=True) + assert host.cache is existing + + +def test_without_cache_responses_the_host_response_cache_is_untouched(make_core: MakeCore, host: ScriptedHost) -> None: + _one(make_core, redis_host="localhost", redis_port=6399) + assert host.cache is None + + +# --- callback registration ---------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("strategy", "callbacks", "input_callbacks"), + [ + ("simple-shuffle", [], []), + ("lar1", [], []), + ("least-busy", ["LeastBusyLoggingHandler"], ["LeastBusyLoggingHandler"]), + ("usage-based-routing", ["LowestTPMLoggingHandler"], []), + ("usage-based-routing-v2", ["LowestTPMLoggingHandler_v2"], []), + ("latency-based-routing", ["LowestLatencyLoggingHandler"], []), + ("cost-based-routing", ["LowestCostLoggingHandler"], []), + ], +) +def test_each_routing_strategy_registers_its_selector_on_the_host( + make_core: MakeCore, host: ScriptedHost, strategy: str, callbacks: list[str], input_callbacks: list[str] +) -> None: + _one(make_core, routing_strategy=strategy) + names: Final = host_callback_names(host) + assert (names["callbacks"], names["input_callback"]) == (callbacks, input_callbacks) + + +def test_the_least_busy_selector_is_the_router_leastbusy_logger(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, routing_strategy="least-busy") + assert host.lists["callbacks"] == [core.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] + + +# --- alerting ---------------------------------------------------------------------------------- + + +def test_alerting_config_registers_slack_alerting_and_its_slow_response_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one( + make_core, alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook", alerting_threshold=10) + ) + names: Final = host_callback_names(host) + assert names["callbacks"] == ["SlackAlerting"] + assert names["success_callback"] == ["sync_deployment_callback_on_success", "response_taking_too_long_callback"] + logger: Final = core.observe().slack_alerting_logger + assert host.lists["callbacks"] == [logger] + assert logger is not None + assert (logger.alerting_threshold, logger.alerting, logger.default_webhook_url) == ( + 10.0, + ["slack"], + "https://example.invalid/hook", + ) + + +def test_observe_reports_no_slack_alerting_without_an_alerting_config(make_core: MakeCore) -> None: + assert _one(make_core).observe().slack_alerting_logger is None + + +# --- optional pre-call checks: registration ----------------------------------------------- + + +@pytest.mark.parametrize( + ("checks", "expected"), + [ + (["prompt_caching"], ["PromptCachingDeploymentCheck"]), + (["enforce_model_rate_limits"], ["ModelRateLimitingCheck"]), + (["deployment_affinity"], ["DeploymentAffinityCheck"]), + (["responses_api_deployment_check"], ["DeploymentAffinityCheck"]), + (["session_affinity"], ["DeploymentAffinityCheck"]), + (["deployment_affinity", "responses_api_deployment_check", "session_affinity"], ["DeploymentAffinityCheck"]), + (["encrypted_content_affinity"], ["EncryptedContentAffinityCheck"]), + (["prompt_caching", "prompt_caching"], ["PromptCachingDeploymentCheck"]), + (["prompt_caching", "enforce_model_rate_limits"], ["PromptCachingDeploymentCheck", "ModelRateLimitingCheck"]), + ], +) +def test_each_optional_pre_call_check_registers_one_callback_on_the_router_and_the_host( + make_core: MakeCore, host: ScriptedHost, checks: list[str], expected: list[str] +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=checks) + assert [type(callback).__name__ for callback in core.optional_callbacks] == expected # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize( + ("checks", "flags"), + [ + (["deployment_affinity"], (True, False, False)), + (["responses_api_deployment_check"], (False, True, False)), + (["session_affinity"], (False, False, True)), + (["deployment_affinity", "session_affinity"], (True, False, True)), + ], +) +def test_affinity_checks_share_one_callback_with_one_flag_each( + make_core: MakeCore, checks: list[str], flags: tuple[bool, bool, bool] +) -> None: + (callback,) = _one( + make_core, optional_pre_call_checks=checks, deployment_affinity_ttl_seconds=60 + ).optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert ( + callback.enable_user_key_affinity, + callback.enable_responses_api_affinity, + callback.enable_session_id_affinity, + ) == flags + assert callback.ttl_seconds == 60 + + +def test_encrypted_content_affinity_is_placed_before_deployment_affinity( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["deployment_affinity", "encrypted_content_affinity"]) + expected: Final = ["EncryptedContentAffinityCheck", "DeploymentAffinityCheck"] + assert [type(callback).__name__ for callback in core.optional_callbacks] == expected # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == expected + + +def test_forward_client_headers_by_model_group_registers_nothing(make_core: MakeCore, host: ScriptedHost) -> None: + """Oddity: a listed optional pre-call check that the router ignores.""" + core: Final = _one(make_core, optional_pre_call_checks=["forward_client_headers_by_model_group"]) + assert core.optional_callbacks is None # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == [] + + +def test_an_empty_optional_pre_call_check_list_registers_nothing(make_core: MakeCore, host: ScriptedHost) -> None: + assert _one(make_core, optional_pre_call_checks=[]).optional_callbacks is None # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == [] + + +def test_router_budget_limiting_cannot_be_built_outside_an_event_loop(make_core: MakeCore) -> None: + """Oddity: the budget limiter starts a background task in its constructor, so a router that needs + one cannot be built from synchronous code with no running loop.""" + with pytest.raises(RuntimeError, match="no running event loop"): + _one(make_core, optional_pre_call_checks=["router_budget_limiting"]) + + +@pytest.mark.asyncio +async def test_router_budget_limiting_registers_the_budget_limiter(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["router_budget_limiting"]) + limiter: Final = core.router_budget_logger # pyright: ignore[reportAttributeAccessIssue] + assert isinstance(limiter, RouterBudgetLimiting) + assert core.registry().plugins("budget_limiter") is limiter + assert host.lists["callbacks"] == [limiter] + core.retire() + + +@pytest.mark.asyncio +async def test_a_provider_budget_config_adds_the_budget_limiter_by_itself(make_core: MakeCore) -> None: + core: Final = _one(make_core, provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}) + assert [type(callback).__name__ for callback in core.optional_callbacks] == ["RouterBudgetLimiting"] # pyright: ignore[reportAttributeAccessIssue] + core.retire() + + +@pytest.mark.asyncio +async def test_a_deployment_budget_adds_the_budget_limiter_by_itself(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_budget=1, budget_duration="1d")]) + assert isinstance(core.registry().plugins("budget_limiter"), RouterBudgetLimiting) + core.retire() + + +def test_model_group_affinity_config_adds_an_affinity_callback_with_every_global_flag_off(make_core: MakeCore) -> None: + (callback,) = _one(make_core, model_group_affinity_config={"g": ["deployment_affinity"]}).optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert type(callback).__name__ == "DeploymentAffinityCheck" + assert ( + callback.enable_user_key_affinity, + callback.enable_responses_api_affinity, + callback.enable_session_id_affinity, + ) == (False, False, False) + assert callback.model_group_affinity_config == {"g": ["deployment_affinity"]} + + +def test_model_group_encrypted_content_affinity_adds_both_affinity_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + _one(make_core, model_group_affinity_config={"g": ["encrypted_content_affinity"]}) + assert host_callback_names(host)["callbacks"] == ["EncryptedContentAffinityCheck", "DeploymentAffinityCheck"] + + +def test_mid_stream_fallback_continuation_registers_a_global_prefill_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, enable_mid_stream_fallback_continuation=True) + assert host_callback_names(host)["callbacks"] == ["ContinuationPrefillDeploymentCheck"] + assert core.optional_callbacks is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_add_optional_pre_call_checks_after_construction_registers_the_new_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.add_optional_pre_call_checks(["enforce_model_rate_limits", "prompt_caching"]) # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == ["PromptCachingDeploymentCheck", "ModelRateLimitingCheck"] + + +def test_add_optional_pre_call_checks_with_none_and_apply_default_settings_change_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + before: Final = list(host.lists["callbacks"]) + core.add_optional_pre_call_checks(None) # pyright: ignore[reportAttributeAccessIssue] + core.apply_default_settings() # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == before + + +def test_runtime_optional_pre_call_checks_drop_unlisted_toggleable_checks_but_keep_affinity( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching", "deployment_affinity"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["enforce_model_rate_limits"]})) + expected: Final = ["DeploymentAffinityCheck", "ModelRateLimitingCheck"] + assert [type(callback).__name__ for callback in core.optional_callbacks] == expected # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == expected + + +def test_runtime_optional_pre_call_checks_none_changes_nothing(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": None})) + assert host_callback_names(host)["callbacks"] == ["PromptCachingDeploymentCheck"] + + +def test_runtime_removal_keeps_the_host_callback_while_another_live_router_uses_that_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the check is left in this host's list because some other live router (here on another + host) still has one of that type.""" + collect_other_routers() + other_host: Final = CallbackHost() + from tests.unit.test_router.core.scripted_host import python_core + + other: Final = python_core( + other_host, model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"] + ) + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": []})) + assert core.optional_callbacks == [] # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == ["PromptCachingDeploymentCheck"] + other.retire() + + +def test_enable_pre_call_checks_is_not_a_runtime_updatable_setting(make_core: MakeCore) -> None: + core: Final = _one(make_core) + core.apply(ChangeSet(settings={"enable_pre_call_checks": True})) + assert core.settings().enable_pre_call_checks is False + + +def test_enable_pre_call_checks_can_be_assigned(make_core: MakeCore) -> None: + core: Final = _one(make_core) + core.apply(ChangeSet(assign={"enable_pre_call_checks": True})) + assert core.settings().enable_pre_call_checks is True + + +# --- retire ------------------------------------------------------------------------------ + + +def test_retire_removes_selectors_optional_checks_and_alerting_but_keeps_foreign_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + host.lists["callbacks"].append("foreign") + host.lists["success_callback"].append("foreign-success") + core: Final = _one( + make_core, + routing_strategy="least-busy", + optional_pre_call_checks=["enforce_model_rate_limits"], + alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook"), + ) + core.retire() + assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == { + "callbacks": ["foreign"], + "success_callback": ["foreign-success"], + } + + +def test_retire_twice_is_harmless(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core) + core.retire() + core.retire() + assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == {} + + +def test_retiring_one_router_keeps_the_callbacks_of_another_on_the_same_host( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + first: Final = _one(make_core, routing_strategy="least-busy") + second: Final = _one(make_core, routing_strategy="least-busy") + first.retire() + assert host.lists["callbacks"] == [second.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["input_callback"] == [second.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["_async_success_callback"] == ["deployment_callback_on_success"] + assert all(callback.__self__ is second for callback in host.lists["_async_success_callback"]) # pyright: ignore[reportAttributeAccessIssue] + + +def test_retire_hands_a_vacated_check_to_another_live_router_on_a_different_host( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: retire re-registers the same-keyed check of any other live router in the process into + this host's list, even one built on a different host.""" + collect_other_routers() + from tests.unit.test_router.core.scripted_host import python_core + + other: Final = python_core( + CallbackHost(), model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"] + ) + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.retire() + assert host.lists["callbacks"] == other.optional_callbacks + other.retire() + + +# --- observe ----------------------------------------------------------------------------- + + +class RecordingLogging: + def __init__(self) -> None: + self.calls: list[tuple[str, object]] = [] + + def add_dynamic_function_callbacks(self, **callbacks: object) -> None: # kwargs-ok: the logging object's API + self.calls.append(("functions", sorted(callbacks))) + + def replace_global_callback(self, callback: object) -> None: + self.calls.append(("replace", type(callback).__name__)) + + def add_dynamic_callback(self, callback: object) -> None: + self.calls.append(("add", type(callback).__name__)) + + +def _bound_outcome_callbacks(core: RouterCore) -> list[tuple[str, object]]: # mutable-ok: a fresh list + return ( + [("functions", ["async_failure", "async_success", "failure"])] + if records_outcomes_through_callbacks(core) + else [] + ) + + +def test_observe_request_scoped_removes_every_callback_the_router_registered_globally( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + observers: Final = core.observe(request_scoped=True) + assert observers.slack_alerting_logger is None + assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == {} + + +def test_observe_with_a_logging_object_binds_the_router_outcome_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The Python core binds its three outcome callbacks; the Rust core records every attempt's + outcome itself and binds none, as it registers none on the host.""" + core: Final = _one(make_core) + before: Final = host_callback_names(host) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportArgumentType] + assert logging_obj.calls == _bound_outcome_callbacks(core) + assert host_callback_names(host) == before + + +def test_observe_request_scoped_then_binding_hands_selectors_and_policies_to_the_request(make_core: MakeCore) -> None: + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + core.observe(request_scoped=True) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportArgumentType] + assert logging_obj.calls == [ + *_bound_outcome_callbacks(core), + ("replace", "LeastBusyLoggingHandler"), + ("add", "PromptCachingDeploymentCheck"), + ] + + +def test_binding_without_request_scoping_hands_the_request_no_selector_or_policy(make_core: MakeCore) -> None: + """Without ``observe(request_scoped=True)`` the selectors and checks stay on the host lists, so a + bound request gets none of them.""" + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core + assert logging_obj.calls == _bound_outcome_callbacks(core) + + +def test_binding_skips_a_policy_the_host_already_runs(make_core: MakeCore, host: ScriptedHost) -> None: + """A request-scoped policy whose key is back on the host's ``callbacks`` list (here re-added by + hand) is not added to the request again; the selector still replaces the host's.""" + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + core.observe(request_scoped=True) # pyright: ignore[reportCallIssue] # see above + host.lists["callbacks"].append(PromptCachingDeploymentCheck(cache=DualCache())) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core + assert logging_obj.calls == [*_bound_outcome_callbacks(core), ("replace", "LeastBusyLoggingHandler")] + + +# --- register_plugin --------------------------------------------------------------------- + + +class PickB(CustomRoutingStrategyBase): + def __init__(self) -> None: + self.asked: list[str] = [] + + @staticmethod + def _b() -> dict[str, object]: + return deployment("g", "b") + + async def async_get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool = False, + request_kwargs: object = None, + ) -> dict[str, object]: + self.asked.append(model) + return self._b() + + def get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool = False, + request_kwargs: object = None, + ) -> dict[str, object]: + self.asked.append(model) + return self._b() + + +@pytest.mark.asyncio +async def test_a_registered_routing_plugin_picks_for_select(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + plugin: Final = PickB() + core.register_plugin(plugin) + picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(5)} # pyright: ignore[reportIndexIssue] + assert picked == {"b"} + assert plugin.asked == ["g"] * 5 + + +@pytest.mark.asyncio +async def test_a_registered_routing_plugin_picks_for_async_and_sync_endpoints( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + core.register_plugin(PickB()) + for _ in range(3): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"b"} + + +@pytest.mark.asyncio +async def test_changing_the_routing_strategy_drops_a_registered_routing_plugin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + plugin: Final = PickB() + core.register_plugin(plugin) + core.apply(ChangeSet(settings={"routing_strategy": "least-busy"})) + for _ in range(5): + await core.select("g") + assert plugin.asked == [] + + +# --- flush_cache / reset ------------------------------------------------------------------ + + +def test_flush_cache_drops_the_response_cache_routing_state_and_prompt_caching_pins( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core) + host.cache = object() + core.health().store.set_cache("k", 1) + core.registry().plugins("prompt_caching_pins").set_cache("p", 1) + before: Final = host_callback_names(host) + core.flush_cache() # pyright: ignore[reportAttributeAccessIssue] + assert host.cache is None + assert core.health().store.get_cache("k") is None + assert core.registry().plugins("prompt_caching_pins").get_cache("p") is None + assert host_callback_names(host) == before + + +def test_reset_empties_every_success_and_failure_list_and_the_retry_policy( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Note: reset clears the host's success/failure lists outright, foreign callbacks included, but + leaves ``callbacks`` and ``input_callback`` alone.""" + core: Final = _one(make_core, retry_policy={"RateLimitErrorRetries": 3}, routing_strategy="least-busy") + host.lists["success_callback"].append("foreign") + core.health().store.set_cache("k", 1) + core.reset() # pyright: ignore[reportAttributeAccessIssue] + assert {name: callbacks for name, callbacks in host_callback_names(host).items() if callbacks} == { + "callbacks": ["LeastBusyLoggingHandler"], + "input_callback": ["LeastBusyLoggingHandler"], + } + assert core.retry_policy is None # pyright: ignore[reportAttributeAccessIssue] + assert core.health().store.get_cache("k") is None + + +# --- attribute writes ---------------------------------------------------------------------- + + +def test_assigning_num_retries_after_construction_is_reported_by_settings(make_core: MakeCore) -> None: + core: Final = _one(make_core) + core.num_retries = 7 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().as_dict()["num_retries"] == 7 + + +def test_assigning_a_setting_attribute_replaces_the_settings_snapshot(make_core: MakeCore) -> None: + core: Final = _one(make_core) + before: Final = core.settings() + core.enable_pre_call_checks = True # pyright: ignore[reportAttributeAccessIssue] + assert core.settings() is not before + assert core.settings().enable_pre_call_checks is True + assert before.enable_pre_call_checks is False + + +def test_assigning_a_plain_setting_keeps_the_settings_snapshot(make_core: MakeCore) -> None: + core: Final = _one(make_core) + before: Final = core.settings() + core.num_retries = 7 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings() is before + + +def test_assigning_an_unrelated_attribute_keeps_the_settings_snapshot(make_core: MakeCore) -> None: + core: Final = _one(make_core) + before: Final = core.settings() + core.some_custom_attribute = 1 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings() is before + assert core.some_custom_attribute == 1 # pyright: ignore[reportAttributeAccessIssue] + + +def test_assigning_a_new_model_list_advances_the_registry_version_once(make_core: MakeCore) -> None: + core: Final = _one(make_core) + version: Final = core.registry().version + core.model_list = [deployment("g", "z")] # pyright: ignore[reportAttributeAccessIssue] + assert core.registry().version == version + 1 + assert [d["model_info"]["id"] for d in core.registry().all()] == ["z"] + + +def test_deleting_a_setting_attribute_makes_it_unset_in_settings(make_core: MakeCore) -> None: + core: Final = _one(make_core) + del core.routing_strategy # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().routing_strategy is None + assert core.settings().is_set("routing_strategy") is False + + +def test_registry_settings_and_health_views_are_built_once(make_core: MakeCore) -> None: + core: Final = _one(make_core) + assert core.registry() is core.registry() + assert core.settings() is core.settings() + assert core.health() is core.health() + + +# --- get_valid_args ----------------------------------------------------------------------- + + +def test_get_valid_args_lists_every_constructor_argument_but_self_and_the_host() -> None: + names: Final = PythonCore.get_valid_args() + assert names[0] == "model_list" + assert "self" not in names + assert "_host" not in names + assert {"routing_strategy", "optional_pre_call_checks", "enable_pre_call_checks", "redis_url"} <= set(names) + assert len(names) == len(set(names)) + + +def test_get_valid_args_matches_the_router_init_args() -> None: + assert tuple(PythonCore.get_valid_args()) == tuple(router_init_args()) + + +def test_get_valid_args_on_a_built_router(make_core: MakeCore) -> None: + assert _one(make_core).get_valid_args() == PythonCore.get_valid_args() # pyright: ignore[reportAttributeAccessIssue] + + +# --- print_deployment ---------------------------------------------------------------------- + +SECRET_DEPLOYMENT: Final = { + "model_name": "g", + "litellm_params": { + "model": "openai/x", + "api_key": "sk-abcdef", + "api_base": "https://h", + "aws_secret_access_key": "secretvalue", + }, + "model_info": {"id": "x"}, +} + + +def test_print_deployment_shows_two_characters_of_the_api_key_then_ten_stars(make_core: MakeCore) -> None: + printed: Final = _one(make_core).print_deployment(SECRET_DEPLOYMENT) # pyright: ignore[reportAttributeAccessIssue] + assert printed["litellm_params"] == { + "model": "openai/x", + "api_key": "sk**********", + "api_base": "https://h", + "aws_secret_access_key": "secretvalue", + } + assert printed["model_info"] == {"id": "x"} + + +def test_print_deployment_leaves_the_given_deployment_unchanged(make_core: MakeCore) -> None: + _one(make_core).print_deployment(SECRET_DEPLOYMENT) # pyright: ignore[reportAttributeAccessIssue] + assert SECRET_DEPLOYMENT["litellm_params"]["api_key"] == "sk-abcdef" # pyright: ignore[reportIndexIssue] + + +def test_print_deployment_without_an_api_key_returns_an_equal_copy(make_core: MakeCore) -> None: + given: Final = {"litellm_params": {"model": "m"}} + printed: Final = _one(make_core).print_deployment(given) # pyright: ignore[reportAttributeAccessIssue] + assert printed == given + assert printed is not given + + +def test_print_deployment_without_litellm_params_raises_key_error(make_core: MakeCore) -> None: + with pytest.raises(KeyError, match="litellm_params"): + _one(make_core).print_deployment({"model_name": "x"}) # pyright: ignore[reportAttributeAccessIssue] + + +def test_print_deployment_with_redaction_masks_every_secret_keeping_its_length( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core) + host.set_dynamic(redact_user_api_key_info=True) + printed: Final = core.print_deployment(SECRET_DEPLOYMENT) # pyright: ignore[reportAttributeAccessIssue] + assert printed["litellm_params"] == { + "model": "openai/x", + "api_key": "sk*******", + "api_base": "https://h", + "aws_secret_access_key": "se*********", + } + + +def test_settings_snapshot_reads_constructor_mappings_as_given(make_core: MakeCore) -> None: + alias: Final[Mapping[str, object]] = {"alias": "g"} + core: Final = _one(make_core, model_group_alias=dict(alias)) + assert core.registry().aliases() == alias diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py new file mode 100644 index 00000000000..bb8d3d30774 --- /dev/null +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -0,0 +1,663 @@ +""" +Characterization of the pre-call checks on a router core: with ``enable_pre_call_checks`` the +context window (token counting over messages, Responses input, instructions, tools and an Anthropic +system block), the allowed region and unsupported params (``response_format``) narrow a group's +deployments before one is picked; image requests skip text-only deployments; budget limits drop +deployments over budget. Covers the errors raised when nothing is left (type and message), which +deployment ``select`` and the endpoints reach, and that no provider call is made then. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import re +from collections.abc import Callable, Generator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RouterRateLimitError +from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, restoring_model_cost, windowed +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, python_core + +MakeCore = Callable[..., RouterCore] +PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once in a thousand runs +SHORT: Final = [{"role": "user", "content": "hi"}] +LONG_TEXT: Final = " ".join(["word"] * 100) +LONG: Final = [{"role": "user", "content": LONG_TEXT}] +IMAGE: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + ], + } +] +UNCOUNTABLE: Final = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"detail": "not-a-real-detail-value"}}]} +] +JSON_MODE: Final = {"response_format": {"type": "json_object"}} +NO_JSON_MODE_MODEL: Final = "openai/gpt-3.5-turbo-16k" +JSON_MODE_MODEL: Final = "openai/gpt-4o" +CONTEXT_WINDOW_PREFIX: Final = ( + "litellm.ContextWindowExceededError: litellm.BadRequestError: litellm._pre_call_checks: Context Window exceeded " + "for given call. No models have context window large enough for this call.\n" +) +UNSUPPORTED_MESSAGE: Final = ( + "litellm.UnsupportedParamsError: no deployment of model=g supports ['response_format']. To drop these, set " + "`litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, or send " + "allowed_openai_params=['response_format'] to pass them through" +) + + +@pytest.fixture +def host() -> ScriptedHost: + return CallbackHost() + + +@pytest.fixture(autouse=True) +def isolated_model_cost() -> Generator[None, None, None]: + """Routers write their deployments' model info into the process cost map; undo it after each test.""" + yield from restoring_model_cost("gpt-3.5-turbo-16k") + + +async def _picked( + core: RouterCore, model: str = "g", times: int = PICKS, **select_arguments: object +) -> set[str]: # kwargs-ok: select's arguments + picked: Final[set[str]] = set() + for _ in range(times): + chosen = await core.select(model, **select_arguments) # pyright: ignore[reportCallIssue] + assert isinstance(chosen, Mapping) + picked.add(str(chosen["model_info"]["id"])) + return picked + + +def _small_and_big( + make_core: MakeCore, small: int = 5, **arguments: object +) -> RouterCore: # kwargs-ok: constructor arguments + return make_core( + model_list=[windowed("g", "small", small), windowed("g", "big", 10000)], + **{"enable_pre_call_checks": True, **arguments}, + ) + + +def _json_mode_group(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL), deployment("g", "new", model=JSON_MODE_MODEL)], + **{"enable_pre_call_checks": True, **arguments}, + ) + + +# --- context window ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_without_enable_pre_call_checks_a_too_small_context_window_is_still_picked(make_core: MakeCore) -> None: + core: Final = _small_and_big(make_core, enable_pre_call_checks=False) + assert await _picked(core, messages=LONG) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_without_enable_pre_call_checks_a_too_long_prompt_reaches_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)]) + await core.acompletion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["small"] + + +@pytest.mark.asyncio +async def test_select_skips_a_deployment_whose_context_window_the_prompt_exceeds(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core), messages=LONG) == {"big"} + + +@pytest.mark.asyncio +async def test_a_short_prompt_may_pick_either_deployment(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core, small=50), messages=SHORT) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_async_completion_reaches_only_a_deployment_whose_context_window_fits( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _small_and_big(make_core) + for _ in range(8): + await core.acompletion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"big"} + + +def test_sync_completion_reaches_only_a_deployment_whose_context_window_fits( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _small_and_big(make_core) + for _ in range(8): + core.completion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"big"} + + +@pytest.mark.asyncio +async def test_pass_through_selection_skips_a_deployment_whose_context_window_the_prompt_exceeds( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[ + windowed("g", "small", 5, use_in_pass_through=True), + windowed("g", "big", 10000, use_in_pass_through=True), + ], + enable_pre_call_checks=True, + ) + assert await _picked(core, kind="pass_through", messages=LONG) == {"big"} + + +@pytest.mark.asyncio +async def test_select_without_messages_or_input_runs_no_context_window_check(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core)) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_when_no_context_window_fits_select_raises_context_window_exceeded(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError) as error: + await core.select("g", messages=LONG) + message: Final = str(error.value) + assert message.startswith(CONTEXT_WINDOW_PREFIX) + assert re.fullmatch( + r"Model=openai/cw-small-[0-9a-f]{8}, Max Input Tokens=5, Got=\d+", message.removeprefix(CONTEXT_WINDOW_PREFIX) + ) + assert error.value.model == "g" + assert error.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_when_no_context_window_fits_async_completion_raises_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError, match="Context Window exceeded for given call"): + await core.acompletion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + assert host.sleeps == [] + + +def test_when_no_context_window_fits_sync_completion_raises_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError): + core.completion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_the_context_window_error_lists_each_deployment_back_to_back(make_core: MakeCore) -> None: + """Oddity: the per-deployment entries are concatenated with no separator between them.""" + core: Final = make_core(model_list=[windowed("g", "one", 5), windowed("g", "two", 6)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError) as error: + await core.select("g", messages=LONG) + assert re.fullmatch( + r"Model=openai/cw-one-[0-9a-f]{8}, Max Input Tokens=5, Got=(\d+)Model=openai/cw-two-[0-9a-f]{8}, Max Input Tokens=6, Got=\1", + str(error.value).removeprefix(CONTEXT_WINDOW_PREFIX), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("where", "expected"), + [("model_info", "azure/gpt-4-mi-base"), ("litellm_params", "azure/gpt-4-lp-base")], +) +async def test_the_context_window_error_names_a_deployment_by_its_base_model( + make_core: MakeCore, where: str, expected: str +) -> None: + extra: Final[dict[str, object]] = ( + {"model_info": {"base_model": expected}} if where == "model_info" else {"base_model": expected} + ) + core: Final = make_core(model_list=[windowed("g", f"base-{where}", 5, **extra)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError) as error: + await core.select("g", messages=LONG) + assert f"Model={expected}, Max Input Tokens=5, Got=" in str(error.value) + + +@pytest.mark.asyncio +async def test_when_no_context_window_fits_the_context_window_fallback_group_answers( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[windowed("small-group", "small", 5), deployment("big-group", "big")], + enable_pre_call_checks=True, + context_window_fallbacks=[{"small-group": ["big-group"]}], + num_retries=0, + ) + await core.acompletion(model="small-group", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["big"] + + +@pytest.mark.asyncio +async def test_a_deployment_without_a_known_context_window_is_kept_for_any_prompt(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[windowed("g", "small", 5), deployment("g", "unknown-window-model")], enable_pre_call_checks=True + ) + assert await _picked(core, messages=LONG) == {"unknown-window-model"} + + +@pytest.mark.asyncio +async def test_the_context_window_comes_from_the_cost_map_when_the_deployment_declares_none( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "gpt4", model="openai/gpt-4"), deployment("g", "gpt4o", model="openai/gpt-4o")], + enable_pre_call_checks=True, + ) + assert await _picked(core, times=10, messages=[{"role": "user", "content": "word " * 10000}]) == {"gpt4o"} + + +@pytest.mark.asyncio +async def test_an_uncountable_prompt_skips_the_context_window_check(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[windowed("g", "tiny", 1)], enable_pre_call_checks=True) + assert await _picked(core, times=2, messages=UNCOUNTABLE) == {"tiny"} + + +@pytest.mark.asyncio +async def test_responses_input_text_is_counted_against_the_context_window(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core), times=10, input=LONG_TEXT) == {"big"} + + +@pytest.mark.asyncio +async def test_a_short_responses_input_may_pick_either_deployment(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core, small=50), input="hi") == {"small", "big"} + + +@pytest.mark.asyncio +async def test_responses_instructions_are_counted_against_the_context_window(make_core: MakeCore) -> None: + core: Final = _small_and_big(make_core, small=50) + assert await _picked(core, request_kwargs={"instructions": LONG_TEXT}, input="hi") == {"big"} + + +@pytest.mark.asyncio +async def test_tool_definitions_are_counted_against_the_context_window(make_core: MakeCore) -> None: + tool: Final = { + "type": "function", + "function": {"name": "f", "description": LONG_TEXT, "parameters": {"type": "object", "properties": {}}}, + } + core: Final = _small_and_big(make_core, small=50) + assert await _picked(core, request_kwargs={"tools": [tool]}, messages=SHORT) == {"big"} + + +@pytest.mark.asyncio +async def test_an_anthropic_system_block_is_counted_against_the_context_window(make_core: MakeCore) -> None: + core: Final = _small_and_big(make_core, small=50) + assert await _picked(core, request_kwargs={"system": LONG_TEXT}, messages=SHORT) == {"big"} + + +@pytest.mark.asyncio +async def test_a_deployment_context_window_is_written_into_the_process_cost_map(make_core: MakeCore) -> None: + """Oddity: a deployment's declared ``max_input_tokens`` replaces the process-wide cost-map entry of its + model, so it outlives the router and narrows another router's deployment of the same model.""" + first: Final = make_core( + model_list=[deployment("g", "declares-five", model=NO_JSON_MODE_MODEL, model_info={"max_input_tokens": 5})] + ) + first.retire() + other: Final = python_core( + CallbackHost(), + model_list=[deployment("g", "declares-nothing", model=NO_JSON_MODE_MODEL)], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError, match="Max Input Tokens=5"): + await other.select("g", messages=SHORT) + other.retire() + + +# --- region ------------------------------------------------------------------------------ + + +def _regional(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return make_core( + model_list=[ + deployment("g", "eu", region_name="eu"), + deployment("g", "us", region_name="us"), + deployment("g", "none"), + ], + **{"enable_pre_call_checks": True, **arguments}, + ) + + +@pytest.mark.asyncio +async def test_allowed_model_region_keeps_only_deployments_in_that_region(make_core: MakeCore) -> None: + core: Final = _regional(make_core) + assert await _picked(core, request_kwargs={"allowed_model_region": "eu"}, messages=SHORT) == {"eu"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_enable_pre_call_checks(make_core: MakeCore) -> None: + core: Final = _regional(make_core, enable_pre_call_checks=False) + assert await _picked(core, request_kwargs={"allowed_model_region": "eu"}, messages=SHORT) == {"eu", "us", "none"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_messages_or_input(make_core: MakeCore) -> None: + core: Final = _regional(make_core) + assert await _picked(core, request_kwargs={"allowed_model_region": "eu"}) == {"eu", "us", "none"} + + +@pytest.mark.asyncio +async def test_no_deployment_in_the_allowed_region_raises_router_rate_limit_error(make_core: MakeCore) -> None: + core: Final = _regional(make_core) + with pytest.raises(RouterRateLimitError) as error: + await core.select("g", {"allowed_model_region": "ap"}, messages=SHORT) + assert str(error.value) == ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=g. " + "pre-call-checks=True, cooldown_list=[]" + ) + + +@pytest.mark.asyncio +async def test_no_deployment_in_the_allowed_region_never_calls_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _regional(make_core) + with pytest.raises(RouterRateLimitError): + await core.acompletion(model="g", messages=SHORT, allowed_model_region="ap") # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_async_completion_with_an_allowed_region_reaches_only_that_region( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _regional(make_core) + for _ in range(6): + await core.acompletion(model="g", messages=SHORT, allowed_model_region="us") # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"us"} + + +# --- unsupported params ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_with_drop_params_off_a_deployment_that_lacks_response_format_is_skipped( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + assert await _picked(_json_mode_group(make_core), request_kwargs=dict(JSON_MODE), messages=SHORT) == {"new"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [None, True]) +async def test_unless_drop_params_is_exactly_false_unsupported_params_are_not_checked( + make_core: MakeCore, host: ScriptedHost, drop_params: bool | None +) -> None: + """Oddity: the unset process default (``None``) skips the check as ``drop_params=True`` does.""" + host.set_dynamic(drop_params=drop_params) + assert await _picked(_json_mode_group(make_core), request_kwargs=dict(JSON_MODE), messages=SHORT) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_unsupported_params_are_not_checked_without_enable_pre_call_checks( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core, enable_pre_call_checks=False) + assert await _picked(core, request_kwargs=dict(JSON_MODE), messages=SHORT) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_unsupported_params_are_not_checked_without_messages_or_input( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + assert await _picked(_json_mode_group(make_core), request_kwargs=dict(JSON_MODE)) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_a_string_response_format_is_not_checked(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core) + assert await _picked(core, request_kwargs={"response_format": "b64_json"}, messages=SHORT) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_when_no_deployment_supports_the_param_select_raises_unsupported_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core(model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True) + with pytest.raises(litellm.UnsupportedParamsError) as error: + await core.select("g", dict(JSON_MODE), messages=SHORT) + assert str(error.value) == UNSUPPORTED_MESSAGE + assert error.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_when_no_deployment_supports_the_param_async_completion_raises_without_calling_or_retrying( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True, num_retries=2 + ) + with pytest.raises(litellm.UnsupportedParamsError, match="allowed_openai_params"): + await core.acompletion(model="g", messages=SHORT, **JSON_MODE) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + assert host.sleeps == [] + + +def test_when_no_deployment_supports_the_param_sync_completion_raises_without_calling( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core(model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True) + with pytest.raises(litellm.UnsupportedParamsError): + core.completion(model="g", messages=SHORT, **JSON_MODE) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +REQUEST_ESCAPES: Final = [ + {"allowed_openai_params": ["response_format"]}, + {"drop_params": True}, + {"additional_drop_params": ["response_format"]}, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("escape", REQUEST_ESCAPES, ids=["allowlist", "drop", "additional-drop"]) +async def test_a_request_that_allows_or_drops_the_param_keeps_the_only_deployment( + make_core: MakeCore, host: ScriptedHost, escape: dict[str, object] +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core(model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True) + assert await _picked(core, times=2, request_kwargs={**JSON_MODE, **escape}, messages=SHORT) == {"old"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("escape", REQUEST_ESCAPES, ids=["allowlist", "drop", "additional-drop"]) +async def test_a_deployment_that_takes_the_param_is_preferred_over_dropping_it( + make_core: MakeCore, host: ScriptedHost, escape: dict[str, object] +) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core) + assert await _picked(core, request_kwargs={**JSON_MODE, **escape}, messages=SHORT) == {"new"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment_escape", + [{"drop_params": True}, {"allowed_openai_params": ["response_format"]}], + ids=["drop", "allowlist"], +) +async def test_a_deployment_that_allows_or_drops_the_param_is_kept( + make_core: MakeCore, host: ScriptedHost, deployment_escape: dict[str, object] +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL, **deployment_escape)], enable_pre_call_checks=True + ) + assert await _picked(core, times=2, request_kwargs=dict(JSON_MODE), messages=SHORT) == {"old"} + + +@pytest.mark.asyncio +async def test_a_request_drop_params_false_overrides_the_deployment_drop_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL, drop_params=True)], enable_pre_call_checks=True + ) + with pytest.raises(litellm.UnsupportedParamsError): + await core.select("g", {**JSON_MODE, "drop_params": False}, messages=SHORT) + + +@pytest.mark.asyncio +async def test_a_context_window_error_wins_over_an_unsupported_param(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "small-no-json", model=NO_JSON_MODE_MODEL, model_info={"max_input_tokens": 5})], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await core.select("g", dict(JSON_MODE), messages=LONG) + core.retire() + + +@pytest.mark.asyncio +async def test_a_capable_deployment_outside_the_region_gives_no_deployments_rather_than_a_400( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[ + deployment("g", "old", model=NO_JSON_MODE_MODEL, region_name="eu"), + deployment("g", "new", model=JSON_MODE_MODEL, region_name="us"), + ], + enable_pre_call_checks=True, + ) + with pytest.raises(RouterRateLimitError, match=r"^No deployments available for selected model"): + await core.select("g", {**JSON_MODE, "allowed_model_region": "eu"}, messages=SHORT) + + +@pytest.mark.asyncio +async def test_dropping_the_param_keeps_the_in_region_deployment_that_lacks_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[ + deployment("g", "old", model=NO_JSON_MODE_MODEL, region_name="eu"), + deployment("g", "new", model=JSON_MODE_MODEL, region_name="us"), + ], + enable_pre_call_checks=True, + ) + request: Final = {**JSON_MODE, "allowed_model_region": "eu", "drop_params": True} + assert await _picked(core, times=5, request_kwargs=request, messages=SHORT) == {"old"} + + +@pytest.mark.asyncio +async def test_a_capable_deployment_in_cooldown_gives_no_deployments_rather_than_a_400( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core) + core.health().cooldowns.add_deployment_to_cooldown( + "new", original_exception=Exception("down"), exception_status=429, cooldown_time=60 + ) + with pytest.raises(RouterRateLimitError) as error: + await core.select("g", dict(JSON_MODE), messages=SHORT) + assert str(error.value).startswith( + "No deployments available for selected model, Try again in 60 seconds. Passed model=g. pre-call-checks=True" + ) + + +# --- images -------------------------------------------------------------------------------- + + +def _vision_group(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment("g", "text", model_info={"supports_vision": False}), + deployment("g", "vision", model_info={"supports_vision": True}), + ] + ) + + +@pytest.mark.asyncio +async def test_an_image_request_skips_a_declared_text_only_deployment_even_without_pre_call_checks( + make_core: MakeCore, +) -> None: + assert await _picked(_vision_group(make_core), messages=IMAGE) == {"vision"} + + +@pytest.mark.asyncio +async def test_an_image_in_the_request_messages_counts_as_an_image_request(make_core: MakeCore) -> None: + assert await _picked(_vision_group(make_core), request_kwargs={"messages": IMAGE}) == {"vision"} + + +@pytest.mark.asyncio +async def test_a_text_request_may_pick_a_text_only_deployment(make_core: MakeCore) -> None: + assert await _picked(_vision_group(make_core), messages=SHORT) == {"text", "vision"} + + +@pytest.mark.asyncio +async def test_an_image_request_keeps_every_deployment_when_all_are_text_only(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "text", model_info={"supports_vision": False}), + deployment("g", "text2", model_info={"supports_vision": False}), + ] + ) + assert await _picked(core, messages=IMAGE) == {"text", "text2"} + + +@pytest.mark.asyncio +async def test_async_completion_with_an_image_reaches_the_vision_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _vision_group(make_core) + for _ in range(6): + await core.acompletion(model="g", messages=IMAGE) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"vision"} + + +# --- budgets ------------------------------------------------------------------------------- + +PROVIDER_BUDGET: Final = {"openai": {"budget_limit": 1, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_a_provider_over_budget_leaves_no_deployment_and_select_raises_value_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + await core.health().store.async_set_cache("provider_spend:openai:1d", 5.0) + with pytest.raises(ValueError, match=r"^No deployments available - crossed budget: ") as error: + await core.select("g") + assert ( + str(error.value) + == "No deployments available - crossed budget: Exceeded budget for provider openai: 5.0 >= 1.0\n" + ) + core.retire() + + +@pytest.mark.asyncio +async def test_a_provider_over_budget_raises_from_async_completion_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + await core.health().store.async_set_cache("provider_spend:openai:1d", 5.0) + with pytest.raises(ValueError, match=r"^No deployments available - crossed budget"): + await core.acompletion(model="g", messages=SHORT) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + core.retire() + + +@pytest.mark.asyncio +async def test_a_provider_under_budget_is_picked(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + await core.health().store.async_set_cache("provider_spend:openai:1d", 0.5) + assert await _picked(core, times=2) == {"a"} + core.retire() + + +@pytest.mark.asyncio +async def test_a_deployment_over_its_own_budget_is_skipped(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_budget=1, budget_duration="1d"), deployment("g", "b")]) + await core.health().store.async_set_cache("deployment_spend:a:1d", 5.0) + assert await _picked(core, times=10) == {"b"} + core.retire() diff --git a/tests/unit/test_router/core/test_deprecated_members.py b/tests/unit/test_router/core/test_deprecated_members.py new file mode 100644 index 00000000000..b73dda78ecd --- /dev/null +++ b/tests/unit/test_router/core/test_deprecated_members.py @@ -0,0 +1,665 @@ +""" +The deprecated Router members (``DeprecatedRouterMembers``, litellm/router_deprecated.py) on a Router +over each core: every old attribute reads what the Python core's ``_core_`` holds, through the +entry point named in its DeprecationWarning, and the old writes have the Python core's effect. +""" + +from collections.abc import Callable +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router import Router +from litellm.router_api import SETTING_FIELDS, ChangeSet +from litellm.router_core import RouterCore +from litellm.router_utils.cooldown_cache import CooldownCache +from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery +from litellm.types.router import AllowedFailsPolicy, Deployment, RouterGeneralSettings, RoutingContext +from tests.unit.test_router.core.registry_helpers import MESSAGES, MakeCore, ids, team_deployment +from tests.unit.test_router.core.scripted_host import deployment + + +def as_router(core: RouterCore) -> Router: + """A ``litellm.Router`` facade over ``core``, as ``Router.from_config`` attaches the core it builds.""" + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach, as from_config calls it + return router + + +def picked_id(picked: object) -> object: + """The ``model_info.id`` of a picked deployment dict.""" + return (as_str_mapping((as_str_mapping(picked) or {}).get("model_info")) or {}).get("id") + + +def read(router: Router, name: str) -> object: + with pytest.warns(DeprecationWarning, match=rf"Router\.{name} is deprecated"): + return getattr(router, name) # pyright: ignore[reportAny] # each old member's own type + + +def write(router: Router, name: str, value: object) -> None: + with pytest.warns(DeprecationWarning, match=rf"Router\.{name} is deprecated"): + setattr(router, name, value) + + +_CONFIGURED: Final[dict[str, object]] = { # mutable-ok: constructor arguments, copied per router + "model_list": [ + deployment("gpt", "a"), + deployment("gpt", "b"), + deployment("claude", "c", model="anthropic/claude-x"), + team_deployment("int-t", "t1", "TA", "pub", model="openai/team-model"), + deployment("openai/*", "w", model="openai/*"), + ], + "fallbacks": [{"gpt": ["claude"]}], + "context_window_fallbacks": [{"gpt": ["claude"]}], + "content_policy_fallbacks": [{"claude": ["gpt"]}], + "cooldown_time": 7, + "allowed_fails": 2, + "allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=4), + "disable_cooldowns": True, + "enable_pre_call_checks": True, + "enable_tag_filtering": True, + "tag_filtering_match_any": False, + "tag_routing_prefix": "x-", + "default_litellm_params": {"metadata": {"k": "v"}}, + "routing_strategy": "least-busy", + "default_max_parallel_requests": 5, + "guardrail_list": [{"guardrail_name": "g", "litellm_params": {"guardrail": "x"}}], + "search_tools": [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily"}}], + "health_check_ignore_transient_errors": True, + "enable_health_check_routing": True, + "background_health_check_model_groups": ["gpt"], + "model_group_alias": {"fast": "gpt"}, + "cache_responses": True, +} + +# The Python core's values for _CONFIGURED, by setting (the objects it builds are compared by type). +_CONFIGURED_SETTINGS: Final[dict[str, object]] = { # mutable-ok: expected values + "routing_strategy": "least-busy", + "fallbacks": [{"gpt": ["claude"]}], + "context_window_fallbacks": [{"gpt": ["claude"]}], + "content_policy_fallbacks": [{"claude": ["gpt"]}], + "default_deployment": None, + "default_litellm_params": {"timeout": None, "max_retries": 0, "metadata": {"k": "v", "caching_groups": None}}, + "default_max_parallel_requests": 5, + "cooldown_time": 7, + "allowed_fails": 2, + "allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=4), + "disable_cooldowns": True, + "enable_pre_call_checks": True, + "enable_health_check_routing": True, + "health_check_ignore_transient_errors": True, + "background_health_check_model_groups": frozenset({"gpt"}), + "enable_tag_filtering": True, + "tag_filtering_match_any": False, + "tag_routing_prefix": "x-", + "router_general_settings": RouterGeneralSettings(), + "provider_budget_config": None, + "provider_default_deployment_ids": ["w"], + "fallback_access_check": None, + "fallback_budget_check": None, + "search_tools": [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily"}}], + "guardrail_list": [{"guardrail_name": "g", "litellm_params": {"guardrail": "x"}}], + "cache_responses": True, +} + + +def _configured(make_core: MakeCore) -> Router: + return as_router(make_core(**{name: _copy(value) for name, value in _CONFIGURED.items()})) + + +def _copy(value: object) -> object: + import copy + + return copy.deepcopy(value) + + +# -- settings ------------------------------------------------------------------------------------ + + +@pytest.mark.parametrize("name", sorted(set(SETTING_FIELDS) - {"kubernetes_pod_discovery"})) +def test_a_setting_attribute_reads_the_configured_setting(make_core: MakeCore, name: str) -> None: + router: Final = _configured(make_core) + value: Final = read(router, name) + assert value == _CONFIGURED_SETTINGS[name] + assert value is getattr(router.settings(), name) + + +def test_the_pod_discovery_setting_is_the_routers_pod_lookup(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + discovery: Final = read(router, "kubernetes_pod_discovery") + assert isinstance(discovery, KubernetesPodDiscovery) + assert discovery is router.settings().kubernetes_pod_discovery + + +def test_default_settings_read_as_the_python_cores_defaults(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + assert read(router, "routing_strategy") == "simple-shuffle" + assert read(router, "cooldown_time") == 5 + assert read(router, "allowed_fails") == 3 + assert read(router, "fallbacks") is None + assert read(router, "default_litellm_params") == { + "timeout": None, + "max_retries": 0, + "metadata": {"caching_groups": None}, + } + assert read(router, "tag_filtering_match_any") is True + assert read(router, "tag_routing_prefix") == "" + assert read(router, "cache_responses") is False + assert read(router, "search_tools") == [] + assert read(router, "guardrail_list") == [] + + +def test_a_setting_attribute_write_is_a_changeset(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + write(router, "cooldown_time", 9) + write(router, "fallbacks", [{"gpt": ["gpt"]}]) + assert read(router, "cooldown_time") == 9 + assert router.settings().cooldown_time == 9 + assert read(router, "fallbacks") == [{"gpt": ["gpt"]}] + + +def test_a_deleted_setting_attribute_no_longer_reads(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with pytest.warns(DeprecationWarning, match=r"Router\.fallbacks is deprecated"): + del router.fallbacks + assert not router.settings().is_set("fallbacks") + with pytest.raises(AttributeError, match="'fallbacks'"): + read(router, "fallbacks") + + +# -- registry state ------------------------------------------------------------------------------ + +_REGISTRY_READS: Final[dict[str, Callable[[Router], object]]] = { + "model_list": lambda r: r.registry().all(), + "model_names": lambda r: r.registry().names(kind="group_set"), + "deployment_names": lambda r: r.registry().names(kind="deployment_names"), + "team_public_model_names": lambda r: r.registry().names(kind="team_public"), + "model_name_to_deployment_indices": lambda r: r.registry().index("model_name"), + "team_model_to_deployment_indices": lambda r: r.registry().index("team_model"), + "cached_deployment_model_info": lambda r: r.registry().deployment_model_info_cache(), + "model_group_alias": lambda r: r.registry().aliases(), + "pattern_router": lambda r: r.registry().patterns(), + "team_pattern_routers": lambda r: r.registry().patterns(team=True), + "routing_groups": lambda r: r.registry().routing_groups(), + "auto_routers": lambda r: r.registry().plugins("auto"), + "complexity_routers": lambda r: r.registry().plugins("complexity"), + "quality_routers": lambda r: r.registry().plugins("quality"), + "adaptive_routers": lambda r: r.registry().plugins("adaptive"), + "routing_plugins": lambda r: r.registry().plugins("routing"), + "prompt_caching_pin_cache": lambda r: r.registry().plugins("prompt_caching_pins"), + "cache": lambda r: r.health().store, + "cooldown_cache": lambda r: r.health().cooldowns, + "health_state_cache": lambda r: r.health().health_checks, + "deployment_latency_map": lambda r: r.health().latency, +} + + +@pytest.mark.parametrize("name", sorted(_REGISTRY_READS)) +def test_a_state_attribute_is_the_live_object_its_entry_point_returns(make_core: MakeCore, name: str) -> None: + router: Final = _configured(make_core) + assert read(router, name) is _REGISTRY_READS[name](router) + + +def test_the_registry_index_attributes_hold_the_python_cores_values(make_core: MakeCore) -> None: + router: Final = _configured(make_core) + assert read(router, "model_names") == {"gpt", "claude", "int-t", "openai/*"} + assert read(router, "deployment_names") == [ + "openai/a", + "openai/b", + "anthropic/claude-x", + "openai/team-model", + "openai/*", + ] + assert read(router, "team_public_model_names") == frozenset({"pub"}) + assert read(router, "model_name_to_deployment_indices") == { + "gpt": [0, 1], + "claude": [2], + "int-t": [3], + "openai/*": [4], + } + assert read(router, "team_model_to_deployment_indices") == {("TA", "pub"): [3]} + assert read(router, "model_group_alias") == {"fast": "gpt"} + assert read(router, "routing_groups") == {} + with pytest.warns(DeprecationWarning, match=r"Router\.routing_groups is deprecated"): + assert router._routing_groups == {} # pyright: ignore[reportPrivateUsage] # the old private alias + assert read(router, "routing_plugins") == [] + assert read(router, "deployment_latency_map") == { + "openai/a": 0, + "openai/b": 0, + "anthropic/claude-x": 0, + "openai/team-model": 0, + "openai/*": 0, + } + assert ids(read(router, "model_list")) == ["a", "b", "c", "t1", "w"] # pyright: ignore[reportArgumentType] # the live model list + + +def test_the_cached_model_info_attribute_is_the_lru_cache(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with pytest.warns(DeprecationWarning, match=r"Router\.cached_deployment_model_info is deprecated"): + cached: Final = router.cached_deployment_model_info + first: Final = cached("a", "gpt") + assert cached("a", "gpt") is first + assert cached.cache_info().hits == 1 + assert first == router.registry().deployment_model_info("a", "gpt") + + +def test_a_configured_routing_plugin_is_listed(make_core: MakeCore) -> None: + class Keep: + async def run(self, context: RoutingContext) -> RoutingContext: + return context + + plugin: Final = Keep() + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], plugins=[plugin])) + assert read(router, "routing_plugins") == [plugin] + + +def test_a_routing_plugin_appended_to_the_live_list_runs(make_core: MakeCore) -> None: + class Narrow: + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [m for m in context.candidate_models if m == "openai/b"] + return context + + class Keep: + async def run(self, context: RoutingContext) -> RoutingContext: + return context + + router: Final = as_router(make_core(model_list=[deployment("g", "a"), deployment("g", "b")], plugins=[Keep()])) + with pytest.warns(DeprecationWarning, match=r"Router\.routing_plugins is deprecated"): + plugins: Final = router.routing_plugins + plugins.append(Narrow()) + + import asyncio + + for _ in range(4): + picked = asyncio.run(router.select("g", {}, messages=MESSAGES)) # pyright: ignore[reportCallIssue] + assert picked_id(picked) == "b" + + +# -- state writes -------------------------------------------------------------------------------- + + +def test_a_registry_state_write_is_what_the_registry_reads(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + write(router, "model_names", {"x"}) + write(router, "deployment_names", ["openai/x"]) + write(router, "model_name_to_deployment_indices", {"x": [0]}) + write(router, "auto_routers", {}) + assert read(router, "model_names") == {"x"} + assert router.registry().names(kind="group_set") == {"x"} + assert router.registry().names(kind="deployment_names") == ["openai/x"] + assert router.registry().index("model_name") == {"x": [0]} + assert router.registry().plugins("auto") == {} + + +def test_a_routing_state_write_is_what_health_reads(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + store: Final = DualCache() + cooldowns: Final = CooldownCache(cache=store, default_cooldown_time=1) + latency: Final[dict[str, float]] = {"openai/a": 0.5} + write(router, "cache", store) + write(router, "cooldown_cache", cooldowns) + write(router, "deployment_latency_map", latency) + assert read(router, "cache") is store + assert router.health().store is store + assert router.health().cooldowns is cooldowns + assert router.health().latency is latency + + +def test_an_assigned_slack_alerting_logger_is_the_routers_observer(make_core: MakeCore) -> None: + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with ( + pytest.raises(AttributeError, match="slack_alerting_logger"), + pytest.warns(DeprecationWarning, match=r"Router\.slack_alerting_logger is deprecated"), + ): + _ = router.slack_alerting_logger + logger: Final = SlackAlerting(alerting=["slack"], default_webhook_url="https://example.invalid/hook") + write(router, "slack_alerting_logger", logger) + assert read(router, "slack_alerting_logger") is logger + assert router.observe().slack_alerting_logger is logger + + +@pytest.mark.asyncio +async def test_an_assigned_async_pick_answers_select(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")])) + calls: Final[list[str]] = [] # mutable-ok: the models the pick was asked for + + async def pick( + model: str, + request_kwargs: dict[str, object], # mutable-ok: the old pick's signature + messages: object = None, + input: object = None, + specific_deployment: object = False, + ) -> dict[str, object]: # mutable-ok: a deployment dict + calls.append(model) + return deployment("gpt", "b") + + write(router, "async_get_available_deployment", pick) + assert read(router, "async_get_available_deployment") is pick + picked: Final = await router.select("gpt", {}) + assert picked_id(picked) == "b" + assert calls == ["gpt"] + + +@pytest.mark.asyncio +async def test_the_async_pick_attribute_picks_like_select(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with pytest.warns(DeprecationWarning, match=r"Router\.async_get_available_deployment is deprecated"): + pick: Final = router.async_get_available_deployment + picked: Final = await pick(model="gpt", request_kwargs={}) + assert picked_id(picked) == "a" + + +# -- methods ------------------------------------------------------------------------------------- + + +def _call(router: Router, name: str, *args: object, **kwargs: object) -> object: # kwargs-ok: the old member's + with pytest.warns(DeprecationWarning, match=rf"Router\.{name} is deprecated"): + return getattr(router, name)(*args, **kwargs) # pyright: ignore[reportAny] # each old member's own type + + +def test_registry_methods_answer_through_the_registry(make_core: MakeCore) -> None: + router: Final = _configured(make_core) + registry: Final = router.registry() + found: Final = _call(router, "get_deployment", "a") + assert isinstance(found, Deployment) + assert found.model_info.id == "a" + assert _call(router, "has_model_id", "a") is True + assert _call(router, "get_model_info", "zzz") is None + assert ids(_call(router, "get_model_list", "gpt")) == ["a", "b"] # pyright: ignore[reportArgumentType] + assert ids(_call(router, "get_all_deployments", "gpt")) == ["a", "b"] # pyright: ignore[reportArgumentType] + assert ids(_call(router, "get_model_list_from_model_alias", "fast")) == ["a", "b"] # pyright: ignore[reportArgumentType] + assert _call(router, "get_model_names") == registry.names() + assert _call(router, "get_model_ids", "gpt") == ["a", "b"] + assert _call(router, "get_model_from_alias", "fast") == "gpt" + assert _call(router, "routable_model_group", "fast") == "gpt" + assert _call(router, "is_recognized_model", "claude") is True + assert _call(router, "get_model_group", "a") == registry.resolve("a", kind="group_of_id") + assert _call(router, "map_team_model", "pub", "TA") == "pub" + assert _call(router, "get_model_access_groups") == registry.access_groups() + assert _call(router, "get_deployment_credentials", "a") == {"api_key": "sk-test"} + assert _call(router, "get_router_deployment_budget_limiter") is None + assert _call(router, "should_include_deployment", "gpt", deployment("gpt", "a")) is True + assert _call(router, "get_router_model_info", None, "gpt", id="a") == registry.router_model_info( + None, "gpt", id="a" + ) + + +def test_settings_methods_answer_through_settings(make_core: MakeCore) -> None: + router: Final = _configured(make_core) + assert _call(router, "get_settings") == dict(router.settings().as_dict()) + assert _call(router, "get_available_guardrail", "g") == { + "guardrail_name": "g", + "litellm_params": {"guardrail": "x"}, + } + with ( + pytest.raises(ValueError, match="No guardrail found with name: nope"), + pytest.warns(DeprecationWarning, match=r"Router\.get_available_guardrail is deprecated"), + ): + router.get_available_guardrail("nope") + + +@pytest.mark.asyncio +async def test_health_and_selection_methods_answer_through_their_entry_points(make_core: MakeCore) -> None: + router: Final = as_router( + make_core(model_list=[deployment("gpt", "a"), deployment("pt", "p", use_in_pass_through=True)]) + ) + assert _call(router, "get_fully_blocked_model_names") == set() + with pytest.warns(DeprecationWarning, match=r"Router\.async_get_fully_unhealthy_model_names is deprecated"): + unhealthy: Final = await router.async_get_fully_unhealthy_model_names() + assert unhealthy == set() + with pytest.warns(DeprecationWarning, match=r"Router\.get_model_group_usage is deprecated"): + usage: Final = await router.get_model_group_usage("gpt") + assert usage == (None, None) + with pytest.warns(DeprecationWarning, match=r"Router\.get_remaining_model_group_usage is deprecated"): + remaining: Final = await router.get_remaining_model_group_usage("gpt") + assert remaining == {} + with pytest.warns(DeprecationWarning, match=r"Router\.async_get_healthy_deployments is deprecated"): + healthy = await router.async_get_healthy_deployments("gpt", {}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the old untyped dicts + assert isinstance(healthy, list) + assert [picked_id(entry) for entry in healthy] == ["a"] # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # see above + assert picked_id(_call(router, "get_available_deployment_for_pass_through", "pt")) == "p" + + +def test_write_methods_go_through_apply(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + _call(router, "update_settings", allowed_fails=6) + assert router.settings().allowed_fails == 6 + router.apply(ChangeSet()) + assert _call(router, "delete_deployment", "a") is not None + assert router.registry().all() == [] + + +def test_the_priority_split_answers_through_health(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], default_priority=2)) + assert _call(router, "split_request_priority", "gpt", {"priority": 3, "x": 1}) == (3, {"x": 1}) + assert _call(router, "split_request_priority", "gpt", {"x": 1}) == (2, {"x": 1}) + assert router.health().priority_header_fields("gpt", {}, "5") == {"priority": 5} + assert router.health().priority_header_fields("gpt", {}, None) == {} + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'bad'"): + router.health().split_priority("gpt", {"priority": "bad"}) + + +# -- plugin context, request logging ------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_pre_routing_hook_answers_through_the_routing_context(make_core: MakeCore) -> None: + router: Final = as_router( + make_core(model_list=[deployment("small", "s"), deployment("big", "b"), _complexity_marker()]) + ) + request: Final[dict[str, object]] = {} + with pytest.warns(DeprecationWarning, match=r"Router\.async_pre_routing_hook is deprecated"): + response = await router.async_pre_routing_hook("smart", request, messages=MESSAGES) # pyright: ignore[reportUnknownMemberType] # the old signature's bare ``list`` + assert response is not None + assert response.model == "small" + decision: Final = as_str_mapping((as_str_mapping(request.get("metadata")) or {}).get("routing_decision")) or {} + assert decision.get("router_model_name") == "smart" + assert decision.get("routed_model") == "small" + plain: Final[dict[str, object]] = {} + with pytest.warns(DeprecationWarning, match="async_pre_routing_hook"): + hooked: Final = await router.async_pre_routing_hook("small", plain) # pyright: ignore[reportUnknownMemberType] # see above + assert hooked is None + assert plain == {} + + +def _complexity_marker() -> dict[str, object]: # mutable-ok: a model_list entry + return { + "model_name": "smart", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "small", "MEDIUM": "big", "COMPLEX": "big"}}, + "complexity_router_default_model": "big", + }, + "model_info": {"id": "marker"}, + } + + +def test_log_retry_records_the_failure_on_the_request_metadata(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a", base_model="bm")])) + kwargs: Final[dict[str, object]] = { + "model": "gpt", + "metadata": {"model_info": {"id": "a"}, "deployment": "openai/a", "attempted_retries": 1}, + } + assert _call(router, "log_retry", kwargs, ValueError("boom")) is kwargs + assert kwargs["metadata"] == { + "model_info": {"id": "a"}, + "deployment": "openai/a", + "attempted_retries": 1, + "previous_models": ( + { + "model_group": "gpt", + "deployment_id": "a", + "model": "openai/a", + "base_model": "bm", + "exception_type": "ValueError", + "exception_string": "boom", + "attempted_retries": 1, + }, + ), + "request_retry_count": 1, + } + _call(router, "log_retry", kwargs, KeyError("again")) + metadata: Final = as_str_mapping(kwargs["metadata"]) or {} + assert metadata["request_retry_count"] == 2 + previous: Final = metadata["previous_models"] + assert isinstance(previous, tuple) + assert len(previous) == 2 # pyright: ignore[reportUnknownArgumentType] # an untyped tuple of breadcrumbs + + +def test_arming_the_read_prefetch_without_a_request_batch_does_nothing(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], routing_strategy="usage-based-routing-v2")) + assert _call(router, "arm_routing_read_prefetch", "gpt", {}) is None + assert _call(router, "arm_routing_read_prefetch", "unknown-model") is None + + +def test_binding_to_request_logging_binds_through_observe(make_core: MakeCore) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + logging_obj: Final = Logging( + model="gpt", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + assert _call(router, "bind_to_request_logging", logging_obj) is None + with ( + pytest.raises(AttributeError, match="add_dynamic_function_callbacks"), + pytest.warns(DeprecationWarning, match=r"Router\.bind_to_request_logging is deprecated"), + ): + router.bind_to_request_logging(None) # pyright: ignore[reportArgumentType] + + +# -- writes with a routing effect ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_assigned_routing_groups_route_from_the_next_call(make_core: MakeCore) -> None: + from litellm.types.router import RoutingGroup + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")])) + with pytest.raises(Exception, match="grp"): + await router.select("grp", {}) + write( + router, + "routing_groups", + {"grp": RoutingGroup(group_name="grp", models=["claude"], routing_strategy="simple-shuffle")}, + ) + assert picked_id(await router.select("grp", {})) == "c" + assert router.registry().routing_groups().keys() == {"grp"} + + +class _NarrowTo: + def __init__(self, keep: str) -> None: + self.keep: Final = keep + + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [m for m in context.candidate_models if m == self.keep] + return context + + +@pytest.mark.asyncio +async def test_a_plugin_added_to_an_empty_plugin_list_runs(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("g", "a"), deployment("g", "b")])) + with pytest.warns(DeprecationWarning, match=r"Router\.routing_plugins is deprecated"): + plugins: Final = router.routing_plugins + plugins.append(_NarrowTo("openai/b")) + for _ in range(4): + assert picked_id(await router.select("g", {}, messages=MESSAGES)) == "b" # pyright: ignore[reportCallIssue] + plugins.clear() + assert {picked_id(await router.select("g", {}, messages=MESSAGES)) for _ in range(40)} == {"a", "b"} # pyright: ignore[reportCallIssue] + + +@pytest.mark.asyncio +async def test_an_assigned_plugin_list_runs(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("g", "a"), deployment("g", "b")])) + write(router, "routing_plugins", [_NarrowTo("openai/a")]) + for _ in range(4): + assert picked_id(await router.select("g", {}, messages=MESSAGES)) == "a" # pyright: ignore[reportCallIssue] + + +def test_an_assigned_prompt_caching_pin_cache_is_the_one_flushed(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + pins: Final = DualCache() + write(router, "prompt_caching_pin_cache", pins) + assert read(router, "prompt_caching_pin_cache") is pins + pins.in_memory_cache.set_cache("k", "v") # pyright: ignore[reportUnknownMemberType] # the in-memory cache is untyped + router.flush_cache() # pyright: ignore[reportAttributeAccessIssue] + assert pins.in_memory_cache.get_cache("k") is None # pyright: ignore[reportUnknownMemberType] # see above + + +# -- pod resolution and cooldown_time -------------------------------------------------------------- + + +class _TaggingDiscovery: + """Stands in for KubernetesPodDiscovery: stamps which discovery resolved the pick.""" + + def __init__(self, name: str) -> None: + self.name: Final = name + + def resolve_deployment(self, deployment: object, request_kwargs: object) -> object: + return {**(as_str_mapping(deployment) or {}), "resolved_by": self.name} + + async def async_resolve_deployment(self, deployment: object, request_kwargs: object) -> object: + return self.resolve_deployment(deployment, request_kwargs) + + +@pytest.mark.asyncio +async def test_a_custom_strategy_resolves_pods_through_the_discovery_bound_when_registered( + make_core: MakeCore, +) -> None: + from litellm.types.router import CustomRoutingStrategyBase + + class Pick(CustomRoutingStrategyBase): + async def async_get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: object = False, + request_kwargs: object = None, + ) -> object: + return deployment("gpt", "b") + + def get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: object = False, + request_kwargs: object = None, + ) -> object: + return deployment("gpt", "b") + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")])) + router.apply(ChangeSet(assign={"kubernetes_pod_discovery": _TaggingDiscovery("first")})) + assert (as_str_mapping(await router.select("gpt", {})) or {}).get("resolved_by") == "first" + router.register_plugin(Pick()) + router.apply(ChangeSet(assign={"kubernetes_pod_discovery": _TaggingDiscovery("second")})) + picked: Final = as_str_mapping(await router.select("gpt", {})) or {} + assert picked_id(picked) == "b" + assert picked.get("resolved_by") == "first" + + +def test_cooldown_time_reads_back_as_given(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], cooldown_time=7)) + value: Final = read(router, "cooldown_time") + assert value == 7 and type(value) is int + assert router.settings().as_dict()["cooldown_time"] == 7 + assert type(router.settings().as_dict()["cooldown_time"]) is int + write(router, "cooldown_time", 2.5) + assert router.settings().cooldown_time == 2.5 + _call(router, "update_settings", cooldown_time="3") + updated: Final = router.settings().cooldown_time + assert updated == 3 and type(updated) is int + default: Final = as_router(make_core(model_list=[deployment("gpt", "a")])).settings().cooldown_time + assert default == 5 and type(default) is int diff --git a/tests/unit/test_router/core/test_endpoints.py b/tests/unit/test_router/core/test_endpoints.py index c0c4c8cf6dd..ce56ae403ba 100644 --- a/tests/unit/test_router/core/test_endpoints.py +++ b/tests/unit/test_router/core/test_endpoints.py @@ -22,6 +22,7 @@ from tests.unit.test_router.core.scripted_host import ( completion_response, deployment, embedding_response, + records_outcomes_through_callbacks, ) MESSAGES: Final = [{"role": "user", "content": "hi"}] @@ -45,11 +46,16 @@ ENDPOINTS: Final[dict[str, tuple[str, EndpointCall]]] = { "anthropic_messages", lambda core: core.aanthropic_messages(model="g", messages=MESSAGES, max_tokens=5), ), + # The proxy's /v1/messages route calls the router by this name (route_request). + "anthropic_messages": ( + "anthropic_messages", + lambda core: core.anthropic_messages(model="g", messages=MESSAGES, max_tokens=5), + ), "call-acompletion": ("acompletion", lambda core: core.call("acompletion", {"model": "g", "messages": MESSAGES})), "call-completion": ("completion", lambda core: core.call("completion", {"model": "g", "messages": MESSAGES})), } SYNC_ENDPOINTS: Final = frozenset({"completion", "embedding", "image_generation"}) -LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages"}) +LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages", "anthropic_messages"}) EMBEDDING_ENDPOINTS: Final = frozenset({"aembedding", "embedding"}) @@ -276,14 +282,15 @@ async def test_request_level_api_key_replaces_the_deployment_key_on_the_host_cal async def test_construction_registers_one_success_and_one_failure_callback_per_mode( make_core: MakeCore, host: ScriptedHost ): - _one_group(make_core) + core: Final = _one_group(make_core) registered: Final = {name: len(callbacks) for name, callbacks in host.lists.items()} + one: Final = 1 if records_outcomes_through_callbacks(core) else 0 assert registered == { "callbacks": 0, - "success_callback": 1, - "_async_success_callback": 1, - "failure_callback": 1, - "_async_failure_callback": 1, + "success_callback": one, + "_async_success_callback": one, + "failure_callback": one, + "_async_failure_callback": one, "input_callback": 0, } assert all(callable(callbacks[0]) for callbacks in host.lists.values() if callbacks) @@ -416,7 +423,6 @@ def _scripted_stream(*texts: str) -> Callable[..., CustomStreamWrapper]: yield ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": text}}]) # A stream that ends without a finish reason fails since #47665. yield ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]) - yield ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]) logging_obj: Final = LiteLLMLogging( model="d1", diff --git a/tests/unit/test_router/core/test_endpoints_characterization_batch.py b/tests/unit/test_router/core/test_endpoints_characterization_batch.py new file mode 100644 index 00000000000..ab3543d7ad9 --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_characterization_batch.py @@ -0,0 +1,516 @@ +""" +Characterization of the fan-out endpoints and request shapers: the files API (one file per +deployment), batches (create, retrieve across deployments, cancel, list), batch completions, +fastest-response completion, the scheduler (priority) and prompt-management model groups. +""" + +import re +import datetime +import io +import json +from collections.abc import Callable, Mapping +from types import SimpleNamespace +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_core import RouterCore +from litellm.types.utils import LiteLLMBatch, ModelResponse +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import ScriptedHost, bad_request, deployment + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] +TWO_DEPLOYMENTS: Final = [deployment("g", "d1"), deployment("g", "d2")] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +def _bucket(kwargs: Mapping[str, object], name: str) -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +def _batch(batch_id: str = "b1") -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + ) + + +def _jsonl() -> io.BytesIO: + return io.BytesIO(b'{"custom_id": "1", "body": {"model": "g", "messages": []}}\n') + + +# --- files ---------------------------------------------------------------------------------- + + +async def test_acreate_file_creates_the_file_on_every_deployment_of_the_group(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + assert sorted(host.deployments_called("acreate_file")) == ["d1", "d2"] + + +async def test_acreate_file_returns_the_first_response_with_the_deployment_file_id_mapping( + make_core: MakeCore, host: EndpointHost +): + first: Final = ModelResponse(id="file-1") + second: Final = ModelResponse(id="file-2") + host.answer("openai/d1", first) + host.answer("openai/d2", second) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + response: Final = await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + assert response in (first, second) + assert response._hidden_params["model_file_id_mapping"] == {"d1": "file-1", "d2": "file-2"} + assert "model_id" not in response._hidden_params + + +async def test_acreate_file_rewrites_the_jsonl_model_to_each_deployment_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + models: Final = {} + for call in host.calls: + sent_file: Final = call.kwargs["file"] + assert isinstance(sent_file, io.BytesIO) + models[call.deployment_id] = json.loads(sent_file.getvalue().decode().strip())["body"]["model"] + assert models == {"d1": "d1", "d2": "d2"} + + +async def test_acreate_file_keeps_a_non_batch_file_as_sent(make_core: MakeCore, host: EndpointHost): + upload: Final = _jsonl() + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_file(model="g", file=upload, purpose="assistants") + assert host.calls[0].kwargs["file"] is not None + sent: Final = host.calls[0].kwargs["file"] + assert isinstance(sent, io.BytesIO) + assert json.loads(sent.getvalue().decode().strip())["body"]["model"] == "g" + + +async def test_acreate_file_puts_the_model_group_in_metadata_and_the_deployment_in_litellm_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + sent: Final = host.calls[0].kwargs + assert _bucket(sent, "metadata")["model_group"] == "g" + assert "deployment" not in _bucket(sent, "metadata") + assert _bucket(sent, "litellm_metadata")["deployment"] == "openai/d1" + assert "model_group" not in _bucket(sent, "litellm_metadata") + assert sent["custom_llm_provider"] == "openai" + + +async def test_acreate_file_without_a_purpose_raises_before_calling_the_host(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(Exception, match="file and file_purpose are required for create_file"): + await core.acreate_file(model="g", file=_jsonl()) + assert host.calls == [] + + +async def test_acreate_file_sends_target_model_names_to_a_litellm_proxy_deployment( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="litellm_proxy/remote-model")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch", extra_body={"keep": 1}) + assert dict(host.calls[0].kwargs["extra_body"]) == {"keep": 1, "target_model_names": "remote-model"} # pyright: ignore[reportArgumentType] # a mapping + + +async def test_acreate_file_sends_the_deployment_gcs_bucket_in_litellm_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", gcs_bucket_name="bucket-1")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + assert _bucket(host.calls[0].kwargs, "litellm_metadata")["gcs_bucket_name"] == "bucket-1" + + +# --- batches -------------------------------------------------------------------------------- + + +async def test_acreate_batch_routes_with_litellm_metadata_only(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.acreate_batch( + model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h" + ) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.kwargs["custom_llm_provider"]) == ( + "acreate_batch", + "openai/d1", + "d1", + "openai", + ) + assert "metadata" not in call.kwargs + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "g" + assert response._hidden_params["model_id"] == "d1" + + +async def test_acreate_batch_is_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1"), _batch()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=1) + await core.acreate_batch(model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h") + assert len(host.calls_to("acreate_batch")) == 2 + + +async def test_aretrieve_batch_asks_every_deployment_and_returns_the_batch_found( + make_core: MakeCore, host: EndpointHost +): + found: Final = _batch() + host.answer("openai/d1", bad_request("d1", "no such batch")) + host.answer("openai/d2", found) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + assert await core.aretrieve_batch(model="g", batch_id="b1") is found + assert sorted(host.deployments_called("aretrieve_batch")) == ["d1", "d2"] + for call in host.calls: + assert (call.kwargs["batch_id"], call.kwargs["custom_llm_provider"], call.model_group) == ("b1", "openai", "g") + assert "metadata" not in call.kwargs + + +async def test_aretrieve_batch_raises_the_first_error_when_no_deployment_has_the_batch( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", bad_request("d1", "first missing")) + host.answer("openai/d2", bad_request("d2", "second missing")) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + with pytest.raises(litellm.BadRequestError, match="missing"): + await core.aretrieve_batch(model="g", batch_id="b1") + assert len(host.calls) == 2 + + +async def test_aretrieve_batch_ignores_a_non_batch_answer_and_raises_a_generic_error( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(Exception, match="Unable to find batch in any model"): + await core.aretrieve_batch(model="g", batch_id="b1") + + +async def test_aretrieve_batch_without_a_model_asks_every_deployment_with_its_own_group( + make_core: MakeCore, host: EndpointHost +): + host.answer("*", bad_request(message="missing")) + core: Final = make_core(model_list=[*TWO_DEPLOYMENTS, deployment("h", "d3")]) + with pytest.raises(litellm.BadRequestError): + await core.aretrieve_batch(batch_id="b1") + assert sorted((call.deployment_id, call.model_group) for call in host.calls) == [ + ("d1", "g"), + ("d2", "g"), + ("d3", "h"), + ] + + +async def test_aretrieve_batch_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + await core.aretrieve_batch(model="g", batch_id="b1") + assert len(host.calls) == 1 + + +async def test_acancel_batch_routes_with_litellm_metadata(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + response: Final = await core.acancel_batch(model="g", batch_id="b1") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["batch_id"], call.kwargs["api_base"]) == ( + "acancel_batch", + "d1", + "b1", + "https://d1.example", + ) + assert "metadata" not in call.kwargs + assert response._hidden_params["model_id"] == "d1" + + +async def test_alist_batches_merges_the_lists_of_every_deployment(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", SimpleNamespace(first_id="a", last_id="b", data=[1, 2], has_more=False)) + host.answer("openai/d2", SimpleNamespace(first_id="c", last_id="d", data=[3], has_more=True)) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + merged: Final = await core.alist_batches(model="g", limit=2) + assert sorted(merged["data"]) == [1, 2, 3] + assert (merged["object"], merged["has_more"]) == ("list", True) + assert merged["first_id"] in ("a", "c") + assert merged["last_id"] in ("b", "d") + + +async def test_alist_batches_sends_the_deployment_params_and_request_without_routing_metadata( + make_core: MakeCore, host: EndpointHost +): + host.answer("*", SimpleNamespace(first_id=None, last_id=None, data=[], has_more=False)) + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + await core.alist_batches(model="g", limit=2, api_base="https://request.example") + call: Final = host.calls[0] + assert (call.model, call.kwargs["limit"], call.kwargs["api_base"], call.deployment_id) == ( + "openai/d1", + 2, + "https://request.example", + None, + ) + assert "metadata" not in call.kwargs and "litellm_metadata" not in call.kwargs + + +async def test_alist_batches_drops_failing_deployments(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", bad_request("d1")) + host.answer("openai/d2", SimpleNamespace(first_id="c", last_id="d", data=[3], has_more=False)) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + merged: Final = await core.alist_batches(model="g") + assert merged == {"object": "list", "data": [3], "first_id": "c", "last_id": "d", "has_more": False} + + +async def test_alist_batches_for_an_unknown_group_returns_an_empty_list(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + assert await core.alist_batches(model="nope") == { + "object": "list", + "data": [], + "first_id": None, + "last_id": None, + "has_more": False, + } + assert host.calls == [] + + +# --- batch completions ------------------------------------------------------------------------ + + +async def test_abatch_completion_sends_one_request_to_each_model_and_returns_errors_in_place( + make_core: MakeCore, host: EndpointHost +): + error: Final = bad_request("d1", "a failed") + host.answer("openai/d1", error) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + results: Final = await core.abatch_completion(models=["a", "b"], messages=MESSAGES) + assert isinstance(results[0], litellm.BadRequestError) + assert isinstance(results[1], ModelResponse) and results[1]._hidden_params["model_id"] == "d2" + + +async def test_abatch_completion_with_message_lists_groups_results_per_request(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")]) + second: Final = [{"role": "user", "content": "second"}] + results: Final = await core.abatch_completion(models=["a", "b"], messages=[MESSAGES, second]) + assert [[response._hidden_params["model_id"] for response in row] for row in results] == [ + ["d1", "d2"], + ["d1", "d2"], + ] + assert sorted(str(call.kwargs["messages"]) for call in host.calls) == sorted( + [str(MESSAGES), str(MESSAGES), str(second), str(second)] + ) + + +async def test_abatch_completion_forwards_extra_kwargs_to_every_request(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")]) + await core.abatch_completion(models=["a", "b"], messages=MESSAGES, max_tokens=15) + assert [call.kwargs["max_tokens"] for call in host.calls] == [15, 15] + + +async def test_abatch_completion_one_model_multiple_requests_returns_one_result_per_request( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", ModelResponse(id="r1"), bad_request("d1"), ModelResponse(id="r3")) + core: Final = make_core(model_list=[deployment("a", "d1")], num_retries=0) + results: Final = await core.abatch_completion_one_model_multiple_requests(model="a", messages=[MESSAGES] * 3) + assert len(results) == 3 + assert sum(isinstance(result, litellm.BadRequestError) for result in results) == 1 + assert len(host.calls) == 3 + + +async def test_fastest_response_returns_the_first_success_marked_in_hidden_params( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", bad_request("d1")) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + response: Final = await core.abatch_completion_fastest_response(model="a, b", messages=MESSAGES) + assert response._hidden_params["model_id"] == "d2" + assert response._hidden_params["fastest_response_batch_completion"] is True + + +async def test_fastest_response_raises_when_every_model_fails(make_core: MakeCore, host: EndpointHost): + host.answer("*", bad_request()) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + with pytest.raises(Exception, match="All tasks failed"): + await core.abatch_completion_fastest_response(model="a,b", messages=MESSAGES) + + +async def test_fastest_response_with_an_unknown_model_counts_it_as_failed(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1")], num_retries=0) + response: Final = await core.abatch_completion_fastest_response(model="nope, a", messages=MESSAGES) + assert response._hidden_params["model_id"] == "d1" + + +# --- scheduler / priority ------------------------------------------------------------------- + + +async def test_schedule_acompletion_marks_the_response_as_prioritized_and_does_not_send_priority( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.schedule_acompletion(model="g", messages=MESSAGES, priority=0) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + assert "priority" not in host.calls[0].kwargs + + +async def test_schedule_acompletion_without_a_priority_or_default_raises(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises( + litellm.BadRequestError, + match=re.escape("priority is required to queue a request, send one or set router_settings.default_priority"), + ): + await core.schedule_acompletion(model="g", messages=MESSAGES) + assert host.calls == [] + + +async def test_schedule_acompletion_without_a_priority_uses_the_default_priority( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_priority=5) + response: Final = await core.schedule_acompletion(model="g", messages=MESSAGES) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +@pytest.mark.parametrize("endpoint", ["schedule_acompletion", "acompletion"]) +async def test_a_non_integer_priority_raises_bad_request(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'high'"): + await core.call(endpoint, {"model": "g", "messages": MESSAGES, "priority": "high"}) + assert host.calls == [] + + +async def test_a_non_integer_priority_with_drop_params_is_served_unscheduled(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES, priority="high", drop_params=True) + assert "x-litellm-request-prioritization-used" not in response._hidden_params["additional_headers"] + assert "priority" not in host.calls[0].kwargs + + +async def test_a_non_integer_priority_with_drop_params_queues_at_the_default_priority( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_priority=5) + response: Final = await core.acompletion(model="g", messages=MESSAGES, priority="high", drop_params=True) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +@pytest.mark.parametrize( + ("endpoint", "request_body"), + [("acompletion", {"messages": MESSAGES}), ("atext_completion", {"prompt": "p"})], +) +async def test_a_request_priority_schedules_the_call( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_body: dict[str, object] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.call(endpoint, {"model": "g", "priority": 3, **request_body}) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True # pyright: ignore[reportAttributeAccessIssue] # the scripted response + assert "priority" not in host.calls[0].kwargs + + +async def test_default_priority_schedules_every_acompletion(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], default_priority=5) + response: Final = await core.acompletion(model="g", messages=MESSAGES) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +async def test_a_scheduled_call_error_carries_its_priority(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", bad_request("d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.BadRequestError) as raised: + await core.acompletion(model="g", messages=MESSAGES, priority=4) + assert raised.value.priority == 4 # pyright: ignore[reportAttributeAccessIssue] # stamped by the scheduler + + +@pytest.mark.parametrize(("header", "expected"), [("7", {"priority": 7}), (None, {}), ("", {})]) +async def test_priority_header_update_parses_the_priority_header( + make_core: MakeCore, host: EndpointHost, header: str | None, expected: dict[str, int] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + assert dict(core.priority_header_update("g", {}, header)) == expected + + +# --- prompt management models ----------------------------------------------------------------- + + +def _logging_obj() -> LiteLLMLogging: + return LiteLLMLogging( + model="pm", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c1", + function_id="f1", + ) + + +async def test_prompt_management_group_without_a_prompt_id_raises_value_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("pm", "p1", model="langfuse/gpt-4o")]) + with pytest.raises(ValueError, match=re.escape("Prompt ID is not set or not a string. Got=None")): + await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") + assert host.calls == [] + + +async def test_prompt_management_group_with_non_dict_prompt_variables_raises_value_error( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="langfuse/gpt-4o", prompt_id="p", prompt_variables="x")] + ) + with pytest.raises(ValueError, match=re.escape("Prompt variables is set but not a dictionary. Got=x")): + await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") + + +async def test_prompt_management_group_without_a_call_id_or_logging_object_raises_key_error( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("pm", "p1", model="langfuse/gpt-4o")]) + with pytest.raises(KeyError, match="litellm_call_id"): + await core.acompletion(model="pm", messages=MESSAGES) + + +async def test_litellm_agent_group_without_a_logging_object_raises_key_error(make_core: MakeCore, host: EndpointHost): + """function_setup consumes original_function, so the direct-call path cannot pop it.""" + core: Final = make_core(model_list=[deployment("pm", "p1", model="litellm_agent/openai/gpt-4o")]) + with pytest.raises(KeyError, match="original_function"): + await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") + assert host.calls == [] + + +async def test_litellm_agent_group_calls_the_stripped_model_directly_with_the_deployment_params( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="litellm_agent/openai/gpt-4o", timeout=9, temperature=0.2)] + ) + response: Final = await core.acompletion(model="pm", messages=MESSAGES, litellm_logging_obj=_logging_obj()) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.model_group) == ("acompletion", "openai/gpt-4o", "p1", "pm") + assert (call.kwargs["timeout"], call.kwargs["temperature"], "original_function" in call.kwargs) == (9.0, 0.2, False) + assert response._hidden_params == {} + + +async def test_litellm_agent_group_naming_another_group_routes_into_it_with_the_agent_timeout( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="litellm_agent/g", timeout=9), deployment("g", "d1")] + ) + response: Final = await core.acompletion(model="pm", messages=MESSAGES, litellm_logging_obj=_logging_obj()) + call: Final = host.calls[0] + assert (call.model, call.model_group, call.kwargs["timeout"]) == ("openai/d1", "pm", 9.0) + # The agent deployment's api_key travels as a request credential, so the target gets a request-scoped id. + assert call.deployment_id != "d1" + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "g" + + +async def test_a_group_with_two_prompt_management_deployments_routes_normally(make_core: MakeCore, host: ScriptedHost): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="langfuse/gpt-4o"), deployment("pm", "p2", model="langfuse/gpt-4o")] + ) + await core.acompletion(model="pm", messages=MESSAGES) + call: Final = host.calls[0] + assert (call.op, call.model, call.model_group) == ("acompletion", "langfuse/gpt-4o", "pm") diff --git a/tests/unit/test_router/core/test_endpoints_characterization_generic.py b/tests/unit/test_router/core/test_endpoints_characterization_generic.py new file mode 100644 index 00000000000..f0f2db3c102 --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_characterization_generic.py @@ -0,0 +1,780 @@ +""" +Characterization of the generic endpoints (Responses API and its get/delete/cancel/compact/input +items, Google GenAI, OCR, search, decisions, video, containers, skills, interactions, managed +agents, realtime, fine-tuning, files, image edit, LLM pass-through, vector stores, assistants, +guardrails): which host op each reaches, the kwargs the core builds, and how an id in the request +pins the deployment. + +The Python core binds litellm's own functions for some of these endpoints instead of going through +``host.invoke``; the ``host`` fixture replaces those functions with recording stand-ins (see +``endpoints_helpers.py``), so every call shows up in ``host.calls``. +""" + +import re +import inspect +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_core import RouterCore +from litellm.types.videos.utils import encode_video_id_with_provider +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] + + +@pytest.fixture +def host(monkeypatch: pytest.MonkeyPatch) -> EndpointHost: + endpoint_host: Final = EndpointHost() + patch_direct_endpoints(monkeypatch, endpoint_host) + return endpoint_host + + +async def _call(core: RouterCore, endpoint: str, request: Mapping[str, object]) -> object: + result: Final = getattr(core, endpoint)(**request) + return await result if inspect.isawaitable(result) else result + + +def _bucket(kwargs: Mapping[str, object], name: str) -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +# endpoint -> (recorded op, request without model) +ROUTED_ASYNC: Final[dict[str, tuple[str, dict[str, object]]]] = { + "aresponses": ("aresponses", {"input": "hi"}), + "agenerate_content": ("agenerate_content", {"contents": []}), + "agenerate_content_stream": ("agenerate_content_stream", {"contents": []}), + "afile_delete": ("afile_delete", {"file_id": "f1"}), + "afile_content": ("afile_content", {"file_id": "f1"}), + "afile_list": ("afile_list", {}), + "aget_responses": ("aget_responses", {"response_id": "resp_plain"}), + "acancel_responses": ("acancel_responses", {"response_id": "resp_plain"}), + "acompact_responses": ("acompact_responses", {"response_id": "resp_plain"}), + "adelete_responses": ("adelete_responses", {"response_id": "resp_plain"}), + "alist_input_items": ("alist_input_items", {"response_id": "resp_plain"}), + "acreate_realtime_client_secret": ("acreate_realtime_client_secret", {}), + "arealtime_calls": ("arealtime_calls", {}), + "acreate_realtime_transcription_session": ("acreate_realtime_transcription_session", {}), + "acreate_fine_tuning_job": ("acreate_fine_tuning_job", {"training_file": "f"}), + "acancel_fine_tuning_job": ("acancel_fine_tuning_job", {"fine_tuning_job_id": "j"}), + "alist_fine_tuning_jobs": ("alist_fine_tuning_jobs", {}), + "aretrieve_fine_tuning_job": ("aretrieve_fine_tuning_job", {"fine_tuning_job_id": "j"}), + "aimage_edit": ("aimage_edit", {"image": b"x", "prompt": "p"}), + "allm_passthrough_route": ("allm_passthrough_route", {"method": "POST", "endpoint": "/v1/x"}), + "aocr": ("aocr", {"document": {}}), + "adecisions": ("adecisions", {}), + "asystemone": ("asystemone", {}), + "aextraction": ("aextraction", {}), + "avideo_generation": ("avideo_generation", {"prompt": "p"}), + "avideo_list": ("avideo_list", {}), + "avideo_status": ("avideo_status", {"video_id": "v"}), + "avideo_content": ("avideo_content", {"video_id": "v"}), + "avideo_remix": ("avideo_remix", {"video_id": "v", "prompt": "p"}), + "avideo_create_character": ("avideo_create_character", {}), + "avideo_get_character": ("avideo_get_character", {}), + "avideo_edit": ("avideo_edit", {"video_id": "v"}), + "avideo_extension": ("avideo_extension", {"video_id": "v"}), + "acreate_container": ("acreate_container", {"name": "c"}), + "alist_containers": ("alist_containers", {}), + "aretrieve_container": ("aretrieve_container", {"container_id": "cntr_native"}), + "adelete_container": ("adelete_container", {"container_id": "cntr_native"}), + "aupload_container_file": ("aupload_container_file", {"container_id": "cntr_native", "file": b"x"}), + "alist_container_files": ("alist_container_files", {"container_id": "cntr_native"}), + "aretrieve_container_file": ("aretrieve_container_file", {"container_id": "cntr_native", "file_id": "f"}), + "adelete_container_file": ("adelete_container_file", {"container_id": "cntr_native", "file_id": "f"}), + "aretrieve_container_file_content": ( + "aretrieve_container_file_content", + {"container_id": "cntr_native", "file_id": "f"}, + ), + "acreate_skill": ("acreate_skill", {}), + "alist_skills": ("alist_skills", {}), + "aget_skill": ("aget_skill", {"skill_id": "s"}), + "adelete_skill": ("adelete_skill", {"skill_id": "s"}), + "acreate_interaction": ("interactions.acreate", {"input": "x"}), + "avector_store_search": ("vs.asearch", {"vector_store_id": "v", "query": "q"}), + "avector_store_retrieve": ("vs.aretrieve", {"vector_store_id": "v"}), + "avector_store_list": ("vs.alist", {}), + "avector_store_update": ("vs.aupdate", {"vector_store_id": "v"}), + "avector_store_delete": ("vs.adelete", {"vector_store_id": "v"}), + "avector_store_file_create": ("vsf.acreate", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_list": ("vsf.alist", {"vector_store_id": "v"}), + "avector_store_file_retrieve": ("vsf.aretrieve", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_content": ("vsf.aretrieve_content", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_update": ("vsf.aupdate", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_delete": ("vsf.adelete", {"vector_store_id": "v", "file_id": "f"}), +} + +ROUTED_SYNC: Final[dict[str, tuple[str, dict[str, object]]]] = { + "responses": ("responses", {"input": "hi"}), + "generate_content": ("generate_content", {"contents": []}), + "generate_content_stream": ("generate_content_stream", {"contents": []}), + "ocr": ("ocr", {"document": {}}), + "decisions": ("decisions", {}), + "systemone": ("systemone", {}), + "extraction": ("extraction", {}), + "video_generation": ("video_generation", {"prompt": "p"}), + "video_list": ("video_list", {}), + "video_status": ("video_status", {"video_id": "v"}), + "video_content": ("video_content", {"video_id": "v"}), + "video_remix": ("video_remix", {"video_id": "v", "prompt": "p"}), + "create_container": ("create_container", {"name": "c"}), + "list_containers": ("list_containers", {}), + "retrieve_container": ("retrieve_container", {"container_id": "c"}), + "delete_container": ("delete_container", {"container_id": "c"}), + "vector_store_search": ("vs.search", {"vector_store_id": "v", "query": "q"}), + "vector_store_create": ("vs.create", {"name": "v"}), + "vector_store_retrieve": ("vs.retrieve", {"vector_store_id": "v"}), + "vector_store_list": ("vs.list", {}), + "vector_store_update": ("vs.update", {"vector_store_id": "v"}), + "vector_store_delete": ("vs.delete", {"vector_store_id": "v"}), +} + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_ASYNC)) +async def test_async_generic_endpoint_routes_to_the_group_deployment_with_litellm_metadata( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + op, request = ROUTED_ASYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + response: Final = await _call(core, endpoint, {"model": "g", **request}) + assert [call.op for call in host.calls] == [op] + call: Final = host.calls[0] + assert (call.model, call.deployment_id, call.kwargs["api_base"], call.kwargs["custom_llm_provider"]) == ( + "openai/d1", + "d1", + "https://d1.example", + "openai", + ) + metadata: Final = _bucket(call.kwargs, "litellm_metadata") + assert (metadata["model_group"], metadata["deployment"]) == ("g", "openai/d1") + assert "metadata" not in call.kwargs + assert response._hidden_params["model_id"] == "d1" # pyright: ignore[reportAttributeAccessIssue] # the scripted response + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "g" # pyright: ignore[reportAttributeAccessIssue] # the scripted response + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_SYNC)) +async def test_sync_generic_endpoint_routes_to_the_group_deployment_and_returns_the_response( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + op, request = ROUTED_SYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + response: Final = getattr(core, endpoint)(model="g", **request) + assert not inspect.isawaitable(response) + assert [call.op for call in host.calls] == [op] + call: Final = host.calls[0] + assert (call.model, call.deployment_id, call.kwargs["api_base"], call.kwargs["custom_llm_provider"]) == ( + "openai/d1", + "d1", + "https://d1.example", + "openai", + ) + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "g" + assert response is not None + + +@pytest.mark.parametrize("endpoint", ["ocr", "generate_content", "video_generation"]) +async def test_sync_generic_endpoints_set_no_routing_hidden_params( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + _, request = ROUTED_SYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = getattr(core, endpoint)(model="g", **request) + assert response._hidden_params == {} + + +@pytest.mark.parametrize("endpoint", ["ocr", "generate_content"]) +async def test_sync_generic_endpoints_are_not_retried(make_core: MakeCore, host: EndpointHost, endpoint: str): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + _, request = ROUTED_SYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + getattr(core, endpoint)(model="g", **request) + assert len(host.calls) == 1 + + +async def test_async_generic_endpoint_is_retried_with_the_router_num_retries(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=2) + with pytest.raises(litellm.InternalServerError): + await core.aocr(model="g", document={}) + assert len(host.calls) == 3 + + +@pytest.mark.parametrize( + "endpoint", + [ + "aadapter_generate_content", + "video_create_character", + "video_get_character", + "video_edit", + "video_extension", + "upload_container_file", + "list_container_files", + "retrieve_container_file", + "delete_container_file", + "retrieve_container_file_content", + "get_interaction", + "delete_interaction", + "cancel_interaction", + ], +) +async def test_endpoints_without_a_dispatch_branch_resolve_to_none_without_calling_anything( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + """These are wrapped by factory_function's async wrapper, which has no branch for them.""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + result: Final = getattr(core, endpoint)(model="g", container_id="c", video_id="v", interaction_id="i") + assert inspect.isawaitable(result) + assert await result is None + assert host.calls == [] + + +async def test_sync_create_interaction_is_awaited_and_fails_with_type_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(TypeError, match="can't be used in 'await' expression"): + await core.create_interaction(model="g", input="x") + assert [call.op for call in host.calls] == ["interactions.create"] + + +# --- request params, provider and timeouts ------------------------------------------------------ + + +@pytest.mark.parametrize("endpoint", ["aocr", "ocr"]) +async def test_generic_request_params_win_over_deployment_params( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1", temperature=0.1)]) + await _call(core, endpoint, {"model": "g", "document": {}, "temperature": 0.9}) + assert host.calls[0].kwargs["temperature"] == 0.9 + + +async def test_generic_call_always_sends_the_deployment_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aocr(model="g", document={}) + assert host.calls[0].model == "openai/d1" + + +async def test_generic_call_sends_the_declared_custom_llm_provider(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", model="my-model", custom_llm_provider="hosted_vllm")]) + await core.aocr(model="g", document={}) + assert (host.calls[0].model, host.calls[0].kwargs["custom_llm_provider"]) == ("my-model", "hosted_vllm") + + +@pytest.mark.parametrize( + ("endpoint", "expected"), + [("aocr", 600.0), ("aresponses", 600.0), ("ocr", 6000.0), ("generate_content", 6000.0)], +) +async def test_generic_default_timeout_is_the_passthrough_default_for_async_and_the_process_setting_for_sync( + make_core: MakeCore, host: EndpointHost, endpoint: str, expected: float +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + request: Final = {**(ROUTED_ASYNC.get(endpoint) or ROUTED_SYNC[endpoint])[1], "model": "g"} + await _call(core, endpoint, request) + assert host.calls[0].kwargs["timeout"] == expected + + +@pytest.mark.parametrize("endpoint", ["aocr", "ocr"]) +async def test_generic_timeout_takes_the_router_timeout(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) + await _call(core, endpoint, {"model": "g", "document": {}}) + assert host.calls[0].kwargs["timeout"] == 42 + + +@pytest.mark.parametrize( + ("endpoint", "request_body"), [("aresponses", {"input": "hi"}), ("agenerate_content", {"contents": []})] +) +async def test_streaming_generic_call_sends_the_deployment_stream_timeout_as_timeout( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_body: dict[str, object] +): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)]) + await _call(core, endpoint, {"model": "g", "stream": True, **request_body}) + sent: Final = host.calls[0].kwargs + assert (sent["timeout"], sent["stream_timeout"]) == (3.0, 3.0) + + +async def test_non_streaming_aresponses_sends_the_deployment_timeout(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)]) + await core.aresponses(model="g", input="hi") + assert host.calls[0].kwargs["timeout"] == 11.0 + + +async def test_realtime_client_secret_session_model_follows_the_routed_deployment( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_realtime_client_secret(model="g", session={"model": "g", "type": "realtime"}) + assert host.calls[0].kwargs["session"] == {"model": "openai/d1", "type": "realtime"} + + +async def test_generic_endpoint_with_an_unknown_group_raises_without_calling_the_host( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(litellm.BadRequestError, match="model=nope"): + await core.agenerate_content(model="nope", contents=[]) + assert host.calls == [] + + +# --- Responses API ids pin the deployment ----------------------------------------------------- + +TWO_DEPLOYMENTS: Final = [deployment("g", "d1"), deployment("g", "d2")] +ENCODED_RESPONSE_ID: Final = ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "d2", "resp_abc") # pyright: ignore[reportPrivateUsage] # litellm's encoder + + +@pytest.mark.parametrize( + "endpoint", ["aget_responses", "acancel_responses", "acompact_responses", "adelete_responses", "alist_input_items"] +) +async def test_a_litellm_response_id_pins_the_deployment_it_encodes( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await _call(core, endpoint, {"response_id": ENCODED_RESPONSE_ID}) + assert host.deployments_called() == ["d2"] * 5 + + +async def test_a_litellm_response_id_overrides_the_request_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await core.adelete_responses(response_id=ENCODED_RESPONSE_ID, model="g") + assert host.deployments_called() == ["d2"] * 5 + assert _bucket(host.calls[0].kwargs, "litellm_metadata")["model_group"] == "d2" + + +async def test_the_response_id_reaches_the_host_still_encoded(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aget_responses(response_id=ENCODED_RESPONSE_ID) + assert host.calls[0].kwargs["response_id"] == ENCODED_RESPONSE_ID + + +async def test_a_plain_response_id_without_a_model_fails_with_type_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + with pytest.raises(TypeError, match="missing 1 required positional argument: 'model'"): + await core.aget_responses(response_id="resp_plain") + assert host.calls == [] + + +# --- container ids ------------------------------------------------------------------------ + +ENCODED_CONTAINER_ID: Final = ResponsesAPIRequestUtils.build_container_id("azure", "d2", "cntr_upstream") + + +async def test_a_litellm_container_id_pins_its_deployment_and_is_sent_decoded(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await core.aretrieve_container(container_id=ENCODED_CONTAINER_ID) + assert host.deployments_called() == ["d2"] * 5 + assert {call.kwargs["container_id"] for call in host.calls} == {"cntr_upstream"} + + +async def test_a_routed_container_call_sends_the_deployment_provider_not_the_decoded_one( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aretrieve_container(container_id=ENCODED_CONTAINER_ID, custom_llm_provider="openai") + assert host.calls[0].kwargs["custom_llm_provider"] == "openai" + + +async def test_a_native_container_id_routes_by_the_forwarded_model_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await core.aretrieve_container(container_id="cntr_native", model_id=" d1 ") + assert host.deployments_called() == ["d1"] * 5 + assert host.calls[0].kwargs["container_id"] == "cntr_native" + + +async def test_a_native_container_id_without_a_model_calls_the_handler_directly( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aretrieve_container(container_id="cntr_native", custom_llm_provider="azure") + call: Final = host.calls[0] + assert (call.deployment_id, call.kwargs["container_id"], call.kwargs["custom_llm_provider"]) == ( + None, + "cntr_native", + "azure", + ) + assert "litellm_metadata" not in call.kwargs + + +async def test_create_container_with_an_unknown_model_passes_through_with_that_model( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_container(name="c", model="not-a-group") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id) == ("acreate_container", "not-a-group", None) + + +async def test_create_container_without_a_model_calls_the_handler_directly(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_container(name="c") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, "model" in call.kwargs) == ("acreate_container", None, False) + + +# --- video ids ----------------------------------------------------------------------------- + +ENCODED_VIDEO_ID: Final = encode_video_id_with_provider("video_up", "openai", model_id="d2", deployment_id="d2") + + +@pytest.mark.parametrize( + "endpoint", ["avideo_status", "avideo_content", "avideo_remix", "avideo_edit", "avideo_extension"] +) +async def test_a_litellm_video_id_pins_the_deployment_that_created_it( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(6): + await _call(core, endpoint, {"model": "g", "video_id": ENCODED_VIDEO_ID, "prompt": "p"}) + assert host.deployments_called() == ["d2"] * 6 + + +async def test_a_video_id_on_a_video_generation_call_does_not_pin(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(30): + await core.avideo_generation(model="g", prompt="p", video_id=ENCODED_VIDEO_ID) + assert set(host.deployments_called()) == {"d1", "d2"} + + +async def test_sync_video_status_pins_the_deployment_of_the_video_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(6): + core.video_status(model="g", video_id=ENCODED_VIDEO_ID) + assert host.deployments_called() == ["d2"] * 6 + + +# --- LLM pass-through ---------------------------------------------------------------------- + + +async def test_llm_passthrough_with_an_unknown_model_calls_the_handler_with_that_model( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.allm_passthrough_route(model="not-a-group", method="POST", endpoint="/model/not-a-group/invoke") + call: Final = host.calls[0] + assert (call.model, call.deployment_id, call.kwargs["endpoint"]) == ( + "not-a-group", + None, + "/model/not-a-group/invoke", + ) + + +async def test_llm_passthrough_rewrites_the_group_name_in_the_endpoint_to_the_provider_model( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="bedrock/us.anthropic.claude")]) + await core.allm_passthrough_route(model="g", method="POST", endpoint="/model/g/invoke") + call: Final = host.calls[0] + assert (call.model, call.kwargs["endpoint"], call.kwargs["custom_llm_provider"]) == ( + "bedrock/us.anthropic.claude", + "/model/us.anthropic.claude/invoke", + "bedrock", + ) + + +async def test_llm_passthrough_forwards_the_request_client(make_core: MakeCore, host: EndpointHost): + client: Final = object() + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.allm_passthrough_route(model="g", method="POST", endpoint="/v1/x", client=client) + assert host.calls[0].kwargs["client"] is client + + +# --- interactions and managed agents ---------------------------------------------------------- + + +async def test_interactions_get_without_a_model_calls_the_handler_with_the_given_provider( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aget_interaction(interaction_id="i", custom_llm_provider="vertex_ai") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["custom_llm_provider"]) == ("interactions.aget", None, "vertex_ai") + + +@pytest.mark.parametrize("endpoint", ["aget_interaction", "adelete_interaction", "acancel_interaction"]) +async def test_interactions_default_to_the_gemini_provider(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await _call(core, endpoint, {"interaction_id": "i"}) + assert host.calls[0].kwargs["custom_llm_provider"] == "gemini" + + +async def test_interaction_with_an_agent_named_as_the_model_is_not_routed(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("ag", "d1")]) + await core.acreate_interaction(agent="ag", model="ag", input="x") + call: Final = host.calls[0] + assert (call.deployment_id, call.kwargs["model"], call.kwargs["custom_llm_provider"]) == (None, None, "gemini") + + +@pytest.mark.parametrize( + ("endpoint", "op"), + [ + ("acreate_agent", "agents.acreate"), + ("alist_agents", "agents.alist"), + ("aget_agent", "agents.aget"), + ("adelete_agent", "agents.adelete"), + ("alist_agent_versions", "agents.alist_versions"), + ("create_agent", "agents.create"), + ("list_agents", "agents.list"), + ("get_agent", "agents.get"), + ("delete_agent", "agents.delete"), + ("list_agent_versions", "agents.list_versions"), + ], +) +async def test_managed_agent_endpoints_call_the_handler_directly_with_gemini_by_default( + make_core: MakeCore, host: EndpointHost, endpoint: str, op: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await _call(core, endpoint, {"name": "a", "model": "g"}) + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["custom_llm_provider"], call.kwargs["model"]) == ( + op, + None, + "gemini", + "g", + ) + + +@pytest.mark.parametrize("endpoint", ["acreate_agent", "create_agent"]) +async def test_managed_agent_endpoints_keep_the_given_provider(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await _call(core, endpoint, {"name": "a", "custom_llm_provider": "vertex_ai"}) + assert host.calls[0].kwargs["custom_llm_provider"] == "vertex_ai" + + +# --- assistants ---------------------------------------------------------------------------- + +ASSISTANT_ENDPOINTS: Final = ( + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", +) + + +@pytest.mark.parametrize("endpoint", ASSISTANT_ENDPOINTS) +async def test_assistants_endpoints_without_a_provider_or_config_raise( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(Exception, match=re.escape("'custom_llm_provider' must be set. Either via:")): + await getattr(core, endpoint)() + assert host.calls == [] + + +@pytest.mark.parametrize("endpoint", ASSISTANT_ENDPOINTS) +async def test_assistants_endpoints_use_the_assistants_config(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core( + model_list=[deployment("g", "d1")], + assistants_config={ + "custom_llm_provider": "azure", + "litellm_params": {"api_key": "k", "api_base": "https://az"}, + }, + ) + await getattr(core, endpoint)(thread_id="t") + call: Final = host.calls[0] + assert (call.op, call.kwargs["custom_llm_provider"], call.kwargs["api_key"], call.kwargs["api_base"]) == ( + endpoint, + "azure", + "k", + "https://az", + ) + assert (call.kwargs["client"], call.kwargs["thread_id"]) == (None, "t") + + +async def test_assistants_endpoint_with_an_explicit_provider_ignores_the_config( + make_core: MakeCore, host: EndpointHost +): + client: Final = object() + core: Final = make_core( + model_list=[deployment("g", "d1")], + assistants_config={"custom_llm_provider": "azure", "litellm_params": {"api_key": "k"}}, + ) + await core.arun_thread(custom_llm_provider="openai", client=client, thread_id="t") + call: Final = host.calls[0] + assert (call.kwargs["custom_llm_provider"], call.kwargs["client"], "api_key" in call.kwargs) == ( + "openai", + client, + False, + ) + + +# --- vector stores --------------------------------------------------------------------------- + + +async def test_vector_store_search_without_a_model_gets_the_router_and_overrides_a_client_router( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.avector_store_search(vector_store_id="v", query="q", router="client-value") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["router"]) == ("vs.asearch", None, core) + assert "_direct_vector_store_embedding_executor" in call.kwargs + + +async def test_sync_vector_store_search_without_a_model_gets_the_router(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + core.vector_store_search(vector_store_id="v", query="q") + assert host.calls[0].kwargs["router"] is core + + +async def test_routed_vector_store_search_gets_the_embedding_executor_but_not_the_router( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.avector_store_search(model="g", vector_store_id="v", query="q") + sent: Final = host.calls[0].kwargs + assert ("router" in sent, "_direct_vector_store_embedding_executor" in sent) == (False, True) + + +async def test_vector_store_retrieve_without_a_model_calls_the_handler_directly( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.avector_store_retrieve(vector_store_id="v", custom_llm_provider="openai") + call: Final = host.calls[0] + assert (call.deployment_id, call.kwargs["custom_llm_provider"], "router" in call.kwargs) == (None, "openai", False) + + +async def test_sync_vector_store_create_without_a_model_passes_the_provider(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + core.vector_store_create(name="v", custom_llm_provider="openai") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["custom_llm_provider"]) == ("vs.create", None, "openai") + + +@pytest.mark.parametrize( + ("endpoint", "op"), + [ + ("vector_store_file_create", "vsf.create"), + ("vector_store_file_list", "vsf.list"), + ("vector_store_file_retrieve", "vsf.retrieve"), + ("vector_store_file_content", "vsf.retrieve_content"), + ("vector_store_file_update", "vsf.update"), + ("vector_store_file_delete", "vsf.delete"), + ], +) +async def test_sync_vector_store_file_endpoints_never_route_even_with_a_model( + make_core: MakeCore, host: EndpointHost, endpoint: str, op: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + getattr(core, endpoint)(model="g", vector_store_id="v", file_id="f") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.kwargs["custom_llm_provider"], call.kwargs["client"]) == ( + op, + "g", + None, + None, + None, + ) + + +async def test_avector_store_create_with_a_model_routes_with_metadata_and_no_model_group( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + await core.avector_store_create(model="g", name="n") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.kwargs["custom_llm_provider"]) == ( + "vs.acreate", + "openai/d1", + "d1", + "openai", + ) + metadata: Final = _bucket(call.kwargs, "metadata") + assert (metadata["deployment"], "model_group" in metadata) == ("openai/d1", False) + + +async def test_avector_store_create_without_a_model_calls_the_handler_directly(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.avector_store_create(model=None, name="n") + call: Final = host.calls[0] + assert (call.op, call.deployment_id) == ("vs.acreate", None) + + +async def test_avector_store_create_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + await core.avector_store_create(model="g", name="n") + assert len(host.calls) == 1 + + +# --- search tools --------------------------------------------------------------------------- + +SEARCH_TOOLS: Final = [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily", "api_key": "tk"}}] + + +async def test_asearch_calls_the_named_search_tool_with_its_litellm_params(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS) + response: Final = await core.asearch(model="s", query="q") + call: Final = host.calls[0] + assert (call.op, call.kwargs["search_provider"], call.kwargs["api_key"], call.kwargs["query"]) == ( + "asearch", + "tavily", + "tk", + "q", + ) + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "s" + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "s" # pyright: ignore[reportAttributeAccessIssue] # the scripted response + + +async def test_asearch_with_an_unknown_tool_raises_value_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS, num_retries=0) + with pytest.raises(ValueError, match=re.escape("Search tool 'g' not found in router.search_tools")): + await core.asearch(model="g", query="q") + assert host.calls == [] + + +# --- guardrails --------------------------------------------------------------------------- + +GUARDRAILS: Final = [{"guardrail_name": "gr", "litellm_params": {"guardrail": "presidio", "mode": "pre_call"}}] + + +async def test_aguardrail_runs_the_function_with_the_selected_guardrail(make_core: MakeCore, host: ScriptedHost): + seen: Final[list[Mapping[str, object]]] = [] + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + seen.append(kwargs) + return "verdict" + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + assert await core.aguardrail(guardrail_name="gr", original_function=guard, data={"x": 1}) == "verdict" + assert sorted(seen[0]) == ["data", "litellm_metadata", "litellm_trace_id", "selected_guardrail"] + assert seen[0]["selected_guardrail"] == GUARDRAILS[0] + assert _bucket(seen[0], "litellm_metadata")["model_group"] == "gr" + assert host.calls == [] + + +async def test_aguardrail_with_an_unknown_name_raises_value_error(make_core: MakeCore, host: ScriptedHost): + async def guard(**_kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + return "verdict" + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS, num_retries=0) + with pytest.raises(ValueError, match="No guardrail found with name: nope"): + await core.aguardrail(guardrail_name="nope", original_function=guard) + + +async def test_aguardrail_does_not_retry_a_failing_guardrail_function(make_core: MakeCore, host: ScriptedHost): + """A guardrail name is no model group (model_group_size 0), so the router num_retries do not apply.""" + attempts: Final[list[int]] = [] + + async def guard(**_kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + attempts.append(1) + raise litellm.InternalServerError(message="boom", llm_provider="", model="gr") + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS, num_retries=2) + with pytest.raises(litellm.InternalServerError, match="boom"): + await core.aguardrail(guardrail_name="gr", original_function=guard) + assert len(attempts) == 1 diff --git a/tests/unit/test_router/core/test_endpoints_characterization_inference.py b/tests/unit/test_router/core/test_endpoints_characterization_inference.py new file mode 100644 index 00000000000..8db0d6bea12 --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_characterization_inference.py @@ -0,0 +1,648 @@ +""" +Characterization of the inference endpoints (completion, text completion, embedding, image, +speech, transcription, rerank, adapter completion, moderation): which host op each invokes and +which keyword arguments the core builds for it. + +Every test builds its core through ``make_core`` on the ``host`` fixture, so it runs once per core. +""" + +import asyncio +import io +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, timeout + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +def _sent(host: ScriptedHost, index: int = -1) -> Mapping[str, object]: + return host.calls[index].kwargs + + +def _bucket(kwargs: Mapping[str, object], name: str = "metadata") -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +# --- host op per endpoint ------------------------------------------------------------------- + + +async def test_aadapter_completion_invokes_the_adapter_host_op_with_the_adapter_id( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aadapter_completion(adapter_id="my-adapter", model="g", messages=MESSAGES) + call: Final = host.calls[0] + assert (call.op, call.model, call.kwargs["adapter_id"], call.deployment_id, call.model_group) == ( + "aadapter_completion", + "openai/d1", + "my-adapter", + "d1", + "g", + ) + + +async def test_amoderation_routes_a_known_group_through_its_deployment(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + await core.amoderation(model="g", input="x") + call: Final = host.calls[0] + assert (call.op, call.model, call.kwargs["api_base"], call.deployment_id, call.model_group) == ( + "amoderation", + "openai/d1", + "https://d1.example", + "d1", + "g", + ) + assert "litellm_metadata" not in call.kwargs + + +async def test_amoderation_calls_the_host_directly_when_the_model_is_not_a_group( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.amoderation(model="omni-moderation-latest", input="x") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id) == ("amoderation", "omni-moderation-latest", None) + assert _bucket(call.kwargs)["model_group"] == "omni-moderation-latest" + + +async def test_amoderation_deployment_params_overwrite_the_request_params(make_core: MakeCore, host: EndpointHost): + """The deployment's litellm_params are applied after the request's (kwargs.update(data)).""" + core: Final = make_core(model_list=[deployment("g", "d1", temperature=0.5)]) + await core.amoderation(model="g", input="x", temperature=1.0) + assert _sent(host)["temperature"] == 0.5 + + +async def test_amoderation_with_a_request_api_key_keeps_the_deployment_key_under_a_request_scoped_id( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.amoderation(model="g", input="x", api_key="sk-request") + sent: Final = _sent(host) + model_info: Final = _bucket(sent, "model_info") + assert sent["api_key"] == "sk-test" + assert (model_info["original_model_id"], model_info["id"] != "d1") == ("d1", True) + + +async def test_amoderation_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + await core.amoderation(model="g", input="x") + assert len(host.calls) == 1 + assert host.sleeps == [] + + +async def test_arerank_sends_the_builtin_input_function_as_input(make_core: MakeCore, host: EndpointHost): + """arerank sets kwargs["input"] = input, the Python builtin, and it reaches the host.""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.arerank(model="g", query="q", documents=["a"]) + sent: Final = _sent(host) + assert (sent["input"], sent["query"], sent["documents"]) == (input, "q", ["a"]) + + +@pytest.mark.parametrize(("request_voice", "expected"), [(None, "deployment-voice"), ("alloy", "alloy")]) +async def test_aspeech_voice_defaults_to_the_deployment_voice( + make_core: MakeCore, host: EndpointHost, request_voice: str | None, expected: str +): + core: Final = make_core(model_list=[deployment("g", "d1", voice="deployment-voice")]) + await core.aspeech(model="g", input="hello", voice=request_voice) + assert (_sent(host)["voice"], _sent(host)["input"]) == (expected, "hello") + + +async def test_aspeech_sends_no_caching_flag(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], cache_responses=True) + await core.aspeech(model="g", input="hello", voice="alloy") + assert "caching" not in _sent(host) + + +async def test_atranscription_forwards_the_same_file_object(make_core: MakeCore, host: EndpointHost): + audio: Final = io.BytesIO(b"audio") + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.atranscription(model="g", file=audio) + assert _sent(host)["file"] is audio + + +@pytest.mark.parametrize( + "endpoint", + ["acompletion", "atext_completion", "aembedding", "aimage_generation", "atranscription", "arerank", "completion"], +) +async def test_cache_responses_reaches_the_host_as_caching(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")], cache_responses=True) + requests: Final[dict[str, dict[str, object]]] = { + "acompletion": {"messages": MESSAGES}, + "completion": {"messages": MESSAGES}, + "atext_completion": {"prompt": "hi"}, + "aembedding": {"input": ["a"]}, + "aimage_generation": {"prompt": "a cat"}, + "atranscription": {"file": io.BytesIO(b"x")}, + "arerank": {"query": "q", "documents": ["a"]}, + } + await core.call(endpoint, {"model": "g", **requests[endpoint]}) + assert _sent(host)["caching"] is True + + +@pytest.mark.parametrize("endpoint", ["acompletion", "completion"]) +async def test_cache_responses_turned_on_after_construction_reaches_the_host( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + """The proxy turns ``cache_responses`` on through ``apply`` once its cache exists.""" + from litellm.router_api import ChangeSet + + core: Final = make_core(model_list=[deployment("g", "d1")]) + _ = core.apply(ChangeSet(assign={"cache_responses": True})) + await core.call(endpoint, {"model": "g", "messages": MESSAGES}) + assert core.settings().cache_responses is True + assert _sent(host)["caching"] is True + + +# --- sync text completion and image generation ---------------------------------------------- + + +async def test_sync_text_completion_sends_the_router_num_retries_and_no_routing_fields( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=4) + core.text_completion(model="g", prompt="p") + sent: Final = _sent(host) + assert sent["num_retries"] == 4 + assert "model_info" not in sent + assert "litellm_trace_id" not in sent + assert sent.get("timeout") is None + metadata: Final = _bucket(sent) + assert metadata["model_group"] == "g" + assert "deployment" not in metadata + + +async def test_sync_text_completion_request_num_retries_wins(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=4) + core.text_completion(model="g", prompt="p", num_retries=1) + assert _sent(host)["num_retries"] == 1 + + +async def test_sync_text_completion_merges_default_metadata_into_request_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"x": 1}, "user": "u"} + ) + core.text_completion(model="g", prompt="p", metadata={"y": 2}) + sent: Final = _sent(host) + metadata: Final = _bucket(sent) + assert (metadata["x"], metadata["y"], metadata["model_group"], sent["user"]) == (1, 2, "g", "u") + + +async def test_sync_text_completion_sends_the_model_group_name_not_the_deployment_model( + make_core: MakeCore, host: EndpointHost +): + """The request's model="g" is spread after the deployment params, so it replaces "openai/d1".""" + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + core.text_completion(model="g", prompt="p") + call: Final = host.calls[0] + assert (call.model, call.kwargs["api_base"]) == ("g", "https://d1.example") + + +async def test_sync_text_completion_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("g", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + core.text_completion(model="g", prompt="p") + assert len(host.calls) == 1 + assert host.sleeps == [] + + +async def test_sync_image_generation_sends_routing_metadata_without_a_trace_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.image_generation(model="g", prompt="p") + sent: Final = _sent(host) + assert "litellm_trace_id" not in sent + assert (_bucket(sent)["model_group"], _bucket(sent)["deployment"], sent["prompt"]) == ("g", "openai/d1", "p") + + +@pytest.mark.parametrize("endpoint", ["acompletion", "atext_completion", "aembedding", "aresponses"]) +async def test_async_endpoints_generate_a_trace_id_unless_the_request_sends_one( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + request: Final[dict[str, object]] = ( + {"messages": MESSAGES} + if endpoint == "acompletion" + else {"prompt": "p"} + if endpoint == "atext_completion" + else {"input": ["a"]} + if endpoint == "aembedding" + else {"input": "hi"} + ) + await core.call(endpoint, {"model": "g", **request}) + await core.call(endpoint, {"model": "g", "litellm_trace_id": "trace-1", **request}) + generated: Final = _sent(host, 0)["litellm_trace_id"] + assert isinstance(generated, str) and len(generated) == 36 + assert _sent(host, 1)["litellm_trace_id"] == "trace-1" + + +# --- routing metadata --------------------------------------------------------------------- + + +async def test_a_model_group_alias_is_recorded_as_the_model_group_and_its_alias( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], model_group_alias={"al": "g"}) + await core.acompletion(model="al", messages=MESSAGES) + metadata: Final = _bucket(_sent(host)) + assert (metadata["model_group"], metadata["model_group_alias"], metadata["deployment_model_name"]) == ( + "al", + "al", + "g", + ) + + +async def test_a_request_for_a_deployment_id_routes_to_that_deployment(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("g", "d2")]) + for _ in range(4): + await core.acompletion(model="d2", messages=MESSAGES) + assert host.deployments_called() == ["d2"] * 4 + assert _bucket(_sent(host))["model_group"] == "d2" + + +async def test_deployment_tags_and_credential_name_join_the_request_tags(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tags=["t1"], litellm_credential_name="cred")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"tags": ["r1"]}) + sent: Final = _sent(host) + assert _bucket(sent)["tags"] == ["r1", "t1", "Credential: cred"] + assert sent["litellm_credential_name"] == "cred" + + +async def test_daily_reset_time_is_not_sent_to_the_host(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, daily_reset_time="00:00") + assert "daily_reset_time" not in _sent(host) + + +async def test_deployment_model_id_param_is_sent_as_model_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", model_id="arn:pinned")]) + await core.acompletion(model="g", messages=MESSAGES) + core.text_completion(model="g", prompt="p") + assert [call.kwargs["model_id"] for call in host.calls] == ["arn:pinned", "arn:pinned"] + + +async def test_deployment_model_id_reaches_the_host_even_when_clientside_auth_allows_model_id( + make_core: MakeCore, host: EndpointHost +): + """The pin is skipped, but the deployment's litellm_params (with model_id) are still spread into the call.""" + core: Final = make_core( + model_list=[deployment("g", "d1", model_id="arn:pinned", configurable_clientside_auth_params=["model_id"])] + ) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["model_id"] == "arn:pinned" + + +async def test_request_model_id_wins_over_a_clientside_configurable_deployment_model_id( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("g", "d1", model_id="arn:pinned", configurable_clientside_auth_params=["model_id"])] + ) + await core.acompletion(model="g", messages=MESSAGES, model_id="arn:request") + assert _sent(host)["model_id"] == "arn:request" + + +# --- default litellm params ------------------------------------------------------------------- + + +async def test_default_litellm_params_metadata_merges_into_request_metadata(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) + await core.acompletion(model="g", messages=MESSAGES, metadata={"trace": "x"}) + metadata: Final = _bucket(_sent(host)) + assert (metadata["team"], metadata["trace"], metadata["model_group"]) == ("t", "x", "g") + + +async def test_default_litellm_params_metadata_overwrites_the_same_request_metadata_key( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) + await core.acompletion(model="g", messages=MESSAGES, metadata={"team": "request"}) + assert _bucket(_sent(host))["team"] == "t" + + +async def test_default_litellm_params_metadata_goes_to_litellm_metadata_for_generic_endpoints( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) + await core.aresponses(model="g", input="hi") + sent: Final = _sent(host) + assert _bucket(sent, "litellm_metadata")["team"] == "t" + assert "metadata" not in sent + + +async def test_default_litellm_params_with_none_values_are_not_sent(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"user": None}) + await core.acompletion(model="g", messages=MESSAGES) + assert "user" not in _sent(host) + + +# --- tools ---------------------------------------------------------------------------------- + +DEPLOYMENT_TOOL: Final = {"type": "function", "function": {"name": "dep", "parameters": {}}} +REQUEST_TOOL: Final = {"type": "function", "function": {"name": "req", "parameters": {}}} + + +async def test_deployment_tools_come_before_request_tools(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL])]) + await core.acompletion(model="g", messages=MESSAGES, tools=[REQUEST_TOOL]) + assert _sent(host)["tools"] == [DEPLOYMENT_TOOL, REQUEST_TOOL] + + +async def test_deployment_tools_are_sent_when_the_request_has_none(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL])]) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["tools"] == [DEPLOYMENT_TOOL] + + +@pytest.mark.parametrize(("request_choice", "expected"), [(None, "required"), ("none", "none")]) +async def test_tool_choice_comes_from_the_request_else_the_deployment( + make_core: MakeCore, host: EndpointHost, request_choice: str | None, expected: str +): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL], tool_choice="required")]) + extra: Final = {} if request_choice is None else {"tool_choice": request_choice} + await core.acompletion(model="g", messages=MESSAGES, **extra) + assert _sent(host)["tool_choice"] == expected + + +async def test_merging_tools_does_not_change_the_deployment_tools(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL])]) + await core.acompletion(model="g", messages=MESSAGES, tools=[REQUEST_TOOL]) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["tools"] == [DEPLOYMENT_TOOL] + + +# --- reasoning effort ----------------------------------------------------------------------- + + +async def test_request_reasoning_effort_drops_the_deployment_thinking_param(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", thinking={"type": "enabled", "budget_tokens": 100})]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + sent: Final = _sent(host) + assert (sent["reasoning_effort"], "thinking" in sent) == ("low", False) + + +async def test_deployment_thinking_is_kept_without_a_request_reasoning_effort(make_core: MakeCore, host: EndpointHost): + thinking: Final = {"type": "enabled", "budget_tokens": 100} + core: Final = make_core(model_list=[deployment("g", "d1", thinking=thinking)]) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["thinking"] == thinking + + +async def test_request_reasoning_effort_strips_effort_from_deployment_extra_body( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[ + deployment( + "g", + "d1", + extra_body={"reasoning_effort": "high", "thinking": {"type": "enabled"}, "other": 1}, + reasoning={"effort": "high", "summary": "auto"}, + ) + ] + ) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + sent: Final = _sent(host) + assert sent["extra_body"] == {"other": 1} + assert sent["reasoning"] == {"summary": "auto"} + + +async def test_request_reasoning_effort_drops_an_extra_body_left_empty(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", extra_body={"reasoning_effort": "high"})]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + assert "extra_body" not in _sent(host) + + +async def test_reasoning_override_does_not_change_the_deployment_config(make_core: MakeCore, host: EndpointHost): + thinking: Final = {"type": "enabled", "budget_tokens": 100} + core: Final = make_core(model_list=[deployment("g", "d1", thinking=thinking)]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["thinking"] == thinking + + +CLASSIFIER_METADATA: Final = {"internal_call_origin": "autorouter_classifier"} + + +async def test_classifier_call_drops_reasoning_effort_the_deployment_does_not_support( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="openai/gpt-3.5-turbo")]) + body: Final[dict[str, object]] = {"reasoning_effort": "low", "keep": 1} + await core.acompletion( + model="g", + messages=MESSAGES, + reasoning_effort="low", + metadata=dict(CLASSIFIER_METADATA), + proxy_server_request={"body": body}, + ) + sent: Final = _sent(host) + assert "reasoning_effort" not in sent + assert sent["proxy_server_request"] == {"body": {"keep": 1}} + + +async def test_non_classifier_call_keeps_reasoning_effort_on_a_deployment_without_support( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="openai/gpt-3.5-turbo")]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + assert _sent(host)["reasoning_effort"] == "low" + + +async def test_classifier_call_keeps_reasoning_effort_on_a_deployment_with_support( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="openai/o3-mini")]) + await core.acompletion( + model="g", messages=MESSAGES, reasoning_effort="low", litellm_metadata=dict(CLASSIFIER_METADATA) + ) + assert _sent(host)["reasoning_effort"] == "low" + + +# --- silent experiments --------------------------------------------------------------------- + + +async def _settle() -> None: + for _ in range(20): + await asyncio.sleep(0) + + +async def test_silent_model_mirrors_acompletion_to_the_silent_group(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"trace": "t"}) + await _settle() + assert sorted(host.deployments_called("acompletion")) == ["d1", "s1"] + primary: Final = next(call for call in host.calls if call.deployment_id == "d1") + shadow: Final = next(call for call in host.calls if call.deployment_id == "s1") + assert "silent_model" not in primary.kwargs + assert "is_silent_experiment" not in _bucket(primary.kwargs) + shadow_metadata: Final = _bucket(shadow.kwargs) + assert (shadow_metadata["is_silent_experiment"], shadow_metadata["model_group"], shadow_metadata["trace"]) == ( + True, + "shadow", + "t", + ) + + +async def test_silent_model_list_mirrors_to_every_listed_group(make_core: MakeCore, host: EndpointHost): + core: Final = make_core( + model_list=[ + deployment("g", "d1", silent_model=["shadow-a", "shadow-b"]), + deployment("shadow-a", "sa"), + deployment("shadow-b", "sb"), + ] + ) + await core.acompletion(model="g", messages=MESSAGES) + await _settle() + assert sorted(host.deployments_called("acompletion")) == ["d1", "sa", "sb"] + + +async def test_a_failing_silent_experiment_does_not_fail_the_primary_call(make_core: MakeCore, host: EndpointHost): + host.answer("openai/s1", litellm.BadRequestError(message="shadow broke", llm_provider="openai", model="s1")) + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES) + await _settle() + assert response is not None + assert "s1" in host.deployments_called("acompletion") + + +async def test_a_malformed_silent_model_is_ignored(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", silent_model=5)]) + await core.acompletion(model="g", messages=MESSAGES) + await _settle() + assert host.deployments_called() == ["d1"] + assert "silent_model" not in _sent(host) + + +# --- timeouts ------------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("stream", "expected"), [(False, (99.0, None)), (True, (99.0, 33.0))]) +async def test_service_tier_timeouts_come_from_the_deployment_tier_maps( + make_core: MakeCore, host: EndpointHost, stream: bool, expected: tuple[float, float | None] +): + core: Final = make_core( + model_list=[ + deployment( + "g", "d1", timeout=11, timeout_by_service_tier={"flex": 99}, stream_timeout_by_service_tier={"flex": 33} + ) + ] + ) + await core.acompletion(model="g", messages=MESSAGES, service_tier="flex", stream=stream) + sent: Final = _sent(host) + assert (sent["timeout"], sent.get("stream_timeout")) == expected + + +async def test_router_stream_timeout_applies_to_streaming_acompletion(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], stream_timeout=5) + await core.acompletion(model="g", messages=MESSAGES, stream=True) + sent: Final = _sent(host) + assert (sent["timeout"], sent["stream_timeout"]) == (11.0, 5.0) + + +async def test_non_streaming_acompletion_sends_no_stream_timeout(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", stream_timeout=3)]) + await core.acompletion(model="g", messages=MESSAGES) + assert "stream_timeout" not in _sent(host) or _sent(host)["stream_timeout"] == 3 + + +async def test_aembedding_uses_the_request_timeout_over_the_deployment_timeout(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)]) + await core.aembedding(model="g", input=["a"], request_timeout=4) + assert _sent(host)["timeout"] == 4 + + +async def test_timeout_message_carries_deployment_timeouts_when_router_debug_is_exposed( + make_core: MakeCore, host: EndpointHost +): + host.set_dynamic(expose_router_debug_in_errors=True) + host.answer("openai/d1", timeout("d1")) + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)], num_retries=0) + with pytest.raises(litellm.Timeout) as raised: + await core.acompletion(model="g", messages=MESSAGES) + assert "Deployment Info: request_timeout: None\ntimeout: 11.0\nstream_timeout: 3.0" in str(raised.value) + + +async def test_timeout_message_has_no_deployment_info_by_default(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", timeout("d1")) + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], num_retries=0) + with pytest.raises(litellm.Timeout) as raised: + await core.acompletion(model="g", messages=MESSAGES) + assert "Deployment Info" not in str(raised.value) + + +# --- clientside credentials ------------------------------------------------------------------- + + +@pytest.mark.parametrize("credential", [{"api_key": "sk-client"}, {"api_base": "https://client.example"}]) +async def test_clientside_credentials_get_a_request_scoped_deployment_id( + make_core: MakeCore, host: EndpointHost, credential: dict[str, str] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, **credential) + sent: Final = _sent(host) + model_info: Final = _bucket(sent, "model_info") + assert model_info["original_model_id"] == "d1" + assert isinstance(model_info["id"], str) and len(model_info["id"]) == 64 + assert _bucket(_bucket(sent), "model_info")["id"] == model_info["id"] + + +async def test_clientside_credential_id_is_stable_for_the_same_credentials(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-a") + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-a") + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-b") + ids: Final = [call.deployment_id for call in host.calls] + assert ids[0] == ids[1] != ids[2] + + +async def test_clientside_credentials_do_not_register_a_deployment(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-client") + await core.acompletion(model="g", messages=MESSAGES) + assert host.deployments_called()[-1] == "d1" + assert [d["model_info"]["id"] for d in core.registry().deployments("g")] == ["d1"] # pyright: ignore[reportIndexIssue, reportUnknownMemberType] # registry rows are deployment dicts + + +# --- embedding batch size -------------------------------------------------------------------- + + +@pytest.mark.parametrize("endpoint", ["aembedding", "embedding"]) +async def test_embedding_over_the_deployment_batch_size_raises_without_calling_the_host( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core( + model_list=[deployment("g", "d1", model_info={"max_embedding_batch_size": 2})], num_retries=0 + ) + with pytest.raises( + litellm.BadRequestError, + match="Embedding request has 3 inputs, which exceeds max_embedding_batch_size=2 for model=g", + ): + await core.call(endpoint, {"model": "g", "input": ["a", "b", "c"]}) + assert host.calls == [] + + +@pytest.mark.parametrize("endpoint", ["aembedding", "embedding"]) +async def test_embedding_within_the_deployment_batch_size_is_sent( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1", model_info={"max_embedding_batch_size": 2})]) + await core.call(endpoint, {"model": "g", "input": ["a", "b"]}) + assert host.calls[0].kwargs["input"] == ["a", "b"] diff --git a/tests/unit/test_router/core/test_fallbacks_characterization.py b/tests/unit/test_router/core/test_fallbacks_characterization.py new file mode 100644 index 00000000000..333d9d73078 --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_characterization.py @@ -0,0 +1,584 @@ +""" +Characterization of a router core's fallbacks (non-streaming): chain sources and formats, +client-side fallback lists, per-fallback params, validation, mock-testing flags, depth bounds, +default fallbacks and wildcard lookup. Complements ``test_fallbacks.py``. +""" + +import pytest + +import litellm +from tests.unit.test_router.core.fallbacks_helpers import ( + TIERED_A, + MakeCore, + build, + complete, + complete_sync, + headers, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + content_policy, + context_window, + deployment, + rate_limit, +) + +# --- validation of the configured chains --------------------------------------------------------- + + +@pytest.mark.parametrize( + ("field", "value", "error", "message"), + [ + ("fallbacks", ["b"], ValueError, "Item 'b' is not a dictionary."), + ( + "fallbacks", + [{"a": ["b"], "c": ["b"]}], + ValueError, + "Dictionary '{'a': ['b'], 'c': ['b']}' must have exactly one key, but has 2 keys.", + ), + ( + "fallbacks", + [{"a": "b"}], + TypeError, + "Fallbacks for 'a' must be a list of model names, got 'b'. Use {'a': ['b']}", + ), + ( + "content_policy_fallbacks", + [{"a": "b"}], + TypeError, + "Fallbacks for 'a' must be a list of model names, got 'b'. Use {'a': ['b']}", + ), + ], + ids=["plain-string-item", "two-keys", "string-target", "content-policy-string-target"], +) +def test_a_malformed_fallbacks_configuration_is_refused_at_construction( + make_core: MakeCore, field: str, value: object, error: type[Exception], message: str +) -> None: + with pytest.raises(error) as raised: + build(make_core, **{field: value}) + + assert str(raised.value) == message + + +@pytest.mark.asyncio +async def test_context_window_fallbacks_are_not_validated_at_construction( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, context_window_fallbacks=[{"a": "b"}]) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_constructor_chain_of_parameter_dicts_passes_validation_and_applies_to_any_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"model": "b"}]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_an_empty_constructor_chain_defers_to_the_process_setting( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(fallbacks=[{"a": ["c"]}]) + core = build(make_core, fallbacks=[]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + + +# --- client-side fallback lists ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_plain_request_list_is_walked_in_order(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + response = await complete(core, fallbacks=["b", "c"]) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + assert headers(response)["x-litellm-model-group"] == "c" + + +@pytest.mark.asyncio +async def test_a_plain_request_list_overrides_the_context_window_chain(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, context_window_fallbacks=[{"a": ["c"]}]) + host.answer("openai/a", context_window()) + + await complete(core, fallbacks=["b"]) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_dict_target_with_params_overrides_the_request_for_that_attempt( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + override = [{"role": "user", "content": "shorter"}] + + await complete(core, fallbacks=[{"model": "b", "temperature": 0.3, "messages": override}]) + + primary, fallback = host.calls_to("acompletion") + assert (primary.deployment_id, primary.kwargs.get("temperature")) == ("a", None) + assert primary.kwargs["messages"] == [{"role": "user", "content": "hi"}] + assert (fallback.deployment_id, fallback.kwargs["temperature"], fallback.kwargs["messages"]) == ("b", 0.3, override) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "arguments", + [ + {"request": [{"model": "b", "temperature": 0.3}, {"model": "c"}]}, + {"constructor": [{"a": [{"model": "b", "temperature": 0.3}, "c"]}]}, + ], + ids=["request-list", "constructor-chain"], +) +async def test_a_dict_targets_params_carry_over_to_the_next_target( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, list[object]] +) -> None: + core = build(make_core, **({"fallbacks": arguments["constructor"]} if "constructor" in arguments else {})) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + await complete(core, **({"fallbacks": arguments["request"]} if "request" in arguments else {})) + + assert [(c.deployment_id, c.kwargs.get("temperature")) for c in host.calls_to("acompletion")] == [ + ("a", None), + ("b", 0.3), + ("c", 0.3), + ] + + +@pytest.mark.asyncio +async def test_a_dict_target_naming_the_failed_group_with_params_retries_that_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit(), None) + + response = await complete(core, fallbacks=[{"model": "a", "temperature": 0}]) + + assert [(c.deployment_id, c.kwargs.get("temperature")) for c in host.calls_to("acompletion")] == [ + ("a", None), + ("a", 0), + ] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +async def test_a_dict_target_with_its_own_api_key_runs_on_a_client_credential_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + + response = await complete(core, fallbacks=[{"model": "b", "api_key": "sk-other"}]) + + primary, fallback = host.calls_to("acompletion") + assert (fallback.model, fallback.kwargs["api_key"]) == ("openai/b", "sk-other") + assert fallback.deployment_id not in {None, "b"} + assert headers(response)["x-litellm-model-group"] == "b" + assert primary.deployment_id == "a" + + +@pytest.mark.asyncio +async def test_a_dict_target_carrying_an_oidc_api_key_is_refused_and_the_primary_error_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, fallbacks=[{"model": "b", "api_key": "oidc/google/x"}]) + + assert host.deployments_called("acompletion") == ["a"] + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to {'model': 'b', 'api_key': 'oidc*********'} also failed: Rejected Request: api_key holds an " + "oidc/ secret reference, which only a deployment or credential a proxy admin configures can carry." + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fallbacks", "expected"), + [(["c", {"a": ["b"]}], ["a", "b", "c"]), (["b", {"x": ["c"]}], ["a", "b"])], + ids=["keyed-entry-wins-then-string-catches-the-next-hop", "string-entry-is-a-chain-for-every-group"], +) +async def test_a_mixed_list_of_names_and_keyed_entries_is_resolved_as_a_chain_map( + make_core: MakeCore, host: ScriptedHost, fallbacks: list[object], expected: list[str] +) -> None: + core = build(make_core) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, fallbacks=fallbacks) + + assert host.deployments_called("acompletion") == expected + + +@pytest.mark.asyncio +async def test_the_request_fallbacks_kwarg_is_not_forwarded_to_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + + await complete(core, fallbacks=["b"]) + + assert [("fallbacks" in c.kwargs) for c in host.calls_to("acompletion")] == [False, False] + + +# --- disable_fallbacks ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_disable_fallbacks_given_as_a_string_does_not_disable_them( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + await complete(core, disable_fallbacks="true") + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_disable_fallbacks_also_stops_the_order_tier_walk(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, model_list=TIERED_A) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, disable_fallbacks=True) + + assert host.deployments_called("acompletion") == ["a1"] + + +@pytest.mark.asyncio +async def test_disable_fallbacks_is_not_forwarded_to_the_provider(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core) + + await complete(core, disable_fallbacks=True) + + assert "disable_fallbacks" not in host.calls_to("acompletion")[0].kwargs + + +# --- depth bounds ----------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_request_max_fallbacks_overrides_the_routers(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, max_fallbacks=1) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_max_fallbacks_zero_also_stops_the_order_tier_walk(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, model_list=TIERED_A, max_fallbacks=0) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a1"] + + +@pytest.mark.asyncio +async def test_max_fallbacks_one_still_tries_every_target_of_the_first_level( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b", "c"]}, {"b": ["c"]}], max_fallbacks=1) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_a_fallback_group_spends_its_own_retries(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}], num_retries=1) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "a", "b", "b"] + assert str(raised.value) == "litellm.RateLimitError: slow down LiteLLM Retried: 1 times, LiteLLM Max Retries: 1" + + +@pytest.mark.asyncio +async def test_a_fallback_hop_resumes_the_original_groups_chain(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b", "c"]}], max_fallbacks=5) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + response = await complete(core, include_fallback_errors=True) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + # c answered inside b's hop (b resumed a's chain), so b's error is listed before a's. + assert headers(response)["x-litellm-fallback-errors"] == ( + '[{"message": "litellm.RateLimitError: slow down", "type": "throttling_error", "param": null, "code": "429"}, ' + '{"message": "litellm.RateLimitError: slow down", "type": "throttling_error", "param": null, "code": "429"}]' + ) + + +# --- order tiers with other chains --------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "arguments", + [{"request": {"fallbacks": ["c"]}}, {"constructor": {"default_fallbacks": ["c"]}}], + ids=["plain-request-list", "default-fallbacks"], +) +async def test_order_tiers_run_before_a_plain_list_or_the_default_chain( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, dict[str, object]] +) -> None: + core = build(make_core, model_list=TIERED_A, **arguments.get("constructor", {})) + host.answer("*", rate_limit()) + host.answer("openai/c", None) + + response = await complete(core, **arguments.get("request", {})) + + assert host.deployments_called("acompletion") == ["a1", "a2", "c"] + assert headers(response)["x-litellm-model-group"] == "c" + + +# --- default fallbacks --------------------------------------------------------------------------- + + +def test_default_fallbacks_are_appended_to_the_process_fallbacks_list_in_place( + make_core: MakeCore, host: ScriptedHost +) -> None: + process_fallbacks: list[object] = [{"a": ["b"]}] + host.set_dynamic(fallbacks=process_fallbacks) + + core = build(make_core, default_fallbacks=["c"]) + + assert core.settings().fallbacks == [{"a": ["b"]}, {"*": ["c"]}] + assert process_fallbacks == [{"a": ["b"]}, {"*": ["c"]}] + + +def test_default_fallbacks_leave_the_constructor_fallbacks_list_unchanged(make_core: MakeCore) -> None: + fallbacks: list[object] = [{"a": ["b"]}] + + core = build(make_core, fallbacks=fallbacks, default_fallbacks=["c"]) + + assert core.settings().fallbacks == [{"a": ["b"]}, {"*": ["c"]}] + assert fallbacks == [{"a": ["b"]}] + + +@pytest.mark.asyncio +async def test_an_unknown_group_is_served_by_the_first_default_fallback_without_a_fallback_hop( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, default_fallbacks=["b", "c"]) + + response = await complete(core, model="unknown") + + assert host.deployments_called("acompletion") == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 0 + assert headers(response)["x-litellm-model-group"] == "unknown" + assert host.calls_to("acompletion")[0].model_group == "unknown" + + +@pytest.mark.asyncio +async def test_an_unknown_group_with_its_own_chain_falls_back_through_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"unknown": ["b"]}]) + + response = await complete(core, model="unknown") + + assert host.deployments_called("acompletion") == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + assert headers(response)["x-litellm-model-group"] == "b" + + +@pytest.mark.asyncio +async def test_default_fallbacks_catch_a_content_policy_error_without_a_content_policy_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, default_fallbacks=["c"]) + host.answer("openai/a", content_policy()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + + +@pytest.mark.asyncio +async def test_a_request_context_window_fallbacks_of_none_falls_through_to_the_general_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, context_window_fallbacks=[{"a": ["c"]}], fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", context_window()) + + await complete(core, context_window_fallbacks=None) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +# --- wildcard groups ----------------------------------------------------------------------------- + +WILDCARD = [deployment("openai/*", "w", model="openai/*"), deployment("b", "b")] + + +@pytest.mark.asyncio +async def test_a_wildcard_routed_request_uses_the_chain_keyed_by_the_bare_model_name( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=WILDCARD, fallbacks=[{"gpt-x": ["b"]}]) + host.answer("openai/gpt-x", rate_limit()) + + response = await complete(core, model="openai/gpt-x") + + assert host.deployments_called("acompletion") == ["w", "b"] + assert [c.model for c in host.calls_to("acompletion")] == ["openai/gpt-x", "openai/b"] + assert headers(response)["x-litellm-model-group"] == "b" + + +@pytest.mark.asyncio +async def test_a_chain_keyed_by_the_wildcard_pattern_does_not_match_a_routed_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=WILDCARD, fallbacks=[{"openai/*": ["b"]}]) + host.answer("openai/gpt-x", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, model="openai/gpt-x") + + assert host.deployments_called("acompletion") == ["w"] + + +# --- mock testing flags ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("flag", "field", "expected"), + [ + ("mock_testing_fallbacks", "fallbacks", "b"), + ("mock_testing_context_fallbacks", "context_window_fallbacks", "c"), + ("mock_testing_content_policy_fallbacks", "content_policy_fallbacks", "c"), + ], + ids=["general", "context-window", "content-policy"], +) +async def test_a_mock_testing_flag_skips_the_primary_and_walks_the_matching_chain( + make_core: MakeCore, host: ScriptedHost, flag: str, field: str, expected: str +) -> None: + core = build(make_core, num_retries=2, **{"fallbacks": [{"a": ["b"]}], field: [{"a": [expected]}]}) + + response = await complete(core, **{flag: True}) + + assert host.deployments_called("acompletion") == [expected] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + assert headers(response)["x-litellm-attempted-retries"] == 0 + assert host.sleeps == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("flag", "error", "message"), + [ + ( + "mock_testing_fallbacks", + litellm.InternalServerError, + "litellm.InternalServerError: This is a mock exception for model=a, to trigger a fallback. Fallbacks=None", + ), + ( + "mock_testing_context_fallbacks", + litellm.ContextWindowExceededError, + "litellm.ContextWindowExceededError: litellm.BadRequestError: This is a mock exception for model=a, " + "to trigger a fallback. Context_Window_Fallbacks=None", + ), + ( + "mock_testing_content_policy_fallbacks", + litellm.ContentPolicyViolationError, + "litellm.BadRequestError: litellm.ContentPolicyViolationError: This is a mock exception for model=a, " + "to trigger a fallback. Context_Policy_Fallbacks=None", + ), + ], + ids=["general", "context-window", "content-policy"], +) +async def test_a_mock_testing_flag_without_a_chain_raises_the_mock_error_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost, flag: str, error: type[Exception], message: str +) -> None: + core = build(make_core) + + with pytest.raises(error) as raised: + await complete(core, **{flag: True}) + + assert str(raised.value) == message + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_the_mock_error_names_the_configured_chain_and_the_fallback_outcome( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.InternalServerError) as raised: + await complete(core, mock_testing_fallbacks=True) + + assert str(raised.value) == ( + "litellm.InternalServerError: This is a mock exception for model=a, to trigger a fallback. " + "Fallbacks=[{'a': ['b']}]\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to b also failed: litellm.RateLimitError: slow down" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "expected"), [("true", ["b"]), ("false", ["a"]), (False, ["a"])]) +async def test_a_mock_testing_flag_given_as_a_string_is_parsed_as_a_bool( + make_core: MakeCore, host: ScriptedHost, value: object, expected: list[str] +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + + await complete(core, mock_testing_fallbacks=value) + + assert host.deployments_called("acompletion") == expected + + +@pytest.mark.asyncio +async def test_the_mock_testing_flag_is_not_forwarded_to_the_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + + await complete(core, mock_testing_fallbacks=True) + + assert "mock_testing_fallbacks" not in host.calls_to("acompletion")[0].kwargs + + +def test_sync_completion_honours_the_mock_testing_flag(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + + response = complete_sync(core, mock_testing_fallbacks=True) + + assert host.deployments_called("completion") == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 diff --git a/tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py b/tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py new file mode 100644 index 00000000000..631b9dcf3c2 --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py @@ -0,0 +1,675 @@ +""" +Characterization of what a router core's fallbacks leave behind (non-streaming): response headers +and hidden params, the metadata each attempt carries, the error text when fallbacks give up, +fallback access and budget checks, fallback events, content-filter responses and weighted failover. +""" + +from collections.abc import Mapping + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.unit.test_router.core.fallbacks_helpers import ( + GROUP_A_OF_THREE, + MakeCore, + build, + complete, + complete_sync, + fails_once_then_answers, + headers, + hidden, + metadata_of, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + completion_response, + context_window, + deployment, + rate_limit, + server_error, + timeout, +) + +ERROR_429 = '{"message": "litellm.RateLimitError: slow down", "type": "throttling_error", "param": null, "code": "429"}' +ERROR_500 = ( + '{"message": "litellm.InternalServerError: boom", "type": "internal_server_error", "param": null, "code": "500"}' +) + +# --- headers and hidden params ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_two_hop_chain_reports_one_attempted_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + assert headers(response)["x-litellm-model-group"] == "c" + + +@pytest.mark.asyncio +async def test_the_answering_fallback_deployment_is_the_responses_model_id( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert hidden(response)["model_id"] == "b" + + +@pytest.mark.asyncio +async def test_success_headers_after_a_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert headers(response) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +@pytest.mark.asyncio +async def test_include_fallback_errors_lists_the_primary_error_in_a_header( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core, include_fallback_errors=True) + + assert headers(response)["x-litellm-fallback-errors"] == f"[{ERROR_429}]" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallbacks", + [[{"a": ["b", "c"]}], [{"a": ["b"]}, {"b": ["c"]}]], + ids=["one-list", "chained"], +) +async def test_include_fallback_errors_lists_the_inner_hops_errors_first( + make_core: MakeCore, host: ScriptedHost, fallbacks: list[object] +) -> None: + core = build(make_core, fallbacks=fallbacks) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", server_error()) + + response = await complete(core, include_fallback_errors=True) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + assert headers(response)["x-litellm-fallback-errors"] == f"[{ERROR_500}, {ERROR_429}]" + + +@pytest.mark.asyncio +async def test_without_include_fallback_errors_no_errors_header_is_set(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert "x-litellm-fallback-errors" not in headers(response) + + +# --- what each attempt carries to the host ------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_each_attempt_carries_the_fallback_count_and_original_group_in_its_metadata( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + await complete(core) + + assert [ + (meta["model_group"], meta["original_model_group"], meta["attempted_fallbacks"]) + for meta in (metadata_of(call.kwargs) for call in host.calls_to("acompletion")) + ] == [("a", "a", 0), ("b", "a", 1), ("c", "a", 2)] + + +@pytest.mark.asyncio +async def test_a_fallback_attempt_also_receives_the_routers_walk_state_as_kwargs( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + primary, fallback = host.calls_to("acompletion") + assert {"fallback_depth", "max_fallbacks", "attempted_targets"} & set(primary.kwargs) == set() + assert (fallback.kwargs["fallback_depth"], fallback.kwargs["max_fallbacks"]) == (1, 5) + assert "attempted_targets" in fallback.kwargs + + +@pytest.mark.asyncio +async def test_the_fallback_attempt_sees_the_failed_attempt_in_previous_models( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + previous = metadata_of(host.calls_to("acompletion")[1].kwargs)["previous_models"] + assert [ + (entry["model_group"], entry["deployment_id"], entry["exception_type"]) + for entry in previous # pyright: ignore[reportGeneralTypeIssues] # a tuple of dicts + ] == [("a", "a", "RateLimitError")] + + +# --- the error when fallbacks give up ------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_router_debug_on_a_context_window_error_without_its_chain_names_the_missing_setting( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("openai/a", context_window()) + + with pytest.raises(litellm.ContextWindowExceededError) as raised: + await complete(core) + + assert str(raised.value) == ( + "litellm.ContextWindowExceededError: litellm.BadRequestError: too long\n" + "model=a. context_window_fallbacks=None. fallbacks=None.\n\n" + "Set 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks\n\n" + "LiteLLM: model group 'a' failed with the error above. No fallback was attempted." + ) + + +@pytest.mark.asyncio +async def test_router_debug_on_a_context_window_error_that_also_fails_its_fallback_nests_the_hint( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("*", context_window()) + + with pytest.raises(litellm.ContextWindowExceededError) as raised: + await complete(core) + + hint = "context_window_fallbacks=None. fallbacks=[{'a': ['b']}].\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + assert str(raised.value) == ( + f"litellm.ContextWindowExceededError: litellm.BadRequestError: too long\nmodel=a. {hint}\n\n" + "LiteLLM: model group 'a' failed with the error above. Fallback to b also failed: " + f"litellm.ContextWindowExceededError: litellm.BadRequestError: too long\nmodel=b. {hint}" + ) + + +@pytest.mark.asyncio +async def test_router_debug_names_a_plain_request_list_that_failed(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, fallbacks=["b", "c"]) + + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to b, c also failed: litellm.RateLimitError: slow down" + ) + + +@pytest.mark.asyncio +async def test_router_debug_with_max_fallbacks_one_names_only_the_first_level( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}], max_fallbacks=1) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert str(raised.value).endswith("Fallback to b also failed: litellm.RateLimitError: slow down") + + +@pytest.mark.asyncio +async def test_router_debug_reports_the_first_fallback_errors_message_not_the_last( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", server_error()) + host.answer("openai/c", timeout()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to b also failed: litellm.InternalServerError: boom" + ) + + +@pytest.mark.asyncio +async def test_router_debug_truncates_a_long_fallback_error(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit(message="x" * 3000)) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + detail = "litellm.RateLimitError: " + "x" * 3000 + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + f"Fallback to b also failed: {detail[:2000]}... [truncated {len(detail) - 2000} characters]" + ) + + +@pytest.mark.asyncio +async def test_router_debug_on_an_unmatched_content_policy_chain_says_no_fallback_was_attempted( + make_core: MakeCore, host: ScriptedHost +) -> None: + from tests.unit.test_router.core.scripted_host import content_policy + + core = build(make_core, content_policy_fallbacks=[{"other": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", content_policy()) + + with pytest.raises(litellm.ContentPolicyViolationError) as raised: + await complete(core) + + assert str(raised.value) == ( + "litellm.BadRequestError: litellm.ContentPolicyViolationError: blocked\n\n" + "LiteLLM: model group 'a' failed with the error above. No fallback was attempted." + ) + + +@pytest.mark.asyncio +async def test_router_debug_no_chain_message_without_any_configured_keys( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, fallbacks=[{"x": ["b"]}, {"y": ["c"]}]) + + assert str(raised.value).endswith( + "LiteLLM: model group 'a' failed with the error above and no fallback model group was found for it, " + "so the request was not retried on another model. Fallbacks are configured for: x, y. " + "Add a fallbacks entry for that model group (Router fallbacks or proxy router_settings.fallbacks) " + "to retry on another model." + ) + + +# --- fallback access and budget checks ------------------------------------------------------------- + + +def denying(denied: str, seen: list[str]): # noqa: ANN201 # a FallbackAccessCheck / FallbackBudgetCheck + async def check(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + seen.append(model) + return model != denied + + return check + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["fallback_access_check", "fallback_budget_check"]) +async def test_a_fallback_target_the_check_rejects_is_skipped( + make_core: MakeCore, host: ScriptedHost, field: str +) -> None: + seen: list[str] = [] + core = build(make_core, fallbacks=[{"a": ["b", "c"]}], **{field: denying("b", seen)}) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + assert headers(response)["x-litellm-model-group"] == "c" + assert seen == ["b", "c"] + + +def test_sync_completion_runs_the_fallback_access_check(make_core: MakeCore, host: ScriptedHost) -> None: + seen: list[str] = [] + core = build(make_core, fallbacks=[{"a": ["b", "c"]}], fallback_access_check=denying("b", seen)) + host.answer("openai/a", rate_limit()) + + complete_sync(core) + + assert host.deployments_called("completion") == ["a", "c"] + assert seen == ["b", "c"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("field", "sentence"), + [ + ("fallback_access_check", "Fallback to b was skipped because the caller is not authorized to call it."), + ("fallback_budget_check", "Fallback to b was skipped because the caller is over budget for it."), + ], +) +async def test_router_debug_explains_a_skipped_fallback_target( + make_core: MakeCore, host: ScriptedHost, field: str, sentence: str +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}], **{field: denying("b", [])}) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a"] + assert str(raised.value) == ( + f"litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. {sentence}" + ) + + +@pytest.mark.asyncio +async def test_router_debug_names_both_the_failed_and_the_skipped_targets( + make_core: MakeCore, host: ScriptedHost +) -> None: + seen: list[str] = [] + core = build( + make_core, + fallbacks=[{"a": ["b", "c"]}], + fallback_access_check=denying("b", seen), + fallback_budget_check=denying("b", seen), + ) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to c also failed: litellm.RateLimitError: slow down " + "Fallback to b was skipped because the caller is not authorized to call it." + ) + # The access check runs before the budget check; a skipped target is asked again in c's hop. + assert seen == ["b", "c", "c", "b"] + + +# --- fallback events --------------------------------------------------------------------------------- + + +class FallbackEvents(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: list[tuple[str, str, object, str]] = [] + + async def log_success_fallback_event( + self, original_model_group: str, kwargs: dict, original_exception: Exception + ) -> None: + self.events.append(("success", original_model_group, kwargs.get("model"), type(original_exception).__name__)) + + async def log_failure_fallback_event( + self, original_model_group: str, kwargs: dict, original_exception: Exception + ) -> None: + self.events.append(("failure", original_model_group, kwargs.get("model"), type(original_exception).__name__)) + + +@pytest.mark.asyncio +async def test_fallback_events_reach_the_process_custom_loggers_innermost_hop_first( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + events = FallbackEvents() + monkeypatch.setattr(litellm, "callbacks", [events]) + core = build(make_core, fallbacks=[{"a": ["b", "c"]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", server_error()) + + await complete(core) + + assert events.events == [ + ("success", "b", "c", "InternalServerError"), + ("success", "a", "b", "RateLimitError"), + ] + + +@pytest.mark.asyncio +async def test_a_fallback_that_fails_everywhere_logs_failure_events( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + events = FallbackEvents() + monkeypatch.setattr(litellm, "callbacks", [events]) + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert events.events == [("failure", "a", "b", "RateLimitError")] + + +# --- content-filter responses ------------------------------------------------------------------------ + + +def content_filtered() -> object: + response = completion_response() + response.choices[0].finish_reason = "content_filter" # pyright: ignore[reportAttributeAccessIssue] # Choices + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ({"content_policy_fallbacks": [{"a": ["b"]}]}, ["a", "b"]), + ({"default_fallbacks": ["c"]}, ["a", "c"]), + ({"fallbacks": [{"a": ["b"]}]}, ["a"]), + ({"content_policy_fallbacks": [{"other": ["b"]}], "default_fallbacks": ["c"]}, ["a"]), + ], + ids=["content-policy-chain", "default-fallbacks", "general-chain-only", "unmatched-content-policy-chain"], +) +async def test_a_content_filtered_response_falls_back_only_when_a_content_policy_fallback_resolves( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], expected: list[str] +) -> None: + core = build(make_core, **arguments) + host.answer("openai/a", content_filtered()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == expected + finish_reason = response.choices[0].finish_reason # pyright: ignore[reportAttributeAccessIssue] # ModelResponse + assert finish_reason == ("content_filter" if expected == ["a"] else "stop") + + +def test_sync_completion_falls_back_on_a_content_filtered_response(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", content_filtered()) + + complete_sync(core) + + assert host.deployments_called("completion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_content_filtered_response_with_fallbacks_disabled_raises_content_policy_violation( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", content_filtered()) + + with pytest.raises(litellm.ContentPolicyViolationError) as raised: + await complete(core, disable_fallbacks=True) + + assert ( + str(raised.value) + == "litellm.BadRequestError: litellm.ContentPolicyViolationError: Response output was blocked." + ) + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_the_last_content_filtered_response_is_returned_when_every_fallback_is_filtered( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + host.answer("*", content_filtered()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert response.choices[0].finish_reason == "content_filter" # pyright: ignore[reportAttributeAccessIssue] # ModelResponse + assert headers(response)["x-litellm-model-group"] == "b" + + +# --- weighted failover --------------------------------------------------------------------------- + + +def a_deployments_fail(host: ScriptedHost, error: object) -> None: + for deployment_id in ("a1", "a2", "a3"): + host.answer(f"openai/{deployment_id}", error) + + +@pytest.mark.asyncio +async def test_weighted_failover_tries_every_deployment_of_the_group_before_the_external_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}]) + a_deployments_fail(host, rate_limit()) + + response = await complete(core) + + called = host.deployments_called("acompletion") + assert sorted(called[:3]) == ["a1", "a2", "a3"] + assert called[3:] == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +async def test_without_weighted_failover_one_deployment_fails_over_to_the_external_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, fallbacks=[{"a": ["b"]}]) + a_deployments_fail(host, rate_limit()) + + await complete(core) + + called = host.deployments_called("acompletion") + assert len(called) == 2 + assert called[0] in {"a1", "a2", "a3"} + assert called[1] == "b" + + +@pytest.mark.asyncio +async def test_a_weighted_failover_that_answers_counts_as_a_fallback_inside_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True) + host.answer("*", fails_once_then_answers(rate_limit())) + + response = await complete(core) + + calls = host.calls_to("acompletion") + assert len(calls) == 2 + assert calls[0].deployment_id != calls[1].deployment_id + assert [call.model_group for call in calls] == ["a", "a"] + assert [metadata_of(call.kwargs)["attempted_fallbacks"] for call in calls] == [0, 1] + assert headers(response)["x-litellm-model-group"] == "a" + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("max_fallbacks", "group_attempts"), [(0, 1), (1, 2), (2, 3)]) +async def test_weighted_failover_is_bounded_by_max_fallbacks( + make_core: MakeCore, host: ScriptedHost, max_fallbacks: int, group_attempts: int +) -> None: + core = build( + make_core, + model_list=GROUP_A_OF_THREE, + enable_weighted_failover=True, + fallbacks=[{"a": ["b"]}], + max_fallbacks=max_fallbacks, + ) + a_deployments_fail(host, rate_limit()) + + if max_fallbacks == 0: + with pytest.raises(litellm.RateLimitError): + await complete(core) + else: + await complete(core) + + called = host.deployments_called("acompletion") + assert len(set(called[:group_attempts])) == group_attempts + assert called[group_attempts:] == ([] if max_fallbacks == 0 else ["b"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "error"), + [ + ({"routing_strategy": "latency-based-routing"}, rate_limit), + ({}, context_window), + ], + ids=["not-simple-shuffle", "context-window-error"], +) +async def test_weighted_failover_does_not_apply( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], error: object +) -> None: + core = build( + make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}], **arguments + ) + a_deployments_fail(host, error()) # pyright: ignore[reportCallIssue] # an error factory + + await complete(core) + + called = host.deployments_called("acompletion") + assert len(called) == 2 + assert called[1] == "b" + + +@pytest.mark.asyncio +async def test_weighted_failover_with_fallbacks_disabled_tries_one_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}]) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, disable_fallbacks=True) + + assert len(host.deployments_called("acompletion")) == 1 + + +def test_sync_completion_also_runs_weighted_failover(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}]) + a_deployments_fail(host, rate_limit()) + + complete_sync(core) + + called = host.deployments_called("completion") + assert sorted(called[:3]) == ["a1", "a2", "a3"] + assert called[3:] == ["b"] + + +@pytest.mark.asyncio +async def test_weighted_failover_on_a_single_deployment_group_goes_straight_to_the_external_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build( + make_core, + model_list=[deployment("a", "a1"), deployment("b", "b")], + enable_weighted_failover=True, + fallbacks=[{"a": ["b"]}], + ) + host.answer("openai/a1", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a1", "b"] diff --git a/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py b/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py new file mode 100644 index 00000000000..d7c893a36f1 --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py @@ -0,0 +1,101 @@ +""" +Characterization of a router core's Anthropic safeguard-refusal fallbacks on /v1/messages +(non-streaming): a ``stop_reason: "refusal"`` response carrying ``stop_details`` re-enters the +fallback chain only when a fallback would resolve for it. + +The core recognises a /v1/messages call by its endpoint's name, so this module runs on +``NamedEndpointHost`` (a ``ScriptedHost`` whose ``anthropic_messages`` endpoint is named so). +""" + +import pytest + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.fallbacks_helpers import ( + MESSAGES, + MakeCore, + NamedEndpointHost, + anthropic_message, + anthropic_refusal, + build, +) + + +@pytest.fixture +def host() -> NamedEndpointHost: + return NamedEndpointHost() + + +async def messages(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=10, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ({"content_policy_fallbacks": [{"a": ["b"]}]}, ["a", "b"]), + ({"fallbacks": [{"a": ["b"]}]}, ["a", "b"]), + ({"default_fallbacks": ["c"]}, ["a", "c"]), + ], + ids=["content-policy-chain", "general-chain", "default-fallbacks"], +) +async def test_a_safeguard_refusal_falls_back_when_a_chain_resolves( + make_core: MakeCore, host: NamedEndpointHost, arguments: dict[str, object], expected: list[str] +) -> None: + core = build(make_core, **arguments) + host.answer("openai/a", anthropic_refusal()) + host.answer(f"openai/{expected[-1]}", anthropic_message()) + + response = await messages(core) + + assert host.deployments_called("anthropic_messages") == expected + assert isinstance(response, dict) + assert response["stop_reason"] == "end_turn" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "request_kwargs", "refusal"), + [ + ({}, {}, anthropic_refusal()), + ({"content_policy_fallbacks": [{"other": ["b"]}], "fallbacks": [{"a": ["b"]}]}, {}, anthropic_refusal()), + ({"fallbacks": [{"a": ["b"]}]}, {"disable_fallbacks": True}, anthropic_refusal()), + ({"fallbacks": [{"a": ["b"]}]}, {}, anthropic_refusal(with_stop_details=False)), + ], + ids=[ + "no-chain", + "unmatched-content-policy-chain-is-authoritative", + "fallbacks-disabled", + "plain-refusal-without-stop-details", + ], +) +async def test_a_refusal_is_returned_to_the_caller_when_no_fallback_would_resolve( + make_core: MakeCore, + host: NamedEndpointHost, + arguments: dict[str, object], + request_kwargs: dict[str, object], + refusal: dict[str, object], +) -> None: + core = build(make_core, **arguments) + host.answer("openai/a", refusal) + + response = await messages(core, **request_kwargs) + + assert host.deployments_called("anthropic_messages") == ["a"] + assert isinstance(response, dict) + assert response["stop_reason"] == "refusal" + assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 0 + + +@pytest.mark.asyncio +async def test_the_last_refusal_is_returned_when_every_fallback_refuses( + make_core: MakeCore, host: NamedEndpointHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("*", anthropic_refusal()) + + response = await messages(core) + + assert host.deployments_called("anthropic_messages") == ["a", "b"] + assert isinstance(response, dict) + assert response["stop_reason"] == "refusal" diff --git a/tests/unit/test_router/core/test_gaps_adapter_timeout.py b/tests/unit/test_router/core/test_gaps_adapter_timeout.py new file mode 100644 index 00000000000..4b895673481 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_adapter_timeout.py @@ -0,0 +1,34 @@ +""" +Characterization gap: a streaming adapter completion splits its timeouts like a streaming chat +completion, sending the request timeout as ``timeout`` and the stream timeout beside it, so a +fake-streaming provider stays bounded by the request timeout rather than the stream timeout. + +Every test builds its core through ``make_core`` on the ``host`` fixture, so it runs once per core. +""" + +from collections.abc import Callable +from typing import Final + +import pytest + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import deployment + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +@pytest.mark.parametrize(("stream", "expected"), [(False, (11.0, None)), (True, (11.0, 5.0))]) +async def test_adapter_completion_sends_the_request_timeout_and_the_stream_timeout_beside_it( + make_core: MakeCore, host: EndpointHost, stream: bool, expected: tuple[float, float | None] +): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], stream_timeout=5) + await core.aadapter_completion(adapter_id="my-adapter", model="g", messages=MESSAGES, stream=stream) + sent: Final = host.calls[-1].kwargs + assert (sent["timeout"], sent.get("stream_timeout")) == expected diff --git a/tests/unit/test_router/core/test_gaps_l0001.py b/tests/unit/test_router/core/test_gaps_l0001.py new file mode 100644 index 00000000000..bb6a4fe6c44 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l0001.py @@ -0,0 +1,445 @@ +""" +Gap tests for the first part of ``litellm/router_python_core.py`` (module helpers, the entry points, +``__init__``, callback registration, strategy setup, routing groups and the strategy dispatch): +behaviors no other characterization test pinned down. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import datetime +import re +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_api import ChangeSet +from litellm.types.router import Deployment, RouterRateLimitError +from litellm.utils import reapply_runtime_model_cost_registrations +from tests.unit.test_router.core.construction_precall_helpers import collect_other_routers, host_callback_names +from tests.unit.test_router.core.registry_helpers import unique_id +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, picked_ids + +NO_DEPLOYMENTS_MESSAGE: Final = ( + "No deployments available for selected model, Try again in 5 seconds. Passed model={model}. " + "pre-call-checks=False, cooldown_list=[]" +) + + +def _group( + name: str, models: list[str], strategy: str, **extra: object +) -> dict[str, object]: # mutable-ok: a routing_groups entry # kwargs-ok: extra RoutingGroup fields + return {"group_name": name, "models": models, "routing_strategy": strategy, **extra} + + +def _logging_obj() -> LiteLLMLogging: + return LiteLLMLogging( + model="g", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c1", + function_id="f1", + ) + + +def _dynamic_callback_names(logging_obj: LiteLLMLogging) -> list[str]: # mutable-ok: a fresh list + return [type(callback).__name__ for callback in logging_obj.dynamic_async_success_callbacks or ()] + + +# --- construction: routing_strategy=None and legacy attributes ------------------------------------- + + +@pytest.mark.asyncio +async def test_routing_strategy_none_is_accepted_but_leaves_every_pick_without_a_deployment( + make_core: MakeCore, +) -> None: + """Oddity: ``routing_strategy=None`` passes validation and is stored as None (not the + ``simple-shuffle`` default), so no strategy arm matches and every pick rate-limits, while the + healthy listing still lists the deployment.""" + core: Final = make_core(model_list=[deployment("g", "a")], routing_strategy=None) + assert core.settings().routing_strategy is None + assert core.settings().as_dict()["routing_strategy"] is None + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g") + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g") + healthy: Final = await core.select("g", kind="healthy") + assert isinstance(healthy, list) + assert [entry["model_info"]["id"] for entry in healthy] == ["a"] + + +LEGACY_ATTRIBUTE_DEFAULTS: Final = { + "total_calls": {}, + "fail_calls": {}, + "success_calls": {}, + "deployment_stats": {}, + "default_cache_time_seconds": 3600, + "tenacity": None, + "lowesttpm_logger_v2": None, + "lowestlatency_logger": None, + "lowestcost_logger": None, +} + + +@pytest.mark.parametrize(("name", "expected"), sorted(LEGACY_ATTRIBUTE_DEFAULTS.items())) +def test_legacy_attribute_default(make_core: MakeCore, name: str, expected: object) -> None: + assert getattr(make_core(model_list=[deployment("g", "a")]), name) == expected + + +def test_healthy_deployments_is_the_live_deployment_list(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.healthy_deployments is core.registry().all() # pyright: ignore[reportAttributeAccessIssue] + + +def test_chat_completions_create_calls_the_router_with_the_default_litellm_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], default_litellm_params={"temperature": 0.3}) + core.chat.completions.create(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert [(call.op, call.deployment_id, call.kwargs.get("temperature")) for call in host.calls] == [ + ("completion", "a", 0.3) + ] + + +@pytest.mark.asyncio +async def test_a_provider_budget_config_adds_the_budget_limiter_after_the_listed_checks(make_core: MakeCore) -> None: + """A provider budget appends ``router_budget_limiting`` to the optional pre-call checks given.""" + core: Final = make_core( + model_list=[deployment("g", "a")], + provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}, + optional_pre_call_checks=["prompt_caching"], + ) + assert [type(callback).__name__ for callback in core.optional_callbacks] == [ # pyright: ignore[reportAttributeAccessIssue] + "PromptCachingDeploymentCheck", + "RouterBudgetLimiting", + ] + core.retire() + + +# --- price data reloads --------------------------------------------------------------------------- + + +def test_a_price_data_reload_restores_a_live_routers_deployment_pricing_until_it_retires( + make_core: MakeCore, +) -> None: + model_id: Final = unique_id("l0001") + core: Final = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + try: + litellm.model_cost.pop(model_id) + reapply_runtime_model_cost_registrations() + assert litellm.model_cost[model_id]["input_cost_per_token"] == 1e-3 + core.retire() + litellm.model_cost.pop(model_id) + reapply_runtime_model_cost_registrations() + assert model_id not in litellm.model_cost + finally: + litellm.model_cost.pop(model_id, None) + + +# --- runtime routing_strategy_args ---------------------------------------------------------------- + + +def test_new_routing_strategy_args_rebuild_the_selector_and_replace_it_on_the_host( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 5} + ) + (before,) = host.lists["callbacks"] + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 9}})) + (after,) = host.lists["callbacks"] + assert after is not before + assert after is core.lowestlatency_logger # pyright: ignore[reportAttributeAccessIssue] + assert after.routing_args.ttl == 9 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().as_dict()["routing_strategy_args"]["ttl"] == 9 + + +def test_invalid_routing_strategy_args_at_runtime_keep_the_previous_selector( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the invalid args are still stored as ``routing_strategy_args``, while the selector (and + so ``as_dict``, which reads the selector) keeps the previous ones.""" + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 5} + ) + (before,) = host.lists["callbacks"] + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": "abc"}})) + assert host.lists["callbacks"] == [before] + assert core.lowestlatency_logger is before # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().as_dict()["routing_strategy_args"]["ttl"] == 5 + assert core.routing_strategy_args == {"ttl": "abc"} # pyright: ignore[reportAttributeAccessIssue] + + +def test_new_least_busy_args_replace_the_selector_in_callbacks_and_input_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], routing_strategy="least-busy") + before: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"routing_strategy_args": {"x": 1}})) + after: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] + assert after is not before + assert (host.lists["callbacks"], host.lists["input_callback"]) == ([after], [after]) + + +def test_new_routing_strategy_args_under_simple_shuffle_register_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 9}})) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 9} + assert (host.lists["callbacks"], host.lists["input_callback"]) == ([], []) + + +# --- routing groups: selectors on the host, shadowing --------------------------------------------- + + +def test_routing_group_selectors_are_registered_on_the_host_and_replaced_when_groups_change( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + routing_strategy="least-busy", + routing_groups=[_group("g1", ["m1"], "latency-based-routing", routing_strategy_args={"ttl": 7})], + ) + assert host_callback_names(host)["callbacks"] == ["LeastBusyLoggingHandler", "LowestLatencyLoggingHandler"] + (latency,) = ( + callback for callback in host.lists["callbacks"] if type(callback).__name__ == "LowestLatencyLoggingHandler" + ) + assert latency.routing_args.ttl == 7 # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"routing_groups": [_group("g2", ["m2"], "usage-based-routing-v2")]})) + assert host_callback_names(host)["callbacks"] == ["LeastBusyLoggingHandler", "LowestTPMLoggingHandler_v2"] + core.apply(ChangeSet(settings={"routing_groups": []})) + assert host_callback_names(host)["callbacks"] == ["LeastBusyLoggingHandler"] + assert host_callback_names(host)["input_callback"] == ["LeastBusyLoggingHandler"] + + +@pytest.mark.asyncio +async def test_a_group_shadowed_by_a_model_name_still_routes_its_members_with_its_strategy( + make_core: MakeCore, +) -> None: + """The shadowed group's name is not callable (calling it reaches the model), but its least-busy + strategy still serves its member: with no traffic least-busy always picks the first deployment.""" + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b"), deployment("m2", "c")], + routing_groups=[_group("m1", ["m2"], "least-busy")], + ) + assert await picked_ids(core, "m2", times=30, messages=MESSAGES) == {"b"} + assert core.registry().routing_group("m1") is None + + +@pytest.mark.asyncio +async def test_a_group_named_like_an_alias_is_shadowed_by_the_alias(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + model_group_alias={"al": "m1"}, + routing_groups=[_group("al", ["m2"], "simple-shuffle")], + ) + assert await picked_ids(core, "al") == {"a"} + assert core.registry().routing_group("al") is None + + +def test_a_priority_group_named_like_an_alias_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=re.escape("Priority routing group names must not shadow a model or alias")): + make_core( + model_list=[deployment("m1", "a")], + model_group_alias={"al": "m1"}, + routing_groups=[_group("al", ["m1"], "priority", model_priorities={"m1": 1})], + ) + + +@pytest.mark.asyncio +async def test_a_deployment_added_under_a_group_name_shadows_the_group(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + routing_groups=[_group("grp", ["m1", "m2"], "simple-shuffle")], + ) + core.apply(ChangeSet(upsert=[Deployment(**deployment("grp", "z"))])) # pyright: ignore[reportArgumentType] + assert await picked_ids(core, "grp") == {"z"} + assert core.registry().routing_group("grp") is None + + +# --- registry().resolve kinds backed by this range ------------------------------------------------ + + +def test_routing_group_alternatives_count_the_deployments_of_every_member(make_core: MakeCore) -> None: + """Only a callable routing group (or an alias to one) can have alternatives; a plain model group + with two deployments has none by this rule.""" + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m1", "a2"), deployment("m2", "b"), deployment("m3", "c")], + model_group_alias={"al": "grp"}, + routing_groups=[ + _group("grp", ["m2", "m3"], "simple-shuffle"), + _group("one", ["m2"], "priority", model_priorities={"m2": 1}), + ], + ) + registry: Final = core.registry() + answers: Final = { + name: registry.resolve(name, kind="routing_group_alternatives") # pyright: ignore[reportCallIssue, reportArgumentType] + for name in (None, "grp", "one", "al", "m1", "zzz") + } + assert answers == {None: False, "grp": True, "one": False, "al": True, "m1": False, "zzz": False} + + +TEAM_MODELS: Final = [ + deployment("pub", "t1", model_info={"team_id": "T", "team_public_model_name": "pubm"}), + deployment("pub2", "t2", model_info={"team_id": "T", "team_public_model_name": "pubm"}), + deployment("solo", "t3", model_info={"team_id": "T", "team_public_model_name": "solo"}), + deployment("plain", "p1"), + deployment("half", "h1", model_info={"team_id": "T"}), +] + + +@pytest.mark.parametrize( + ("deployment_id", "expected"), + [("t1", True), ("t3", False), ("p1", False), ("h1", False), ("zzz", False)], + ids=["two-siblings", "only-one", "no-team", "no-public-name", "unknown-id"], +) +def test_team_model_alternatives_need_another_routable_deployment_under_the_same_team_name( + make_core: MakeCore, deployment_id: str, expected: bool +) -> None: + core: Final = make_core(model_list=list(TEAM_MODELS)) + assert core.registry().resolve(deployment_id, kind="team_model_alternatives") is expected # pyright: ignore[reportCallIssue, reportArgumentType] + + +def test_a_blocked_team_sibling_is_no_alternative(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("pub", "t1", model_info={"team_id": "T", "team_public_model_name": "pubm"}), + deployment("pub2", "t2", model_info={"team_id": "T", "team_public_model_name": "pubm", "blocked": True}), + ] + ) + assert core.registry().resolve("t1", kind="team_model_alternatives") is False # pyright: ignore[reportCallIssue, reportArgumentType] + + +def test_search_tool_support_comes_from_model_info_else_an_anthropic_request_shape(make_core: MakeCore) -> None: + """Per routable deployment: a configured ``supports_search_tool`` wins, a deployment whose chat + config builds an Anthropic request answers True, anything else None; an unknown model has none.""" + core: Final = make_core( + model_list=[ + deployment("g", "o1"), + deployment("g", "an1", model="anthropic/claude-3-5-sonnet-20240620"), + deployment( + "g", "cfg", model="anthropic/claude-3-5-sonnet-20240620", model_info={"supports_search_tool": False} + ), + ] + ) + registry: Final = core.registry() + assert registry.resolve("g", kind="search_tool_support") == (None, True, False) # pyright: ignore[reportCallIssue, reportArgumentType] + assert registry.resolve("zzz", kind="search_tool_support") == () # pyright: ignore[reportCallIssue, reportArgumentType] + + +# --- per-request strategy overrides --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_override_selector_is_bound_to_the_request_logging_object_and_reused(make_core: MakeCore) -> None: + """An override strategy the router does not register globally is bound to the request's logging + object; one selector per strategy serves every request. An override matching the router's own + strategy reuses the global selector and binds nothing.""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="least-busy") + own: Final = _logging_obj() + await core.select("g", {"routing_strategy": "least-busy", "litellm_logging_obj": own}) + first: Final = _logging_obj() + await core.select("g", {"routing_strategy": "usage-based-routing-v2", "litellm_logging_obj": first}) + second: Final = _logging_obj() + await core.select("g", {"routing_strategy": "usage-based-routing-v2", "litellm_logging_obj": second}) + assert _dynamic_callback_names(own) == [] + assert _dynamic_callback_names(first) == ["LowestTPMLoggingHandler_v2"] + assert (first.dynamic_async_success_callbacks or [None])[0] is (second.dynamic_async_success_callbacks or [])[0] + + +@pytest.mark.asyncio +async def test_an_override_matching_a_routing_group_strategy_is_not_bound_to_the_request(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("m", "b")], + routing_groups=[_group("gg", ["m"], "usage-based-routing-v2")], + ) + logging_obj: Final = _logging_obj() + await core.select("g", {"routing_strategy": "usage-based-routing-v2", "litellm_logging_obj": logging_obj}) + assert _dynamic_callback_names(logging_obj) == [] + + +def test_a_sync_usage_based_v2_override_counts_rpm_before_the_call(make_core: MakeCore, host: ScriptedHost) -> None: + """The override selector is on no callback list, so the sync pick runs its pre-call check itself: + the first call is counted against ``rpm=1`` and the next pick finds no deployment. Oddity: the sync + path raises a bare ``ValueError`` where the async one raises ``RateLimitError``.""" + core: Final = make_core(model_list=[deployment("g", "a", rpm=1)], num_retries=0) + core.completion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(ValueError, match=re.escape("No deployments available for selected model. Passed model=g.")): + core.completion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + + +@pytest.mark.asyncio +async def test_an_async_usage_based_v2_override_counts_rpm_before_the_call( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=1)], num_retries=0) + await core.acompletion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(litellm.RateLimitError, match=re.escape("No deployments available for selected model.")): + await core.acompletion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + + +@pytest.mark.asyncio +async def test_request_scoped_routing_still_runs_the_strategy_rpm_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + """After ``observe(request_scoped=True)`` the selector is off the host's lists, but the router's + own pre-call checks still reach it, so ``rpm=1`` still holds.""" + collect_other_routers() + core: Final = make_core( + model_list=[deployment("g", "a", rpm=1)], routing_strategy="usage-based-routing-v2", num_retries=0 + ) + core.observe(request_scoped=True) + assert host.lists["callbacks"] == [] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + + +# --- strategy dispatch arms ----------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_async_pick_under_usage_based_routing_v1_picks_the_first_configured(make_core: MakeCore) -> None: + """The async pick calls v1's sync selector inline; with no traffic it picks the first deployment.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="usage-based-routing" + ) + assert await picked_ids(core, "g", times=10, messages=MESSAGES) == {"a"} + + +def test_sync_completion_under_latency_based_routing_reaches_the_group(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="latency-based-routing" + ) + for _ in range(10): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) <= {"a", "b"} + assert len(host.calls) == 10 + + +def test_sync_completion_under_cost_based_routing_finds_no_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + """Oddity: the cost-based selector has no sync pick, so every sync pick under cost-based routing + rate-limits without calling a provider.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="cost-based-routing", num_retries=0 + ) + with pytest.raises(RouterRateLimitError) as raised: + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g") + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_pass_through_pick_under_cost_based_routing_finds_no_deployment(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", use_in_pass_through=True)], routing_strategy="cost-based-routing" + ) + with pytest.raises(RouterRateLimitError): + await core.select("g", kind="pass_through") diff --git a/tests/unit/test_router/core/test_gaps_l11100_registry.py b/tests/unit/test_router/core/test_gaps_l11100_registry.py new file mode 100644 index 00000000000..bd9fa3a4e8c --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l11100_registry.py @@ -0,0 +1,657 @@ +""" +Gap tests for the deployment registry of ``router_python_core.py`` (about lines 11550-14060): +provider validation in ``_add_deployment``, the os.environ references of Azure ``dataSources``, +vertex pass-through registration, config-sourced rows, index repair after a delete, cost-map +registration (``responses/`` keys, strategy-marker pricing, replay), deployment budgets kept in +step with upserts and deletes, credentials with provider, router / deployment model info, and how +model group info merges its deployments. +""" + +import os +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.types.router import Deployment +from tests.unit.test_router.core.gaps_l11100_helpers import MakeCore, ModelsEndpoint, complexity_marker, model_ids +from tests.unit.test_router.core.registry_helpers import team_deployment +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +def _unique(prefix: str = "l11100") -> str: + return f"{prefix}-{uuid.uuid4().hex}" + + +@pytest.fixture +def data_source_env(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("L11100_ENDPOINT", "https://search.example") + monkeypatch.delenv("L11100_MISSING", raising=False) + yield + + +# --- provider validation ------------------------------------------------------------------------------ + + +def test_a_provider_missing_from_the_hosts_provider_list_is_unsupported( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(provider_list=["anthropic"]) + + with pytest.raises(Exception, match=r"^Unsupported provider - openai$"): + make_core(model_list=[deployment("gpt", "a")]) + + +def test_ignore_invalid_deployments_drops_a_deployment_of_an_unsupported_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(provider_list=["anthropic"]) + + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("claude", "c", model="anthropic/claude-opus-4-6")], + ignore_invalid_deployments=True, + ) + + assert model_ids(core) == ["c"] + + +def test_a_prompt_management_model_skips_provider_validation(make_core: MakeCore, host: ScriptedHost) -> None: + """A ``/...`` model (here langfuse) is resolved from the prompt at call time, so it is + accepted even when no provider is supported at all.""" + host.set_dynamic(provider_list=[]) + + core: Final = make_core(model_list=[deployment("prompted", "p", model="langfuse/my-prompt")]) + + assert model_ids(core) == ["p"] + assert core.registry().names(kind="deployment_names") == ["langfuse/my-prompt"] + + +# --- Azure dataSources and vertex pass-through --------------------------------------------------------- + + +@pytest.mark.usefixtures("data_source_env") +def test_data_source_endpoint_and_key_env_references_are_resolved(make_core: MakeCore) -> None: + """Only ``endpoint`` and ``key`` are read from the environment; a missing variable becomes ``""``.""" + core: Final = make_core( + model_list=[ + deployment( + "gpt", + "a", + dataSources=[ + { + "type": "AzureCognitiveSearch", + "parameters": { + "endpoint": "os.environ/L11100_ENDPOINT", + "key": "os.environ/L11100_MISSING", + "indexName": "os.environ/L11100_ENDPOINT", + }, + } + ], + ) + ] + ) + + stored: Final = core.registry().all()[0]["litellm_params"]["dataSources"] + assert stored == [ + { + "type": "AzureCognitiveSearch", + "parameters": { + "endpoint": "https://search.example", + "key": "", + "indexName": "os.environ/L11100_ENDPOINT", + }, + } + ] + + +def test_a_vertex_pass_through_deployment_without_project_and_location_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="vertex_project") as raised: + make_core(model_list=[deployment("gem", "g", model="vertex_ai/gemini-pro", use_in_pass_through=True)]) + assert str(raised.value) == ( + "vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints." + ) + + +def test_pass_through_needs_no_registration_for_other_providers(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a", use_in_pass_through=True)]) + + assert model_ids(core) == ["a"] + + +# --- config rows and index repair ----------------------------------------------------------------------- + + +def test_config_deployments_leave_out_rows_that_came_from_the_db(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("gpt", "a"), + deployment("gpt", "b", model_info={"db_model": True}), + deployment("other", "c", model_info={"db_model": False}), + ] + ) + + assert [entry["model_info"]["id"] for entry in core.registry().deployments(source="config")] == ["a", "c"] + + +def test_deleting_a_later_deployment_keeps_earlier_team_indices_and_shifts_later_ones(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + team_deployment("internal-1", "t1", "TA", "pub"), + deployment("gpt", "b"), + team_deployment("internal-2", "t2", "TA", "other"), + ] + ) + + core.apply(ChangeSet(delete=["b"])) + + assert dict(core.registry().index("team_model")) == {("TA", "pub"): [0], ("TA", "other"): [1]} + + +def test_deleting_a_teams_only_wildcard_deployment_drops_its_team_pattern_router(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[team_deployment("internal", "tw", "TA", "openai/*"), deployment("gpt", "b")], + ) + assert list(core.registry().patterns(team=True)) == ["TA"] + + core.apply(ChangeSet(delete=["tw"])) + + assert dict(core.registry().patterns(team=True)) == {} + + +# --- cost-map registration --------------------------------------------------------------------------------- + + +def test_a_responses_model_registers_its_backend_key_with_and_without_the_responses_segment( + make_core: MakeCore, +) -> None: + """The shared backend keys carry no custom pricing: that stays under the deployment id.""" + name: Final = _unique() + deployment_id: Final = _unique("l11100-id") + make_core(model_list=[deployment("gpt", deployment_id, model=f"openai/responses/{name}", input_cost_per_token=0.1)]) + + assert litellm.model_cost[f"openai/responses/{name}"].get("input_cost_per_token") is None + assert litellm.model_cost[f"openai/{name}"].get("input_cost_per_token") is None + assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.1 + + +def test_custom_pricing_on_a_strategy_marker_is_not_registered_as_a_price(make_core: MakeCore) -> None: + """A marker is never the deployment that is billed, so even an explicit zero price stays out of its + cost-map entry.""" + marker_id: Final = _unique() + make_core( + model_list=[ + deployment("small", "s"), + deployment("big", "b"), + complexity_marker("smart", marker_id, input_cost_per_token=0.0, output_cost_per_token=0.0), + ] + ) + + entry: Final = litellm.model_cost[marker_id] + assert "input_cost_per_token" not in entry + assert "output_cost_per_token" not in entry + + +def test_replaying_cost_registrations_restores_the_entries_of_deployments_still_served(make_core: MakeCore) -> None: + kept: Final = _unique() + deleted: Final = _unique() + core: Final = make_core( + model_list=[ + deployment("gpt", kept, input_cost_per_token=0.25), + deployment("gpt", deleted, input_cost_per_token=0.5), + ] + ) + core.apply(ChangeSet(delete=[deleted])) + litellm.model_cost.pop(kept, None) + litellm.model_cost.pop(deleted, None) + + core.replay_model_cost_registrations() # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert litellm.model_cost[kept]["input_cost_per_token"] == 0.25 + assert deleted not in litellm.model_cost + + +# --- deployment budgets -------------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deployment_budgets_from_the_model_list_are_registered_at_construction(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("gpt", "a", max_budget=1.0, budget_duration="1d"), + deployment("gpt", "b", max_budget=2.0, budget_duration="1d"), + deployment("gpt", "c"), + ] + ) + + limiter: Final = core.registry().plugins("budget_limiter") + assert limiter is not None + assert limiter.deployment_budget_config is not None + assert {key: config.max_budget for key, config in limiter.deployment_budget_config.items()} == {"a": 1.0, "b": 2.0} + + +@pytest.mark.asyncio +async def test_deleting_and_upserting_deployments_keep_the_budget_limiter_in_step(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("gpt", "a", max_budget=1.0, budget_duration="1d"), + deployment("gpt", "b", max_budget=2.0, budget_duration="1d"), + ] + ) + limiter: Final = core.registry().plugins("budget_limiter") + assert limiter is not None + + core.apply(ChangeSet(delete=["a"])) + assert list(limiter.deployment_budget_config or {}) == ["b"] + + core.apply(ChangeSet(upsert=[Deployment(**deployment("gpt", "b"))])) # pyright: ignore[reportArgumentType] + assert limiter.deployment_budget_config is None + + core.apply(ChangeSet(upsert=[Deployment(**deployment("gpt", "b", max_budget=3.0, budget_duration="1d"))])) # pyright: ignore[reportArgumentType] + assert {key: config.max_budget for key, config in (limiter.deployment_budget_config or {}).items()} == {"b": 3.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("params", [{"max_budget": 1.0}, {"budget_duration": "1d"}], ids=["no-duration", "no-budget"]) +async def test_adding_a_deployment_with_half_a_budget_installs_no_budget_limiter( + make_core: MakeCore, params: dict[str, object] +) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")]) + + core.add_deployment(Deployment(**deployment("gpt", "half", **params))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert core.registry().plugins("budget_limiter") is None + + +# --- credentials with provider ------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("entry", "expected_provider"), + [ + (deployment("gpt", "a", model="gpt-4o"), "openai"), + (deployment("gpt", "a", model="gpt-4o", custom_llm_provider="azure", api_base="https://x"), "azure"), + (deployment("gpt", "a", model="anthropic/claude-opus-4-6"), "anthropic"), + ], + ids=["bare-model-defaults-to-openai", "explicit-provider-wins", "provider-prefix"], +) +def test_credentials_with_provider_name_the_provider( + make_core: MakeCore, entry: dict[str, object], expected_provider: str +) -> None: + found: Final = make_core(model_list=[entry]).registry().credentials("a", form="with_provider") + + assert found is not None + assert found["custom_llm_provider"] == expected_provider + + +def test_credentials_with_provider_drop_an_unknown_credential_name(make_core: MakeCore) -> None: + found: Final = ( + make_core(model_list=[deployment("gpt", "a", litellm_credential_name="l11100-no-such-credential")]) + .registry() + .credentials("a", form="with_provider") + ) + + assert found == {"api_key": "sk-test", "model": "openai/a", "custom_llm_provider": "openai"} + + +# --- router and deployment model info ------------------------------------------------------------------------ + + +@pytest.mark.parametrize( + ("model", "model_info", "expected_key"), + [ + ("azure/gpt-4o", {}, "azure/gpt-4o"), + ("azure/my-deployment", {"base_model": "azure/gpt-4o"}, "azure/gpt-4o"), + ], + ids=["mapped-deployment-name", "base-model"], +) +def test_router_model_info_of_an_azure_deployment( + make_core: MakeCore, model: str, model_info: dict[str, object], expected_key: str +) -> None: + core: Final = make_core( + model_list=[ + deployment("az", "az", model=model, api_base="https://x", api_version="2024", model_info=model_info) + ] + ) + + info: Final = core.registry().router_model_info(None, "az", id="az") + builtin: Final = litellm.get_model_info("azure/gpt-4o") + assert info["key"] == expected_key + assert info["max_input_tokens"] == builtin["max_input_tokens"] + assert info["input_cost_per_token"] == builtin["input_cost_per_token"] + + +def test_router_model_info_of_an_unmapped_azure_deployment_is_its_zeroed_stub(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("az", "az", model="azure/l11100-unmapped", api_base="https://x", api_version="2024")] + ) + + info: Final = core.registry().router_model_info(None, "az", id="az") + assert info["key"] == "azure/l11100-unmapped" + assert info["max_input_tokens"] is None + assert info["input_cost_per_token"] == 0 + + +def test_deployment_model_info_takes_defaults_from_a_custom_base_model(make_core: MakeCore) -> None: + """The deployment's own pricing wins; the rest comes from its ``base_model``'s built-in entry.""" + name: Final = _unique() + core: Final = make_core( + model_list=[ + deployment( + "gpt", name, model=f"openai/{name}", input_cost_per_token=0.5, model_info={"base_model": "gpt-4o"} + ) + ] + ) + + info: Final = core.registry().deployment_model_info(name, f"openai/{name}") + builtin: Final = litellm.get_model_info("gpt-4o") + assert info is not None + assert info["input_cost_per_token"] == 0.5 + assert info["output_cost_per_token"] == builtin["output_cost_per_token"] + assert info["max_input_tokens"] == builtin["max_input_tokens"] + + +def test_deployment_model_info_of_an_unpriced_unmapped_model_reports_no_prices(make_core: MakeCore) -> None: + """``get_model_info`` reads a missing price as 0; the deployment reports ``None`` instead.""" + name: Final = _unique() + core: Final = make_core(model_list=[deployment("gpt", name, model=f"openai/{name}")]) + + info: Final = core.registry().deployment_model_info(name, f"openai/{name}") + assert info is not None + assert (info["input_cost_per_token"], info["output_cost_per_token"]) == (None, None) + + +# --- model group info ------------------------------------------------------------------------------------------ + + +def test_group_info_takes_the_widest_limits_highest_prices_and_any_capability(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="openai/gpt-3.5-turbo"), + deployment("g", "b", model="openai/gpt-4o"), + deployment("g", "c", model="openai/o3-mini"), + deployment("g", "d", model="anthropic/claude-opus-4-6"), + ] + ) + + info: Final = core.registry().group("g") + opus: Final = litellm.get_model_info("anthropic/claude-opus-4-6") + assert info is not None + assert info.providers == ["openai", "anthropic"] + assert (info.max_input_tokens, info.max_output_tokens) == (opus["max_input_tokens"], opus["max_output_tokens"]) + assert (info.input_cost_per_token, info.output_cost_per_token) == ( + opus["input_cost_per_token"], + opus["output_cost_per_token"], + ) + assert info.supports_vision is True + assert info.supports_function_calling is True + assert info.supports_parallel_function_calling is True + assert info.supports_web_search is True + assert info.supports_reasoning is True + assert info.supports_url_context is False + assert info.supports_fast_mode is False + + +def test_group_info_supports_url_context_when_any_deployment_does(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", model="openai/gpt-4o"), deployment("g", "b", model="gemini/gemini-2.5-flash")] + ) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supports_url_context is True + + +@pytest.mark.parametrize( + ("models", "fast"), + [ + (["anthropic/claude-opus-4-8"], True), + (["anthropic/claude-opus-4-8", "anthropic/claude-opus-4-6"], False), + ], + ids=["every-deployment-fast", "one-not-fast"], +) +def test_group_info_supports_fast_mode_only_when_every_deployment_does( + make_core: MakeCore, models: list[str], fast: bool +) -> None: + core: Final = make_core(model_list=[deployment("g", f"d{i}", model=m) for i, m in enumerate(models)]) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supports_fast_mode is fast + + +@pytest.mark.parametrize( + ("models", "efforts"), + [ + (["anthropic/claude-opus-4-8"], ("none", "minimal", "low", "medium", "high", "xhigh", "max")), + ( + ["anthropic/claude-opus-4-8", "anthropic/claude-opus-4-6"], + ("none", "minimal", "low", "medium", "high", "max"), + ), + (["anthropic/claude-sonnet-4-5-20250929", "anthropic/claude-opus-4-6"], None), + (["anthropic/claude-opus-4-6", "anthropic/claude-sonnet-4-5-20250929"], None), + (["openai/gpt-4o", "openai/gpt-4o"], ()), + ], + ids=["one", "intersection", "unknown-first-stays-unknown", "unknown-later-clears", "none-supported"], +) +def test_group_info_reasoning_efforts_are_the_intersection_or_unknown( + make_core: MakeCore, models: list[str], efforts: tuple[str, ...] | None +) -> None: + core: Final = make_core(model_list=[deployment("g", f"d{i}", model=m) for i, m in enumerate(models)]) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supported_reasoning_efforts == efforts + + +def test_group_info_supported_endpoints_are_the_union_in_deployment_order(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="openai/gpt-5"), + deployment("g", "b", model="openai/gpt-4o", model_info={"supported_endpoints": ["/v1/embeddings"]}), + ] + ) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supported_endpoints == ("/v1/chat/completions", "/v1/batch", "/v1/responses", "/v1/embeddings") + + +def test_group_info_of_a_prompt_management_model_names_its_logger_as_provider(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("prompted", "p", model="langfuse/my-prompt")]) + + info: Final = core.registry().group("prompted") + assert info is not None + assert info.providers == ["langfuse"] + + +def test_group_info_of_an_alias_to_neither_a_name_nor_a_dict_is_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")]) + core.apply(ChangeSet(assign={"model_group_alias": {"odd": 5}})) + + assert core.registry().group("odd") is None + + +def test_group_info_of_an_unpriced_unmapped_model_has_no_prices_or_limits(make_core: MakeCore) -> None: + name: Final = _unique() + core: Final = make_core(model_list=[deployment("g", name, model=f"openai/{name}")]) + + info: Final = core.registry().group("g") + assert info is not None + assert (info.input_cost_per_token, info.output_cost_per_token, info.max_input_tokens) == (None, None, None) + + +# --- lookups --------------------------------------------------------------------------------------------------- + + +def test_from_id_of_a_deployment_id_is_its_group_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")]) + + assert core.registry().resolve("a", kind="from_id") == "gpt" + + +def test_from_id_on_an_empty_core_or_of_an_empty_id_is_none(make_core: MakeCore) -> None: + assert make_core(model_list=[]).registry().resolve("x", kind="from_id") is None + assert make_core(model_list=[deployment("gpt", "a")]).registry().resolve("", kind="from_id") is None + + +def test_route_candidates_of_a_routing_group_are_its_members_deployments(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("h", "b"), deployment("other", "c")], + routing_groups=[{"group_name": "rg", "models": ["g", "h"], "routing_strategy": "simple-shuffle"}], + ) + + assert core.registry().names("rg", kind="route_candidates") == frozenset({"a", "b"}) + + +def test_search_tool_support_reads_model_info_then_the_anthropic_route(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="anthropic/claude-opus-4-6"), + deployment("g", "b", model="openai/gpt-4o", model_info={"supports_search_tool": False}), + deployment("g", "c", model="openai/gpt-4o"), + ] + ) + + assert core.registry().resolve("g", kind="search_tool_support") == (True, False, None) + + +# --- discovered model info (aapply refresh_model_info) ---------------------------------------------------- + +VLLM_CARDS: Final = {"data": [{"id": "my-model", "max_model_len": 8192}, {"id": "other", "max_model_len": 4096}]} + + +def _vllm( + deployment_id: str = "v", **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: litellm_params + return deployment( + "vllm", deployment_id, model="hosted_vllm/my-model", api_base="http://vllm.local", api_key="sk-vllm", **params + ) + + +@pytest.mark.asyncio +async def test_a_refresh_reads_the_upstream_models_endpoint_with_the_deployments_key(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[_vllm(extra_headers={"X-Team": "a"})]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert endpoint.requests == [("http://vllm.local/v1/models", {"authorization": "Bearer sk-vllm", "x-team": "a"})] + assert dict(core.registry().discovered_model_info("v")) == { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + } + + +@pytest.mark.asyncio +async def test_discovered_limits_never_override_configured_ones_and_feed_the_group_limits(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[_vllm(model_info={"max_output_tokens": 1000})]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + registry: Final = core.registry() + assert dict(registry.discovered_model_info("v")) == {"max_tokens": 8192, "max_input_tokens": 8192} + assert registry.group("vllm", field="token_limits") == (8192, 1000) + listing: Final = registry.group("vllm", field="listing_info") + assert listing is not None + assert (listing.max_input_tokens, listing.max_output_tokens) == (8192, 1000) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entry", + [ + deployment("plain", "p", model="anthropic/claude-opus-4-6", api_base="http://vllm.local"), + deployment("wild", "w", model="hosted_vllm/*", api_base="http://vllm.local"), + deployment("nobase", "n", model="openai/gpt-4o"), + ], + ids=["provider-without-discovery", "wildcard-model", "no-api-base"], +) +async def test_a_refresh_skips_deployments_it_cannot_discover(make_core: MakeCore, entry: dict[str, object]) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[entry]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert endpoint.requests == [] + + +@pytest.mark.asyncio +async def test_an_unreachable_models_endpoint_discovers_nothing(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(error=ConnectionError("down")) + core: Final = make_core(model_list=[_vllm()]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert len(endpoint.requests) == 1 + assert dict(core.registry().discovered_model_info("v")) == {} + + +@pytest.mark.asyncio +async def test_discovered_limits_are_forgotten_when_the_deployment_is_replaced(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[_vllm()]) + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + core.apply(ChangeSet(upsert=[Deployment(**_vllm(rpm=5))])) # pyright: ignore[reportArgumentType] + + assert dict(core.registry().discovered_model_info("v")) == {} + + +# --- failed upserts and replays under ignore_invalid_deployments --------------------------------------------- + + +def test_a_failing_upsert_of_a_new_deployment_under_ignore_invalid_deployments_adds_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")], ignore_invalid_deployments=True) + host.set_dynamic(provider_list=["anthropic"]) + + result: Final = core.apply(ChangeSet(upsert=[Deployment(**deployment("gpt", "new"))])) # pyright: ignore[reportArgumentType] + + assert list(result.upserted) == [None] + assert model_ids(core) == ["a"] + + +def test_a_replay_retries_deployments_whose_provider_did_not_resolve_and_keeps_them_out_while_it_still_fails( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("odd", "u", model="l11100-no-provider/x")], + ignore_invalid_deployments=True, + ) + assert model_ids(core) == ["a"] + + core.replay_model_cost_registrations() # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert model_ids(core) == ["a"] + + +# --- group info without any model info ------------------------------------------------------------------------ + + +def test_group_info_without_a_cost_map_entry_falls_back_to_the_configured_mode_and_costs(make_core: MakeCore) -> None: + """With the entries construction registered withdrawn from the cost map (as a price data reload + does before the replay) and a model litellm does not map, the group is described from the + deployment's own model_info: its mode, its costs (a string cost read as a float, a missing one as + None) and the provider's supported params, with no token limits.""" + name: Final = _unique() + core: Final = make_core( + model_list=[ + deployment( + "g", name, model=f"openai/{name}", model_info={"mode": "embedding", "input_cost_per_token": "0.5"} + ) + ] + ) + litellm.model_cost.pop(name, None) + litellm.model_cost.pop(f"openai/{name}", None) + + info: Final = core.registry().group("g") + assert info is not None + assert (info.mode, info.input_cost_per_token, info.output_cost_per_token) == ("embedding", 0.5, None) + assert (info.max_input_tokens, info.max_output_tokens) == (None, None) + assert info.supported_openai_params diff --git a/tests/unit/test_router/core/test_gaps_l11100_strategy.py b/tests/unit/test_router/core/test_gaps_l11100_strategy.py new file mode 100644 index 00000000000..2459c577b9c --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l11100_strategy.py @@ -0,0 +1,438 @@ +""" +Gap tests for the strategy-router (pre-routing) deployments of ``router_python_core.py`` (about lines +11114-11550, plus the strategy hooks of set_model_list / upsert / delete): construction-time +validation of auto, adaptive and quality router markers, the (model_name, tags) slot each marker +holds and when it is released, the licensed capability ceiling, and the adaptive router's derived +inputs and post-call hooks. +""" + +from typing import Final + +import pytest + +from litellm.router_api import ChangeSet +from litellm.types.router import Deployment +from tests.unit.test_router.core.gaps_l11100_helpers import ( + TIERS, + MakeCore, + RemovingHost, + adaptive_marker, + complexity_marker, + marker, + model_ids, + registered_tags, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +@pytest.fixture +def host() -> ScriptedHost: + """The adaptive router's hook sync removes callbacks from every list, so the host must support it.""" + return RemovingHost() + + +PLAIN: Final = [deployment("small", "s"), deployment("big", "b")] + +HEURISTIC_V2: Final = {"classifier_type": "heuristic_v2", "tiers": dict(TIERS)} + +HEURISTIC_V2_LIMIT_MESSAGE: Final = ( + "At most 1 auto-router(s) using the heuristic_v2 classifier can be registered but this would make 2. " + "Choose the heuristic v1 classifier for this router or remove an existing heuristic_v2 router." +) + + +# --- marker validation at construction ------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("params", "message"), + [ + ( + {}, + "auto_router_config_path or auto_router_config is required for auto-router deployments. " + "Please set it in the litellm_params", + ), + ( + {"auto_router_config": "{}"}, + "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params", + ), + ( + {"auto_router_config": "{}", "auto_router_default_model": "big"}, + "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params", + ), + ], + ids=["no-config", "no-default-model", "no-embedding-model"], +) +def test_an_auto_router_marker_missing_a_required_param_is_rejected( + make_core: MakeCore, params: dict[str, object], message: str +) -> None: + with pytest.raises(ValueError, match="is required") as raised: + make_core(model_list=[*PLAIN, marker("auto", "m", "auto_router/semantic", **params)]) + assert str(raised.value) == message + + +def test_an_adaptive_router_marker_without_its_config_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="is required") as raised: + make_core(model_list=[*PLAIN, marker("ad", "m", "auto_router/adaptive_router")]) + assert str(raised.value) == "adaptive_router_config is required for adaptive-router deployments." + + +def test_a_quality_router_marker_without_a_default_model_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="is required") as raised: + make_core(model_list=[*PLAIN, marker("q", "m", "auto_router/quality_router")]) + assert str(raised.value) == ( + "quality_router_default_model is required for quality-router deployments, or set default_model in " + "quality_router_config. Please configure it in the litellm_params" + ) + + +def test_a_quality_router_takes_its_default_model_from_its_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + marker("q", "m", "auto_router/quality_router", quality_router_config={"default_model": "big"}), + ] + ) + registry: Final = core.registry().plugins("quality") + + assert registered_tags(core, "quality") == {"q": [()]} + assert registry["q"][0].strategy.config.default_model == "big" + + +def test_quality_router_default_model_param_beats_the_one_in_its_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + marker( + "q", + "m", + "auto_router/quality_router", + quality_router_default_model="small", + quality_router_config={"default_model": "big"}, + ), + ] + ) + + assert core.registry().plugins("quality")["q"][0].strategy.config.default_model == "small" + + +# --- the (model_name, tags) slot -------------------------------------------------------------------- + + +@pytest.mark.parametrize(("tags", "shown"), [(None, "[]"), (["x"], "['x']")], ids=["untagged", "tagged"]) +def test_a_second_marker_with_the_same_name_and_tags_is_rejected( + make_core: MakeCore, tags: list[str] | None, shown: str +) -> None: + extra: Final = {} if tags is None else {"tags": tags} + with pytest.raises(ValueError, match="already exists") as raised: + make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", **extra), complexity_marker("smart", "m2", **extra)] + ) + assert str(raised.value) == ( + f"Complexity-router deployment smart with tags {shown} already exists. " + "Please use a different model name or set different tags." + ) + + +def test_ignore_invalid_deployments_drops_the_repeated_marker_and_keeps_the_first(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1"), complexity_marker("smart", "m2")], + ignore_invalid_deployments=True, + ) + + assert model_ids(core) == ["s", "b", "m1"] + assert registered_tags(core, "complexity") == {"smart": [()]} + + +def test_markers_of_one_name_with_different_tags_each_register(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", tags=["a"]), complexity_marker("smart", "m2", tags=["b"])] + ) + + assert registered_tags(core, "complexity") == {"smart": [("a",), ("b",)]} + + +def test_adding_a_marker_whose_slot_is_taken_raises_and_adds_nothing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1")]) + + with pytest.raises(ValueError, match="already exists"): + core.add_deployment(Deployment(**complexity_marker("smart", "m2"))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + assert model_ids(core) == ["s", "b", "m1"] + assert registered_tags(core, "complexity") == {"smart": [()]} + + +def test_deleting_a_marker_frees_its_slot_for_a_new_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1")]) + + core.apply(ChangeSet(delete=["m1"])) + assert registered_tags(core, "complexity") == {} + result: Final = core.apply(ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m2"))])) # pyright: ignore[reportArgumentType] + + assert result.upserted[0] is not None + assert registered_tags(core, "complexity") == {"smart": [()]} + assert model_ids(core) == ["s", "b", "m2"] + + +def test_deleting_one_tagged_marker_keeps_the_other_tags_of_its_name(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", tags=["a"]), complexity_marker("smart", "m2", tags=["b"])] + ) + + core.apply(ChangeSet(delete=["m1"])) + + assert registered_tags(core, "complexity") == {"smart": [("b",)]} + + +def test_deleting_a_plain_deployment_sharing_a_markers_name_keeps_the_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1"), deployment("smart", "p")]) + + core.apply(ChangeSet(delete=["p"])) + + assert registered_tags(core, "complexity") == {"smart": [()]} + + +def test_an_upsert_that_retags_a_marker_moves_it_to_the_new_slot(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1")]) + + core.apply(ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m1", tags=["t"]))])) # pyright: ignore[reportArgumentType] + + assert registered_tags(core, "complexity") == {"smart": [("t",)]} + assert model_ids(core) == ["s", "b", "m1"] + + +def test_set_model_list_drops_the_strategy_routers_of_the_old_list(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("smart", "m1"), + marker("q", "m2", "auto_router/quality_router", quality_router_default_model="big"), + ] + ) + + core.set_model_list(list(PLAIN)) # pyright: ignore[reportAttributeAccessIssue] + + assert registered_tags(core, "complexity") == {} + assert registered_tags(core, "quality") == {} + + +# --- the licensed capability ceiling ----------------------------------------------------------------- + + +def test_a_router_beyond_the_capability_limit_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="At most 1") as raised: + make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: 1, + ) + assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE + + +def test_routers_within_the_capability_limit_or_without_the_capability_register(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2"), + complexity_marker("three", "m3"), + ], + auto_router_capability_limit=lambda: 1, + ) + + assert registered_tags(core, "complexity") == {"one": [()], "two": [()], "three": [()]} + + +def test_without_a_capability_limit_any_number_of_gated_routers_register(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ] + ) + + assert sorted(registered_tags(core, "complexity")) == ["one", "two"] + + +def test_the_capability_limit_is_read_again_on_every_registration(make_core: MakeCore) -> None: + limit: Final = {"value": 2} + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: limit["value"], + ) + limit["value"] = 1 + + with pytest.raises(ValueError, match="At most 1") as raised: + core.add_deployment(Deployment(**complexity_marker("three", "m3", config=HEURISTIC_V2))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE.replace("would make 2", "would make 3") + + +def test_a_failed_edit_rolls_back_to_the_serving_router_even_above_a_tightened_limit(make_core: MakeCore) -> None: + """The edited router is popped first, so the edit itself counts 1 other + itself = 2 > 1 and fails; + the rollback re-adds the previous router with the ceiling switched off.""" + limit: Final = {"value": 2} + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: limit["value"], + ignore_invalid_deployments=True, + ) + limit["value"] = 1 + edited: Final = complexity_marker("two", "m2", config=HEURISTIC_V2, tags=["edited"]) + + result: Final = core.apply(ChangeSet(upsert=[Deployment(**edited)])) # pyright: ignore[reportArgumentType] + + assert list(result.upserted) == [None] + assert model_ids(core) == ["s", "b", "m1", "m2"] + assert registered_tags(core, "complexity") == {"one": [()], "two": [()]} + restored: Final = core.registry().deployment("m2") + assert restored is not None + assert restored.litellm_params.tags is None + + +def test_a_failed_edit_without_ignore_invalid_deployments_raises_and_leaves_the_router_removed( + make_core: MakeCore, +) -> None: + limit: Final = {"value": 2} + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: limit["value"], + ) + limit["value"] = 1 + + with pytest.raises(ValueError, match="At most 1") as raised: + core.apply( + ChangeSet(upsert=[Deployment(**complexity_marker("two", "m2", config=HEURISTIC_V2, tags=["edited"]))]) # pyright: ignore[reportArgumentType] + ) + assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE + assert model_ids(core) == ["s", "b", "m1"] + assert registered_tags(core, "complexity") == {"one": [()]} + + +# --- the adaptive router ------------------------------------------------------------------------------- + + +def _adaptive_list() -> list[dict[str, object]]: # mutable-ok: a model_list + return [ + adaptive_marker("ar", "ar1", ["x", "y", "z"]), + deployment( + "x", + "x", + input_cost_per_token=0.5, + model_info={"input_cost_per_token": 0.9, "adaptive_router_preferences": {"quality_tier": 2}}, + ), + deployment("y", "y", model_info={"input_cost_per_token": 0.25}), + ] + + +def _hook_counts(host: ScriptedHost) -> tuple[int, int]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, + ) + + callbacks: Final = host.lists["callbacks"] + return ( + sum(isinstance(callback, AdaptiveRouterPostCallHook) for callback in callbacks), + sum(isinstance(callback, AdaptiveRouterResponseHeaderHook) for callback in callbacks), + ) + + +def test_an_adaptive_router_is_built_after_the_whole_list_from_the_models_it_names(make_core: MakeCore) -> None: + """The marker is listed before x and y, yet sees them: litellm_params pricing wins over model_info + pricing, preferences come from model_info, and z (not configured) gets neither.""" + core: Final = make_core(model_list=_adaptive_list()) + registry: Final = core.registry().plugins("adaptive") + + assert registered_tags(core, "adaptive") == {"ar": [()]} + strategy: Final = registry["ar"][0].strategy + assert strategy.model_to_cost == {"x": 0.5, "y": 0.25} + assert {name: prefs.quality_tier for name, prefs in strategy.model_to_prefs.items()} == {"x": 2} + + +def test_an_adaptive_router_registers_one_post_call_hook_and_a_header_hook( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_adaptive_list()) + if not registered_tags(core, "adaptive"): + pytest.fail("no adaptive router was registered") + + post_call, header = _hook_counts(host) + assert post_call == 1 + assert header >= 1 + + +def test_deleting_an_adaptive_router_releases_it_and_its_post_call_hook( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_adaptive_list()) + + core.apply(ChangeSet(delete=["ar1"])) + + assert registered_tags(core, "adaptive") == {} + assert _hook_counts(host)[0] == 0 + + +def test_a_complexity_router_with_adaptive_enabled_also_registers_an_adaptive_router(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True})] + ) + + assert registered_tags(core, "complexity") == {"smart": [()]} + assert registered_tags(core, "adaptive") == {"smart": [()]} + + +def test_deleting_an_adaptive_complexity_router_releases_both_registrations(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True})] + ) + + core.apply(ChangeSet(delete=["m1"])) + + assert registered_tags(core, "complexity") == {} + assert registered_tags(core, "adaptive") == {} + + +def test_upserting_an_adaptive_complexity_router_registers_its_adaptive_companion(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PLAIN)) + + core.apply( + ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True}))]) # pyright: ignore[reportArgumentType] + ) + + assert registered_tags(core, "adaptive") == {"smart": [()]} + + +def test_an_upsert_that_turns_adaptive_off_releases_the_adaptive_companion(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True})] + ) + + core.apply(ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m1"))])) # pyright: ignore[reportArgumentType] + + assert registered_tags(core, "complexity") == {"smart": [()]} + assert registered_tags(core, "adaptive") == {} + + +def test_editing_an_adaptive_router_rebuilds_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_adaptive_list()) + + core.apply(ChangeSet(upsert=[Deployment(**adaptive_marker("ar", "ar1", ["y"]))])) # pyright: ignore[reportArgumentType] + + registry: Final = core.registry().plugins("adaptive") + assert registered_tags(core, "adaptive") == {"ar": [()]} + assert registry["ar"][0].strategy.model_to_cost == {"y": 0.25} diff --git a/tests/unit/test_router/core/test_gaps_l14000.py b/tests/unit/test_router/core/test_gaps_l14000.py new file mode 100644 index 00000000000..51c47bbf802 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l14000.py @@ -0,0 +1,299 @@ +""" +Gaps in the characterization of ``router_python_core.py`` lines 14000-17669 that the views answer: +``settings().allowed_fails_for`` under an allowed_fails_policy, ``settings().as_dict()`` for the +latency strategy and after a ``routing_strategy_args`` update, the ``registry()`` deployment +sources ``routable`` / ``alias`` / ``routing_groups`` / ``routed_group`` / ``for_request`` and team +aliases of ``all``, cached clients of a pod-discovery deployment, and PTU-share ceilings. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_api import ChangeSet +from litellm.types.router import AllowedFailsPolicy +from tests.unit.test_router.core.scripted_host import ( + auth_error, + bad_request, + content_policy, + context_window, + deployment, + rate_limit, + server_error, + timeout, +) +from tests.unit.test_router.core.selection_helpers import MakeCore, picked_ids + +# --- settings().allowed_fails_for ---------------------------------------------------------------- + + +def _bad_gateway() -> litellm.BadGatewayError: + return litellm.BadGatewayError(message="bad gateway", llm_provider="openai", model="gpt-4o") + + +def _connection_error() -> litellm.APIConnectionError: + return litellm.APIConnectionError(message="no route", llm_provider="openai", model="gpt-4o") + + +FULL_POLICY: Final = AllowedFailsPolicy( + ContentPolicyViolationErrorAllowedFails=1, + BadRequestErrorAllowedFails=2, + BadGatewayErrorAllowedFails=3, + RateLimitErrorAllowedFails=4, +) + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + (content_policy(), 1), + (bad_request(), 2), + # A context-window error is a BadRequestError, so it takes the bad-request allowance. + (context_window(), 2), + (_bad_gateway(), 3), + (rate_limit(), 4), + # Types the policy names no allowance for answer None, never the router's allowed_fails. + (server_error(), None), + (timeout(), None), + (auth_error(), None), + (_connection_error(), None), + ], + ids=[ + "content-policy", + "bad-request", + "context-window", + "bad-gateway", + "rate-limit", + "server-error", + "timeout", + "auth", + "connection", + ], +) +def test_allowed_fails_for_reads_the_policy_field_of_the_error_type( + make_core: MakeCore, error: Exception, expected: int | None +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], allowed_fails_policy=FULL_POLICY) + assert core.settings().allowed_fails_for(error) == expected + + +def test_allowed_fails_for_a_content_policy_error_falls_back_to_the_bad_request_allowance(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=2) + ) + assert core.settings().allowed_fails_for(content_policy()) == 2 + + +def test_allowed_fails_for_without_a_policy_is_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], allowed_fails=7) + assert core.settings().allowed_fails_for(rate_limit()) is None + + +# --- settings().as_dict() and routing_strategy_args ---------------------------------------------- + +LATENCY_ARGS_DEFAULTS: Final = {"lowest_latency_buffer": 0, "max_latency_list_size": 10, "ttft_percentile": None} + + +def test_latency_routing_reports_the_strategys_full_routing_args(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 30} + ) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 30.0, **LATENCY_ARGS_DEFAULTS} + + +def test_updating_routing_strategy_args_reaches_the_latency_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 30} + ) + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 99}})) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 99.0, **LATENCY_ARGS_DEFAULTS} + + +def test_routing_strategy_args_under_simple_shuffle_are_reported_as_given(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.settings().as_dict()["routing_strategy_args"] == {} + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 99}})) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 99} + + +def test_switching_to_latency_routing_takes_the_args_sent_with_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + core.apply(ChangeSet(settings={"routing_strategy": "latency-based-routing", "routing_strategy_args": {"ttl": 7}})) + settings: Final = core.settings().as_dict() + assert settings["routing_strategy"] == "latency-based-routing" + assert settings["routing_strategy_args"] == {"ttl": 7.0, **LATENCY_ARGS_DEFAULTS} + + +# --- registry() deployment sources --------------------------------------------------------------- + +TEAMS_AND_GROUPS: Final = [ + deployment("internal-x", "tx", model_info={"team_id": "team-x", "team_public_model_name": "pub"}), + deployment("internal-y", "ty", model_info={"team_id": "team-y", "team_public_model_name": "pub"}), + deployment("g", "a"), + deployment("g", "b", model_info={"access_groups": ["ag"]}), +] + +ROUTING_GROUP: Final = {"group_name": "grp", "models": ["g"], "routing_strategy": "simple-shuffle"} + + +def _ids(deployments: object) -> list[str]: # mutable-ok: a fresh list + return [deployment_id for _, deployment_id in _rows(deployments)] + + +def _rows(deployments: object) -> list[tuple[str, str]]: # mutable-ok: a fresh list + assert isinstance(deployments, (list, tuple)) + return [(str(row["model_name"]), str(row["model_info"]["id"])) for row in deployments] + + +def test_all_deployments_of_a_team_public_name_are_renamed_to_the_alias(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + rows: Final = core.registry().deployments("pub", source="all", team_id="team-x", model_alias="nick") + assert _rows(rows) == [("nick", "tx")] + + +def test_an_alias_to_a_routing_group_lists_the_members_under_the_alias_without_access_groups( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=list(TEAMS_AND_GROUPS), routing_groups=[dict(ROUTING_GROUP)], model_group_alias={"al": "grp"} + ) + rows: Final = core.registry().deployments("al") + assert _rows(rows) == [("al", "a"), ("al", "b")] + assert isinstance(rows, list) + assert all("access_groups" not in row["model_info"] for row in rows) + + +def test_the_alias_source_lists_every_alias_row(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(TEAMS_AND_GROUPS), + routing_groups=[dict(ROUTING_GROUP)], + model_group_alias={"al": "grp", "al2": "g"}, + ) + rows: Final = core.registry().deployments(source="alias") + assert _rows(rows) == [("al", "a"), ("al", "b"), ("al2", "a"), ("al2", "b")] + # A plain alias row keeps the member's access groups; only routing-group rows drop them. + assert isinstance(rows, list) + assert rows[3]["model_info"].get("access_groups") == ["ag"] + + +def test_the_routing_groups_source_lists_members_under_the_group_without_access_groups(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS), routing_groups=[dict(ROUTING_GROUP)]) + rows: Final = core.registry().deployments(source="routing_groups") + assert _rows(rows) == [("grp", "a"), ("grp", "b")] + assert isinstance(rows, tuple) + assert all("access_groups" not in row["model_info"] for row in rows) + + +def test_the_routed_group_of_a_team_public_name_is_the_teams_own_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + assert _rows(core.registry().deployments("pub", source="routed_group", team_id="team-y")) == [("internal-y", "ty")] + assert core.registry().deployments("pub", source="routed_group") == [] + + +def test_the_routable_source_lists_the_deployments_routing_would_consider(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + registry: Final = core.registry() + assert _ids(registry.deployments("g", source="routable")) == ["a", "b"] + assert registry.deployments("nope", source="routable") == () + assert _ids(registry.deployments("a", source="routable")) == ["a"] + team_request: Final = {"metadata": {"user_api_key_team_id": "team-x"}} + assert _ids(registry.deployments("pub", source="routable", request=team_request)) == ["tx"] + + +def test_the_routable_source_answers_nothing_for_a_name_several_teams_share_for_an_admin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + admin: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")}} # pyright: ignore[reportCallIssue] + assert core.registry().deployments("pub", source="routable", request=admin) == () + + +def test_the_for_request_source_finds_one_teams_public_name_for_a_proxy_admin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[TEAMS_AND_GROUPS[0], *TEAMS_AND_GROUPS[2:]]) + admin: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")}} # pyright: ignore[reportCallIssue] + assert _ids(core.registry().deployments("pub", source="for_request", request=admin)) == ["tx"] + assert core.registry().deployments("pub", source="for_request", request={}) == [] + + +def test_the_for_request_source_refuses_a_name_several_teams_share_for_a_proxy_admin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + admin: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")}} # pyright: ignore[reportCallIssue] + with pytest.raises(litellm.BadRequestError) as raised: + core.registry().deployments("pub", source="for_request", request=admin) + assert str(raised.value) == ( + "litellm.BadRequestError: Model name 'pub' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ) + + +@pytest.mark.parametrize("client_type", ["async", None]) +def test_a_pod_discovery_deployment_has_no_cached_client(make_core: MakeCore, client_type: str | None) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + pod_discovered: Final = deployment("g", "k", kubernetes_pod_discovery=True) + assert core.registry().client(pod_discovered, client_type=client_type) is None + + +# --- PTU shares ---------------------------------------------------------------------------------- + + +def _ptu_shared(model_name: str, deployment_id: str) -> dict[str, object]: # mutable-ok: a model_list entry + return deployment( + model_name, + deployment_id, + model=f"azure/{deployment_id}", + api_base="https://example.openai.azure.com", + model_info={ + "ptu_shares": {"TA": 50}, + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "base_model": "gpt-4o", + }, + ) + + +PTU_MODELS: Final = [_ptu_shared("gpt", "shared"), _ptu_shared("solo", "s2"), deployment("solo", "open")] + + +def test_a_share_holding_team_has_a_ptu_ceiling_on_its_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PTU_MODELS)) + expected: Final = PTUTeamCeiling( + model_group="gpt", tpm_limit=125000, output_to_input_ratio=4.0, cached_input_ratio=0.0 + ) + assert core.health().team_ptu_ceiling("TA", "gpt") == expected + assert core.health().team_ptu_ceiling("TA", "solo") == PTUTeamCeiling( + model_group="solo", tpm_limit=125000, output_to_input_ratio=4.0, cached_input_ratio=0.0 + ) + assert core.health().team_ptu_ceiling("TB", "gpt") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "group", "expected"), + [ + ({"metadata": {"user_api_key_team_id": "TA"}}, "gpt", "shared"), + ({"metadata": {"user_api_key_team_id": "TA"}}, "solo", "s2"), + ({"metadata": {"user_api_key_team_id": "TB"}}, "solo", "open"), + ({}, "solo", "open"), + ], + ids=["holder-shared", "holder-prefers-its-share", "other-team-open", "no-team-open"], +) +async def test_ptu_shares_route_holders_to_their_share_and_others_to_open_deployments( + make_core: MakeCore, request_kwargs: dict[str, object], group: str, expected: str +) -> None: + core: Final = make_core(model_list=list(PTU_MODELS)) + assert await picked_ids(core, group, request_kwargs=request_kwargs) == {expected} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", [{"metadata": {"user_api_key_team_id": "TB"}}, {}], ids=["other", "none"]) +async def test_a_group_only_share_holders_may_use_refuses_everyone_else( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=list(PTU_MODELS)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("gpt", dict(request_kwargs)) + assert str(raised.value) == ( + "litellm.BadRequestError: Deployment gpt is reserved for the teams holding a PTU share of it" + ) diff --git a/tests/unit/test_router/core/test_gaps_l14000_selection.py b/tests/unit/test_router/core/test_gaps_l14000_selection.py new file mode 100644 index 00000000000..8cff6bcdbe5 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l14000_selection.py @@ -0,0 +1,491 @@ +""" +Gaps in the characterization of ``router_python_core.py`` lines 14000-17669 on the selection path: +health-check routing (the cooldown bypass under an allowed_fails_policy, health-check probes, the +sync filter the pass-through pick uses), pass-through picks (paused deployments by id, priority +groups, cooling, other strategies, the async entry point), the failure dispatch to a request's +logging object, ``model_alias_map``, access groups and litellm model names, tier ``litellm_params`` +(dropped params, effort carriers, restored output ceilings), routing-decision redaction and the +Claude Code session binding's guards. +""" + +import asyncio +import time +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.router import AllowedFailsPolicy, RouterRateLimitError +from tests.unit.test_router.core.scripted_host import ScriptedHost, auth_error, deployment, rate_limit +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, ids_of, picked_ids +from tests.unit.test_router.core.test_selection_characterization_prerouting import ( + HARD_MESSAGES, + SESSION_ID, + TIERS, + _complexity_router, + _session_request, + _tiered, +) + +SPREAD_PICKS: Final = 40 + + +def _unhealthy(*ids: str) -> dict[str, dict[str, object]]: # mutable-ok: health states + return {i: {"is_healthy": False, "timestamp": time.time(), "reason": "probe failed"} for i in ids} + + +def _metadata(request: Mapping[str, object], bucket: str = "metadata") -> Mapping[str, object]: + metadata: Final = request.get(bucket) + assert isinstance(metadata, Mapping) + return metadata + + +# --- health-check routing ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_health_check_routing_with_an_allowed_fails_policy_ignores_cooldowns_when_all_cool( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + enable_health_check_routing=True, + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + cooldown_time=60, + ) + assert [core.health().record_failure(i, rate_limit(), 429) for i in ("a", "b")] == [True, True] + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_health_check_routing_without_a_policy_keeps_honoring_cooldowns(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], enable_health_check_routing=True, cooldown_time=60 + ) + for deployment_id in ("a", "b"): + core.health().record_failure(deployment_id, rate_limit(), 429) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g") + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 60 seconds. All deployments for selected " + "model are in cooldown. Passed model=g. pre-call-checks=False, cooldown_list=['a', 'b']" + ) + + +@pytest.mark.asyncio +async def test_a_health_check_probe_gets_no_deployments_when_all_are_unhealthy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(_unhealthy("a", "b")) # pyright: ignore[reportArgumentType] # health states + assert ids_of(await core.select("g", kind="healthy")) == ["a", "b"] + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", kind="healthy", health_check_probe=True) # pyright: ignore[reportCallIssue] + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=g. " + "pre-call-checks=False, cooldown_list=[]" + ) + + +@pytest.mark.asyncio +async def test_a_health_check_probe_still_drops_only_the_unhealthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(_unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + assert ids_of(await core.select("g", kind="healthy", health_check_probe=True)) == ["b"] # pyright: ignore[reportCallIssue] + + +PASS_THROUGH_PAIR: Final = [ + deployment("g", "pt", use_in_pass_through=True), + deployment("g", "pt2", use_in_pass_through=True), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ({}, {"pt2"}), + ({"allowed_fails_policy": {"RateLimitErrorAllowedFails": 1}}, {"pt", "pt2"}), + ({"background_health_check_model_groups": ["other"]}, {"pt", "pt2"}), + ( + {"background_health_check_model_groups": ["g"], "allowed_fails_policy": {"RateLimitErrorAllowedFails": 1}}, + {"pt2"}, + ), + ], + ids=["filtered", "policy-disables", "group-not-listed", "listed-group-with-policy"], +) +async def test_pass_through_selection_applies_the_health_check_filter( + make_core: MakeCore, arguments: dict[str, object], expected: set[str] +) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH_PAIR), enable_health_check_routing=True, **arguments) + core.health().health_checks.set_deployment_health_states(_unhealthy("pt")) # pyright: ignore[reportArgumentType] # health states + assert await picked_ids(core, "g", times=SPREAD_PICKS, kind="pass_through") == expected + + +@pytest.mark.asyncio +async def test_pass_through_selection_uses_every_deployment_when_all_are_unhealthy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH_PAIR), enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(_unhealthy("pt", "pt2")) # pyright: ignore[reportArgumentType] # health states + assert await picked_ids(core, "g", times=SPREAD_PICKS, kind="pass_through") == {"pt", "pt2"} + + +# --- pass-through selection ---------------------------------------------------------------------- + +PAUSED_PASS_THROUGH: Final = [ + deployment("g", "pt", use_in_pass_through=True, model_info={"blocked": True}), + deployment("g", "pt2", use_in_pass_through=True), +] + +PAUSED_MESSAGE: Final = ( + # Oddity: the message names the deployment's litellm model, not the id that was asked for. + "litellm.ServiceUnavailableError: Model 'openai/pt' is currently paused and cannot accept requests." +) + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_paused_deployment_by_id_is_unavailable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PAUSED_PASS_THROUGH)) + with pytest.raises(litellm.ServiceUnavailableError) as raised: + await core.select("pt", kind="pass_through") + assert str(raised.value) == PAUSED_MESSAGE + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_of_a_paused_deployment_by_id_is_unavailable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PAUSED_PASS_THROUGH)) + with pytest.raises(litellm.ServiceUnavailableError) as raised: + await core.async_get_available_deployment_for_pass_through("pt", {}) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == PAUSED_MESSAGE + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_by_id_returns_a_pass_through_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PAUSED_PASS_THROUGH)) + chosen: Final = await core.async_get_available_deployment_for_pass_through("pt2", {}) # pyright: ignore[reportAttributeAccessIssue] + assert chosen["model_info"]["id"] == "pt2" + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_by_id_refuses_a_plain_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.async_get_available_deployment_for_pass_through("plain", {}) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == ( + "litellm.BadRequestError: Deployment plain does not support pass-through endpoint (use_in_pass_through=False)" + ) + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_picks_only_pass_through_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")]) + for _ in range(10): + chosen = await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + assert chosen["model_info"]["id"] == "pt" + + +@pytest.mark.asyncio +async def test_pass_through_selection_with_a_selector_strategy_picks_through_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH_PAIR), routing_strategy="least-busy") + assert await picked_ids(core, "g", kind="pass_through") == {"pt"} + chosen: Final = await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + assert chosen["model_info"]["id"] == "pt" + + +PRIORITY_PASS_THROUGH: Final = [ + deployment("m1", "a", use_in_pass_through=True), + deployment("m2", "b", use_in_pass_through=True), +] +PRIORITY_GROUP: Final = { + "group_name": "pg", + "models": ["m1", "m2"], + "routing_strategy": "priority", + "model_priorities": {"m1": 1, "m2": 2}, +} + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_priority_group_keeps_the_lowest_priority(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PRIORITY_PASS_THROUGH), routing_groups=[dict(PRIORITY_GROUP)]) + assert await picked_ids(core, "pg", kind="pass_through") == {"a"} + assert await picked_ids(core, "pg", request_kwargs={"_target_order": 2}, kind="pass_through") == {"b"} + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_priority_order_nothing_has_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PRIORITY_PASS_THROUGH), routing_groups=[dict(PRIORITY_GROUP)]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("pg", {"_target_order": 3}, kind="pass_through") + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=pg. " + "pre-call-checks=False, cooldown_list=[]" + ) + + +@pytest.mark.asyncio +async def test_pass_through_selection_with_every_pass_through_deployment_cooling_rate_limits( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")], cooldown_time=60 + ) + assert core.health().record_failure("pt", auth_error(), 401) is True + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", kind="pass_through") + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 60 seconds. All deployments for selected " + "model are in cooldown. Passed model=g. pre-call-checks=False, cooldown_list=['pt']" + ) + + +# --- failure dispatch to the request's logging object -------------------------------------------- + + +class _RecordingLoggingObject: + """Stands in for a request's ``litellm_logging_obj``: records the failure dispatches.""" + + def __init__(self) -> None: + self.dispatched: list[dict[str, object]] = [] # mutable-ok: the dispatches, in order + + async def dispatch_failure_handlers(self, **kwargs: object) -> None: # kwargs-ok: the dispatch arguments + self.dispatched.append(kwargs) + + +async def _let_tasks_run() -> None: + for _ in range(3): + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_refused_pick_is_dispatched_to_the_requests_logging_object(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + logging_obj: Final = _RecordingLoggingObject() + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("nope", {"litellm_logging_obj": logging_obj}) + await _let_tasks_run() + (dispatched,) = logging_obj.dispatched + assert dispatched["exception"] is raised.value + assert dispatched["prefer_async_handlers"] is True + assert "Traceback" in str(dispatched["traceback_exception"]) + + +@pytest.mark.asyncio +async def test_a_refused_async_pass_through_pick_is_dispatched_to_the_requests_logging_object( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + logging_obj: Final = _RecordingLoggingObject() + with pytest.raises(litellm.BadRequestError) as raised: + await core.async_get_available_deployment_for_pass_through("g", {"litellm_logging_obj": logging_obj}) # pyright: ignore[reportAttributeAccessIssue] + await _let_tasks_run() + assert [entry["exception"] for entry in logging_obj.dispatched] == [raised.value] + + +# --- model_alias_map, access groups and litellm model names -------------------------------------- + + +@pytest.mark.asyncio +async def test_the_process_model_alias_map_renames_the_model_in_a_no_deployments_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(model_alias_map={"g": "renamed"}) + core: Final = make_core(model_list=[deployment("g", "a")], cooldown_time=60) + assert (await core.select("g"))["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + core.health().record_failure("a", auth_error(), 401) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g") + assert "Passed model=renamed." in str(raised.value) + + +ACCESS_GROUPED: Final = [ + deployment("g", "in_a", model_info={"access_groups": ["grpA"]}), + deployment("g", "in_b", model_info={"access_groups": ["grpB"]}), + deployment("h", "hh"), +] + + +@pytest.mark.asyncio +async def test_an_access_group_only_caller_reaches_any_deployment_by_its_litellm_model(make_core: MakeCore) -> None: + # Oddity: the access-group narrowing looks up access groups by the requested name, and a + # litellm model string names no group, so a grpA-only caller reaches the grpB deployment. + core: Final = make_core(model_list=list(ACCESS_GROUPED), default_fallbacks=["h"]) + auth: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(models=["grpA"])}} # pyright: ignore[reportCallIssue] + assert (await core.select("openai/in_b", dict(auth)))["model_info"]["id"] == "in_b" # pyright: ignore[reportIndexIssue] + assert (await core.select("openai/in_a", dict(auth)))["model_info"]["id"] == "in_a" # pyright: ignore[reportIndexIssue] + assert (await core.select("unknown", dict(auth)))["model_info"]["id"] == "hh" # pyright: ignore[reportIndexIssue] + + +# --- tier litellm_params ------------------------------------------------------------------------- + + +def _pinned(**tier_params: object) -> dict[str, object]: # mutable-ok: tiers # kwargs-ok: tier litellm_params + return {**TIERS, "SIMPLE": {"model_name": "small", "litellm_params": tier_params}} + + +@pytest.mark.asyncio +async def test_a_tier_param_no_deployment_supports_is_dropped(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s", model="anthropic/claude-3-5-sonnet-20240620"), + deployment("big", "b"), + _complexity_router(tiers=_pinned(logit_bias={"1": 2}, temperature=0.2)), + ] + ) + request: Final[dict[str, object]] = {} + chosen: Final = await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + assert request["temperature"] == 0.2 + assert "logit_bias" not in request + + +@pytest.mark.asyncio +async def test_a_tier_param_the_request_allowlists_is_kept(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s", model="anthropic/claude-3-5-sonnet-20240620"), + deployment("big", "b"), + _complexity_router(tiers=_pinned(logit_bias={"1": 2})), + ] + ) + request: Final[dict[str, object]] = {"allowed_openai_params": ["logit_bias"]} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["logit_bias"] == {"1": 2} + + +@pytest.mark.asyncio +async def test_a_tier_reasoning_effort_removes_the_callers_other_effort_carriers(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=_pinned(reasoning_effort="low"))) + request: Final[dict[str, object]] = { + "thinking": {"type": "enabled"}, + "output_config": {"effort": "high", "format": "json"}, + "reasoning": {"effort": "high"}, + "max_completion_tokens": 100, + } + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["reasoning_effort"] == "low" + assert "thinking" not in request + assert request["output_config"] == {"format": "json"} + assert "reasoning" not in request + assert request["max_completion_tokens"] == 100 + + +@pytest.mark.asyncio +async def test_a_tier_without_an_output_ceiling_leaves_the_callers_ceiling(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=_pinned(temperature=0.1))) + request: Final[dict[str, object]] = {"max_completion_tokens": 100} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["temperature"] == 0.1 + assert request["max_completion_tokens"] == 100 + assert "max_tokens" not in request + # The caller's ceiling is stamped on every routed pick, pinned or not. + assert _metadata(request)["_client_output_ceiling"] == {"max_completion_tokens": 100} + + +@pytest.mark.asyncio +async def test_a_later_pick_restores_the_callers_ceiling_a_tier_pin_replaced(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=_pinned(max_tokens=7))) + request: Final[dict[str, object]] = {"max_completion_tokens": 100} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["max_tokens"] == 7 + assert "max_completion_tokens" not in request + assert _metadata(request)["_client_output_ceiling"] == {"max_completion_tokens": 100} + + chosen: Final = await core.select("smart", request, messages=HARD_MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + assert request["max_completion_tokens"] == 100 + assert "max_tokens" not in request + + +@pytest.mark.asyncio +async def test_a_second_pick_leaves_the_first_picks_forwarded_marker_params_untracked(make_core: MakeCore) -> None: + # Oddity: forwarded marker params already on the request are not "newly forwarded" on the + # next attempt, so the tracking key is dropped while the params themselves stay behind. + core: Final = make_core(model_list=_tiered(temperature=0.3)) + request: Final[dict[str, object]] = {} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + forwarded: Final = request["_alias_marker_forwarded_params"] + assert isinstance(forwarded, tuple) + assert "temperature" in forwarded + assert request["temperature"] == 0.3 + + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert "_alias_marker_forwarded_params" not in request + assert request["temperature"] == 0.3 + + +# --- routing decision redaction ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_the_routing_decision_keeps_its_prompt_signals_by_default(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + decision: Final = _metadata(request)["routing_decision"] + assert isinstance(decision, Mapping) + assert "signals" in decision + + +@pytest.mark.asyncio +async def test_redacted_message_logging_drops_the_prompt_quoting_routing_decision_fields(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {"standard_callback_dynamic_params": {"turn_off_message_logging": True}} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + decision: Final = _metadata(request)["routing_decision"] + assert isinstance(decision, Mapping) + assert "signals" not in decision + assert (decision["router_model_name"], decision["tier"], decision["routed_model"]) == ("smart", "SIMPLE", "small") + + +@pytest.mark.asyncio +async def test_routing_metadata_goes_to_litellm_metadata_when_the_request_has_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {"litellm_metadata": {}} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert "routing_decision" in _metadata(request, "litellm_metadata") + assert "metadata" not in request + + +# --- the Claude Code session binding's guards ---------------------------------------------------- + + +async def _subagent_pick(core: object) -> str: + chosen: Final = await core.select("big", _session_request(agent_id="agent-1"), messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + return str(chosen["model_info"]["id"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("depth", "expected"), [(1, "b"), (0, "s")], ids=["fallback-attempt", "first-attempt"]) +async def test_only_a_first_attempt_binds_its_session(make_core: MakeCore, depth: int, expected: str) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["fallback_depth"] = depth + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == expected + + +@pytest.mark.asyncio +async def test_a_request_without_a_caller_hash_never_binds(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["metadata"] = {} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == "b" + + +@pytest.mark.asyncio +async def test_the_caller_hash_is_read_from_litellm_metadata_when_present(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["metadata"] = {} + first["litellm_metadata"] = {"user_api_key_hash": "caller-hash"} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == "s" + + +@pytest.mark.asyncio +async def test_the_session_header_name_is_matched_case_insensitively(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["proxy_server_request"] = {"headers": {"X-Claude-Code-Session-Id": SESSION_ID}} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == "s" diff --git a/tests/unit/test_router/core/test_gaps_l2800.py b/tests/unit/test_router/core/test_gaps_l2800.py new file mode 100644 index 00000000000..78e237f0696 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l2800.py @@ -0,0 +1,431 @@ +""" +Gap tests for ``router_python_core.py`` lines 2800-5600: optional pre-call check re-registration, +silent experiments, per-deployment ``num_retries`` stamped on errors, and the request kwargs +``_update_kwargs_with_deployment`` builds (Kubernetes pod routing, baggage stamps, stream usage +injection, malformed request tools, per-user GitHub Copilot credentials). + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import asyncio +import threading +import time +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.constants import BAGGAGE_MAPPED_TAGS_METADATA_KEY, KUBERNETES_POD_ROUTING_KEY +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD +from litellm.types.router import BaggageMappedTagsStamp +from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, host_callback_names +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +@pytest.fixture +def host() -> ScriptedHost: + return CallbackHost() + + +def _metadata(kwargs: Mapping[str, object], name: str = "metadata") -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +async def _settle() -> None: + for _ in range(30): + await asyncio.sleep(0) + + +# --- optional pre-call checks added again at runtime (add_optional_pre_call_checks) ------------ + + +def test_runtime_affinity_checks_merge_into_the_existing_affinity_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], optional_pre_call_checks=["deployment_affinity"]) + (existing,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["deployment_affinity", "session_affinity"]})) + (merged,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert merged is existing + assert ( + merged.enable_user_key_affinity, + merged.enable_responses_api_affinity, + merged.enable_session_id_affinity, + ) == ( + True, + False, + True, + ) + assert host_callback_names(host)["callbacks"] == ["DeploymentAffinityCheck"] + + +def test_runtime_affinity_flags_are_only_ever_turned_on(make_core: MakeCore) -> None: + """Oddity: a runtime list without ``deployment_affinity`` leaves the flag it set on, since the + existing callback's flags are OR-ed with the new list's.""" + core: Final = make_core(model_list=[deployment("g", "a")], optional_pre_call_checks=["deployment_affinity"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["responses_api_deployment_check"]})) + (merged,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert ( + merged.enable_user_key_affinity, + merged.enable_responses_api_affinity, + merged.enable_session_id_affinity, + ) == ( + True, + True, + False, + ) + + +def test_runtime_affinity_merge_takes_the_model_group_affinity_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + optional_pre_call_checks=["deployment_affinity"], + model_group_affinity_config={"g": ["session_affinity"]}, + ) + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["session_affinity"]})) + (merged,) = [cb for cb in core.optional_callbacks if type(cb).__name__ == "DeploymentAffinityCheck"] # pyright: ignore[reportAttributeAccessIssue] + assert merged.model_group_affinity_config == {"g": ["session_affinity"]} + assert merged.enable_session_id_affinity is True + + +def test_runtime_encrypted_content_affinity_reuses_the_existing_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], optional_pre_call_checks=["encrypted_content_affinity"]) + (existing,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["encrypted_content_affinity"]})) + (kept,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert kept is existing + assert kept.enable_global_affinity is True + assert host_callback_names(host)["callbacks"] == ["EncryptedContentAffinityCheck"] + + +def test_runtime_encrypted_content_affinity_turns_global_affinity_on_for_a_group_scoped_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], model_group_affinity_config={"g": ["encrypted_content_affinity"]} + ) + (existing,) = [cb for cb in core.optional_callbacks if type(cb).__name__ == "EncryptedContentAffinityCheck"] # pyright: ignore[reportAttributeAccessIssue] + assert existing.enable_global_affinity is False + core.add_optional_pre_call_checks(["encrypted_content_affinity"]) # pyright: ignore[reportAttributeAccessIssue] + assert existing.enable_global_affinity is True + assert host_callback_names(host)["callbacks"] == ["EncryptedContentAffinityCheck", "DeploymentAffinityCheck"] + + +async def test_router_budget_limiting_is_not_added_twice_beside_a_provider_budget_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}, + optional_pre_call_checks=["router_budget_limiting"], + ) + limiters: Final = [cb for cb in core.optional_callbacks if isinstance(cb, RouterBudgetLimiting)] # pyright: ignore[reportAttributeAccessIssue] + assert len(limiters) == 1 + assert core.registry().plugins("budget_limiter") is limiters[0] + core.retire() + + +# --- silent experiments ----------------------------------------------------------------------- + + +def _wait_for_calls(host: ScriptedHost, count: int) -> None: + deadline: Final = time.monotonic() + 5 + while len(host.calls) < count and time.monotonic() < deadline: + time.sleep(0.01) + + +def test_sync_completion_mirrors_to_the_silent_model_on_a_background_thread( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + core.completion(model="g", messages=MESSAGES, metadata={"trace": "t"}) # pyright: ignore[reportAttributeAccessIssue] + _wait_for_calls(host, 2) + # the thread starts before the primary call, so the two calls can land in either order + assert sorted((call.op, call.deployment_id) for call in host.calls) == [("acompletion", "s1"), ("completion", "d1")] + primary: Final = next(call for call in host.calls if call.deployment_id == "d1") + shadow: Final = _metadata(next(call for call in host.calls if call.deployment_id == "s1").kwargs) + assert (shadow["is_silent_experiment"], shadow["model_group"], shadow["trace"]) == (True, "shadow", "t") + assert "is_silent_experiment" not in _metadata(primary.kwargs) + + +async def test_a_request_already_marked_silent_does_not_start_another_experiment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"is_silent_experiment": True}) # pyright: ignore[reportAttributeAccessIssue] + await _settle() + assert host.deployments_called() == ["d1"] + + +def test_a_sync_request_already_marked_silent_does_not_start_another_experiment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + core.completion(model="g", messages=MESSAGES, metadata={"is_silent_experiment": True}) # pyright: ignore[reportAttributeAccessIssue] + time.sleep(0.2) + assert host.deployments_called() == ["d1"] + + +async def test_silent_experiment_marks_its_own_metadata_when_the_request_metadata_cannot_be_deep_copied( + make_core: MakeCore, host: ScriptedHost +) -> None: + lock: Final = threading.Lock() + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"lock": lock}) # pyright: ignore[reportAttributeAccessIssue] + await _settle() + primary: Final = next(call for call in host.calls if call.deployment_id == "d1") + shadow: Final = next(call for call in host.calls if call.deployment_id == "s1") + assert "is_silent_experiment" not in _metadata(primary.kwargs) + assert _metadata(primary.kwargs)["model_group"] == "g" + assert (_metadata(shadow.kwargs)["is_silent_experiment"], _metadata(shadow.kwargs)["lock"]) == (True, lock) + + +async def test_a_streaming_request_mirrors_a_streaming_silent_experiment( + make_core: MakeCore, host: ScriptedHost +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + host.answer("openai/d1", chat_stream("A", model="d1")) + host.answer("openai/s1", chat_stream("S", model="s1")) + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + assert await collect_texts(stream) == "A" + await _settle() + shadow: Final = next(call for call in host.calls if call.deployment_id == "s1") + assert shadow.kwargs["stream"] is True + + +# --- per-deployment num_retries stamped on the error ------------------------------------------ + + +async def test_a_num_retries_already_on_the_error_wins_over_the_deployment_num_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + error: Final = rate_limit() + error.num_retries = 0 # pyright: ignore[reportAttributeAccessIssue] + host.answer("openai/d1", error) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries=2)]) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 1 + + +async def test_a_string_deployment_num_retries_is_used_as_an_int(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/d1", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries="1")]) + with pytest.raises(litellm.RateLimitError) as raised: + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 2 + assert getattr(raised.value, "num_retries") == 1 + + +async def test_a_non_numeric_deployment_num_retries_is_ignored(make_core: MakeCore, host: ScriptedHost) -> None: + """The router's own default (2 retries) applies when the deployment value cannot be read as an int.""" + host.answer("openai/d1", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries="abc")]) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 3 + + +# --- request kwargs built for the deployment -------------------------------------------------- + + +async def test_responses_replace_kubernetes_pod_routing_in_the_request_metadata( + make_core: MakeCore, host: ScriptedHost +) -> None: + """On the Responses path the routing goes to ``litellm_metadata``; a pod routing key the caller put + in ``metadata`` is overwritten with the serving deployment's (none here).""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] + model="g", input="hi", metadata={KUBERNETES_POD_ROUTING_KEY: {"pod": "stale"}, "keep": 1} + ) + sent: Final = host.calls[0].kwargs + assert _metadata(sent) == {KUBERNETES_POD_ROUTING_KEY: None, "keep": 1} + assert _metadata(sent, "litellm_metadata")[KUBERNETES_POD_ROUTING_KEY] is None + + +async def test_a_baggage_stamp_records_the_serving_deployment_tags(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", tags=["t1", "t2"])]) + stamp: Final = BaggageMappedTagsStamp(tags=("b1",)) + await core.acompletion( # pyright: ignore[reportAttributeAccessIssue] + model="g", messages=MESSAGES, metadata={BAGGAGE_MAPPED_TAGS_METADATA_KEY: stamp, "tags": ["b1"]} + ) + metadata: Final = _metadata(host.calls[0].kwargs) + assert metadata[BAGGAGE_MAPPED_TAGS_METADATA_KEY] == BaggageMappedTagsStamp( + tags=("b1",), serving_deployment_tags=("t1", "t2") + ) + assert metadata["tags"] == ["t1", "t2", "b1"] + + +@pytest.mark.parametrize( + ("request_extra", "injected"), + [ + ({"stream": True}, True), + ({"stream": True, "stream_options": {"include_usage": False}}, False), + ({}, False), + ], +) +async def test_a_deployment_stream_options_include_usage_marks_the_streaming_request_as_injected( + make_core: MakeCore, host: ScriptedHost, request_extra: dict[str, object], injected: bool +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + host.answer("openai/d1", chat_stream("A")) + core: Final = make_core(model_list=[deployment("g", "d1", stream_options={"include_usage": True})]) + response: Final = await core.acompletion(model="g", messages=MESSAGES, **request_extra) # pyright: ignore[reportAttributeAccessIssue] + if request_extra.get("stream"): + await collect_texts(response) + assert host.calls[0].kwargs.get(STREAM_USAGE_INJECTED_FIELD, False) is injected + + +async def test_malformed_request_tools_are_dropped_for_the_deployment_tools_when_drop_params_is_set( + make_core: MakeCore, host: ScriptedHost +) -> None: + tool: Final = {"type": "function", "function": {"name": "dep_tool", "parameters": {}}} + core: Final = make_core(model_list=[deployment("g", "d1", tools=[tool], drop_params=True)]) + await core.acompletion(model="g", messages=MESSAGES, tools="not-a-list", tool_choice="auto") # pyright: ignore[reportAttributeAccessIssue] + sent: Final = host.calls[0].kwargs + assert sent["tools"] == [tool] + assert "tool_choice" not in sent + + +async def test_a_per_user_github_copilot_deployment_rejects_another_credential_name( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[ + deployment( + "g", + "d1", + model="github_copilot/gpt-4o", + litellm_credential_name="copilot-cred", + github_copilot_auth_type="per_user_oauth", + ) + ], + num_retries=0, + ) + with pytest.raises(litellm.BadRequestError) as raised: + await core.acompletion(model="g", messages=MESSAGES, litellm_credential_name="someone-else") # pyright: ignore[reportAttributeAccessIssue] + assert "litellm_credential_name cannot be overridden on a deployment that uses per-user GitHub OAuth" in str( + raised.value + ) + assert host.calls == [] + + +# --- timeouts --------------------------------------------------------------------------------- + + +async def test_the_process_request_timeout_outranks_the_router_timeout_once_a_router_timeout_is_set( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + """Oddity: this per-attempt timeout is read from ``litellm.request_timeout`` itself (not the host's + process settings), and only when the router has a ``timeout`` of its own.""" + monkeypatch.setattr(litellm, "request_timeout", 33.0) + core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[0].kwargs["timeout"] == 33.0 + + +async def test_the_process_request_timeout_is_the_stream_timeout_without_a_stream_timeout( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + monkeypatch.setattr(litellm, "request_timeout", 33.0) + host.answer("openai/d1", chat_stream("A")) + core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) + await collect_texts(await core.acompletion(model="g", messages=MESSAGES, stream=True)) # pyright: ignore[reportAttributeAccessIssue] + assert (host.calls[0].kwargs["stream_timeout"], host.calls[0].kwargs["timeout"]) == (33.0, 33.0) + + +async def test_default_litellm_params_timeout_applies_when_no_router_or_process_timeout_is_set( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(request_timeout=None) + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"timeout": 9}) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[0].kwargs["timeout"] == 9 + + +async def test_default_litellm_params_stream_timeout_is_the_last_stream_timeout_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + host.answer("openai/d1", chat_stream("A")) + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"stream_timeout": 4}) + await collect_texts(await core.acompletion(model="g", messages=MESSAGES, stream=True)) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[0].kwargs["stream_timeout"] == 4 + + +# --- sync completion runs the routing pre-call checks for a model group only ------------------- + + +def test_sync_completion_runs_routing_pre_call_checks_for_a_group_but_not_a_deployment_id( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.integrations.custom_logger import CustomLogger + + checked: Final[list[object]] = [] + + class Recorder(CustomLogger): + def pre_call_check(self, deployment: dict) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + checked.append(deployment["model_info"]["id"]) + return None + + host.lists["callbacks"].append(Recorder()) + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + core.completion(model="d1", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert checked == ["d1"] + assert host.deployments_called() == ["d1", "d1"] + + +def test_a_sync_routing_pre_call_check_error_stops_the_call(make_core: MakeCore, host: ScriptedHost) -> None: + from litellm.integrations.custom_logger import CustomLogger + + class OverLimit(CustomLogger): + def pre_call_check(self, deployment: dict) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + raise litellm.RateLimitError(message="over the rpm limit", llm_provider="openai", model="d1") + + host.lists["callbacks"].append(OverLimit()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.RateLimitError, match="over the rpm limit"): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +# --- per-model call counters ------------------------------------------------------------------ + + +async def test_acompletion_counts_total_success_and_failed_calls_per_deployment_model( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + host.answer("openai/d1", rate_limit()) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + counts: Final = ( + core.total_calls["openai/d1"], # pyright: ignore[reportAttributeAccessIssue] + core.success_calls["openai/d1"], # pyright: ignore[reportAttributeAccessIssue] + core.fail_calls["openai/d1"], # pyright: ignore[reportAttributeAccessIssue] + ) + assert counts == (2, 1, 1) + + +def test_sync_completion_does_not_count_calls(make_core: MakeCore, host: ScriptedHost) -> None: + """Oddity: only the async path updates ``total_calls``/``success_calls``/``fail_calls``.""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert core.total_calls["openai/d1"] == 0 # pyright: ignore[reportAttributeAccessIssue] diff --git a/tests/unit/test_router/core/test_gaps_l2800_streaming.py b/tests/unit/test_router/core/test_gaps_l2800_streaming.py new file mode 100644 index 00000000000..42c95c6f894 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l2800_streaming.py @@ -0,0 +1,177 @@ +""" +Gap tests for the stream wrappers of ``router_python_core.py`` (lines 2800-5600): what a mid-stream +fallback yields when the fallback answers without a stream, the usage on fallback chunks, and the +Responses stream's handling of an error that is not a mid-stream fallback error. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +from collections.abc import AsyncIterator, Callable +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.utils import ModelResponseStream, Usage +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + bad_request, + completion_response, + deployment, + rate_limit, +) +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGES, + Event, + chat_stream, + drain_into, + finish_chunk, + responses_stream, + text_chunk, +) + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_AB, "fallbacks": [{"a": ["b"]}], "num_retries": 0, **arguments}) + + +def stream_with_usage(sync: bool = False) -> Callable[..., object]: + """A bare chunk iterator (not a ``CustomStreamWrapper``, which would strip the usage) of one text + chunk and a finish chunk that carries usage.""" + + def answer(**_request: object) -> object: # kwargs-ok: the host passes the request kwargs + chunks: Final = [ + text_chunk("B"), + finish_chunk(usage={"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}), + ] + + async def achunks() -> AsyncIterator[ModelResponseStream]: + for chunk in chunks: + yield chunk + + return iter(chunks) if sync else achunks() + + return answer + + +# --- a fallback that answers without a stream -------------------------------------------------- + + +async def test_a_mid_stream_fallback_that_answers_without_a_stream_yields_one_none_item( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the fallback's whole response is dropped and the stream yields ``None`` in its place.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", completion_response("not streamed", model="b")) + core: Final = build(make_core) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + await drain_into(stream, seen) + assert seen[-1] is None + assert host.deployments_called() == ["a", "b"] + + +def test_a_sync_mid_stream_fallback_that_answers_without_a_stream_yields_the_response_fields( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: a ``ModelResponse`` is iterable (pydantic yields ``(field, value)`` pairs), so the sync + wrapper streams the fallback response's fields instead of chunks.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True, sync=True)) + host.answer("openai/b", completion_response("not streamed", model="b")) + core: Final = build(make_core) + stream: Final = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + items: Final = list(stream) # pyright: ignore[reportArgumentType] + fields: Final = [item for item in items if isinstance(item, tuple)] + assert ("model", "b") in fields + assert {name for name, _ in fields} >= {"id", "choices", "usage"} + assert host.deployments_called() == ["a", "b"] + + +# --- usage on fallback chunks ------------------------------------------------------------------ + + +async def test_a_mid_stream_fallback_chunk_dict_usage_becomes_an_all_zero_usage( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: a fallback chunk's usage is rebuilt through ``combine_usage_objects``, which reads token + counts as attributes, so a usage given as a plain dict comes out as an all-zero ``Usage``.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", stream_with_usage()) + core: Final = build(make_core) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + await drain_into(stream, seen) + usages: Final = [ + getattr(item, "usage") + for item in seen + if isinstance(item, ModelResponseStream) and getattr(item, "usage", None) is not None + ] + assert usages + assert isinstance(usages[-1], Usage) + assert (usages[-1].prompt_tokens, usages[-1].completion_tokens, usages[-1].total_tokens) == (0, 0, 0) + + +def test_a_sync_mid_stream_fallback_chunk_dict_usage_becomes_an_all_zero_usage( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: as in the async wrapper, a plain-dict usage comes out as an all-zero ``Usage``.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True, sync=True)) + host.answer("openai/b", stream_with_usage(sync=True)) + core: Final = build(make_core) + stream: Final = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + usages: Final = [ + getattr(item, "usage") + for item in stream # pyright: ignore[reportGeneralTypeIssues] + if isinstance(item, ModelResponseStream) and getattr(item, "usage", None) is not None + ] + assert usages + assert isinstance(usages[-1], Usage) + assert (usages[-1].prompt_tokens, usages[-1].completion_tokens, usages[-1].total_tokens) == (0, 0, 0) + + +# --- Responses streams ------------------------------------------------------------------------- + + +async def test_a_responses_error_that_is_not_a_mid_stream_error_flushes_held_events_then_raises( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + responses_stream(Event("response.created"), Event("response.in_progress"), error=bad_request(message="plain")), + ) + host.answer("openai/b", responses_stream(Event("response.created"), Event("response.completed"))) + core: Final = build(make_core) + stream: Final = await core.aresponses(model="a", input="hi", stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + with pytest.raises(litellm.BadRequestError, match="plain"): + await drain_into(stream, seen) + assert [getattr(event, "type") for event in seen] == ["response.created", "response.in_progress"] + assert host.deployments_called() == ["a"] + + +async def test_a_responses_stream_fallback_that_answers_without_a_stream_yields_that_response( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.exceptions import MidStreamFallbackError + + cut: Final = MidStreamFallbackError( + message="stream cut", + model="a", + llm_provider="anthropic", + original_exception=rate_limit(message="stream cut"), + generated_content="", + is_pre_first_chunk=True, + ) + plain: Final = completion_response("not streamed", model="b") + host.answer("openai/a", responses_stream(Event("response.created"), error=cut)) + host.answer("openai/b", plain) + core: Final = build(make_core) + stream: Final = await core.aresponses(model="a", input="hi", stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + await drain_into(stream, seen) + assert seen == [plain] + assert host.deployments_called() == ["a", "b"] diff --git a/tests/unit/test_router/core/test_gaps_l5600.py b/tests/unit/test_router/core/test_gaps_l5600.py new file mode 100644 index 00000000000..83ab8e943b9 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l5600.py @@ -0,0 +1,479 @@ +""" +Gap characterization for ``router_python_core.py`` lines 5600-8400: batch completions, the scheduler, +prompt-management groups, the image/audio/rerank/text/adapter endpoints, generic calls and their +pass-through endpoint rewrite, guardrail selection, embeddings, files and batches. + +Each test pins a behavior the earlier characterization files leave unasserted; oddities are pinned as +Python does them today and called out in the test's docstring. +""" + +import datetime +import io +from collections.abc import Callable, Mapping +from types import SimpleNamespace +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_core import RouterCore +from litellm.types.utils import CredentialItem, LiteLLMBatch +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + deployment, + rate_limit, + server_error, +) + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] +TWO_DEPLOYMENTS: Final = [deployment("g", "d1"), deployment("g", "d2")] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +def _batch() -> LiteLLMBatch: + return LiteLLMBatch( + id="b1", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + ) + + +def _hidden(response: object) -> Mapping[str, object]: + return getattr(response, "_hidden_params") + + +def _headers(response: object) -> Mapping[str, object]: + found: Final = _hidden(response)["additional_headers"] + assert isinstance(found, Mapping) + return found + + +# The routed endpoints that answer through ``async_function_with_fallbacks`` with their own attempt helper. +ROUTED_ENDPOINTS: Final[dict[str, dict[str, object]]] = { + "aimage_generation": {"prompt": "p"}, + "atranscription": {"file": b"audio"}, + "aspeech": {"input": "x", "voice": "alloy"}, + "arerank": {"query": "q", "documents": ["a"]}, + "aadapter_completion": {"adapter_id": "ad", "messages": MESSAGES}, + "atext_completion": {"prompt": "p"}, + "aembedding": {"input": ["a"]}, + "acreate_batch": {"input_file_id": "f", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + "acancel_batch": {"batch_id": "b1"}, +} + + +# --- batch completions (5565-5790) ---------------------------------------------------------- + + +async def test_abatch_completion_with_message_lists_keeps_a_failed_request_in_its_own_row( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", rate_limit("d1", "a failed")) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + rows: Final = await core.abatch_completion(models=["a", "b"], messages=[MESSAGES, MESSAGES]) + assert [[type(result).__name__ for result in row] for row in rows] == [ + ["RateLimitError", "ModelResponse"], + ["RateLimitError", "ModelResponse"], + ] + + +async def test_abatch_completion_with_mixed_message_shapes_returns_none_without_calling( + make_core: MakeCore, host: EndpointHost +): + """Neither every element a dict nor every element a list: no branch matches and None comes back.""" + core: Final = make_core(model_list=[deployment("a", "d1")]) + assert await core.abatch_completion(models=["a"], messages=[MESSAGES, {"role": "user", "content": "x"}]) is None + assert host.calls == [] + + +async def test_fastest_response_forwards_the_stream_flag(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1")]) + await core.abatch_completion_fastest_response(model="a", messages=MESSAGES, stream=True) + assert host.calls[0].kwargs["stream"] is True + + +# --- scheduler (5792-5912) ------------------------------------------------------------------ + + +async def test_schedule_acompletion_forwards_the_stream_flag(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.schedule_acompletion(model="g", messages=MESSAGES, priority=1, stream=True) + assert host.calls[0].kwargs["stream"] is True + + +async def test_a_prioritized_request_for_a_deployment_id_is_queued_and_served_by_it( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(3): + response = await core.acompletion(model="d1", messages=MESSAGES, priority=2) + assert _headers(response)["x-litellm-request-prioritization-used"] is True + assert host.deployments_called() == ["d1", "d1", "d1"] + + +async def test_a_bad_priority_for_a_deployment_id_follows_that_deployments_drop_params( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", drop_params=True), deployment("g", "d2")]) + response: Final = await core.acompletion(model="d1", messages=MESSAGES, priority="high") + assert "x-litellm-request-prioritization-used" not in _headers(response) + assert "priority" not in host.calls[0].kwargs + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'high'"): + await core.acompletion(model="d2", messages=MESSAGES, priority="high") + assert host.deployments_called() == ["d1"] + + +async def test_a_bad_priority_is_dropped_when_any_deployment_of_the_group_drops_params( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", drop_params=True), deployment("g", "d2")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES, priority="high") + assert "x-litellm-request-prioritization-used" not in _headers(response) + assert len(host.calls) == 1 + + +# --- prompt-management groups (5893-6034) --------------------------------------------------- + + +class _PromptLogging(LiteLLMLogging): + """A logging object whose prompt template renders to ``rendered_model`` with a system message.""" + + rendered_model: str = "openai/rendered" + seen: Mapping[str, object] = {} + + def get_chat_completion_prompt(self, **kwargs: object): # kwargs-ok: Logging's own keyword signature + self.seen = kwargs + return self.rendered_model, [{"role": "system", "content": "rendered"}], {"temperature": 0.7} + + +def _prompt_logging(rendered_model: str = "openai/rendered") -> _PromptLogging: + logging_obj: Final = _PromptLogging( + model="pm", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c1", + function_id="f1", + ) + logging_obj.rendered_model = rendered_model + return logging_obj + + +async def test_a_prompt_management_group_renders_the_deployment_prompt_and_calls_the_rendered_model( + make_core: MakeCore, host: EndpointHost +): + """The deployment's prompt_id/variables/label reach the template; the prompt-management config keys + (bitbucket_config, ...) are stripped from what is sent; the rendered model, messages and optional + params are sent directly (no routing hidden params) since the rendered model is no group.""" + core: Final = make_core( + model_list=[ + deployment( + "pm", + "p1", + model="langfuse/gpt-4o", + prompt_id="dep-pid", + prompt_variables={"a": 1}, + prompt_label="lbl", + bitbucket_config={"x": 1}, + timeout=7, + ) + ] + ) + logging_obj: Final = _prompt_logging() + response: Final = await core.acompletion(model="pm", messages=MESSAGES, litellm_logging_obj=logging_obj) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.model_group) == ("acompletion", "openai/rendered", "p1", "pm") + sent: Final = call.kwargs + assert (sent["prompt_id"], sent["prompt_variables"], sent["prompt_label"]) == ("dep-pid", {"a": 1}, "lbl") + assert "bitbucket_config" not in sent + assert (sent["temperature"], sent["timeout"], sent["messages"]) == ( + 0.7, + 7.0, + [{"role": "system", "content": "rendered"}], + ) + assert sent["litellm_logging_obj"] is logging_obj + assert (logging_obj.seen["model"], logging_obj.seen["messages"], logging_obj.seen["prompt_id"]) == ( + "langfuse/gpt-4o", + MESSAGES, + "dep-pid", + ) + assert logging_obj.seen["injected_for_every_deployment"] is True + assert _hidden(response) == {} + + +async def test_a_prompt_rendered_to_another_group_routes_into_it_with_the_request_prompt_id_and_timeout( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[ + deployment("pm", "p1", model="langfuse/gpt-4o", prompt_id="dep-pid", timeout=9), + deployment("rendered-g", "r1"), + ] + ) + response: Final = await core.acompletion( + model="pm", + messages=MESSAGES, + litellm_logging_obj=_prompt_logging("rendered-g"), + prompt_id="req-pid", + timeout=3, + ) + call: Final = host.calls[0] + assert (call.model, call.model_group, call.kwargs["prompt_id"], call.kwargs["timeout"]) == ( + "openai/r1", + "pm", + "req-pid", + 3, + ) + assert _headers(response)["x-litellm-model-group"] == "rendered-g" + + +async def test_a_single_deployment_without_a_provider_prefix_is_not_a_prompt_management_group( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="gpt-4o")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES) + assert (host.calls[0].model, host.calls[0].deployment_id) == ("gpt-4o", "d1") + assert _hidden(response)["model_id"] == "d1" + + +# --- routed endpoints: retries and counters (6036-6603, 7646-7789, 8026-8315) ------------- + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_ENDPOINTS)) +async def test_a_routed_endpoint_retries_a_401_on_the_other_deployment( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + for _ in range(4): + host.calls.clear() + host.answer("*", auth_error(), None) + core = make_core(model_list=TWO_DEPLOYMENTS, num_retries=1, disable_cooldowns=True) + await core.call(endpoint, {"model": "g", **ROUTED_ENDPOINTS[endpoint]}) + assert sorted(host.deployments_called()) == ["d1", "d2"] + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_ENDPOINTS)) +async def test_a_routed_endpoint_ignores_the_deployment_num_retries( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + """Unlike acompletion and the generic endpoints, these attempt helpers never stamp the deployment's + num_retries on the error, so the router's num_retries=0 governs.""" + host.answer("*", server_error()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries=2)], num_retries=0) + with pytest.raises(litellm.InternalServerError): + await core.call(endpoint, {"model": "g", **ROUTED_ENDPOINTS[endpoint]}) + assert len(host.calls) == 1 + assert host.raised[-1].num_retries is None # pyright: ignore[reportAttributeAccessIssue] # stamped by the retry loop + + +async def test_a_generic_endpoint_retries_with_the_deployment_num_retries( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + patch_direct_endpoints(monkeypatch, host) + host.answer("*", server_error()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries=2)], num_retries=0) + with pytest.raises(litellm.InternalServerError) as raised: + await core.aresponses(model="g", input="hi") + assert len(host.calls) == 3 + assert raised.value.num_retries == 2 # pyright: ignore[reportAttributeAccessIssue] # stamped from the deployment + + +async def test_sync_image_generation_is_retried_with_the_router_num_retries(make_core: MakeCore, host: EndpointHost): + host.answer("*", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=2) + with pytest.raises(litellm.RateLimitError): + core.image_generation(model="g", prompt="p") + assert len(host.calls_to("image_generation")) == 3 + + +@pytest.mark.parametrize("endpoint", ["aimage_generation", "image_generation"]) +async def test_image_generation_for_an_unknown_group_raises_without_calling_the_host( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.BadRequestError, match="There are no healthy deployments for this model"): + await core.call(endpoint, {"model": "nope", "prompt": "p"}) + assert host.calls == [] + + +@pytest.mark.parametrize( + ("endpoint", "total_key", "fail_key"), + [ + ("aimage_generation", "openai/d1", "openai/d1"), + ("atranscription", "g", "g"), + ("aspeech", "g", "g"), + ("arerank", "openai/d1", "openai/d1"), + ("aadapter_completion", "openai/d1", "g"), + ("atext_completion", "openai/d1", "g"), + ("aembedding", "openai/d1", "openai/d1"), + ("acreate_batch", "openai/d1", "g"), + ("acancel_batch", "openai/d1", "g"), + ("image_generation", "", ""), + ], +) +async def test_call_counters_are_keyed_differently_per_endpoint( + make_core: MakeCore, host: EndpointHost, endpoint: str, total_key: str, fail_key: str +): + """total_calls/fail_calls are keyed by the deployment model for some endpoints and by the model group + for others (fail_calls even differs from total_calls for some); sync image_generation counts under "".""" + host.answer("*", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=1) + request: Final = ROUTED_ENDPOINTS.get(endpoint, {"prompt": "p"}) + with pytest.raises(litellm.RateLimitError): + await core.call(endpoint, {"model": "g", **request}) + assert dict(core.total_calls) == {total_key: 2} # pyright: ignore[reportAttributeAccessIssue] # a public counter + assert dict(core.fail_calls) == {fail_key: 2} # pyright: ignore[reportAttributeAccessIssue] # a public counter + assert dict(core.success_calls) == {} # pyright: ignore[reportAttributeAccessIssue] # a public counter + + +async def test_success_counters_count_the_deployment_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aspeech(model="g", input="x", voice="alloy") + await core.aembedding(model="g", input=["a"]) + assert dict(core.total_calls) == {"g": 1, "openai/d1": 1} # pyright: ignore[reportAttributeAccessIssue] # a public counter + assert dict(core.success_calls) == {"g": 1, "openai/d1": 1} # pyright: ignore[reportAttributeAccessIssue] # a public counter + + +# --- guardrails (6692-6721) ----------------------------------------------------------------- + + +async def test_aguardrail_picks_among_same_named_guardrails_by_weight(make_core: MakeCore, host: ScriptedHost): + picked: Final[list[object]] = [] + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + selected: Final = kwargs["selected_guardrail"] + assert isinstance(selected, Mapping) + picked.append(selected["litellm_params"]["guardrail"]) + return "verdict" + + guardrails: Final = [ + {"guardrail_name": "gr", "litellm_params": {"guardrail": "never", "weight": 0}}, + {"guardrail_name": "gr", "litellm_params": {"guardrail": "always", "weight": 1}}, + ] + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=guardrails) + for _ in range(10): + await core.aguardrail(guardrail_name="gr", original_function=guard) + assert set(picked) == {"always"} + + +async def test_aguardrail_without_weights_spreads_over_same_named_guardrails(make_core: MakeCore, host: ScriptedHost): + picked: Final[set[object]] = set() + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + selected: Final = kwargs["selected_guardrail"] + assert isinstance(selected, Mapping) + picked.add(selected["litellm_params"]["guardrail"]) + return "verdict" + + guardrails: Final = [ + {"guardrail_name": "gr", "litellm_params": {"guardrail": "a"}}, + {"guardrail_name": "gr", "litellm_params": {"guardrail": "b"}}, + ] + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=guardrails) + for _ in range(60): + await core.aguardrail(guardrail_name="gr", original_function=guard) + assert picked == {"a", "b"} + + +# --- generic calls: pass-through endpoint rewrite (6754-6781, 7561-7627) -------------------- + + +async def test_async_passthrough_rewrites_only_whole_path_segments_naming_the_group( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + patch_direct_endpoints(monkeypatch, host) + core: Final = make_core(model_list=[deployment("g", "d1", model="bedrock/us.anthropic.claude")]) + await core.allm_passthrough_route(model="g", method="POST", endpoint="/model/g/invoke/gg/g:stream") + assert host.calls[0].kwargs["endpoint"] == "/model/us.anthropic.claude/invoke/gg/us.anthropic.claude:stream" + + +async def test_sync_generic_call_replaces_every_substring_of_the_group_name_in_the_endpoint( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + """The sync path uses str.replace with the full deployment model (provider prefix kept), so a group + name inside other path segments is rewritten too, unlike the async path's segment-bounded rewrite.""" + patch_direct_endpoints(monkeypatch, host) + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.responses(model="g", input="hi", endpoint="/gg/g") + call: Final = host.calls[0] + assert (call.model, call.kwargs["endpoint"]) == ("openai/d1", "/openai/d1openai/d1/openai/d1") + + +# --- batches (8112-8362) -------------------------------------------------------------------- + + +async def test_aretrieve_batch_sends_the_deployment_provider_over_the_request_one( + make_core: MakeCore, host: EndpointHost +): + host.answer("*", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("h", "d2", custom_llm_provider="azure")]) + await core.aretrieve_batch(model="g", batch_id="b1", custom_llm_provider="azure") + await core.aretrieve_batch(model="h", batch_id="b1") + assert [call.kwargs["custom_llm_provider"] for call in host.calls] == ["openai", "azure"] + + +async def test_alist_batches_takes_the_first_set_first_id_and_the_last_answer_last_id( + make_core: MakeCore, host: EndpointHost +): + """last_id is overwritten by every answer in deployment order, even with None.""" + host.answer("openai/d1", SimpleNamespace(first_id=None, last_id="b", data=[1], has_more=False)) + host.answer("openai/d2", SimpleNamespace(first_id="c", last_id=None, data=[2], has_more=False)) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + assert await core.alist_batches(model="g") == { + "object": "list", + "data": [1, 2], + "first_id": "c", + "last_id": None, + "has_more": False, + } + + +async def test_acancel_batch_resolves_the_deployment_credential_but_acreate_batch_does_not( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="cred1", + credential_values={"api_key": "sk-cred", "api_base": "https://cred.example"}, + credential_info={}, + ) + ], + ) + host.answer("*", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1", litellm_credential_name="cred1")]) + await core.acancel_batch(model="g", batch_id="b1") + await core.acreate_batch(model="g", **ROUTED_ENDPOINTS["acreate_batch"]) + cancel, create = host.calls[0].kwargs, host.calls[1].kwargs + assert (cancel["api_key"], cancel["api_base"], "litellm_credential_name" in cancel) == ( + "sk-cred", + "https://cred.example", + False, + ) + assert (create["api_key"], create.get("api_base"), create["litellm_credential_name"]) == ("sk-test", None, "cred1") + + +# --- files (7814-7934) ---------------------------------------------------------------------- + + +async def test_acreate_file_passthrough_keeps_the_batch_file_model(make_core: MakeCore, host: EndpointHost): + upload: Final = io.BytesIO(b'{"custom_id": "1", "body": {"model": "g", "messages": []}}\n') + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_file(model="g", file=upload, purpose="batch", passthrough=True) + sent: Final = host.calls[0].kwargs["file"] + assert isinstance(sent, io.BytesIO) + assert b'"model": "g"' in sent.getvalue() diff --git a/tests/unit/test_router/core/test_gaps_l5600_anthropic.py b/tests/unit/test_router/core/test_gaps_l5600_anthropic.py new file mode 100644 index 00000000000..e2b19058a50 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l5600_anthropic.py @@ -0,0 +1,103 @@ +""" +Gap characterization for /v1/messages streaming recovery (``router_python_core.py`` 7106-7455): a +``MidStreamFallbackError`` raised by the source stream itself, a same-group stream retry whose own +attempt raises before opening, and a request-level retry policy that is not a ``RetryPolicy``. +""" + +from collections.abc import Callable +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit, server_error +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGE_START, + MESSAGE_STOP, + MESSAGES, + anthropic_stream, + collect_bytes, + drain_into, + error_frame, + install_orjson_stub_if_missing, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] + + +async def _open(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=5, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def _mid_stream(*, pre_first_chunk: bool) -> MidStreamFallbackError: + return MidStreamFallbackError( + message="bridge dropped", + model="a", + llm_provider="openai", + original_exception=rate_limit("a", "bridge original"), + is_pre_first_chunk=pre_first_chunk, + ) + + +async def test_a_raised_mid_stream_error_not_before_the_first_chunk_flushes_held_frames_and_raises_the_original( + make_core: MakeCore, host: ScriptedHost +): + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=_mid_stream(pre_first_chunk=False))) + host.answer("openai/b", anthropic_stream(MESSAGE_START, text_delta("B"), MESSAGE_STOP)) + core: Final = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream: Final = await _open(core) + seen: Final[list[object]] = [] + with pytest.raises(litellm.RateLimitError, match="bridge original"): + await drain_into(stream, seen) + assert seen == [MESSAGE_START] + assert host.deployments_called() == ["a"] + + +async def test_a_raised_mid_stream_error_before_the_first_chunk_falls_back(make_core: MakeCore, host: ScriptedHost): + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=_mid_stream(pre_first_chunk=True))) + host.answer("openai/b", anthropic_stream(MESSAGE_START, text_delta("B"), MESSAGE_STOP)) + core: Final = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames: Final = await collect_bytes(await _open(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_stream_retry_that_raises_before_opening_spends_a_retry_and_the_next_one_serves( + make_core: MakeCore, host: ScriptedHost +): + host.answer( + "openai/a", + anthropic_stream(MESSAGE_START, error_frame()), + server_error("a"), + anthropic_stream(MESSAGE_START, text_delta("A"), MESSAGE_STOP), + ) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=2) + stream: Final = await _open(core) + frames: Final = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a", "a"] + assert len(host.sleeps) == 2 + headers: Final = getattr(stream, "_hidden_params")["additional_headers"] + assert (headers["x-litellm-attempted-retries"], headers["x-litellm-max-retries"]) == (2, 2) + + +async def test_a_request_retry_policy_that_is_not_a_retry_policy_streams_with_the_plain_budget( + make_core: MakeCore, host: ScriptedHost +): + host.answer( + "openai/a", + anthropic_stream(MESSAGE_START, error_frame("rate_limit_error", "slow")), + anthropic_stream(MESSAGE_START, text_delta("A"), MESSAGE_STOP), + ) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=1) + stream: Final = await _open(core, model_group_retry_policy={"a": 5}) + frames: Final = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a"] + assert getattr(stream, "_hidden_params")["additional_headers"]["x-litellm-max-retries"] == 1 diff --git a/tests/unit/test_router/core/test_gaps_l8400_precall.py b/tests/unit/test_router/core/test_gaps_l8400_precall.py new file mode 100644 index 00000000000..1e64257bac8 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l8400_precall.py @@ -0,0 +1,178 @@ +""" +Gaps in the routing-strategy pre-call checks and the Responses-websocket turn admission of ``PythonCore`` +(router_python_core.py ``routing_strategy_pre_call_checks``, ``async_routing_strategy_pre_call_checks``, +``async_admit_deployment_turn``). A ``CustomLogger`` on the host's ``callbacks`` list (``litellm.callbacks``) +takes part in those checks through ``pre_call_check`` / ``async_pre_call_check``. +""" + +from collections.abc import Callable +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.retries_health_helpers import attempt, attempt_sync, cooled_down +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit + +MakeCore = Callable[..., RouterCore] + + +class PreCallGate(CustomLogger): + """Records the deployment ids its pre-call checks saw and raises ``error`` (when set) from them.""" + + def __init__(self, error: Exception | None = None) -> None: + super().__init__() + self.error = error + self.seen: Final[list[tuple[str, str]]] = [] + + async def async_pre_call_check(self, deployment: dict, parent_otel_span: object) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + self.seen.append(("async", deployment["model_info"]["id"])) + if self.error is not None: + raise self.error + return None + + def pre_call_check(self, deployment: dict) -> dict | None: + self.seen.append(("sync", deployment["model_info"]["id"])) + if self.error is not None: + raise self.error + return None + + +def gated(host: ScriptedHost, error: Exception | None = None) -> PreCallGate: + gate: Final = PreCallGate(error) + host.lists["callbacks"].append(gate) + return gate + + +GROUP: Final = [deployment("g", "a"), deployment("g", "b")] + + +# --- async pre-call checks ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_rate_limit_from_a_pre_call_check_cools_the_deployment_down_and_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUP, num_retries=0, cooldown_time=33) + gate: Final = gated(host, rate_limit(message="gate says no")) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert str(error) == "litellm.RateLimitError: gate says no" + assert host.calls == [] + ((kind, checked),) = gate.seen + assert kind == "async" + cooled: Final = cooled_down(core, ["a", "b"]) + assert list(cooled) == [checked] + assert cooled[checked]["cooldown_time"] == 33 + + +@pytest.mark.asyncio +async def test_any_other_error_from_a_pre_call_check_is_raised_without_a_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUP, num_retries=0, cooldown_time=33) + gate: Final = gated(host, ValueError("gate broke")) + + error: Final = await attempt(core) + + assert type(error) is ValueError + assert str(error) == "gate broke" + assert host.calls == [] + assert len(gate.seen) == 1 + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.asyncio +async def test_a_passing_pre_call_check_sees_the_deployment_that_is_then_called( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUP, num_retries=0) + gate: Final = gated(host) + + await attempt(core) + + assert gate.seen == [("async", host.deployments_called("acompletion")[0])] + + +def test_a_rate_limit_from_a_sync_pre_call_check_is_raised_without_a_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the sync twin raises the gate's error but, unlike the async path, sets no cooldown.""" + core: Final = make_core(model_list=GROUP, num_retries=0, cooldown_time=33) + gate: Final = gated(host, rate_limit(message="gate says no")) + + error: Final = attempt_sync(core) + + assert isinstance(error, litellm.RateLimitError) + assert host.calls == [] + assert [kind for kind, _ in gate.seen] == ["sync"] + assert cooled_down(core, ["a", "b"]) == {} + + +# --- Responses websocket turn admission ---------------------------------------------------------- + + +class TestResponsesWebsocketTurnAdmission: + """The core hands the websocket handler a ``turn_admission`` callable that re-runs the deployment's + admission (budget limiter, routing strategy, pre-call checks) before every later turn.""" + + @pytest.fixture + def host(self) -> EndpointHost: + return EndpointHost() + + @staticmethod + async def open_socket(core: RouterCore, host: EndpointHost) -> Callable[[], object]: + await core.call("_aresponses_websocket", {"model": "g"}) + (call,) = host.calls_to("_aresponses_websocket") + admission: Final = call.kwargs["turn_admission"] + assert callable(admission) + return admission + + @pytest.mark.asyncio + async def test_simple_shuffle_admits_every_turn_and_runs_the_pre_call_checks( + self, make_core: MakeCore, host: EndpointHost + ) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=1)]) + gate: Final = gated(host) + admit: Final = await self.open_socket(core, host) + + admitted: Final = [dict(await admit()) for _ in range(3)] # pyright: ignore[reportGeneralTypeIssues] + + assert admitted == [{}, {}, {}] + assert gate.seen == [("async", "a")] * 4 # the socket opening, then each turn + + @pytest.mark.asyncio + async def test_a_rate_limit_from_a_pre_call_check_refuses_the_turn( + self, make_core: MakeCore, host: EndpointHost + ) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + gate: Final = gated(host) + admit: Final = await self.open_socket(core, host) + gate.error = rate_limit(message="turn refused") + + with pytest.raises(litellm.RateLimitError, match="turn refused"): + await admit() # pyright: ignore[reportGeneralTypeIssues] + + @pytest.mark.asyncio + async def test_usage_based_routing_admits_turns_until_the_deployment_rpm_is_used_up( + self, make_core: MakeCore, host: EndpointHost + ) -> None: + """Opening the socket already counts two requests (its pre-call check and its success), so an + rpm of 5 leaves room for three turns.""" + core: Final = make_core(model_list=[deployment("g", "a", rpm=5)], routing_strategy="usage-based-routing-v2") + admit: Final = await self.open_socket(core, host) + + admitted: Final = [dict(await admit()) for _ in range(3)] # pyright: ignore[reportGeneralTypeIssues] + with pytest.raises(litellm.RateLimitError) as refused: + await admit() # pyright: ignore[reportGeneralTypeIssues] + + assert admitted == [{}, {}, {}] + assert str(refused.value).startswith( + "litellm.RateLimitError: No deployments available for selected model. Passed model=g. Deployments=" + ) diff --git a/tests/unit/test_router/core/test_gaps_l8400_pricing.py b/tests/unit/test_router/core/test_gaps_l8400_pricing.py new file mode 100644 index 00000000000..4c0df5a17b6 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l8400_pricing.py @@ -0,0 +1,240 @@ +""" +Gaps in the built-in pricing a custom-priced deployment inherits when the core registers it in +``litellm.model_cost`` (router_python_core.py ``_inherit_builtin_cache_pricing``, ``_cost_map_backend_model``, +``_inherit_builtin_service_tier_pricing``, ``_inherit_builtin_base_rates_for_off_peak``, +``_inherit_builtin_tiered_output_rate``), and the deployment checks of ``_create_deployment`` +(access windows, PTU configuration). Expected rates are read from the built-in cost map, not hardcoded. +""" + +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from tests.unit.test_router.core.registry_helpers import MakeCore, unique_id +from tests.unit.test_router.core.scripted_host import deployment + +OFF_PEAK: Final = {"hours_utc": "01:00-02:00", "input_cost_per_token": 1e-9} + + +def registered(make_core: MakeCore, model: str, **params: object) -> Mapping[str, object]: # kwargs-ok: litellm_params + """The ``litellm.model_cost`` entry registered under a fresh deployment id.""" + model_id: Final = unique_id("l8400") + make_core(model_list=[deployment("g", model_id, model=model, **params)]) + return litellm.model_cost[model_id] + + +def builtin(key: str) -> Mapping[str, object]: + return litellm.model_cost[key] + + +# --- service-tier rates --------------------------------------------------------------------------- + + +def test_custom_input_pricing_inherits_the_backend_service_tier_rates(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/gpt-4o", input_cost_per_token=1e-3) + + for field in ( + "input_cost_per_token_priority", + "output_cost_per_token_priority", + "cache_read_input_token_cost_priority", + ): + assert entry[field] == builtin("gpt-4o")[field] + assert entry.get("input_cost_per_token_batches") is None # batch rates are not a service tier + + +def test_a_user_set_service_tier_rate_wins_over_the_inherited_one(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/gpt-4o", input_cost_per_token=1e-3, input_cost_per_token_priority=9e-3) + + assert entry["input_cost_per_token_priority"] == 9e-3 + assert entry["output_cost_per_token_priority"] == builtin("gpt-4o")["output_cost_per_token_priority"] + + +def test_output_only_custom_pricing_inherits_service_tiers_but_not_cache_pricing(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/gpt-4o", output_cost_per_token=1e-3) + + assert entry["input_cost_per_token_priority"] == builtin("gpt-4o")["input_cost_per_token_priority"] + assert entry.get("cache_read_input_token_cost") is None + + +def test_a_backend_missing_from_the_cost_map_contributes_nothing(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/l8400-not-a-real-model", input_cost_per_token=1e-3) + + assert entry["input_cost_per_token"] == 1e-3 + assert entry.get("cache_read_input_token_cost") is None + assert entry.get("input_cost_per_token_priority") is None + + +# --- which backend model the rates come from ------------------------------------------------------ + + +@pytest.mark.parametrize( + "params", + [{"model_info": {"base_model": "gpt-5"}}, {"base_model": "gpt-5"}], + ids=["model-info-base-model", "litellm-params-base-model"], +) +def test_the_base_model_names_the_backend_the_rates_are_inherited_from( + make_core: MakeCore, params: dict[str, object] +) -> None: + entry: Final = registered(make_core, "openai/my-fine-tune", input_cost_per_token=1e-3, **params) + + assert entry["cache_read_input_token_cost"] == builtin("gpt-5")["cache_read_input_token_cost"] + assert entry["input_cost_per_token_flex"] == builtin("gpt-5")["input_cost_per_token_flex"] + + +# --- explicit cache-read rate ----------------------------------------------------------------------- + + +def test_explicit_cache_read_rate_is_inherited_only_without_a_user_cache_read_rate(make_core: MakeCore) -> None: + backend: Final = builtin("dashscope/qwen3.7-max") + inherited: Final = registered(make_core, "dashscope/qwen3.7-max", input_cost_per_token=1e-3) + user_set: Final = registered( + make_core, "dashscope/qwen3.7-max", input_cost_per_token=1e-3, cache_read_input_token_cost=1e-8 + ) + + assert inherited["explicit_cache_read_input_token_cost"] == backend["explicit_cache_read_input_token_cost"] + assert inherited["cache_read_input_token_cost"] == backend["cache_read_input_token_cost"] + assert user_set.get("explicit_cache_read_input_token_cost") is None + assert user_set["cache_read_input_token_cost"] == 1e-8 + + +# --- off-peak only ---------------------------------------------------------------------------------- + + +def test_an_off_peak_only_deployment_copies_the_whole_backend_entry(make_core: MakeCore) -> None: + backend: Final = builtin("deepseek-v4-flash") + entry: Final = registered(make_core, "deepseek/deepseek-v4-flash", model_info={"off_peak_pricing": OFF_PEAK}) + + assert entry["off_peak_pricing"] == OFF_PEAK + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost", "max_input_tokens"): + assert entry[field] == backend[field] + assert entry["source"] == backend["source"] + + +def test_an_off_peak_deployment_keeps_its_own_fields(make_core: MakeCore) -> None: + entry: Final = registered( + make_core, "deepseek/deepseek-v4-flash", model_info={"off_peak_pricing": OFF_PEAK, "max_input_tokens": 7} + ) + + assert entry["max_input_tokens"] == 7 + assert entry["input_cost_per_token"] == builtin("deepseek-v4-flash")["input_cost_per_token"] + + +def test_an_off_peak_deployment_with_a_base_rate_copies_nothing_more(make_core: MakeCore) -> None: + entry: Final = registered( + make_core, + "deepseek/deepseek-v4-flash", + model_info={"off_peak_pricing": OFF_PEAK}, + input_cost_per_token=5e-3, + ) + + assert entry["input_cost_per_token"] == 5e-3 + assert entry.get("output_cost_per_token") is None + assert entry.get("max_input_tokens") is None + + +def test_off_peak_pricing_given_in_litellm_params_is_not_registered(make_core: MakeCore) -> None: + """Oddity: only ``model_info.off_peak_pricing`` counts; the same key in litellm_params is dropped.""" + entry: Final = registered(make_core, "deepseek/deepseek-v4-flash", off_peak_pricing=OFF_PEAK) + + assert entry.get("off_peak_pricing") is None + assert entry.get("input_cost_per_token") is None + + +# --- tiered pricing output rate --------------------------------------------------------------------- + + +def test_a_tier_table_without_output_rates_inherits_the_backend_flat_output_rate(make_core: MakeCore) -> None: + entry: Final = registered( + make_core, "openai/gpt-4o", tiered_pricing=[{"input_cost_per_token": 1e-9, "range": [0, 1000]}] + ) + + assert entry["output_cost_per_token"] == builtin("gpt-4o")["output_cost_per_token"] + + +@pytest.mark.parametrize( + ("model", "params"), + [ + ("dashscope/qwen-flash", {"tiered_pricing": [{"input_cost_per_token": 1e-9, "range": [0, 1000]}]}), + ( + "openai/gpt-4o", + {"tiered_pricing": [{"input_cost_per_token": 1e-9, "output_cost_per_token": 2e-9, "range": [0, 1000]}]}, + ), + ( + "openai/gpt-4o", + {"tiered_pricing": [{"input_cost_per_token": 1e-9, "range": [0, 1000]}], "output_cost_per_token": 3e-9}, + ), + ], + ids=["tiered-only-backend", "every-tier-has-an-output-rate", "user-output-rate"], +) +def test_the_tiered_output_rate_is_not_inherited(make_core: MakeCore, model: str, params: dict[str, object]) -> None: + entry: Final = registered(make_core, model, **params) + + assert entry.get("output_cost_per_token") == params.get("output_cost_per_token") + + +# --- _create_deployment checks ------------------------------------------------------------------------ + + +def test_an_invalid_access_windows_config_raises_naming_the_model(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=r"^model 'gpt': invalid model_info\.access_windows: "): + make_core(model_list=[deployment("gpt", "a", model_info={"access_windows": "bad"})]) + + +@pytest.mark.parametrize( + "model_info", + [ + {"access_windows": "bad"}, + {"ptu_shares": {"team-a": 1}}, + ], + ids=["access-windows", "ptu-shares"], +) +def test_ignore_invalid_deployments_drops_a_deployment_with_an_invalid_config( + make_core: MakeCore, model_info: dict[str, object] +) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a", model_info=model_info), deployment("gpt", "b")], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +def test_ptu_shares_without_a_count_and_rate_raise(make_core: MakeCore) -> None: + with pytest.raises( + ValueError, + match=r"^PTU configuration on model 'gpt' is invalid: ptu_count and cost_per_ptu_per_hour are required " + r"when ptu_shares is set", + ): + make_core(model_list=[deployment("gpt", "a", model_info={"ptu_shares": {"team-a": 1}})]) + + +def test_an_inverted_ptu_window_with_shares_raises(make_core: MakeCore) -> None: + with pytest.raises( + ValueError, + match=r"^PTU configuration on model 'gpt' is invalid: ptu_effective_to must be after ptu_effective_from", + ): + make_core( + model_list=[ + deployment( + "gpt", + "a", + model_info={ + "ptu_effective_from": "2026-01-02T00:00:00Z", + "ptu_effective_to": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 1}, + "ptu_count": 1, + "cost_per_ptu_per_hour": 1, + }, + ) + ] + ) + + +def test_a_db_sourced_deployment_skips_the_ptu_config_check(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a", model_info={"ptu_shares": {"team-a": 1}, "db_model": True})] + ) + + assert core.registry().names(kind="ids") == ["a"] diff --git a/tests/unit/test_router/core/test_gaps_l8400_retries.py b/tests/unit/test_router/core/test_gaps_l8400_retries.py new file mode 100644 index 00000000000..34373f6e03e --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l8400_retries.py @@ -0,0 +1,138 @@ +""" +Gaps in the retry loop and the per-attempt bookkeeping of ``PythonCore`` (router_python_core.py +lines 8400-11100): the ``UnsupportedEndpointError`` skip-and-stop rule, the retry breadcrumb cap and +``base_model``, and the request the async failure callback counts toward group usage. +""" + +from typing import Final + +import pytest + +import litellm +from tests.unit.test_router.core.retries_health_helpers import ( + BackgroundSuccessLoggingHost, + MakeCore, + MetadataRecorder, + attempt, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit, server_error + + +def unsupported_endpoint() -> litellm.UnsupportedEndpointError: + return litellm.UnsupportedEndpointError(message="endpoint not supported", llm_provider="openai", model="gpt-4o") + + +# --- UnsupportedEndpointError: retried on the other deployments, stopped once every one refused --------- + + +@pytest.mark.asyncio +async def test_unsupported_endpoint_is_retried_once_per_deployment_then_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + """A 400 is normally never retried; ``UnsupportedEndpointError`` is, but each deployment that refused + is skipped and the loop stops as soon as every deployment of the group refused (num_retries=3 unused).""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=3) + host.answer("*", lambda **_: unsupported_endpoint()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.UnsupportedEndpointError) + called: Final = host.deployments_called("acompletion") + assert sorted(called) == ["a", "b"] + assert host.sleeps == [0] + + +@pytest.mark.asyncio +async def test_unsupported_endpoint_on_a_lone_deployment_is_raised_without_a_retry( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3) + host.answer("*", lambda **_: unsupported_endpoint()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.UnsupportedEndpointError) + assert host.deployments_called("acompletion") == ["a"] + assert host.sleeps == [] + + +# --- retry breadcrumbs ---------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_previous_models_keeps_only_the_four_most_recent_failures( + make_core: MakeCore, host: ScriptedHost +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=6) + host.answer("*", *([recorder.then(server_error)] * 6), recorder.then(None)) + + response: Final = await attempt(core) + + assert not isinstance(response, Exception) + trails: Final = [snapshot["previous_models"] for snapshot in recorder.snapshots] + assert trails[0] is None + assert [len(trail) for trail in trails[1:]] == [1, 2, 3, 4, 4, 4] # pyright: ignore[reportArgumentType] + assert [entry["attempted_retries"] for entry in trails[-1]] == [2, 3, 4, 5] # pyright: ignore[reportGeneralTypeIssues, reportOptionalIterable] + assert [snapshot["request_retry_count"] for snapshot in recorder.snapshots] == [None, 1, 2, 3, 4, 5, 6] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("params", "expected"), + [ + ({"base_model": "azure/gpt-4o"}, "azure/gpt-4o"), + ({"model_info": {"base_model": "azure/gpt-4o-mini"}}, "azure/gpt-4o-mini"), + ({}, None), + ], + ids=["litellm-params-base-model", "model-info-base-model", "none"], +) +async def test_a_retry_breadcrumb_names_the_deployments_base_model( + make_core: MakeCore, host: ScriptedHost, params: dict[str, object], expected: str | None +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=[deployment("g", "a", **params)], num_retries=1) + host.answer("*", recorder.then(server_error), recorder.then(None)) + + await attempt(core) + + (entry,) = recorder.snapshots[1]["previous_models"] # pyright: ignore[reportGeneralTypeIssues] + assert entry["base_model"] == expected # pyright: ignore[reportIndexIssue] + + +# --- the async failure callback counts the failed request --------------------------------------- + + +@pytest.mark.asyncio +async def test_a_failed_call_counts_a_request_but_no_tokens_in_group_usage( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=5)], num_retries=0) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert await core.health().usage("g") == (None, 1) + + +class TestUsageWhenTheHostLogsSuccessAfterTheResponse: + """litellm runs the async success callbacks after the endpoint returned (``ScriptedHost`` runs them + before, which lets the success callback and the router's own count both add the same response).""" + + @pytest.fixture + def host(self) -> BackgroundSuccessLoggingHost: + return BackgroundSuccessLoggingHost() + + @pytest.mark.asyncio + @pytest.mark.parametrize(("failures", "requests"), [(0, 1), (1, 2), (2, 3)]) + async def test_every_failed_attempt_and_the_final_success_count_as_requests( + self, make_core: MakeCore, host: BackgroundSuccessLoggingHost, failures: int, requests: int + ) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=50)], num_retries=2) + host.answer("*", *([server_error()] * failures), None) + + await attempt(core) + await host.drain() + + assert await core.health().usage("g") == (10, requests) diff --git a/tests/unit/test_router/core/test_gaps_request_overrides.py b/tests/unit/test_router/core/test_gaps_request_overrides.py new file mode 100644 index 00000000000..3793dd5f898 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_request_overrides.py @@ -0,0 +1,227 @@ +""" +Per-request routing overrides a router core reads (or does not read) on each call: + +- ``default_fallbacks`` and ``retry_policy`` in a request are no routing input: the router neither + falls back on the request's ``default_fallbacks`` nor retries on its ``retry_policy``. Both reach + the deployment call as given (litellm's own ``completion`` reads ``retry_policy`` there), on every + attempt, sync and async, streamed or not. +- ``model_group_retry_policy`` that is not a mapping of policies, in a request or as the router's + setting: the retry loop looks the model group up in it (``group in value``, then ``value[group]``), + so a value that does not hold the group leaves the router's ``retry_policy`` in charge, a ``None`` + entry is no policy at all, and an explicit request ``None`` replaces the router's per-group policies. +""" + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import RetryPolicy +from tests.unit.test_router.core.retries_health_helpers import attempt, attempt_sync +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts, collect_texts_sync + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] +GROUPS: Final = [deployment("a", "a"), deployment("b", "b"), deployment("c", "c")] +OVERRIDES: Final[Mapping[str, object]] = {"default_fallbacks": ["c"], "retry_policy": {"RateLimitErrorRetries": 3}} + + +def sent_overrides(host: ScriptedHost) -> list[tuple[str | None, object, object]]: # mutable-ok: a fresh list + return [ + (call.deployment_id, call.kwargs.get("default_fallbacks"), call.kwargs.get("retry_policy")) + for call in host.calls_to() + ] + + +# --- default_fallbacks and retry_policy: passed to the deployment, never routed on ---------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("endpoint", "request_body"), + [ + ("acompletion", {"messages": MESSAGES}), + ("atext_completion", {"prompt": "p"}), + ("aembedding", {"input": ["x"]}), + ("aresponses", {"input": "hi"}), + ], +) +async def test_request_default_fallbacks_and_retry_policy_reach_the_deployment_as_given( + make_core: MakeCore, host: ScriptedHost, endpoint: str, request_body: Mapping[str, object] +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + + await core.call(endpoint, {"model": "a", **request_body, **OVERRIDES}) + + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +@pytest.mark.asyncio +async def test_a_failed_call_neither_falls_back_on_nor_retries_by_the_request_overrides( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + host.answer("openai/a", rate_limit()) + + error: Final = await attempt(core, "a", **OVERRIDES) + + assert isinstance(error, litellm.RateLimitError) + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +def test_a_failed_sync_call_neither_falls_back_on_nor_retries_by_the_request_overrides( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + host.answer("openai/a", rate_limit()) + + error: Final = attempt_sync(core, "a", **OVERRIDES) + + assert isinstance(error, litellm.RateLimitError) + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +@pytest.mark.asyncio +async def test_the_router_retries_and_fallbacks_carry_the_request_overrides_on_every_attempt( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=1, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit(), rate_limit()) + + response: Final = await attempt(core, "a", **OVERRIDES) + + assert not isinstance(response, Exception) + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] * 2 + [ + ("b", ["c"], {"RateLimitErrorRetries": 3}) + ] + + +@pytest.mark.asyncio +async def test_a_stream_that_fails_before_content_falls_back_with_the_request_overrides( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0)) + host.answer("openai/b", chat_stream("B")) + + stream: Final = await core.call("acompletion", {"model": "a", "messages": MESSAGES, "stream": True, **OVERRIDES}) + + assert await collect_texts(stream) == "B" + assert sent_overrides(host) == [ + ("a", ["c"], {"RateLimitErrorRetries": 3}), + ("b", ["c"], {"RateLimitErrorRetries": 3}), + ] + + +def test_a_sync_stream_carries_the_request_overrides(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + host.answer("openai/a", chat_stream("A", sync=True)) + + completion: Final[Callable[..., object]] = getattr(core, "completion") # pyright: ignore[reportAny] # an endpoint method + stream: Final = completion(model="a", messages=MESSAGES, stream=True, **OVERRIDES) + + assert collect_texts_sync(stream) == "A" + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +# --- model_group_retry_policy that is not a mapping of policies -------------------------------------- + + +def calls_after_a_rate_limit( + core: RouterCore, host: ScriptedHost, **request: object +) -> int: # kwargs-ok: request kwargs + host.answer("openai/a", rate_limit()) + error: Final = attempt_sync(core, "a", **request) + assert isinstance(error, litellm.RateLimitError) + return len(host.calls_to()) + + +@pytest.mark.parametrize( + "policies", + [None, {"a": None}, "xyz", [], ("b",)], + ids=["none", "none-entry", "string-without-the-group", "empty-list", "tuple-without-the-group"], +) +def test_a_request_policy_that_holds_no_policy_for_the_group_keeps_the_plain_retries( + make_core: MakeCore, host: ScriptedHost, policies: object +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=2) + + assert calls_after_a_rate_limit(core, host, model_group_retry_policy=policies) == 3 + + +@pytest.mark.parametrize( + ("policies", "calls"), + [ + ("xyz", 2), # the group is not in it: the router's retry_policy applies + ({"a": None}, 3), # a None entry is no policy: the plain retries apply + (None, 2), # an explicit None replaces the router's per-group policies + ], + ids=["string-without-the-group", "none-entry", "none"], +) +def test_a_request_policy_without_the_group_leaves_the_router_retry_policy_in_charge( + make_core: MakeCore, host: ScriptedHost, policies: object, calls: int +) -> None: + core: Final = make_core( + model_list=GROUPS, + num_retries=2, + retry_policy=RetryPolicy(RateLimitErrorRetries=1), + model_group_retry_policy={"a": RetryPolicy(RateLimitErrorRetries=0)}, + ) + + assert calls_after_a_rate_limit(core, host, model_group_retry_policy=policies) == calls + + +@pytest.mark.parametrize("policies", [5, "abc", {"a": 3}], ids=["int", "string-holding-the-group", "int-entry"]) +def test_a_request_that_opts_out_of_retries_never_reads_its_policy( + make_core: MakeCore, host: ScriptedHost, policies: object +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=2) + + assert calls_after_a_rate_limit(core, host, model_group_retry_policy=policies, num_retries=0) == 1 + + +@pytest.mark.parametrize( + "policies", [{"a": None}, "xyz", []], ids=["none-entry", "string-without-the-group", "empty-list"] +) +@pytest.mark.parametrize("how", ["constructor", "apply"]) +def test_a_router_policy_that_holds_no_policy_for_the_group_keeps_the_plain_retries( + make_core: MakeCore, host: ScriptedHost, policies: object, how: str +) -> None: + if how == "constructor": + core = make_core(model_list=GROUPS, num_retries=2, model_group_retry_policy=policies) + else: + core = make_core(model_list=GROUPS, num_retries=2) + core.apply(ChangeSet(settings={"model_group_retry_policy": policies})) + + stored: Final[object] = getattr(core, "model_group_retry_policy") # pyright: ignore[reportAny] # the router attribute + assert stored == policies + assert calls_after_a_rate_limit(core, host) == 3 + + +@pytest.mark.parametrize("how", ["constructor", "apply"]) +def test_a_router_policy_that_is_no_mapping_is_not_read_when_the_request_opts_out( + make_core: MakeCore, host: ScriptedHost, how: str +) -> None: + if how == "constructor": + core = make_core(model_list=GROUPS, num_retries=2, model_group_retry_policy=5) + else: + core = make_core(model_list=GROUPS, num_retries=2) + core.apply(ChangeSet(settings={"model_group_retry_policy": 5})) + + assert calls_after_a_rate_limit(core, host, num_retries=0) == 1 + + +def test_get_num_retries_from_retry_policy_reads_a_router_policy_that_is_no_mapping(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=GROUPS, retry_policy=RetryPolicy(RateLimitErrorRetries=4), model_group_retry_policy="xyz" + ) + error: Final = rate_limit() + retries_for: Final[Callable[..., object]] = getattr(core, "get_num_retries_from_retry_policy") # pyright: ignore[reportAny] # the router member + + assert retries_for(error, model_group="a") == 4 + with pytest.raises(TypeError, match="string indices must be integers"): + retries_for(error, model_group="x") diff --git a/tests/unit/test_router/core/test_gaps_router_core_crate.py b/tests/unit/test_router/core/test_gaps_router_core_crate.py new file mode 100644 index 00000000000..7f00d26c944 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_router_core_crate.py @@ -0,0 +1,283 @@ +""" +Gaps the Rust core closes inside router-core, pinned to ``PythonCore``'s behaviour (oddities included): +``default_deployment``, a ``model_group_retry_policy`` whose lookup raises, the deprecated registry +index writes, a raw ``routing_groups`` write against the stale member map, and ``cooldown_time=0`` +assigned after construction. +""" + +import warnings +from typing import Final + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router import Router +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import RoutingGroup +from tests.unit.test_router.core.registry_helpers import team_deployment +from tests.unit.test_router.core.retries_health_helpers import MakeCore, attempt, cooled_down, headers_of +from tests.unit.test_router.core.scripted_host import ScriptedHost, completion_response, deployment, rate_limit + +DEFAULT: Final[dict[str, object]] = { # mutable-ok: a deployment dict, copied per test + "model_name": "dm", + "litellm_params": {"model": "openai/other-model"}, + "model_info": {"id": "dd"}, +} + + +def as_router(core: RouterCore) -> Router: + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach, as from_config calls it + return router + + +def write(router: Router, name: str, value: object) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + setattr(router, name, value) + + +async def picked(router: Router, model: str) -> str: + """The id ``select`` picks, or the error's class and message.""" + try: + chosen: Final[object] = await router.select(model, {}) + except Exception as error: + return f"{type(error).__name__}: {error}" + return str((as_str_mapping((as_str_mapping(chosen) or {}).get("model_info")) or {}).get("id")) + + +# --- default_deployment -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_unknown_model_routes_to_the_default_deployment_without_retries_or_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the default deployment skips cooldown and the strategy, and the retry count reads no healthy + deployment, so ``num_retries`` is unused; its id is never cooled while it names no configured deployment.""" + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1, allowed_fails=0) + _ = core.apply(ChangeSet(assign={"default_deployment": dict(DEFAULT)})) + host.answer("other-model", rate_limit(), completion_response()) + + failed: Final = await attempt(core, model="other-model") + assert isinstance(failed, litellm.RateLimitError) + assert cooled_down(core, ["dd", "a"]) == {} + served: Final = await attempt(core, model="other-model") + + assert headers_of(served)["x-litellm-model-group"] == "other-model" + assert [call.deployment_id for call in host.calls] == ["dd", "dd"] + assert host.calls[0].kwargs["model"] == "other-model" + assert host.sleeps == [] + assert core.settings().default_deployment == DEFAULT + + +@pytest.mark.asyncio +async def test_a_known_model_never_reaches_the_default_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + _ = core.apply(ChangeSet(assign={"default_deployment": dict(DEFAULT)})) + + _ = await attempt(core) + + assert host.deployments_called("acompletion") == ["a"] + + +# --- model_group_retry_policy lookups that raise ----------------------------------------------------- + + +RAISING_POLICIES: Final = [ + (3, TypeError, "argument of type 'int' is not iterable"), + ("g", TypeError, "string indices must be integers"), + (["g"], TypeError, "list indices must be integers or slices"), + ({"g": 3}, AttributeError, "'int' object has no attribute 'DefaultRetries'"), + ({"g": {"DefaultRetries": "x"}}, ValidationError, "validation error for RetryPolicy"), +] +"""Each policy, the error its lookup raises for group ``g``, and a fragment of that error's message.""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("policy", "raised", "message"), RAISING_POLICIES) +async def test_a_raising_retry_policy_lookup_replaces_the_error_and_stops_the_hop( + make_core: MakeCore, host: ScriptedHost, policy: object, raised: type[Exception], message: str +) -> None: + """Oddity: the lookup runs on the hop's first failure; its error replaces the attempt's, and no retry + follows. The error is raised while the attempt's is handled, so it chains to it, except pydantic's + ``ValidationError``, which pydantic raises with no context.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=2, model_group_retry_policy=policy + ) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, raised) + assert message in str(error) + if raised is ValidationError: + assert error.__context__ is None + else: + assert isinstance(error.__context__, litellm.RateLimitError) + assert len(host.calls) == 1 + assert host.sleeps == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("policy", "raised", "message"), RAISING_POLICIES) +async def test_a_raising_retry_policy_lookup_takes_the_generic_fallbacks( + make_core: MakeCore, host: ScriptedHost, policy: object, raised: type[Exception], message: str +) -> None: + del raised, message + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("h", "c")], + num_retries=2, + model_group_retry_policy=policy, + fallbacks=[{"g": ["h"]}], + ) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + served: Final = await attempt(core) + + assert headers_of(served)["x-litellm-attempted-fallbacks"] == 1 + assert host.deployments_called("acompletion")[-1] == "c" + assert len(host.calls) == 2 + + +@pytest.mark.asyncio +async def test_a_policy_naming_another_group_leaves_this_group_retrying( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=2, model_group_retry_policy={"h": 3} + ) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert len(host.calls) == 2 + + +# --- deprecated registry index writes ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_written_model_name_index_decides_which_names_resolve(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")])) + write(router, "model_name_to_deployment_indices", {"gpt": [0], "claude": []}) + + assert await picked(router, "gpt") == "a" + assert (await picked(router, "claude")).startswith("BadRequestError") + + +@pytest.mark.asyncio +async def test_a_model_name_index_entry_pointing_at_another_name_resolves_nothing(make_core: MakeCore) -> None: + """Each indexed deployment is still checked against the requested name.""" + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")])) + write(router, "model_name_to_deployment_indices", {"gpt": [1], "claude": [1]}) + + assert (await picked(router, "gpt")).startswith("BadRequestError") + assert await picked(router, "claude") == "c" + + +@pytest.mark.asyncio +async def test_a_written_model_names_set_decides_which_names_reach_the_default_deployment( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")]) + _ = core.apply(ChangeSet(assign={"default_deployment": dict(DEFAULT)})) + router: Final = as_router(core) + + write(router, "model_names", set[str]()) + assert [await picked(router, name) for name in ("gpt", "claude", "other-model")] == ["dd", "dd", "dd"] + + write(router, "model_names", {"gpt", "claude", "other-model"}) + assert await picked(router, "gpt") == "a" + assert (await picked(router, "other-model")).startswith("BadRequestError") + + +@pytest.mark.asyncio +async def test_an_emptied_team_index_still_finds_team_deployments_by_scanning(make_core: MakeCore) -> None: + router: Final = as_router( + make_core(model_list=[deployment("gpt", "a"), team_deployment("int-t", "t1", "TA", "pub")]) + ) + write(router, "team_model_to_deployment_indices", {}) + write(router, "model_name_to_deployment_indices", {"gpt": [0]}) + + chosen: Final = await router.select("pub", {"metadata": {"user_api_key_team_id": "TA"}}) + assert chosen["model_info"]["id"] == "t1" + assert (await picked(router, "int-t")).startswith("BadRequestError") + + +# --- a raw routing_groups write: the member map stays as constructed ------------------------------------ + + +def _grouped(make_core: MakeCore) -> Router: + return as_router( + make_core( + model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("claude", "c")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "least-busy"}], + ) + ) + + +@pytest.mark.asyncio +async def test_dropping_a_routing_group_raw_leaves_its_members_raising_a_key_error(make_core: MakeCore) -> None: + """Oddity: the member-to-group map is built once, so a member of a dropped group looks it up and fails.""" + router: Final = _grouped(make_core) + write(router, "routing_groups", {}) + + assert await picked(router, "gpt") == "KeyError: 'G'" + assert await picked(router, "claude") == "c" + + +@pytest.mark.asyncio +async def test_a_routing_group_added_raw_leaves_its_members_on_the_default_strategy(make_core: MakeCore) -> None: + router: Final = _grouped(make_core) + write( + router, + "routing_groups", + { + "G": RoutingGroup(group_name="G", models=["gpt"], routing_strategy="least-busy"), + "K": RoutingGroup(group_name="K", models=["claude"], routing_strategy="latency-based-routing"), + }, + ) + + assert await picked(router, "gpt") == "a" + assert await picked(router, "claude") == "c" + assert (await picked(router, "K")).startswith("RouterRateLimitError: No deployments available") + + +# --- cooldown_time assigned after construction ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_assigned_zero_cooldown_time_never_keeps_a_deployment_cooled( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: an assigned 0 is a zero cooldown ttl (only the constructor reads 0 as the default).""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0) + _ = core.apply(ChangeSet(assign={"cooldown_time": 0})) + host.answer("*", rate_limit()) + + for _ in range(3): + assert isinstance(await attempt(core), litellm.RateLimitError) + assert cooled_down(core, ["a", "b"]) == {} + + assert core.settings().cooldown_time == 0 + assert len(host.calls) == 3 + + +@pytest.mark.asyncio +async def test_an_assigned_none_cooldown_time_cools_for_the_default(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0) + _ = core.apply(ChangeSet(assign={"cooldown_time": None})) + host.answer("*", rate_limit()) + + _ = await attempt(core) + _ = await attempt(core) + + assert sorted(cooled_down(core, ["a", "b"])) == ["a", "b"] + assert {entry["cooldown_time"] for entry in cooled_down(core, ["a", "b"]).values()} == {5} + assert isinstance(await attempt(core), litellm.RouterRateLimitError) diff --git a/tests/unit/test_router/core/test_gaps_router_core_crate_picks.py b/tests/unit/test_router/core/test_gaps_router_core_crate_picks.py new file mode 100644 index 00000000000..255d731142d --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_router_core_crate_picks.py @@ -0,0 +1,372 @@ +""" +Pick errors the Rust core serves inside router-core, pinned to ``PythonCore``'s behaviour (oddities +included): a pick's error retried by the retry loop, a ``default_deployment`` Python cannot serve, a +written index position past the model list, a stale routing group at a deeper fallback level or in +weighted failover, and a strategy change through ``update_settings`` after a raw ``routing_groups`` write. +""" + +import warnings +from typing import Final + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router import Router +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import RoutingGroup +from tests.unit.test_router.core.registry_helpers import team_deployment +from tests.unit.test_router.core.retries_health_helpers import MakeCore, attempt, attempt_sync +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.test_fallbacks_characterization_outcomes import FallbackEvents + +TEAM: Final = {"metadata": {"user_api_key_team_id": "TA"}} + + +def as_router(core: RouterCore) -> Router: + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach, as from_config calls it + return router + + +def write(router: Router, name: str, value: object) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + setattr(router, name, value) + + +async def selected(router: Router, model: str, **request: object) -> object: # kwargs-ok: the request kwargs + """The deployment ``select`` picks, or the error it raised.""" + try: + return await router.select(model, dict(request)) + except Exception as error: + return error + + +def picked_id(picked: object) -> object: + return (as_str_mapping((as_str_mapping(picked) or {}).get("model_info")) or {}).get("id") + + +# --- a pick's error in the retry loop ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_retry_policy_retries_a_group_whose_deployments_all_cool_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The pick's ``RouterRateLimitError`` is a plain ``ValueError``: ``DefaultRetries`` applies to it.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + num_retries=2, + allowed_fails=0, + model_group_retry_policy={"g": {"DefaultRetries": 1}}, + ) + host.answer("*", rate_limit()) + for _ in range(2): + _ = await attempt(core, num_retries=0) + host.calls.clear() + host.sleeps.clear() + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RouterRateLimitError) + assert isinstance(error.__context__, litellm.RouterRateLimitError) + assert host.calls == [] + assert len(host.sleeps) == 2 + assert all(sleep > 0 for sleep in host.sleeps) + + +@pytest.mark.asyncio +async def test_a_retry_policy_retries_an_unknown_model_and_stamps_the_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2, retry_policy={"DefaultRetries": 1}) + + error: Final = await attempt(core, model="zz") + + assert isinstance(error, litellm.BadRequestError) + assert (getattr(error, "num_retries", None), getattr(error, "max_retries", None)) == (1, 1) + assert isinstance(error.__context__, litellm.BadRequestError) + assert len(host.sleeps) == 2 + + +@pytest.mark.asyncio +async def test_invalid_router_weights_are_retried_while_a_deployment_is_healthy( + make_core: MakeCore, host: ScriptedHost +) -> None: + """No policy: ``should_retry_this_error`` retries a plain exception while a deployment is healthy.""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=2) + + error: Final = await attempt(core, _router_weights={"a": "x"}) + + assert isinstance(error, ValidationError) + assert isinstance(error.__context__, ValidationError) + assert host.calls == [] + assert host.sleeps == [0, 0, 0] + assert all(type(sleep) is int for sleep in host.sleeps) + + +@pytest.mark.asyncio +async def test_without_a_policy_an_unknown_model_is_raised_at_once(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2) + + error: Final = await attempt(core, model="zz") + + assert isinstance(error, litellm.BadRequestError) + assert error.__context__ is None + assert host.sleeps == [] + + +@pytest.mark.asyncio +async def test_the_last_retry_failure_chains_to_the_hops_first_failure(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + num_retries=2, + model_group_retry_policy={"g": {"DefaultRetries": 1}}, + ) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert isinstance(error.__context__, litellm.RateLimitError) + assert error.__context__ is not error + assert host.sleeps[0] == 0 + assert type(host.sleeps[0]) is int + + +# --- a default_deployment Python cannot serve --------------------------------------------------------- + + +UNCOPYABLE_DEFAULTS: Final[list[tuple[object, type[Exception], str]]] = [ # mutable-ok: pytest parameters + ({}, KeyError, "'litellm_params'"), + ({"model_name": "dm"}, KeyError, "'litellm_params'"), + ({"model_name": "dm", "litellm_params": "s", "model_info": {"id": "dd"}}, AttributeError, "no attribute 'copy'"), + ("str", AttributeError, "'str' object has no attribute 'copy'"), + (5, AttributeError, "'int' object has no attribute 'copy'"), + (["a"], TypeError, "list indices must be integers or slices, not str"), +] +"""Each value, and the error its copy raises for any name.""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "raised", "message"), UNCOPYABLE_DEFAULTS) +async def test_a_default_deployment_python_cannot_copy_fails_each_pick_of_an_unknown_name( + make_core: MakeCore, host: ScriptedHost, value: object, raised: type[Exception], message: str +) -> None: + """Oddity: assigning it succeeds; every request for a name nothing else serves raises the copy's error.""" + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1) + _ = core.apply(ChangeSet(assign={"default_deployment": value})) + + for error in ( + await attempt(core, model="zz"), + attempt_sync(core, model="zz"), + await selected(as_router(core), "zz"), + ): + assert isinstance(error, raised) + assert message in str(error) + assert host.deployments_called("acompletion") == [] + _ = await attempt(core) + assert host.deployments_called("acompletion") == ["a"] + + +FAILING_DEFAULTS: Final[list[tuple[object, str, object]]] = [ # mutable-ok: pytest parameters + ({"litellm_params": {"model": "x"}}, "'model_name'", None), + ({"model_name": "dm", "litellm_params": {"model": "x"}}, "'model_info'", None), + ({"model_name": "dm", "litellm_params": {"model": "x"}, "model_info": {}}, "'id'", {}), +] +"""Each value, the ``KeyError`` every call to it raises, and the ``model_info`` its pick returns.""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "missing", "model_info"), FAILING_DEFAULTS) +async def test_a_default_deployment_missing_a_field_is_picked_but_every_call_fails( + make_core: MakeCore, host: ScriptedHost, value: object, missing: str, model_info: object +) -> None: + """Oddity: the pick returns Python's copy; the call's first reads of it raise ``KeyError``, with no retry.""" + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1) + _ = core.apply(ChangeSet(assign={"default_deployment": value})) + + called: Final = await attempt(core, model="zz") + called_sync: Final = attempt_sync(core, model="zz") + chosen: Final = as_str_mapping(await selected(as_router(core), "zz")) or {} + + assert isinstance(called, KeyError) + assert isinstance(called_sync, KeyError) + assert str(called) == missing + assert chosen.get("model_info") == model_info + assert (as_str_mapping(chosen.get("litellm_params")) or {}).get("model") == "zz" + assert host.calls == [] + assert host.sleeps == [] + + +# --- a written index position past the model list ----------------------------------------------------- + + +def _teams(make_core: MakeCore) -> tuple[RouterCore, Router]: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("claude", "c"), team_deployment("int-t", "t1", "TA", "pub")], + num_retries=1, + fallbacks=[{"gpt": ["claude"]}, {"pub": ["claude"]}], + ) + return core, as_router(core) + + +@pytest.mark.asyncio +async def test_a_name_index_position_past_the_list_raises_index_error_for_that_name( + make_core: MakeCore, host: ScriptedHost +) -> None: + core, router = _teams(make_core) + write(router, "model_name_to_deployment_indices", {"gpt": [0, 7], "claude": [1], "int-t": [2]}) + + picked: Final = await selected(router, "gpt") + called: Final = await attempt(core, model="gpt") + + assert isinstance(picked, IndexError) + assert isinstance(called, IndexError) + assert str(called) == "list index out of range" + assert picked_id(await selected(router, "claude")) == "c" + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_a_team_index_position_past_the_list_raises_from_the_fallback_handling( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the level's fallback handling lists the group again and raises a second ``IndexError`` + (chained to the pick's), before any fallback.""" + core, router = _teams(make_core) + write(router, "team_model_to_deployment_indices", {("TA", "pub"): [9]}) + + picked: Final = await selected(router, "pub", **TEAM) + called: Final = await attempt(core, model="pub", **TEAM) + + assert isinstance(picked, IndexError) + assert isinstance(called, IndexError) + assert isinstance(called.__context__, IndexError) + assert host.calls == [] + assert picked_id(await selected(router, "gpt")) == "a" + + +# --- a stale routing group deeper in the walk --------------------------------------------------------- + + +def _stale(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: extra router arguments + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("claude", "c"), deployment("x", "x1")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "simple-shuffle"}], + num_retries=0, + fallbacks=[{"x": ["gpt"]}, {"gpt": ["claude"]}], + **arguments, + ) + write(as_router(core), "routing_groups", {}) + return core + + +@pytest.mark.asyncio +async def test_a_stale_routing_group_at_a_fallback_level_is_that_levels_error( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final = FallbackEvents() + monkeypatch.setattr(litellm, "callbacks", [events]) + core: Final = _stale(make_core) + host.answer("*", rate_limit()) + + error: Final = await attempt(core, model="x") + + assert isinstance(error, litellm.RateLimitError) + assert events.events == [("failure", "gpt", "claude", "KeyError"), ("failure", "x", "gpt", "RateLimitError")] + + +@pytest.mark.asyncio +async def test_weighted_failover_raises_the_stale_groups_key_error_before_any_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: ``_maybe_run_weighted_failover`` reads the routing context first, so the stale member map's + ``KeyError`` escapes the fallback handling (chained to the pick's) and the regular fallback never runs.""" + core: Final = _stale(make_core, enable_weighted_failover=True) + + error: Final = await attempt(core, model="gpt") + + assert isinstance(error, KeyError) + assert str(error) == "'G'" + assert isinstance(error.__context__, KeyError) + assert host.calls == [] + + +# --- update_settings(routing_strategy) after a raw routing_groups write ----------------------------------- + + +@pytest.mark.asyncio +async def test_a_strategy_change_re_inits_the_routing_groups_from_their_last_input(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("claude", "c")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "least-busy"}], + ) + router: Final = as_router(core) + write( + router, "routing_groups", {"K": RoutingGroup(group_name="K", models=["claude"], routing_strategy="least-busy")} + ) + assert isinstance(await selected(router, "gpt"), KeyError) + + _ = core.apply(ChangeSet(settings={"routing_strategy": "usage-based-routing-v2"})) + + assert sorted(core.registry().routing_groups()) == ["G"] + assert picked_id(await selected(router, "G")) == "a" + assert picked_id(await selected(router, "gpt")) == "a" + assert isinstance(await selected(router, "K"), litellm.BadRequestError) + + +@pytest.mark.asyncio +async def test_an_unchanged_strategy_keeps_the_raw_routing_groups(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("claude", "c")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "least-busy"}], + ) + write( + as_router(core), + "routing_groups", + {"K": RoutingGroup(group_name="K", models=["claude"], routing_strategy="least-busy")}, + ) + + _ = core.apply(ChangeSet(settings={"routing_strategy": "simple-shuffle"})) + + assert sorted(core.registry().routing_groups()) == ["K"] + + +@pytest.mark.asyncio +async def test_a_negative_index_position_counts_from_the_end_of_the_list( + make_core: MakeCore, host: ScriptedHost +) -> None: + """As Python's ``model_list[index]`` reads it; one before the start raises ``IndexError``.""" + core, router = _teams(make_core) + write(router, "model_name_to_deployment_indices", {"gpt": [-3], "claude": [-2], "int-t": [-4]}) + + assert picked_id(await selected(router, "gpt")) == "a" + assert picked_id(await selected(router, "claude")) == "c" + before_the_start: Final = await attempt(core, model="int-t") + assert isinstance(before_the_start, IndexError) + assert str(before_the_start) == "list index out of range" + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_a_default_deployment_whose_id_is_not_a_string_is_picked_and_called_as_python_copies_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + value: Final = { + "model_name": "dm", + "litellm_params": {"model": "openai/dm", "api_key": "k"}, + "model_info": {"id": 5}, + } + _ = core.apply(ChangeSet(assign={"default_deployment": value})) + + served: Final = await attempt(core, model="zz") + chosen: Final = await selected(as_router(core), "zz") + + assert isinstance(served, litellm.ModelResponse) + assert host.calls[0].kwargs["model"] == "zz" + assert chosen == {"model_name": "dm", "litellm_params": {"model": "zz", "api_key": "k"}, "model_info": {"id": 5}} diff --git a/tests/unit/test_router/core/test_health_registry_views.py b/tests/unit/test_router/core/test_health_registry_views.py new file mode 100644 index 00000000000..0180f991ccb --- /dev/null +++ b/tests/unit/test_router/core/test_health_registry_views.py @@ -0,0 +1,174 @@ +""" +The ``health()`` priority helpers and the ``registry()`` members that read router state beyond the +deployment list: ``client(client_type="max_parallel_requests")``, ``plugins("routing")`` and +``zero_cost()``. Each test pins what the Python core does; every core must answer the same. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit +from litellm.types.router import Deployment, LiteLLM_Params, RoutingContext +from tests.unit.test_router.core.scripted_host import deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore + +# --- health().split_priority ---------------------------------------------------------------------- + + +def test_split_priority_gives_the_request_priority_and_the_kwargs_without_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().split_priority("g", {"priority": 5, "user": "u"}) == (5, {"user": "u"}) + + +def test_split_priority_falls_back_to_the_default_priority_when_the_request_sets_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], default_priority=2) + assert core.health().split_priority("g", {"user": "u"}) == (2, {"user": "u"}) + + +def test_split_priority_without_a_default_leaves_the_request_unqueued(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().split_priority("g", {"priority": None}) == (None, {}) + + +def test_split_priority_refuses_a_priority_that_is_no_integer(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], default_priority=2) + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'high'"): + core.health().split_priority("g", {"priority": "high"}) + + +def test_split_priority_drops_a_bad_priority_for_a_group_that_drops_params(make_core: MakeCore) -> None: + """A deployment of the group with ``drop_params=True`` turns the bad priority into the default.""" + core: Final = make_core(model_list=[deployment("g", "a", drop_params=True)], default_priority=2) + assert core.health().split_priority("g", {"priority": "high"}) == (2, {}) + + +def test_split_priority_reads_a_pinned_deployment_s_params_for_its_id(make_core: MakeCore) -> None: + """A deployment id as the model reads that deployment's ``drop_params``, not a group's.""" + core: Final = make_core( + model_list=[deployment("g", "a", drop_params=True), deployment("g", "b")], default_priority=1 + ) + assert core.health().split_priority("a", {"priority": "high"}) == (1, {}) + with pytest.raises(litellm.BadRequestError): + core.health().split_priority("b", {"priority": "high"}) + + +# --- health().priority_header_fields -------------------------------------------------------------- + + +def test_priority_header_fields_read_an_integer_header(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {}, "3") == {"priority": 3} + + +@pytest.mark.parametrize("header", [None, ""]) +def test_priority_header_fields_are_empty_without_a_header(make_core: MakeCore, header: str | None) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {}, header) == {} + + +def test_priority_header_fields_pass_a_bad_header_on_for_the_call_to_refuse(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {}, "urgent") == {"priority": "urgent"} + + +def test_priority_header_fields_drop_a_bad_header_when_the_request_drops_params(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {"drop_params": True}, "urgent") == {} + + +# --- registry().client(client_type="max_parallel_requests") --------------------------------------- + + +def test_the_max_parallel_requests_client_is_the_deployment_s_cached_limit(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_parallel_requests=2)]) + entry: Final = deployment("g", "a", max_parallel_requests=2) + limit: Final = core.registry().client(entry, client_type="max_parallel_requests") + assert isinstance(limit, MaxParallelRequestsLimit) + assert (limit.max_parallel_requests, limit.model_id, limit.model_group, limit.in_flight) == (2, "a", "g", 0) + assert core.registry().client(entry, client_type="max_parallel_requests") is limit + + +def test_the_max_parallel_requests_limit_comes_from_rpm_or_the_router_default(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", rpm=7), deployment("g", "b")], default_max_parallel_requests=4 + ) + by_rpm: Final = core.registry().client(deployment("g", "a", rpm=7), client_type="max_parallel_requests") + by_default: Final = core.registry().client(deployment("g", "b"), client_type="max_parallel_requests") + assert isinstance(by_rpm, MaxParallelRequestsLimit) + assert isinstance(by_default, MaxParallelRequestsLimit) + assert (by_rpm.max_parallel_requests, by_default.max_parallel_requests) == (7, 4) + + +def test_a_deployment_without_a_limit_has_no_max_parallel_requests_client(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.registry().client(deployment("g", "a"), client_type="max_parallel_requests") is None + + +def test_the_client_lookup_reads_the_deployment_id_first(make_core: MakeCore) -> None: + """``deployment["model_info"]["id"]``: a deployment without model_info raises, whatever the client type.""" + core: Final = make_core(model_list=[deployment("g", "a")]) + entry: Final = {"model_name": "g", "litellm_params": {"model": "openai/a"}} + for client_type in ("max_parallel_requests", "async", None): + with pytest.raises(KeyError, match="model_info"): + core.registry().client(entry, client_type=client_type) + + +@pytest.mark.asyncio +async def test_a_full_max_parallel_requests_client_turns_the_next_call_away(make_core: MakeCore) -> None: + """The limit ``client`` returns is the one calls take their slot from.""" + core: Final = make_core(model_list=[deployment("g", "a", max_parallel_requests=1)], num_retries=0) + entry: Final = deployment("g", "a", max_parallel_requests=1) + limit: Final = core.registry().client(entry, client_type="max_parallel_requests") + assert isinstance(limit, MaxParallelRequestsLimit) + with limit, pytest.raises(litellm.RateLimitError, match="already has max_parallel_requests=1"): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every core serves acompletion + assert limit.in_flight == 0 + + +# --- registry().plugins("routing") and zero_cost() ------------------------------------------------ + + +class _Keep: + async def run(self, context: RoutingContext) -> RoutingContext: + return context + + +def test_the_routing_plugins_are_the_ones_the_router_was_built_with(make_core: MakeCore) -> None: + plugin: Final = _Keep() + core: Final = make_core(model_list=[deployment("g", "a")], plugins=[plugin]) + assert core.registry().plugins("routing") == [plugin] + + +def test_a_router_without_plugins_lists_no_routing_plugins(make_core: MakeCore) -> None: + assert make_core(model_list=[deployment("g", "a")]).registry().plugins("routing") == [] + + +def test_zero_cost_is_an_empty_lookup_the_auth_checks_fill(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + lookup: Final = core.registry().zero_cost() + assert lookup == {} + assert isinstance(lookup, dict) + lookup["g"] = True + assert core.registry().zero_cost() == {"g": True} + + +def test_a_deployment_change_empties_the_zero_cost_lookup(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + lookup: Final = core.registry().zero_cost() + assert isinstance(lookup, dict) + lookup["g"] = True + core.apply( + ChangeSet( + upsert=[ + Deployment( + model_name="g", + litellm_params=LiteLLM_Params(model="openai/b", api_key="sk-test"), + model_info={"id": "b"}, + ) + ] + ) + ) + assert core.registry().zero_cost() == {} diff --git a/tests/unit/test_router/core/test_host_settings_changes_characterization.py b/tests/unit/test_router/core/test_host_settings_changes_characterization.py new file mode 100644 index 00000000000..fa1446aa51a --- /dev/null +++ b/tests/unit/test_router/core/test_host_settings_changes_characterization.py @@ -0,0 +1,428 @@ +""" +Characterization of the settings only the host reads, changed after construction: ``search_tools``, +``guardrail_list``, ``default_max_parallel_requests``, ``provider_budget_config``, +``health_check_ignore_transient_errors`` and ``router_general_settings``. + +``apply(ChangeSet(assign=...))`` stores each as given, ``settings()`` reports it, and whatever reads +it afterwards sees it from then on. None of them is runtime-updatable, so ``ChangeSet(settings=...)`` +(``update_settings``) ignores them. +""" + +import asyncio +import time +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery +from litellm.types.router import Deployment, RouterGeneralSettings +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.retries_health_helpers import cooled_down +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.selection_helpers import ids_of +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] +SEARCH_TOOLS: Final = [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily", "api_key": "tk"}}] +GUARDRAILS: Final = [{"guardrail_name": "gr", "litellm_params": {"guardrail": "presidio", "mode": "pre_call"}}] +PROVIDER_BUDGET: Final = {"openai": {"budget_limit": 1, "time_period": "1d"}} + +ASSIGNED: Final[list[tuple[str, object]]] = [ # mutable-ok: pytest parameters + ("search_tools", SEARCH_TOOLS), + ("guardrail_list", GUARDRAILS), + ("default_max_parallel_requests", 3), + ("provider_budget_config", PROVIDER_BUDGET), + ("health_check_ignore_transient_errors", True), + ("router_general_settings", RouterGeneralSettings(pass_through_all_models=True)), + ("provider_default_deployment_ids", ["d1"]), + ("enable_health_check_routing", True), + ("background_health_check_model_groups", ["g"]), +] + + +@pytest.fixture +def host(monkeypatch: pytest.MonkeyPatch) -> EndpointHost: + endpoint_host: Final = EndpointHost() + patch_direct_endpoints(monkeypatch, endpoint_host) + return endpoint_host + + +def _endpoint(core: RouterCore, name: str) -> Callable[..., Awaitable[object]]: + """An async endpoint method, which the core serves by name.""" + method: Final[Callable[..., Awaitable[object]]] = getattr(core, name) # pyright: ignore[reportAny] # endpoint methods are served by name + return method + + +async def _open_stream(core: RouterCore) -> object: + return await _endpoint(core, "acompletion")(model="a", messages=MESSAGES, stream=True) + + +# --- settings() --------------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("name", "value"), ASSIGNED) +def test_an_assigned_host_setting_is_reported_by_settings(make_core: MakeCore, name: str, value: object) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + before: Final = core.settings() + _ = core.apply(ChangeSet(assign={name: value})) + assert getattr(core.settings(), name) == value + assert core.settings() is not before + + +@pytest.mark.parametrize(("name", "value"), [("search_tools", SEARCH_TOOLS), ("guardrail_list", GUARDRAILS)]) +def test_an_assigned_host_setting_is_copied_from_the_callers_list( + make_core: MakeCore, name: str, value: list[dict[str, object]] +) -> None: + given: Final = [dict(entry) for entry in value] + core: Final = make_core(model_list=[deployment("g", "d1")]) + _ = core.apply(ChangeSet(assign={name: given})) + given.clear() + assert getattr(core.settings(), name) == value # pyright: ignore[reportAny] # a settings field named by the parameter + + +@pytest.mark.parametrize(("name", "value"), ASSIGNED) +def test_update_settings_ignores_a_host_setting(make_core: MakeCore, name: str, value: object) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + before: Final[object] = getattr(core.settings(), name) # pyright: ignore[reportAny] # a settings field named by the parameter + _ = core.apply(ChangeSet(settings={name: value})) + assert getattr(core.settings(), name) == before + + +def test_a_router_attribute_write_assigns_a_host_setting(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.search_tools = SEARCH_TOOLS # pyright: ignore[reportAttributeAccessIssue] # the deprecated attribute write + assert core.settings().search_tools == SEARCH_TOOLS + + +# --- router_general_settings as a constructor argument ------------------------------------------ + + +def test_router_general_settings_passed_to_the_constructor_are_reported(make_core: MakeCore) -> None: + general: Final = RouterGeneralSettings(async_only_mode=True, pass_through_all_models=True) + core: Final = make_core(model_list=[deployment("g", "d1")], router_general_settings=general) + assert core.settings().router_general_settings == general + + +def test_router_general_settings_of_none_become_the_defaults(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], router_general_settings=None) + assert core.settings().router_general_settings == RouterGeneralSettings() + + +def test_routers_built_without_general_settings_share_one_default_object(make_core: MakeCore) -> None: + """The Python core's default argument is one ``RouterGeneralSettings()`` every router shares.""" + first: Final = make_core(model_list=[deployment("g", "d1")]) + second: Final = make_core(model_list=[deployment("g", "d1")]) + assert first.settings().router_general_settings is second.settings().router_general_settings + + +async def test_router_general_settings_leave_routing_unchanged(make_core: MakeCore, host: EndpointHost) -> None: + general: Final = RouterGeneralSettings(async_only_mode=True, pass_through_all_models=True) + core: Final = make_core(model_list=[deployment("g", "d1")], router_general_settings=general) + _ = await _endpoint(core, "acompletion")(model="g", messages=MESSAGES) + completion: Final[Callable[..., object]] = getattr(core, "completion") # pyright: ignore[reportAny] # see _endpoint + _ = completion(model="g", messages=MESSAGES) + assert host.deployments_called() == ["d1", "d1"] + + +# --- readers ------------------------------------------------------------------------------------ + + +async def test_asearch_reads_search_tools_assigned_after_construction(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + _ = core.apply(ChangeSet(assign={"search_tools": SEARCH_TOOLS})) + _ = await _endpoint(core, "asearch")(model="s", query="q") + call: Final = host.calls[0] + assert (call.op, call.kwargs["search_provider"], call.kwargs["api_key"]) == ("asearch", "tavily", "tk") + + +async def test_asearch_no_longer_finds_a_tool_dropped_by_a_changeset(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS, num_retries=0) + _ = core.apply(ChangeSet(assign={"search_tools": []})) + with pytest.raises(ValueError, match=r"Search tool 's' not found in router\.search_tools"): + _ = await _endpoint(core, "asearch")(model="s", query="q") + assert host.calls == [] + + +async def test_aguardrail_reads_a_guardrail_list_assigned_after_construction(make_core: MakeCore, host: ScriptedHost): + seen: Final[list[Mapping[str, object]]] = [] + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + seen.append(kwargs) + return "verdict" + + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + _ = core.apply(ChangeSet(assign={"guardrail_list": GUARDRAILS})) + assert await _endpoint(core, "aguardrail")(guardrail_name="gr", original_function=guard) == "verdict" + assert seen[0]["selected_guardrail"] == GUARDRAILS[0] + + +async def test_a_default_max_parallel_requests_assigned_later_limits_the_next_requests( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x")) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0) + _ = core.apply(ChangeSet(assign={"default_max_parallel_requests": 1})) + first: Final = await _open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(_open_stream(core), 5) + assert await collect_texts(first) == "x" + + +async def test_a_deployment_keeps_the_slot_limit_it_was_given_before_a_change( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The limit is cached per deployment on first use: a later default does not replace it.""" + host.answer("openai/a", chat_stream("x")) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0, default_max_parallel_requests=1) + assert await collect_texts(await _open_stream(core)) == "x" + _ = core.apply(ChangeSet(assign={"default_max_parallel_requests": 5})) + first: Final = await _open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(_open_stream(core), 5) + assert await collect_texts(first) == "x" + + +async def test_a_budget_config_assigned_beside_the_budget_check_builds_the_limiter_with_it(make_core: MakeCore) -> None: + """``apply_changes`` assigns before ``update_settings`` adds the check.""" + core: Final = make_core(model_list=[deployment("g", "a")]) + _ = core.apply( + ChangeSet( + assign={"provider_budget_config": PROVIDER_BUDGET}, + settings={"optional_pre_call_checks": ["router_budget_limiting"]}, + ) + ) + limiter: Final = core.registry().plugins("budget_limiter") + assert isinstance(limiter, RouterBudgetLimiting) + assert limiter.provider_budget_config is core.settings().provider_budget_config + assert list(limiter.provider_budget_config or {}) == ["openai"] + core.retire() + + +async def test_a_budget_limiter_built_already_keeps_its_config(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + limiter: Final = core.registry().plugins("budget_limiter") + assert isinstance(limiter, RouterBudgetLimiting) + _ = core.apply(ChangeSet(assign={"provider_budget_config": {"azure": {"budget_limit": 2, "time_period": "1d"}}})) + assert core.registry().plugins("budget_limiter") is limiter + assert list(limiter.provider_budget_config or {}) == ["openai"] + assert list(core.settings().provider_budget_config or {}) == ["azure"] + core.retire() + + +# --- settings().guardrail(name) ----------------------------------------------------------------- + + +def test_settings_guardrail_returns_the_one_configured_under_the_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + assert core.settings().guardrail("gr") == GUARDRAILS[0] + + +def test_settings_guardrail_with_an_unknown_name_raises_value_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + with pytest.raises(ValueError, match="No guardrail found with name: nope"): + _ = core.settings().guardrail("nope") + + +def test_settings_guardrail_picks_among_same_named_guardrails_by_weight(make_core: MakeCore) -> None: + guardrails: Final = [ + {"guardrail_name": "gr", "litellm_params": {"guardrail": "never", "weight": 0}}, + {"guardrail_name": "gr", "litellm_params": {"guardrail": "always", "weight": 1}}, + ] + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=guardrails) + assert {core.settings().guardrail("gr")["litellm_params"]["guardrail"] for _ in range(20)} == {"always"} + + +def test_settings_guardrail_reads_a_guardrail_list_assigned_after_construction(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + _ = core.apply(ChangeSet(assign={"guardrail_list": GUARDRAILS})) + assert core.settings().guardrail("gr") == GUARDRAILS[0] + + +# --- a default_max_parallel_requests that is no int --------------------------------------------- + + +@pytest.mark.parametrize("assigned", [False, True]) +async def test_a_float_default_max_parallel_requests_limits_the_slots( + make_core: MakeCore, host: ScriptedHost, assigned: bool +) -> None: + host.answer("openai/a", chat_stream("x")) + given: Final = {} if assigned else {"default_max_parallel_requests": 1.0} + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0, **given) + if assigned: + _ = core.apply(ChangeSet(assign={"default_max_parallel_requests": 1.0})) + first: Final = await _open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(_open_stream(core), 5) + assert await collect_texts(first) == "x" + + +# --- health-check routing assigned after construction ------------------------------------------- + + +async def _healthy_ids(core: RouterCore, model: str) -> list[str]: # mutable-ok: a fresh list + select: Final[Callable[..., Awaitable[object]]] = getattr(core, "select") # pyright: ignore[reportAny] # the core's select, which takes kind + return ids_of(await select(model, kind="healthy")) + + +def _unhealthy(*ids: str) -> dict[str, dict[str, object]]: # mutable-ok: health states + return {i: {"is_healthy": False, "timestamp": time.time(), "reason": "probe failed"} for i in ids} + + +async def test_health_check_routing_turned_on_later_drops_unhealthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + core.health().health_checks.set_deployment_health_states(_unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + assert await _healthy_ids(core, "g") == ["a", "b"] + _ = core.apply(ChangeSet(assign={"enable_health_check_routing": True})) + assert await _healthy_ids(core, "g") == ["b"] + _ = core.apply(ChangeSet(assign={"enable_health_check_routing": False})) + assert await _healthy_ids(core, "g") == ["a", "b"] + + +async def test_background_health_check_model_groups_assigned_later_scope_the_filter(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("h", "c"), deployment("h", "d")], + enable_health_check_routing=True, + ) + core.health().health_checks.set_deployment_health_states(_unhealthy("a", "c")) # pyright: ignore[reportArgumentType] # health states + assert await _healthy_ids(core, "h") == ["d"] + _ = core.apply(ChangeSet(assign={"background_health_check_model_groups": ["g"]})) + assert await _healthy_ids(core, "g") == ["b"] + assert await _healthy_ids(core, "h") == ["c", "d"] + + +# --- provider_default_deployment_ids assigned after construction -------------------------------- + + +@pytest.mark.parametrize(("default_ids", "cooled"), [([], ["a"]), (["a"], [])]) +async def test_an_assigned_provider_default_deployment_never_cools_down( + make_core: MakeCore, host: ScriptedHost, default_ids: list[str], cooled: list[str] +) -> None: + host.answer("openai/a", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0) + _ = core.apply(ChangeSet(assign={"provider_default_deployment_ids": default_ids})) + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "a", "messages": MESSAGES}) + assert sorted(cooled_down(core, ["a", "b"])) == cooled + + +async def test_wildcard_deployments_never_cool_down(make_core: MakeCore, host: ScriptedHost) -> None: + """Each is a provider default deployment, which the cooldown check skips.""" + host.answer("openai/x", rate_limit()) + core: Final = make_core( + model_list=[deployment("openai/*", "w", model="openai/*"), deployment("openai/*", "w2", model="openai/*")], + num_retries=0, + ) + for _ in range(3): + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "openai/x", "messages": MESSAGES}) + assert cooled_down(core, ["w", "w2"]) == {} + assert len(host.calls) == 3 + + +async def test_an_added_wildcard_deployment_never_cools_down(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/x", rate_limit()) + core: Final = make_core(model_list=[deployment("openai/*", "w", model="openai/*")], num_retries=0) + _ = core.apply(ChangeSet(upsert=[Deployment.model_validate(deployment("openai/*", "w2", model="openai/*"))])) + assert core.settings().provider_default_deployment_ids == ["w", "w2"] + for _ in range(3): + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "openai/x", "messages": MESSAGES}) + assert cooled_down(core, ["w", "w2"]) == {} + + +# --- Kubernetes pod discovery ------------------------------------------------------------------- + + +class _RecordingDiscovery(KubernetesPodDiscovery): + """Resolves every pod-discovery deployment to one pod IP, and records the retry counts seen.""" + + def __init__(self) -> None: + super().__init__(proxy_environment=lambda: {}) + self.seen: list[object] = [] # mutable-ok: filled as picks resolve + + def _pod(self, deployment: object) -> object: + record: Final = dict(as_str_mapping(deployment) or {}) + params: Final = as_str_mapping(record.get("litellm_params")) or {} + if params.get("kubernetes_pod_discovery") is not True: + return deployment + record["litellm_params"] = {**params, "api_base": "http://10.0.0.7:4000"} + return record + + def resolve_deployment(self, deployment: object, request_kwargs: Mapping[str, object] | None = None) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # the resolver's generic signature, read as objects + self.seen.append(self._request_retry_count(request_kwargs)) + return self._pod(deployment) + + async def async_resolve_deployment( # pyright: ignore[reportIncompatibleMethodOverride] # see resolve_deployment + self, deployment: object, request_kwargs: Mapping[str, object] | None = None + ) -> object: + self.seen.append(self._request_retry_count(request_kwargs)) + return self._pod(deployment) + + +def _pod_deployment(deployment_id: str) -> dict[str, object]: # mutable-ok: a model_list entry + return deployment("g", deployment_id, api_base="http://svc.ns:4000", kubernetes_pod_discovery=True) + + +def test_settings_report_the_routers_pod_discovery(make_core: MakeCore) -> None: + settings: Final = make_core(model_list=[deployment("g", "d1")]).settings() + assert settings.is_set("kubernetes_pod_discovery") + assert isinstance(settings.kubernetes_pod_discovery, KubernetesPodDiscovery) + + +async def test_an_assigned_pod_discovery_resolves_each_attempt(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", rate_limit()) + discovery: Final = _RecordingDiscovery() + core: Final = make_core(model_list=[_pod_deployment("a")], num_retries=1) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": discovery})) + assert core.settings().kubernetes_pod_discovery is discovery + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "g", "messages": MESSAGES}) + assert [call.kwargs["api_base"] for call in host.calls] == ["http://10.0.0.7:4000"] * 2 + assert discovery.seen == [0, 1] + + +def test_an_assigned_pod_discovery_resolves_a_sync_attempt(make_core: MakeCore, host: ScriptedHost) -> None: + discovery: Final = _RecordingDiscovery() + core: Final = make_core(model_list=[_pod_deployment("a")]) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": discovery})) + completion: Final[Callable[..., object]] = getattr(core, "completion") # pyright: ignore[reportAny] # see _endpoint + _ = completion(model="g", messages=MESSAGES) + assert host.calls[0].kwargs["api_base"] == "http://10.0.0.7:4000" + + +def _api_base(record: object) -> object: + return (as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {}).get("api_base") + + +async def test_a_pick_is_resolved_to_a_pod_and_a_healthy_listing_is_not(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[_pod_deployment("a")]) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": _RecordingDiscovery()})) + assert _api_base(await core.select("g")) == "http://10.0.0.7:4000" + select: Final[Callable[..., Awaitable[object]]] = getattr(core, "select") # pyright: ignore[reportAny] # see _healthy_ids + listed: Final[object] = await select("g", kind="healthy") + assert isinstance(listed, list) + assert [_api_base(record) for record in listed] == ["http://svc.ns:4000"] # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the listing + + +def test_a_sync_pick_is_resolved_to_a_pod(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[_pod_deployment("a")]) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": _RecordingDiscovery()})) + pick: Final[Callable[..., object]] = getattr(core, "get_available_deployment") # pyright: ignore[reportAny] # the deprecated sync pick + assert _api_base(pick(model="g")) == "http://10.0.0.7:4000" + + +async def test_the_default_pod_discovery_resolves_a_service_host(make_core: MakeCore, host: ScriptedHost) -> None: + """``localhost`` resolves to a loopback address, which takes the service host's place.""" + core: Final = make_core( + model_list=[deployment("g", "a", api_base="http://localhost:4000", kubernetes_pod_discovery=True)] + ) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": KubernetesPodDiscovery(proxy_environment=lambda: {})})) + await core.call("acompletion", {"model": "g", "messages": MESSAGES}) + assert str(host.calls[0].kwargs["api_base"]) in ("http://127.0.0.1:4000", "http://[::1]:4000") diff --git a/tests/unit/test_router/core/test_registry_changesets.py b/tests/unit/test_router/core/test_registry_changesets.py index d4aa4d14a75..567b082607f 100644 --- a/tests/unit/test_router/core/test_registry_changesets.py +++ b/tests/unit/test_router/core/test_registry_changesets.py @@ -16,6 +16,7 @@ from tests.unit.test_router.core.scripted_host import ( ScriptedHost, context_window, deployment, + records_outcomes_through_callbacks, server_error, ) @@ -637,6 +638,8 @@ def test_retire_releases_every_callback_the_core_registered_on_its_host( make_core: MakeCore, host: ScriptedHost ) -> None: core = make_core(model_list=_model_list()) + if not records_outcomes_through_callbacks(core): + pytest.skip("this core registers no callbacks on its host") assert _registered(host) != {} core.retire() @@ -646,6 +649,8 @@ def test_retire_releases_every_callback_the_core_registered_on_its_host( def test_retire_leaves_another_cores_callbacks_registered(make_core: MakeCore, host: ScriptedHost) -> None: kept = make_core(model_list=_model_list()) + if not records_outcomes_through_callbacks(kept): + pytest.skip("this core registers no callbacks on its host") retired = make_core(model_list=[deployment("other", "o")]) registered_by_both = _registered(host) diff --git a/tests/unit/test_router/core/test_registry_characterization_construction.py b/tests/unit/test_router/core/test_registry_characterization_construction.py new file mode 100644 index 00000000000..7b80f6c3b71 --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_construction.py @@ -0,0 +1,379 @@ +""" +How the core builds its registry from ``model_list``: deployment ids, provider checks, environment +gating, ``os.environ/`` references, per-organization fan-out, wildcard and team registration, and +the pricing each deployment writes into ``litellm.model_cost``. +""" + +import pytest + +import litellm +from litellm.router_api import deployment_id +from tests.unit.test_router.core.registry_helpers import MakeCore, ids, team_deployment, unique_id +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +def _bare( + model_name: str, model: str, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: litellm_params + """A model_list entry with no ``model_info`` at all.""" + return {"model_name": model_name, "litellm_params": {"model": model, "api_key": "sk-test", **params}} + + +# -- deployment ids ----------------------------------------------------------------------------- + + +def test_a_deployment_without_an_id_gets_the_stable_hash_of_its_group_and_params(make_core: MakeCore) -> None: + core = make_core(model_list=[_bare("gpt", "openai/gpt-4o")]) + + expected = deployment_id("gpt", {"model": "openai/gpt-4o", "api_key": "sk-test"}) + assert core.registry().names(kind="ids") == [expected] + assert core.registry().has_deployment(expected) + + +def test_generated_ids_are_the_same_across_cores_built_from_the_same_list(make_core: MakeCore) -> None: + first = make_core(model_list=[_bare("gpt", "openai/gpt-4o")]) + second = make_core(model_list=[_bare("gpt", "openai/gpt-4o")]) + + assert first.registry().names(kind="ids") == second.registry().names(kind="ids") + + +def test_generated_ids_differ_when_the_group_or_params_differ(make_core: MakeCore) -> None: + core = make_core( + model_list=[_bare("gpt", "openai/gpt-4o"), _bare("other", "openai/gpt-4o"), _bare("gpt", "openai/gpt-4o-mini")] + ) + + assert len(set(core.registry().names(kind="ids"))) == 3 + + +def test_a_declared_non_string_id_is_kept_as_a_string(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + {"model_name": "gpt", "litellm_params": {"model": "openai/x", "api_key": "k"}, "model_info": {"id": 1}} + ] + ) + + assert core.registry().names(kind="ids") == ["1"] + found = core.registry().deployment("1") + assert found is not None + assert found.model_info.id == "1" + + +def test_two_deployments_declaring_one_id_are_both_kept_and_the_id_finds_the_last(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "d"), deployment("emb", "d")]) + + assert [entry["model_name"] for entry in core.registry().all()] == ["gpt", "emb"] + assert core.registry().index("model_name") == {"gpt": [0], "emb": [1]} + found = core.registry().deployment("d") + assert found is not None + assert found.model_name == "emb" + + +# -- the stored entries ------------------------------------------------------------------------- + + +def test_the_stored_entry_carries_litellm_param_defaults_and_model_info_flags(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + (entry,) = core.registry().all() + assert entry["model_name"] == "gpt" + assert entry["litellm_params"]["model"] == "openai/a" + assert entry["litellm_params"]["api_key"] == "sk-test" + assert entry["model_info"] == {"id": "a", "db_model": False, "member_auto_router": False} + + +def test_top_level_rpm_and_tpm_are_copied_into_litellm_params(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "openai/x", "api_key": "k"}, + "model_info": {"id": "r"}, + "rpm": 7, + "tpm": 9, + } + ] + ) + + (entry,) = core.registry().all() + assert (entry["litellm_params"]["rpm"], entry["litellm_params"]["tpm"]) == (7, 9) + assert (entry["rpm"], entry["tpm"]) == (7, 9) + + +def test_os_environ_references_in_litellm_params_are_resolved( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REGISTRY_CHAR_BASE", "https://env.example") + monkeypatch.setenv("REGISTRY_CHAR_KEY", "sk-from-env") + core = make_core( + model_list=[ + deployment("gpt", "a", api_key="os.environ/REGISTRY_CHAR_KEY", api_base="os.environ/REGISTRY_CHAR_BASE") + ] + ) + + params = core.registry().all()[0]["litellm_params"] + assert (params["api_key"], params["api_base"]) == ("sk-from-env", "https://env.example") + + +@pytest.mark.asyncio +async def test_resolved_environment_credentials_are_sent_to_the_provider( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REGISTRY_CHAR_BASE", "https://env.example") + core = make_core(model_list=[deployment("gpt", "a", api_base="os.environ/REGISTRY_CHAR_BASE")]) + + await core.acompletion(model="gpt", messages=[{"role": "user", "content": "hi"}]) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert [call.kwargs.get("api_base") for call in host.calls_to("acompletion")] == ["https://env.example"] + + +def test_a_list_of_organizations_fans_out_into_one_deployment_per_organization(make_core: MakeCore) -> None: + core = make_core(model_list=[_bare("gpt", "openai/x", organization=["org-1", "org-2"])]) + + entries = core.registry().all() + assert [entry["litellm_params"]["organization"] for entry in entries] == ["org-1", "org-2"] + assert len(set(ids(entries))) == 1 + assert core.registry().index("model_name") == {"gpt": [0, 1]} + + +def test_deployment_names_list_each_deployments_litellm_model(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("emb", "e")]) + + assert core.registry().names(kind="deployment_names") == ["openai/a", "openai/b", "openai/e"] + + +def test_every_configured_litellm_model_starts_with_zero_latency(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")]) + + assert core.health().latency == {"openai/a": 0, "openai/b": 0} + + +@pytest.mark.parametrize("model_list", [None, []], ids=["none", "empty"]) +def test_a_core_without_deployments_has_an_empty_registry(make_core: MakeCore, model_list: object) -> None: + core = make_core(model_list=model_list) + + assert core.registry().all() == [] + assert core.registry().names(kind="ids") == [] + assert core.registry().names(kind="group_set") == set() + assert core.registry().deployments() == [] + + +# -- invalid deployments ------------------------------------------------------------------------ + + +def test_an_unknown_provider_raises_bad_request_at_construction(make_core: MakeCore) -> None: + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + make_core(model_list=[deployment("gpt", "a", model="notaprovider/model")]) + + +def test_ignore_invalid_deployments_drops_the_invalid_one_and_keeps_the_rest(make_core: MakeCore) -> None: + core = make_core( + model_list=[deployment("gpt", "a", model="notaprovider/model"), deployment("gpt", "b")], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +def test_an_invalid_tag_regex_raises_value_error_naming_the_model(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=r"Invalid regex in tag_regex for model 'gpt': '\('"): + make_core(model_list=[deployment("gpt", "a", tag_regex=["("])]) + + +def test_ignore_invalid_deployments_drops_a_deployment_with_an_invalid_tag_regex(make_core: MakeCore) -> None: + core = make_core( + model_list=[deployment("gpt", "a", tag_regex=["("]), deployment("gpt", "b", tag_regex=["^ok$"])], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +# -- environment-gated deployments -------------------------------------------------------------- + + +def test_a_deployment_for_another_environment_is_left_out(make_core: MakeCore, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_ENVIRONMENT", "staging") + core = make_core( + model_list=[ + deployment("prod-only", "a", model_info={"supported_environments": ["production"]}), + deployment("both", "b", model_info={"supported_environments": ["staging", "production"]}), + ] + ) + + assert core.registry().names(kind="ids") == ["b"] + assert core.registry().deployments("prod-only") == [] + + +def test_a_left_out_environment_deployment_still_names_its_group( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + """Odd but current: the group set is read from the given model_list, not from what was added.""" + monkeypatch.setenv("LITELLM_ENVIRONMENT", "staging") + core = make_core( + model_list=[ + deployment("prod-only", "a", model_info={"supported_environments": ["production"]}), + deployment("h", "b"), + ] + ) + + assert core.registry().names(kind="group_set") == {"prod-only", "h"} + assert core.registry().index("model_name") == {"h": [0]} + + +def test_supported_environments_without_litellm_environment_raises( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("LITELLM_ENVIRONMENT", raising=False) + + with pytest.raises(ValueError, match="Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set"): + make_core(model_list=[deployment("gpt", "a", model_info={"supported_environments": ["production"]})]) + + +@pytest.mark.parametrize( + ("environment", "supported", "message"), + [ + ("bogus", ["production"], "LITELLM_ENVIRONMENT must be one of"), + ("staging", ["prod"], "supported_environments must be one of"), + ("staging", "production", "supported_environments must be a list of"), + ], + ids=["unknown_environment", "unknown_supported_value", "not_a_list"], +) +def test_invalid_environment_configuration_raises( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, environment: str, supported: object, message: str +) -> None: + monkeypatch.setenv("LITELLM_ENVIRONMENT", environment) + + with pytest.raises(ValueError, match=message): + make_core(model_list=[deployment("gpt", "a", model_info={"supported_environments": supported})]) + + +def test_ignore_invalid_deployments_drops_an_invalid_environment_deployment( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("LITELLM_ENVIRONMENT", raising=False) + core = make_core( + model_list=[ + deployment("gpt", "a", model_info={"supported_environments": ["production"]}), + deployment("gpt", "b"), + ], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +@pytest.mark.parametrize( + ("model_info", "active"), + [ + ({}, True), + ({"supported_environments": ["development"]}, True), + ({"supported_environments": ["production"]}, False), + ], + ids=["unset", "matching", "other"], +) +def test_deployment_is_active_for_environment_follows_litellm_environment( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], active: bool +) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + monkeypatch.setenv("LITELLM_ENVIRONMENT", "development") + candidate = litellm.types.router.Deployment(**deployment("gpt", "z", model_info=model_info)) # pyright: ignore[reportArgumentType] # a model_list entry + + assert core.deployment_is_active_for_environment(candidate) is active # pyright: ignore[reportAttributeAccessIssue] # public on the core + + +# -- wildcard and team registration ------------------------------------------------------------- + + +def test_a_wildcard_group_registers_a_pattern_and_a_provider_default_deployment(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("openai/*", "w", model="openai/*"), + deployment("anthropic/*", "x", model="anthropic/*"), + deployment("gpt", "a"), + ] + ) + + assert core.settings().provider_default_deployment_ids == ["w", "x"] + assert ids(core.registry().deployments("openai/gpt-4o")) == ["w"] + assert ids(core.registry().deployments("anthropic/claude-x")) == ["x"] + + +def test_team_deployments_are_indexed_by_team_and_public_name(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + team_deployment("internal-1", "t1", "team-a", "pub"), + team_deployment("internal-2", "t2", "team-a", "pub"), + team_deployment("internal-3", "t3", "team-b", "other"), + deployment("gpt", "g"), + ] + ) + + assert core.registry().index("team_model") == {("team-a", "pub"): [0, 1], ("team-b", "other"): [2]} + assert core.registry().names(kind="team_public") == frozenset({"pub", "other"}) + assert core.registry().names(kind="team_public", team_id="team-a") == frozenset({"pub"}) + assert core.registry().names(kind="team_public", team_id="nobody") == frozenset() + + +def test_a_team_deployment_needs_both_team_id_and_public_name_to_be_team_indexed(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("gpt", "a", model_info={"team_id": "team-a"}), + deployment("gpt", "b", model_info={"team_public_model_name": "pub"}), + ] + ) + + assert core.registry().index("team_model") == {} + assert core.registry().names(kind="team_public") == frozenset() + + +# -- pricing written into litellm.model_cost ----------------------------------------------------- + + +def test_custom_pricing_is_registered_under_the_deployment_id(make_core: MakeCore) -> None: + model_id = unique_id() + make_core( + model_list=[ + deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3, output_cost_per_token=2e-3) + ] + ) + + entry = litellm.model_cost[model_id] + assert (entry["input_cost_per_token"], entry["output_cost_per_token"]) == (1e-3, 2e-3) + + +def test_custom_input_pricing_inherits_the_backend_models_cache_read_price(make_core: MakeCore) -> None: + model_id = unique_id() + make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + + builtin = litellm.get_model_info("openai/gpt-4o") + assert litellm.model_cost[model_id]["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"] + + +def test_a_user_set_cache_read_price_is_not_overwritten_by_inheritance(make_core: MakeCore) -> None: + model_id = unique_id() + make_core( + model_list=[ + deployment( + "gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3, cache_read_input_token_cost=7e-9 + ) + ] + ) + + assert litellm.model_cost[model_id]["cache_read_input_token_cost"] == 7e-9 + + +def test_a_deployment_without_custom_pricing_registers_no_prices_under_its_id(make_core: MakeCore) -> None: + model_id = unique_id() + make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o")]) + + entry = litellm.model_cost[model_id] + assert entry.get("input_cost_per_token") is None + assert entry.get("output_cost_per_token") is None + + +def test_model_info_fields_are_registered_under_the_deployment_id(make_core: MakeCore) -> None: + model_id = unique_id() + make_core(model_list=[deployment("gpt", model_id, model_info={"max_input_tokens": 1000, "mode": "embedding"})]) + + entry = litellm.model_cost[model_id] + assert (entry["max_input_tokens"], entry["mode"]) == (1000, "embedding") diff --git a/tests/unit/test_router/core/test_registry_characterization_lookups.py b/tests/unit/test_router/core/test_registry_characterization_lookups.py new file mode 100644 index 00000000000..19d58ff10d4 --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_lookups.py @@ -0,0 +1,682 @@ +""" +Registry reads: deployments by group, alias, wildcard and team; names and ids; name resolution; +access groups; blocked and unhealthy groups; credentials lookups. +""" + +import time + +import pytest + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.registry_helpers import MakeCore, ids, names_and_ids, team_deployment, typed +from tests.unit.test_router.core.scripted_host import deployment + + +def _core(make_core: MakeCore) -> RouterCore: + """ + gpt (a, b), emb (e), the openai/* wildcard (w), team TA's ``pub`` (t1) and ``openai/*`` (tw), + team TB's ``pub`` (t2). Aliases: fast -> gpt, hid -> gpt (hidden), gpt -> emb. + """ + return make_core( + model_list=[ + deployment("gpt", "a", model_info={"access_groups": ["ag1", "ag2"]}), + deployment("gpt", "b", model_info={"access_groups": ["ag1"]}), + deployment("emb", "e", model="openai/text-embedding-3-small", model_info={"access_groups": ["ag2"]}), + deployment("openai/*", "w", model="openai/*", model_info={"access_groups": ["wild"]}), + team_deployment("int-a", "t1", "TA", "pub", model="openai/team-model"), + team_deployment("int-b", "t2", "TB", "pub", model="anthropic/claude-x"), + team_deployment("team-wild", "tw", "TA", "openai/*", model="openai/*"), + ], + model_group_alias={"fast": "gpt", "hid": {"model": "gpt", "hidden": True}, "gpt": "emb"}, + ) + + +# -- registry().deployments(): the group listing ------------------------------------------------ + + +def test_listing_every_deployment_puts_visible_alias_rows_before_the_model_list(make_core: MakeCore) -> None: + listed = _core(make_core).registry().deployments() + + assert names_and_ids(listed) == [ + ("fast", "a"), + ("fast", "b"), + ("gpt", "e"), + ("gpt", "a"), + ("gpt", "b"), + ("emb", "e"), + ("openai/*", "w"), + ("int-a", "t1"), + ("int-b", "t2"), + ("team-wild", "tw"), + ] + + +def test_an_alias_lists_its_targets_deployments_under_the_alias_name(make_core: MakeCore) -> None: + assert names_and_ids(_core(make_core).registry().deployments("fast")) == [("fast", "a"), ("fast", "b")] + + +def test_a_hidden_alias_is_listed_when_named(make_core: MakeCore) -> None: + assert names_and_ids(_core(make_core).registry().deployments("hid")) == [("hid", "a"), ("hid", "b")] + + +def test_a_group_that_is_also_an_alias_lists_its_own_deployments_then_the_alias_targets(make_core: MakeCore) -> None: + assert names_and_ids(_core(make_core).registry().deployments("gpt")) == [("gpt", "a"), ("gpt", "b"), ("gpt", "e")] + + +def test_an_alias_row_is_a_copy_and_leaves_the_configured_entry_named_as_configured(make_core: MakeCore) -> None: + core = _core(make_core) + + core.registry().deployments("fast") + + assert [entry["model_name"] for entry in core.registry().all()[:2]] == ["gpt", "gpt"] + + +def test_an_unknown_name_lists_nothing(make_core: MakeCore) -> None: + assert _core(make_core).registry().deployments("nope") == [] + + +def test_a_name_matching_only_a_wildcard_lists_the_wildcard_deployment_under_the_requested_name( + make_core: MakeCore, +) -> None: + assert names_and_ids(_core(make_core).registry().deployments("openai/gpt-x")) == [("openai/gpt-x", "w")] + + +def test_a_team_caller_gets_its_own_wildcard_deployment_after_the_global_one(make_core: MakeCore) -> None: + listed = _core(make_core).registry().deployments("openai/gpt-x", team_id="TA") + + assert names_and_ids(listed) == [("openai/gpt-x", "w"), ("openai/gpt-x", "tw")] + + +@pytest.mark.parametrize(("team_id", "expected"), [("TA", [("int-a", "t1")]), ("TB", [("int-b", "t2")]), (None, [])]) +def test_a_team_public_name_lists_only_that_teams_deployment( + make_core: MakeCore, team_id: str | None, expected: list[tuple[str, str]] +) -> None: + assert names_and_ids(_core(make_core).registry().deployments("pub", team_id=team_id)) == expected + + +@pytest.mark.parametrize(("team_id", "expected"), [(None, ["t1"]), ("TA", ["t1"]), ("TB", [])]) +def test_a_team_deployments_internal_name_lists_it_for_no_team_or_its_own_team_only( + make_core: MakeCore, team_id: str | None, expected: list[str] +) -> None: + assert ids(_core(make_core).registry().deployments("int-a", team_id=team_id)) == expected + + +def test_all_deployments_renames_rows_to_the_given_model_alias(make_core: MakeCore) -> None: + listed = _core(make_core).registry().deployments("gpt", source="all", model_alias="renamed") + + assert names_and_ids(listed) == [("renamed", "a"), ("renamed", "b")] + + +def test_all_deployments_follow_no_alias_and_no_wildcard(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.deployments("fast", source="all") == [] + assert registry.deployments("openai/gpt-x", source="all") == [] + + +def test_all_deployments_for_a_team_public_name_reads_the_team_index(make_core: MakeCore) -> None: + assert ids(_core(make_core).registry().deployments("pub", source="all", team_id="TA")) == ["t1"] + + +def test_named_deployments_are_the_group_index_entry_or_none(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert ids(registry.deployments("gpt", source="named")) == ["a", "b"] + assert registry.deployments("fast", source="named") is None + + +def test_first_in_group_is_the_first_configured_deployment_of_the_group(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + first = registry.deployments("gpt", source="first_in_group") + assert first is not None + assert first.model_info.id == "a" + assert registry.deployments("fast", source="first_in_group") is None + + +@pytest.mark.parametrize("source", ["all", "named", "wildcard", "routable", "first_in_group"]) +def test_a_source_that_needs_a_model_name_raises_without_one(make_core: MakeCore, source: str) -> None: + with pytest.raises(TypeError, match="needs a model_name"): + _core(make_core).registry().deployments(source=source) # pyright: ignore[reportCallIssue, reportArgumentType] # source is one of the literals + + +def test_an_unknown_deployment_source_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown deployment source 'bogus'"): + _core(make_core).registry().deployments("gpt", source="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid source on purpose + + +# -- wildcards ------------------------------------------------------------------------------------ + + +@pytest.mark.parametrize(("team_id", "expected"), [("TA", ["tw"]), ("TB", ["w"]), (None, ["w"])]) +def test_wildcard_route_deployments_prefer_the_teams_own_pattern_over_the_global_one( + make_core: MakeCore, team_id: str | None, expected: list[str] +) -> None: + assert ids(_core(make_core).registry().deployments("openai/gpt-x", source="wildcard", team_id=team_id)) == expected + + +def test_wildcard_routed_models_name_the_requested_model(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names("openai/gpt-x", kind="wildcard") == frozenset({"openai/gpt-x"}) + assert registry.names("anthropic/x", kind="wildcard") == frozenset() + + +# -- names ------------------------------------------------------------------------------------------ + + +def test_group_names_repeat_per_deployment_and_include_visible_aliases(make_core: MakeCore) -> None: + assert _core(make_core).registry().names() == ["fast", "fast", "gpt", "gpt", "gpt", "emb", "openai/*"] + + +def test_group_names_for_a_team_add_its_public_names(make_core: MakeCore) -> None: + assert _core(make_core).registry().names(team_id="TA") == [ + "fast", + "fast", + "gpt", + "gpt", + "gpt", + "emb", + "openai/*", + "pub", + "openai/*", + ] + + +def test_ids_can_leave_out_team_deployments(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names(kind="ids") == ["a", "b", "e", "w", "t1", "t2", "tw"] + assert registry.names(kind="ids", exclude_team_models=True) == ["a", "b", "e", "w"] + + +def test_ids_of_a_group_read_only_the_group_index(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names("fast", kind="ids") == [] + assert registry.names("pub", kind="ids") == [] + assert registry.names("int-a", kind="ids") == ["t1"] + assert registry.names("int-a", kind="ids", exclude_team_models=True) == [] + + +def test_litellm_models_follow_aliases_and_wildcards(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names("fast", kind="litellm_models") == ("openai/a", "openai/b") + assert registry.names("openai/gpt-x", kind="litellm_models") == ("openai/gpt-x",) + assert registry.names("nope", kind="litellm_models") == () + + +def test_litellm_base_models_use_a_deployments_base_model(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment( + "gpt", "a", model="azure/my-deployment", api_base="https://x", model_info={"base_model": "azure/gpt-4o"} + ), + deployment("gpt", "b", model="openai/gpt-4o-mini"), + ] + ) + + assert core.registry().names("gpt", kind="litellm_base_models") == ("azure/gpt-4o", "openai/gpt-4o-mini") + assert core.registry().names("gpt", kind="litellm_models") == ("azure/my-deployment", "openai/gpt-4o-mini") + + +@pytest.mark.parametrize( + ("name", "team_id", "expected"), + [("fast", None, {"a", "b"}), ("openai/gpt-x", None, {"w"}), ("pub", "TA", {"t1"}), ("nope", None, set())], +) +def test_route_candidates_are_the_ids_a_request_for_the_name_may_reach( + make_core: MakeCore, name: str, team_id: str | None, expected: set[str] +) -> None: + assert _core(make_core).registry().names(name, kind="route_candidates", team_id=team_id) == expected + + +@pytest.mark.parametrize("kind", ["litellm_models", "wildcard", "route_candidates"]) +def test_a_name_kind_that_needs_a_model_name_raises_without_one(make_core: MakeCore, kind: str) -> None: + with pytest.raises(TypeError, match="needs a model_name"): + _core(make_core).registry().names(kind=kind) # pyright: ignore[reportCallIssue, reportArgumentType] # kind is one of the literals + + +# -- resolve ---------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "expected"), + [("fast", "gpt"), ("hid", "gpt"), ("gpt", "emb"), ("emb", "emb"), ("nope", "nope")], +) +def test_resolving_a_group_follows_one_alias_hop(make_core: MakeCore, name: str, expected: str) -> None: + assert _core(make_core).registry().resolve(name) == expected + + +@pytest.mark.parametrize(("name", "expected"), [("fast", "gpt"), ("hid", "gpt"), ("emb", None), ("pub", None)]) +def test_resolving_an_alias_reads_plain_and_hidden_aliases( + make_core: MakeCore, name: str, expected: str | None +) -> None: + assert _core(make_core).registry().resolve(name, kind="alias") == expected + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("fast", True), + ("hid", True), + ("gpt", True), + ("a", True), + ("pub", False), + ("openai/gpt-x", False), + ("nope", False), + ], +) +def test_recognized_names_are_groups_aliases_and_deployment_ids(make_core: MakeCore, name: str, expected: bool) -> None: + assert _core(make_core).registry().resolve(name, kind="recognized") is expected + + +@pytest.mark.parametrize( + ("name", "team_id", "expected"), + [ + ("fast", None, "openai/a"), + ("gpt", None, "openai/text-embedding-3-small"), + ("pub", "TA", "openai/team-model"), + ("pub", None, None), + ("openai/gpt-x", None, None), + ("a", None, None), + ], +) +def test_upstream_is_the_first_routable_deployments_litellm_model( + make_core: MakeCore, name: str, team_id: str | None, expected: str | None +) -> None: + assert _core(make_core).registry().resolve(name, kind="upstream", team_id=team_id) == expected + + +def test_upstream_skips_a_blocked_deployment(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("half", "x2", model_info={"blocked": True}), deployment("half", "x3")]) + + assert core.registry().resolve("half", kind="upstream") == "openai/x3" + + +@pytest.mark.parametrize( + ("served", "expected"), [("gpt", "gpt"), ("emb", "fast"), (None, "fast")], ids=["target", "other", "none"] +) +def test_the_served_name_is_the_served_model_only_when_the_name_routes_to_it( + make_core: MakeCore, served: str | None, expected: str +) -> None: + assert _core(make_core).registry().resolve("fast", kind="served", served_model=served) == expected + + +def test_group_of_id_lists_the_deployments_group_including_alias_rows(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert names_and_ids(registry.resolve("a", kind="group_of_id")) == [("gpt", "a"), ("gpt", "b"), ("gpt", "e")] + assert registry.resolve("gpt", kind="group_of_id") is None + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("gpt", "gpt"), + ("a", "gpt"), + ("fast", "fast"), + ("nope", None), + (None, None), + ("", None), + ], +) +def test_from_id_resolves_a_group_name_or_deployment_id( + make_core: MakeCore, name: str | None, expected: str | None +) -> None: + assert _core(make_core).registry().resolve(name, kind="from_id") == expected + + +def test_from_id_matches_a_litellm_model_suffix_and_the_first_matching_row_wins(make_core: MakeCore) -> None: + """Odd but current: the alias row ``gpt -> emb`` is listed first, so ``emb``'s model resolves to ``gpt``.""" + registry = _core(make_core).registry() + + assert registry.resolve("text-embedding-3-small", kind="from_id") == "gpt" + assert registry.resolve("openai/text-embedding-3-small", kind="from_id") == "gpt" + + +def test_from_id_suffix_match_on_a_plain_core(make_core: MakeCore) -> None: + core = make_core( + model_list=[deployment("veo", "v", model="vertex_ai/veo-2.0", vertex_project="p", vertex_location="l")] + ) + + assert core.registry().resolve("veo-2.0", kind="from_id") == "veo" + + +def test_from_id_with_a_provider_returns_the_provider_prefixed_name_a_wildcard_serves(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.resolve("zzz", kind="from_id", custom_llm_provider="openai") == "openai/zzz" + assert registry.resolve("zzz", kind="from_id", custom_llm_provider="anthropic") is None + + +@pytest.mark.parametrize( + ("name", "team_id", "provider", "expected"), + [ + ("pub", "TA", None, "pub"), + ("gpt", "TA", None, None), + (None, "TA", None, "pub"), + (None, "TA", "openai", "pub"), + (None, "TB", "openai", None), + (None, "TB", "anthropic", "pub"), + (None, "nobody", None, None), + ], +) +def test_team_model_maps_a_teams_public_name_or_picks_its_first_deployment( + make_core: MakeCore, name: str | None, team_id: str, provider: str | None, expected: str | None +) -> None: + registry = _core(make_core).registry() + + assert registry.resolve(name, kind="team_model", team_id=team_id, custom_llm_provider=provider) == expected + + +def test_team_model_without_a_team_id_raises(make_core: MakeCore) -> None: + with pytest.raises(TypeError, match="needs a team_id"): + _core(make_core).registry().resolve("pub", kind="team_model") + + +def test_an_unknown_resolve_kind_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown resolve kind 'bogus'"): + _core(make_core).registry().resolve("gpt", kind="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid kind on purpose + + +# -- includes ------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("team_id", "expected"), [(None, False), ("TA", True), ("TB", False)]) +def test_a_team_deployment_is_usable_only_by_its_own_team( + make_core: MakeCore, team_id: str | None, expected: bool +) -> None: + registry = _core(make_core).registry() + team_entry = registry.all()[4] + + assert registry.includes(team_entry, team_id=team_id) is expected + + +def test_a_shared_deployment_is_usable_by_any_team(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert all(registry.includes(registry.all()[0], team_id=team) for team in (None, "TA", "TB")) + + +@pytest.mark.parametrize( + ("model_name", "team_id", "expected"), + [("pub", "TA", True), ("pub", "TB", False), ("int-a", None, True), ("int-a", "TB", False), ("gpt", "TA", False)], +) +def test_includes_with_a_model_name_matches_the_public_or_internal_name( + make_core: MakeCore, model_name: str, team_id: str | None, expected: bool +) -> None: + registry = _core(make_core).registry() + + assert registry.includes(registry.all()[4], model_name=model_name, team_id=team_id) is expected + + +# -- access groups -------------------------------------------------------------------------------- + + +def test_access_groups_map_each_group_to_the_model_names_that_carry_it(make_core: MakeCore) -> None: + assert _core(make_core).registry().access_groups() == { + "ag1": ["fast", "fast", "gpt", "gpt"], + "ag2": ["fast", "gpt", "gpt", "emb"], + "wild": ["openai/*"], + } + + +def test_access_groups_of_one_model_read_its_listing(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert dict(registry.access_groups(model_name="gpt")) == {"ag1": ["gpt", "gpt"], "ag2": ["gpt", "gpt"]} + assert dict(registry.access_groups(model_name="openai/gpt-x")) == {"wild": ["openai/gpt-x"]} + assert dict(registry.access_groups(model_name="nope")) == {} + + +def test_access_groups_filtered_to_one_group(make_core: MakeCore) -> None: + assert dict(_core(make_core).registry().access_groups(model_access_group="ag2")) == { + "ag2": ["fast", "gpt", "gpt", "emb"] + } + + +def test_access_groups_follow_a_deployment_added_later(make_core: MakeCore) -> None: + core = _core(make_core) + assert "late" not in core.registry().access_groups() + + core.add_deployment( # pyright: ignore[reportAttributeAccessIssue] # public on the core + typed("new", "n", model_info={"access_groups": ["late"]}) + ) + + assert core.registry().access_groups()["late"] == ["new"] + + +@pytest.mark.parametrize(("group", "expected"), [("wild", True), ("ag1", False), ("nope", False)]) +def test_an_access_group_is_a_wildcard_route_when_one_of_its_models_is_a_pattern( + make_core: MakeCore, group: str, expected: bool +) -> None: + assert _core(make_core).registry().resolve(group, kind="wildcard_access_group") is expected + + +# -- blocked and unhealthy groups ----------------------------------------------------------------- + + +def _blocked_core(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment("blk", "x1", model_info={"blocked": True}), + deployment("half", "x2", model_info={"blocked": True}), + deployment("half", "x3"), + deployment("ok", "x4"), + ], + model_group_alias={"blk-alias": "blk"}, + ) + + +def test_fully_blocked_names_are_groups_whose_every_deployment_is_blocked(make_core: MakeCore) -> None: + assert _blocked_core(make_core).health().fully_blocked_model_names() == {"blk", "blk-alias"} + + +def test_is_blocked_and_all_blocked_read_model_info_blocked(make_core: MakeCore) -> None: + core = _blocked_core(make_core) + health = core.health() + entries = core.registry().all() + + blocked = core.registry().deployment("x1") + open_ = core.registry().deployment("x3") + assert blocked is not None + assert open_ is not None + assert (health.is_blocked(blocked), health.is_blocked(open_)) == (True, False) + assert health.all_blocked(entries[:1]) is True + assert health.all_blocked(entries[1:3]) is False + assert health.all_blocked([]) is False + + +@pytest.mark.asyncio +async def test_fully_unhealthy_names_are_groups_whose_every_deployment_failed_its_health_check( + make_core: MakeCore, +) -> None: + core = make_core( + model_list=[ + deployment("sick", "s1"), + deployment("half", "h1"), + deployment("half", "h2"), + team_deployment("int", "t1", "TA", "team-sick"), + ] + ) + now = time.time() + core.health().health_checks.set_deployment_health_states( + { + "s1": {"is_healthy": False, "timestamp": now, "reason": "down"}, + "h1": {"is_healthy": False, "timestamp": now, "reason": "down"}, + "h2": {"is_healthy": True, "timestamp": now, "reason": ""}, + "t1": {"is_healthy": False, "timestamp": now, "reason": "down"}, + } + ) + + assert await core.health().fully_unhealthy_model_names() == {"sick", "int", "team-sick"} + + +@pytest.mark.asyncio +async def test_no_health_state_hides_nothing(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert await core.health().fully_unhealthy_model_names() == set() + + +@pytest.mark.asyncio +async def test_a_stale_unhealthy_state_hides_nothing(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + stale = time.time() - 10 * core.health().health_checks.staleness_threshold + core.health().health_checks.set_deployment_health_states( + {"a": {"is_healthy": False, "timestamp": stale, "reason": "x"}} + ) + + assert await core.health().fully_unhealthy_model_names() == set() + + +@pytest.mark.asyncio +async def test_an_allowed_fails_policy_hides_no_unhealthy_group(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")], allowed_fails_policy={"RateLimitErrorAllowedFails": 2}) + core.health().health_checks.set_deployment_health_states( + {"a": {"is_healthy": False, "timestamp": time.time(), "reason": "x"}} + ) + + assert await core.health().fully_unhealthy_model_names() == set() + + +# -- credentials ------------------------------------------------------------------------------------ + + +def _credential_core(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment("gpt", "a", api_base="https://a.example", temperature=0.2), + deployment("plain", "p", model="gpt-4o-mini"), + deployment("claude", "c", model="claude-3-haiku", custom_llm_provider="anthropic"), + deployment("blk", "x", model_info={"blocked": True}), + deployment("openai/*", "w", model="openai/*"), + team_deployment("shared", "ta", "TA", "team-name", api_key="sk-team-a"), + team_deployment("team-wild", "tw", "TA", "openai/*", model="openai/*", api_key="sk-team-wild"), + ] + ) + + +def test_credentials_are_the_credential_params_only(make_core: MakeCore) -> None: + assert _credential_core(make_core).registry().credentials("a") == { + "api_key": "sk-test", + "api_base": "https://a.example", + } + + +def test_credentials_of_an_unknown_or_blocked_deployment_are_none(make_core: MakeCore) -> None: + registry = _credential_core(make_core).registry() + + assert registry.credentials("missing") is None + assert registry.credentials("x") is None + assert registry.credentials("x", form="with_provider") is None + assert registry.credentials("x", form="deployment") is None + + +@pytest.mark.parametrize( + ("lookup", "provider"), + [("a", "openai"), ("p", "openai"), ("c", "anthropic")], + ids=["prefix", "default", "explicit"], +) +def test_credentials_with_provider_add_the_model_and_its_provider( + make_core: MakeCore, lookup: str, provider: str +) -> None: + registry = _credential_core(make_core).registry() + + found = registry.credentials(lookup, form="with_provider") + deployment_found = registry.deployment(lookup) + assert found is not None + assert deployment_found is not None + assert (found["model"], found["custom_llm_provider"]) == (deployment_found.litellm_params.model, provider) + + +@pytest.mark.parametrize( + ("lookup", "team_id", "expected"), + [ + ("a", None, "a"), + ("gpt", None, "a"), + ("team-name", "TA", "ta"), + ("team-name", None, None), + ("shared", "TB", None), + ("shared", "TA", "ta"), + ("openai/gpt-x", None, "w"), + ("openai/gpt-x", "TA", "tw"), + ("openai/gpt-x", "TB", "w"), + ("nope", None, None), + ], +) +def test_the_credential_deployment_is_found_by_id_group_team_name_then_wildcard( + make_core: MakeCore, lookup: str, team_id: str | None, expected: str | None +) -> None: + found = _credential_core(make_core).registry().credentials(lookup, form="deployment", team_id=team_id) + + assert (None if found is None else found.model_info.id) == expected + + +def test_credentials_with_provider_for_a_wildcard_name_carry_the_requested_model(make_core: MakeCore) -> None: + found = _credential_core(make_core).registry().credentials("openai/gpt-x", form="with_provider", team_id="TA") + + assert found is not None + assert (found["api_key"], found["model"], found["custom_llm_provider"]) == ( + "sk-team-wild", + "openai/gpt-x", + "openai", + ) + + +def test_an_unknown_credential_form_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown credential form 'bogus'"): + _credential_core(make_core).registry().credentials("a", form="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid form on purpose + + +# -- single deployments and the raw list ------------------------------------------------------------ + + +def test_deployment_dict_is_the_live_model_list_entry(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.deployment_dict("a") is registry.all()[0] + assert registry.deployment_dict("missing") is None + + +def test_a_looked_up_deployment_is_a_copy(make_core: MakeCore) -> None: + core = _core(make_core) + found = core.registry().deployment("a") + assert found is not None + + found.litellm_params.api_base = "https://mutated.example" + + assert core.registry().all()[0]["litellm_params"].get("api_base") is None + + +def test_a_visible_dict_alias_is_listed_like_a_plain_alias(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")], model_group_alias={"vis": {"model": "gpt", "hidden": False}}) + + assert names_and_ids(core.registry().deployments()) == [("vis", "a"), ("gpt", "a")] + assert core.registry().names() == ["vis", "gpt"] + + +def test_the_deprecated_generate_model_id_is_the_registry_id_hash(make_core: MakeCore) -> None: + from litellm.router_api import deployment_id + + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match="generate_model_id is deprecated"): + generated = core.generate_model_id("g", {"model": "x"}) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert generated == deployment_id("g", {"model": "x"}) + + +def test_assigning_model_list_swaps_the_list_without_reindexing(make_core: MakeCore) -> None: + """Odd but current: the id index still maps ``a`` to position 0, which now holds the new entry.""" + core = make_core(model_list=[deployment("gpt", "a")]) + version = core.registry().version + + core.model_list = [deployment("n", "n")] # pyright: ignore[reportAttributeAccessIssue] # deprecated attribute, still assignable + + assert core.registry().version == version + 1 + assert core.registry().names(kind="ids") == ["a"] + found = core.registry().deployment("a") + assert found is not None + assert found.model_name == "n" diff --git a/tests/unit/test_router/core/test_registry_characterization_model_info.py b/tests/unit/test_router/core/test_registry_characterization_model_info.py new file mode 100644 index 00000000000..d6ded63339d --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_model_info.py @@ -0,0 +1,540 @@ +""" +What the registry knows about models: group info, deployment and router model info, configured +limits/mode/display name/service tiers/listing info, group usage this minute, remaining usage and +the rate limit headers ``set_response_headers`` attaches. +""" + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import Deployment, RouterCacheEnum +from litellm.utils import get_utc_datetime +from tests.unit.test_router.core.registry_helpers import MakeCore, typed, unique_id +from tests.unit.test_router.core.scripted_host import completion_response, deployment + + +def _usage_core(make_core: MakeCore) -> RouterCore: + """gpt: a (tpm 100, rpm 10) and b (top-level tpm 1000, model_info tpm 50 / rpm 5); io: i (itpm 300, otpm 40).""" + return make_core( + model_list=[ + deployment("gpt", "a", model="openai/gpt-4o", tpm=100, rpm=10), + { + "model_name": "gpt", + "litellm_params": {"model": "anthropic/claude-3-haiku-20240307", "api_key": "k"}, + "model_info": {"id": "b", "tpm": 50, "rpm": 5, "description": "desc b"}, + "tpm": 1000, + }, + deployment("io", "i", model="openai/gpt-4o", itpm=300, otpm=40), + deployment("plain", "p", model="openai/gpt-4o"), + ], + model_group_alias={"hid": {"model": "gpt", "hidden": True}, "al": "gpt"}, + ) + + +def _minute() -> str: + return get_utc_datetime().strftime("%H-%M") + + +async def _count(core: RouterCore, kind: RouterCacheEnum, deployment_id: str, model: str, value: int) -> None: + key = kind.value.format(id=deployment_id, model=model, current_minute=_minute()) + await core.health().store.async_set_cache(key, value) + + +# -- model group info ----------------------------------------------------------------------------- + + +def test_group_info_merges_providers_limits_and_description_across_deployments(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("gpt") + + assert info is not None + assert info.model_group == "gpt" + assert info.providers == ["openai", "anthropic"] + assert info.description == "desc b" + assert info.mode == "chat" + builtin = litellm.get_model_info("openai/gpt-4o") + assert info.max_input_tokens == builtin["max_input_tokens"] + assert info.input_cost_per_token == builtin["input_cost_per_token"] + + +def test_group_info_sums_tpm_and_rpm_with_top_level_values_before_model_info(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("gpt") + + assert info is not None + assert (info.tpm, info.rpm) == (1100, 15) + + +def test_group_info_sums_itpm_and_otpm(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("io") + + assert info is not None + assert (info.itpm, info.otpm, info.tpm, info.rpm) == (300, 40, None, None) + + +def test_group_info_through_an_alias_is_reported_under_the_alias(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("al") + + assert info is not None + assert (info.model_group, info.tpm) == ("al", 1100) + + +def test_group_info_of_a_hidden_alias_needs_include_hidden(make_core: MakeCore) -> None: + registry = _usage_core(make_core).registry() + + assert registry.group("hid") is None + shown = registry.group("hid", include_hidden=True) + assert shown is not None + assert shown.model_group == "hid" + + +def test_group_info_of_an_unknown_name_is_none(make_core: MakeCore) -> None: + assert _usage_core(make_core).registry().group("nope") is None + + +def test_group_info_of_a_name_a_wildcard_serves_reads_the_requested_model(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("openai/*", "w", model="openai/*")]) + + info = core.registry().group("openai/gpt-4o-mini") + + assert info is not None + assert (info.model_group, info.providers) == ("openai/gpt-4o-mini", ["openai"]) + assert info.max_input_tokens == litellm.get_model_info("openai/gpt-4o-mini")["max_input_tokens"] + + +def test_group_info_of_an_unmapped_model_uses_its_configured_mode_and_costs(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment( + "unm", + unique_id(), + model="openai/totally-unmapped-registry-model", + model_info={"mode": "embedding", "input_cost_per_token": "0.5"}, + ) + ] + ) + + info = core.registry().group("unm") + + assert info is not None + assert (info.mode, info.input_cost_per_token, info.providers) == ("embedding", 0.5, ["openai"]) + assert info.max_input_tokens is None + + +def test_group_info_reports_configurable_clientside_auth_params(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a", configurable_clientside_auth_params=["api_base"])]) + + info = core.registry().group("gpt") + + assert info is not None + assert info.configurable_clientside_auth_params == ["api_base"] + + +def test_cached_group_info_equals_group_info(make_core: MakeCore) -> None: + registry = _usage_core(make_core).registry() + + assert registry.group("gpt", field="cached_info") == registry.group("gpt") + + +def test_cached_group_info_follows_a_deployment_added_later(make_core: MakeCore) -> None: + core = _usage_core(make_core) + before = core.registry().group("gpt", field="cached_info") + assert before is not None + + core.add_deployment(typed("gpt", "c", tpm=1)) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + after = core.registry().group("gpt", field="cached_info") + assert after is not None + assert (before.tpm, after.tpm) == (1100, 1101) + + +def test_an_unknown_group_field_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown group field 'bogus'"): + _usage_core(make_core).registry().group("gpt", field="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid field on purpose + + +# -- configured values of a group ----------------------------------------------------------------- + + +def _configured_core(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment( + "cfg", + "c1", + model="openai/gpt-4o", + model_info={ + "max_input_tokens": 1000, + "mode": "embedding", + "display_name": "Nice", + "service_tiers": ["flex"], + }, + ), + deployment( + "cfg", "c2", model="openai/gpt-4o-mini", model_info={"max_input_tokens": 5000, "max_output_tokens": 77} + ), + deployment("blank", "b1", model_info={"mode": " ", "display_name": ""}), + deployment("bad", "x1", model_info={"max_input_tokens": "lots"}), + deployment("openai/*", "w", model="openai/*"), + ], + model_group_alias={"cfg-alias": "cfg"}, + ) + + +def test_token_limits_come_from_the_first_deployment_of_the_group(make_core: MakeCore) -> None: + assert _configured_core(make_core).registry().group("cfg", field="token_limits") == (1000, None) + + +@pytest.mark.parametrize("name", ["nope", "openai/gpt-x", "cfg-alias"]) +def test_token_limits_of_an_unknown_wildcard_or_alias_name_are_none(make_core: MakeCore, name: str) -> None: + assert _configured_core(make_core).registry().group(name, field="token_limits") == (None, None) + + +def test_a_malformed_token_limit_reads_as_absent(make_core: MakeCore) -> None: + assert _configured_core(make_core).registry().group("bad", field="token_limits") == (None, None) + + +@pytest.mark.parametrize( + ("name", "mode", "display_name"), + [("cfg", "embedding", "Nice"), ("blank", None, None), ("nope", None, None), ("cfg-alias", None, None)], +) +def test_mode_and_display_name_are_the_first_deployments_non_blank_values( + make_core: MakeCore, name: str, mode: str | None, display_name: str | None +) -> None: + registry = _configured_core(make_core).registry() + + assert registry.group(name, field="mode") == mode + assert registry.group(name, field="display_name") == display_name + + +def test_listing_info_reports_the_widest_limits_and_each_backend_model_once(make_core: MakeCore) -> None: + info = _configured_core(make_core).registry().group("cfg", field="listing_info") + + assert info is not None + assert info.cost_map_keys == ("openai/gpt-4o", "openai/gpt-4o-mini") + assert (info.max_input_tokens, info.max_output_tokens) == (5000, 77) + + +@pytest.mark.parametrize("name", ["nope", "openai/gpt-x", "cfg-alias"]) +def test_listing_info_of_an_unknown_wildcard_or_alias_name_is_none(make_core: MakeCore, name: str) -> None: + assert _configured_core(make_core).registry().group(name, field="listing_info") is None + + +def test_service_tiers_are_each_routable_deployments_value_in_order(make_core: MakeCore) -> None: + registry = _configured_core(make_core).registry() + + assert registry.group("cfg", field="service_tiers") == (["flex"], None) + assert registry.group("cfg-alias", field="service_tiers") == (["flex"], None) + assert registry.group("nope", field="service_tiers") == () + + +def test_service_tiers_leave_out_blocked_deployments_and_other_teams(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("st", "s1", model_info={"service_tiers": ["a"], "blocked": True}), + deployment("st", "s2", model_info={"service_tiers": ["b"]}), + deployment("st", "s3", model_info={"service_tiers": ["c"], "team_id": "TA"}), + ] + ) + + assert core.registry().group("st", field="service_tiers") == (["b"],) + assert core.registry().group("st", field="service_tiers", team_id="TA") == (["b"], ["c"]) + assert core.registry().group("st", field="service_tiers", team_id="TB") == (["b"],) + + +# -- deployment model info -------------------------------------------------------------------------- + + +def test_deployment_model_info_merges_custom_pricing_over_the_builtin_entry(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + + info = core.registry().deployment_model_info(model_id, "openai/gpt-4o") + + assert info is not None + builtin = litellm.get_model_info("openai/gpt-4o") + assert info["input_cost_per_token"] == 1e-3 + assert info["max_input_tokens"] == builtin["max_input_tokens"] + + +def test_deployment_model_info_without_custom_pricing_is_the_builtin_entry(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o")]) + + info = core.registry().deployment_model_info(model_id, "openai/gpt-4o") + + assert info is not None + builtin = litellm.get_model_info("openai/gpt-4o") + assert (info["input_cost_per_token"], info["output_cost_per_token"]) == ( + builtin["input_cost_per_token"], + builtin["output_cost_per_token"], + ) + + +def test_deployment_model_info_of_an_unmapped_model_is_its_custom_entry(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core( + model_list=[ + deployment( + "unm", + model_id, + model="openai/unmapped-registry-thing", + input_cost_per_token=5e-6, + model_info={"max_input_tokens": 1000}, + ) + ] + ) + + info = core.registry().deployment_model_info(model_id, "openai/unmapped-registry-thing") + + assert info is not None + assert (info["input_cost_per_token"], info["max_input_tokens"]) == (5e-6, 1000) + + +def test_deployment_model_info_of_an_unknown_id_and_model_is_none(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert core.registry().deployment_model_info("missing-id", "nope/nothing-at-all") is None + + +def test_cached_deployment_model_info_equals_the_uncached_one(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + + assert core.registry().deployment_model_info(model_id, "openai/gpt-4o", cached=True) == ( + core.registry().deployment_model_info(model_id, "openai/gpt-4o") + ) + + +def test_discovered_model_info_is_empty_without_a_refresh(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert dict(core.registry().discovered_model_info("a")) == {} + assert dict(core.registry().discovered_model_info(None)) == {} + + +# -- router model info ------------------------------------------------------------------------------ + + +def test_router_model_info_by_id_is_the_builtin_entry_with_the_deployments_model_info_on_top( + make_core: MakeCore, +) -> None: + model_id = unique_id() + core = make_core( + model_list=[ + deployment( + "gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3, model_info={"max_input_tokens": 99} + ) + ] + ) + + info = core.registry().router_model_info(None, "gpt", id=model_id) + + builtin = litellm.get_model_info("openai/gpt-4o") + assert info["key"] == builtin["key"] + assert info["max_input_tokens"] == 99 + assert info["max_output_tokens"] == builtin["max_output_tokens"] + + +def test_router_model_info_of_a_passed_deployment_dict(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + info = core.registry().router_model_info(deployment("x", "x", model="openai/gpt-4o-mini"), "x") + + assert info["key"] == litellm.get_model_info("openai/gpt-4o-mini")["key"] + + +def test_router_model_info_reads_the_base_model(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment( + "az", "z", model="azure/my-deployment", api_base="https://x", model_info={"base_model": "azure/gpt-4o"} + ) + ] + ) + + info = core.registry().router_model_info(None, "az", id="z") + + assert info["max_input_tokens"] == litellm.get_model_info("azure/gpt-4o")["max_input_tokens"] + + +def test_router_model_info_of_a_wildcard_deployment_reads_the_requested_model(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("openai/*", "w", model="openai/*")]) + + info = core.registry().router_model_info(None, "openai/gpt-4o-mini", id="w") + + assert info["key"] == litellm.get_model_info("openai/gpt-4o-mini")["key"] + + +def test_router_model_info_for_an_unknown_id_and_no_deployment_raises(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + with pytest.raises(ValueError, match="Deployment not found"): + core.registry().router_model_info(None, "gpt", id="missing") + + +def test_router_model_info_of_a_deployment_without_litellm_params_raises(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + with pytest.raises(ValueError, match="Deployment missing valid litellm_params"): + core.registry().router_model_info({"model_name": "x", "litellm_params": {}, "model_info": {"id": "q"}}, "x") + + +def test_router_model_info_of_an_unmapped_model_is_the_zeroed_stub_construction_registered(make_core: MakeCore) -> None: + """Construction registers every backend model in the cost map, so an unmapped one reads as free, not missing.""" + core = make_core(model_list=[deployment("gpt", "a", model="openai/not-a-known-registry-model")]) + + info = core.registry().router_model_info(None, "gpt", id="a") + + assert info["key"] == "openai/not-a-known-registry-model" + assert (info["input_cost_per_token"], info["output_cost_per_token"], info["max_input_tokens"]) == (0, 0, None) + assert info["litellm_provider"] == "openai" + + +# -- usage this minute ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_group_usage_is_none_until_something_is_counted(make_core: MakeCore) -> None: + core = _usage_core(make_core) + + assert await core.health().usage("gpt") == (None, None) + assert await core.health().usage("nope") == (None, None) + + +@pytest.mark.asyncio +async def test_group_usage_sums_each_deployments_counters_for_this_minute(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.TPM, "a", "openai/gpt-4o", 30) + await _count(core, RouterCacheEnum.TPM, "b", "anthropic/claude-3-haiku-20240307", 4) + await _count(core, RouterCacheEnum.RPM, "b", "anthropic/claude-3-haiku-20240307", 2) + + assert await core.health().usage("gpt") == (34, 2) + assert await core.health().usage("al") == (34, 2) + + +@pytest.mark.asyncio +async def test_io_token_usage_sums_input_and_output_counters(make_core: MakeCore) -> None: + core = _usage_core(make_core) + assert await core.get_model_group_io_token_usage("io") == (None, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core + await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + + assert await core.get_model_group_io_token_usage("io") == (7, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + +@pytest.mark.asyncio +async def test_remaining_usage_is_each_limit_minus_this_minutes_usage(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.TPM, "a", "openai/gpt-4o", 30) + await _count(core, RouterCacheEnum.RPM, "b", "anthropic/claude-3-haiku-20240307", 2) + + assert await core.health().usage("gpt", kind="remaining") == { + "x-ratelimit-remaining-tokens": 1070, + "x-ratelimit-limit-tokens": 1100, + "x-ratelimit-remaining-requests": 13, + "x-ratelimit-limit-requests": 15, + } + + +@pytest.mark.asyncio +async def test_remaining_usage_of_an_io_limited_group_reports_input_and_output_token_headers( + make_core: MakeCore, +) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + + remaining = await core.health().usage("io", kind="remaining") + + assert remaining["x-ratelimit-limit-input-tokens"] == 300 + assert remaining["x-ratelimit-remaining-input-tokens"] == 293 + assert remaining["x-ratelimit-limit-output-tokens"] == 40 + assert remaining["x-ratelimit-remaining-output-tokens"] == 40 + assert "x-ratelimit-limit-tokens" not in remaining + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group", ["plain", "nope"]) +async def test_remaining_usage_of_a_group_without_limits_is_empty(make_core: MakeCore, group: str) -> None: + assert await _usage_core(make_core).health().usage(group, kind="remaining") == {} + + +@pytest.mark.asyncio +async def test_an_unknown_usage_kind_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown usage kind 'bogus'"): + await _usage_core(make_core).health().usage("gpt", kind="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid kind on purpose + + +# -- set_response_headers --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_response_headers_name_the_group_and_its_remaining_limits(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.TPM, "a", "openai/gpt-4o", 30) + response = completion_response() + + returned = await core.set_response_headers(response, model_group="gpt") # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert returned is response + assert response._hidden_params["additional_headers"] == { + "x-litellm-model-group": "gpt", + "x-ratelimit-remaining-tokens": 1070, + "x-ratelimit-limit-tokens": 1100, + "x-ratelimit-remaining-requests": 15, + "x-ratelimit-limit-requests": 15, + } + + +@pytest.mark.asyncio +async def test_response_headers_without_a_group_carry_only_a_none_group(make_core: MakeCore) -> None: + returned = await _usage_core(make_core).set_response_headers(completion_response(), model_group=None) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert returned._hidden_params["additional_headers"] == {"x-litellm-model-group": None} + + +@pytest.mark.asyncio +async def test_response_headers_of_a_group_without_limits_carry_only_the_group(make_core: MakeCore) -> None: + returned = await _usage_core(make_core).set_response_headers(completion_response(), model_group="plain") # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert returned._hidden_params["additional_headers"] == {"x-litellm-model-group": "plain"} + + +@pytest.mark.asyncio +async def test_response_headers_of_no_response_are_none(make_core: MakeCore) -> None: + assert await _usage_core(make_core).set_response_headers(None, model_group="gpt") is None # pyright: ignore[reportAttributeAccessIssue] # public on the core + + +@pytest.mark.asyncio +async def test_response_headers_stamp_the_model_id_of_the_deployment_in_the_request(make_core: MakeCore) -> None: + returned = await _usage_core(make_core).set_response_headers( # pyright: ignore[reportAttributeAccessIssue] # public on the core + completion_response(), model_group="plain", request_kwargs={"metadata": {"model_info": {"id": "p"}}} + ) + + assert returned._hidden_params["model_id"] == "p" + + +# -- budgets per deployment ------------------------------------------------------------------------- + + +def test_no_deployment_budget_means_no_budget_limiter(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert core.registry().plugins("budget_limiter") is None + + +@pytest.mark.asyncio +async def test_adding_a_deployment_with_a_budget_installs_the_budget_limiter(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + core.add_deployment( # pyright: ignore[reportAttributeAccessIssue] # public on the core + Deployment(**deployment("gpt", "budgeted", max_budget=1.0, budget_duration="1d")) # pyright: ignore[reportArgumentType] # a model_list entry + ) + + limiter = core.registry().plugins("budget_limiter") + assert limiter is not None + assert type(limiter).__name__ == "RouterBudgetLimiting" + + +def test_a_team_without_ptu_shares_has_no_ptu_ceiling(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert core.health().team_ptu_ceiling("TA", "gpt") is None diff --git a/tests/unit/test_router/core/test_registry_characterization_writes.py b/tests/unit/test_router/core/test_registry_characterization_writes.py new file mode 100644 index 00000000000..0558a16fe9d --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_writes.py @@ -0,0 +1,516 @@ +""" +Registry and settings writes beyond the version bookkeeping in ``test_registry_changesets.py``: +what ``add_deployment``, ``set_model_list``, upserts and deletes leave behind in every index, how +``update_settings`` / ``ChangeSet(settings=...)`` coerce and filter values, and the ``settings()`` +and ``registry()`` views. +""" + +import warnings + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import Deployment, RetryPolicy +from tests.unit.test_router.core.registry_helpers import ( + MESSAGES, + MakeCore, + ids, + names_and_ids, + team_deployment, + typed, + unique_id, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +def _core(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + """gpt (a, b), emb (e), the openai/* wildcard (w), team TA's ``pub`` (t1, internal name ``int``).""" + return make_core( + model_list=[ + deployment("gpt", "a"), + deployment("gpt", "b"), + deployment("emb", "e"), + deployment("openai/*", "w", model="openai/*"), + team_deployment("int", "t1", "TA", "pub"), + ], + **arguments, + ) + + +# -- add_deployment --------------------------------------------------------------------------------- + + +def test_add_deployment_returns_the_deployment_it_serves(make_core: MakeCore) -> None: + core = _core(make_core) + + added = core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert isinstance(added, Deployment) + assert (added.model_name, added.model_info.id, added.litellm_params.model) == ("new", "n", "openai/n") + + +def test_an_added_deployment_is_appended_and_names_a_new_group(make_core: MakeCore) -> None: + core = _core(make_core) + + core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1", "n"] + assert "new" in core.registry().names(kind="group_set") + assert core.registry().index("model_name")["new"] == [5] + assert core.registry().names(kind="deployment_names")[-1] == "openai/n" + + +def test_an_added_deployment_without_an_id_gets_the_id_its_model_info_generated(make_core: MakeCore) -> None: + core = _core(make_core) + candidate = Deployment(model_name="new", litellm_params={"model": "openai/q", "api_key": "k"}) # pyright: ignore[reportArgumentType] # litellm_params as a dict + + added = core.add_deployment(candidate) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert added is not None + assert added.model_info.id == candidate.model_info.id + assert core.registry().has_deployment(str(candidate.model_info.id)) + + +@pytest.mark.asyncio +async def test_an_added_deployment_serves_requests_for_its_group(make_core: MakeCore, host: ScriptedHost) -> None: + core = _core(make_core) + + core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + await core.acompletion(model="new", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert host.deployments_called("acompletion") == ["n"] + + +def test_an_added_wildcard_deployment_serves_matching_names(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + core.add_deployment(typed("anthropic/*", "x", model="anthropic/*")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert names_and_ids(core.registry().deployments("anthropic/claude-x")) == [("anthropic/claude-x", "x")] + assert core.settings().provider_default_deployment_ids == ["x"] + + +def test_an_added_team_deployment_is_team_indexed(make_core: MakeCore) -> None: + core = _core(make_core) + + core.add_deployment(Deployment(**team_deployment("int-2", "t2", "TB", "pub-b"))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] # public on the core + + assert core.registry().index("team_model")[("TB", "pub-b")] == [5] + assert core.registry().names(kind="team_public") == frozenset({"pub", "pub-b"}) + assert ids(core.registry().deployments("pub-b", team_id="TB")) == ["t2"] + + +def test_adding_a_deployment_with_an_unknown_provider_raises_and_adds_nothing(make_core: MakeCore) -> None: + core = _core(make_core) + version = core.registry().version + + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + core.add_deployment(typed("bad", "bad", model="notaprovider/x")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1"] + assert "bad" not in core.registry().names(kind="group_set") + assert core.registry().version == version + + +def test_an_added_deployment_registers_its_custom_pricing(make_core: MakeCore) -> None: + core = _core(make_core) + model_id = unique_id() + + core.add_deployment(typed("priced", model_id, model="openai/gpt-4o", input_cost_per_token=3e-3)) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert litellm.model_cost[model_id]["input_cost_per_token"] == 3e-3 + + +# -- upsert ----------------------------------------------------------------------------------------- + + +def test_a_changed_upsert_moves_the_deployment_to_the_end_of_the_list(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(upsert=[typed("gpt", "a", api_base="https://changed.example")])) + + assert core.registry().names(kind="ids") == ["b", "e", "w", "t1", "a"] + assert core.registry().index("model_name") == {"gpt": [0, 4], "emb": [1], "openai/*": [2], "int": [3]} + assert ids(core.registry().deployments("gpt")) == ["b", "a"] + + +def test_an_upsert_that_renames_the_group_moves_the_deployment_between_groups(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(upsert=[typed("emb", "a")])) + + assert ids(core.registry().deployments("gpt")) == ["b"] + assert ids(core.registry().deployments("emb")) == ["e", "a"] + + +def test_an_upsert_that_empties_a_group_drops_the_group_name(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(upsert=[typed("emb2", "e")])) + + assert core.registry().names(kind="group_set") == {"gpt", "openai/*", "int", "emb2"} + assert core.registry().deployments("emb") == [] + + +def test_a_failing_upsert_of_an_existing_deployment_raises_and_leaves_it_removed(make_core: MakeCore) -> None: + """Odd but current: without ignore_invalid_deployments the old deployment is popped and not restored.""" + core = _core(make_core) + + with pytest.raises(litellm.BadRequestError): + core.apply(ChangeSet(upsert=[typed("gpt", "a", model="notaprovider/x")])) + + assert core.registry().names(kind="ids") == ["b", "e", "w", "t1"] + + +def test_a_failing_upsert_under_ignore_invalid_deployments_restores_the_previous_deployment( + make_core: MakeCore, +) -> None: + core = _core(make_core, ignore_invalid_deployments=True) + + result = core.apply(ChangeSet(upsert=[typed("gpt", "a", model="notaprovider/x")])) + + assert list(result.upserted) == [None] + restored = core.registry().deployment("a") + assert restored is not None + assert restored.litellm_params.model == "openai/a" + assert core.registry().names(kind="ids") == ["b", "e", "w", "t1", "a"] + + +def test_the_deprecated_upsert_deployment_warns_and_returns_the_deployment(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match=r"upsert_deployment is deprecated; use .*apply\(ChangeSet\(upsert="): + upserted = core.upsert_deployment(typed("gpt", "c")) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert upserted is not None + assert upserted.model_info.id == "c" + + +# -- delete ----------------------------------------------------------------------------------------- + + +def test_delete_returns_the_removed_model_list_entry(make_core: MakeCore) -> None: + core = _core(make_core) + + result = core.apply(ChangeSet(delete=["e"])) + + (removed,) = result.deleted + assert isinstance(removed, dict) + assert (removed["model_name"], removed["model_info"]["id"]) == ("emb", "e") + + +def test_deleting_the_last_deployment_of_a_group_drops_the_group(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["e"])) + + assert "emb" not in core.registry().names(kind="group_set") + assert core.registry().deployments("emb") == [] + assert core.registry().index("model_name") == {"gpt": [0, 1], "openai/*": [2], "int": [3]} + + +def test_deleting_a_deployment_reindexes_the_ones_after_it(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["a"])) + + for deployment_id in ("b", "e", "w", "t1"): + found = core.registry().deployment(deployment_id) + assert found is not None + assert found.model_info.id == deployment_id + assert core.registry().index("team_model") == {("TA", "pub"): [3]} + + +def test_deleting_a_team_deployment_drops_its_team_index_and_public_name(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["t1"])) + + assert core.registry().index("team_model") == {} + assert core.registry().names(kind="team_public") == frozenset() + assert core.registry().deployments("pub", team_id="TA") == [] + + +def test_deleting_a_wildcard_deployment_stops_it_serving_matching_names(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["w"])) + + assert core.registry().deployments("openai/gpt-x") == [] + + +def test_deleting_a_wildcard_deployment_keeps_it_among_provider_default_deployment_ids(make_core: MakeCore) -> None: + """Odd but current: the provider default id list is append-only.""" + core = _core(make_core) + + core.apply(ChangeSet(delete=["w"])) + + assert core.settings().provider_default_deployment_ids == ["w"] + + +@pytest.mark.asyncio +async def test_a_deleted_deployment_is_no_longer_called(make_core: MakeCore, host: ScriptedHost) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["a"])) + for _ in range(3): + await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert host.deployments_called("acompletion") == ["b", "b", "b"] + + +def test_the_deprecated_delete_deployment_warns_and_returns_none_for_an_unknown_id(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match=r"delete_deployment is deprecated; use .*apply\(ChangeSet\(delete="): + assert core.delete_deployment("missing") is None # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + +# -- set_model_list --------------------------------------------------------------------------------- + + +def test_set_model_list_replaces_every_deployment_and_group(make_core: MakeCore) -> None: + core = _core(make_core) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="ids") == ["n"] + assert core.registry().names(kind="group_set") == {"new"} + assert core.registry().index("model_name") == {"new": [0]} + assert core.registry().index("team_model") == {} + assert core.registry().names(kind="team_public") == frozenset() + + +def test_set_model_list_keeps_the_aliases(make_core: MakeCore) -> None: + core = _core(make_core, model_group_alias={"fast": "new"}) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert ids(core.registry().deployments("fast")) == ["n"] + + +def test_set_model_list_leaves_the_given_list_unchanged(make_core: MakeCore) -> None: + core = _core(make_core) + given = [{"model_name": "new", "litellm_params": {"model": "openai/n", "api_key": "k"}}] + + core.set_model_list(given) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert given == [{"model_name": "new", "litellm_params": {"model": "openai/n", "api_key": "k"}}] + + +def test_set_model_list_keeps_the_old_wildcard_routes(make_core: MakeCore) -> None: + """Odd but current: the wildcard pattern router is not reset, so a replaced wildcard still matches.""" + core = _core(make_core) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert names_and_ids(core.registry().deployments("openai/gpt-x")) == [("openai/gpt-x", "w")] + assert core.settings().provider_default_deployment_ids == ["w"] + + +def test_set_model_list_appends_to_the_deployment_names(make_core: MakeCore) -> None: + """Odd but current: deployment_names keep the replaced deployments' models.""" + core = make_core(model_list=[deployment("gpt", "a")]) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="deployment_names") == ["openai/a", "openai/n"] + + +@pytest.mark.asyncio +async def test_after_set_model_list_requests_route_to_the_new_deployments( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = _core(make_core) + + core.set_model_list([deployment("gpt", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert host.deployments_called("acompletion") == ["n"] + + +# -- update_settings / ChangeSet(settings=...) ------------------------------------------------------ + + +@pytest.mark.parametrize( + ("name", "given", "stored"), + [ + ("timeout", "5", 5), + ("num_retries", "3", 3), + ("retry_after", 2.9, 2), + ("allowed_fails", "4", 4), + ("cooldown_time", 7.5, 7), + ], +) +def test_integer_settings_are_cast_to_int(make_core: MakeCore, name: str, given: object, stored: int) -> None: + core = _core(make_core) + + core.apply(ChangeSet(settings={name: given})) + + assert core.settings().as_dict()[name] == stored + + +def test_a_non_integer_value_for_an_integer_setting_raises(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.raises(ValueError, match="invalid literal for int"): + core.apply(ChangeSet(settings={"timeout": "abc"})) + + +def test_a_setting_that_is_not_runtime_updatable_is_ignored(make_core: MakeCore) -> None: + core = _core(make_core) + before = dict(core.settings().as_dict()) + + core.apply(ChangeSet(settings={"bogus": 1, "model_list": []})) + + assert dict(core.settings().as_dict()) == before + assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1"] + + +def test_a_retry_policy_dict_is_stored_as_a_retry_policy(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(settings={"retry_policy": {"RateLimitErrorRetries": 2}})) + + stored = core.settings().as_dict()["retry_policy"] + assert isinstance(stored, RetryPolicy) + assert stored.RateLimitErrorRetries == 2 + + +@pytest.mark.parametrize(("given", "stored"), [("true", True), (True, True), ("false", False), ("yes-ish", False)]) +def test_mid_stream_fallback_continuation_reads_flag_strings(make_core: MakeCore, given: object, stored: bool) -> None: + core = _core(make_core) + + core.apply(ChangeSet(settings={"enable_mid_stream_fallback_continuation": given})) + + assert core.settings().as_dict()["enable_mid_stream_fallback_continuation"] is stored + + +def test_model_group_alias_from_settings_is_visible_to_group_info(make_core: MakeCore) -> None: + core = _core(make_core) + assert core.registry().group("quick") is None + + core.apply(ChangeSet(settings={"model_group_alias": {"quick": "emb"}})) + + info = core.registry().group("quick") + assert info is not None + assert info.model_group == "quick" + + +def test_the_deprecated_update_settings_warns_and_applies(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match=r"update_settings is deprecated"): + core.update_settings(cooldown_time=13) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert core.settings().cooldown_time == 13 + + +# -- the settings() and registry() views ------------------------------------------------------------ + + +def test_settings_as_dict_lists_the_runtime_updatable_settings(make_core: MakeCore) -> None: + core = _core(make_core, model_group_alias={"fast": "gpt"}, num_retries=2, cooldown_time=5) + + settings = core.settings().as_dict() + + assert set(settings) == { + "routing_strategy_args", + "routing_strategy", + "allowed_fails", + "cooldown_time", + "num_retries", + "timeout", + "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", + "routing_groups", + } + assert (settings["routing_strategy"], settings["num_retries"], settings["cooldown_time"]) == ( + "simple-shuffle", + 2, + 5, + ) + assert settings["model_group_alias"] == {"fast": "gpt"} + assert settings["routing_groups"] == [] + + +def test_settings_defaults_on_a_fresh_core(make_core: MakeCore) -> None: + settings = _core(make_core).settings() + + assert settings.routing_strategy == "simple-shuffle" + assert settings.allowed_fails == 3 + assert settings.fallbacks is None + assert settings.enable_tag_filtering is False + assert settings.provider_default_deployment_ids == ["w"] + + +@pytest.mark.parametrize(("name", "expected"), [("cooldown_time", True), ("bogus", False), ("num_retries", False)]) +def test_settings_is_set_names_only_settings_fields(make_core: MakeCore, name: str, expected: bool) -> None: + assert _core(make_core).settings().is_set(name) is expected + + +def test_registry_view_is_reused(make_core: MakeCore) -> None: + core = _core(make_core) + + assert core.registry() is core.registry() + + +def test_registry_all_is_the_live_model_list(make_core: MakeCore) -> None: + core = _core(make_core) + listed = core.registry().all() + + core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert ids(listed) == ["a", "b", "e", "w", "t1", "n"] + + +def test_an_unknown_index_name_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown index 'bogus'"): + _core(make_core).registry().index("bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid name on purpose + + +# -- aapply(refresh_model_info) --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_refreshing_model_info_with_nothing_to_discover_changes_nothing(make_core: MakeCore) -> None: + core = _core(make_core) + version = core.registry().version + + result = await core.aapply(ChangeSet(refresh_model_info=True)) # pyright: ignore[reportAttributeAccessIssue] # aapply is on every core + + assert (result.version, list(result.upserted), list(result.deleted)) == (version, [], []) + assert core.registry().version == version + assert dict(core.registry().discovered_model_info("a")) == {} + + +@pytest.mark.asyncio +async def test_aapply_applies_upserts_and_deletes_like_apply(make_core: MakeCore) -> None: + core = _core(make_core) + + result = await core.aapply(ChangeSet(upsert=[typed("gpt", "c")], delete=["e"])) # pyright: ignore[reportAttributeAccessIssue] # aapply is on every core + + assert [None if d is None else d.model_info.id for d in result.upserted] == ["c"] + assert core.registry().names(kind="ids") == ["a", "b", "w", "t1", "c"] + + +@pytest.mark.asyncio +async def test_the_deprecated_arefresh_model_info_warns(make_core: MakeCore) -> None: + core = _core(make_core) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + await core.arefresh_model_info() # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert any("arefresh_model_info is deprecated" in str(w.message) for w in caught) diff --git a/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py b/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py new file mode 100644 index 00000000000..4aeb11c58ec --- /dev/null +++ b/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py @@ -0,0 +1,845 @@ +""" +Characterization of deployment health in a router core: when a failure cools a deployment down (by status, +group size, allowed_fails and allowed_fails_policy at router and deployment level), what the failure and +success callbacks count, blocked deployments and health-check-driven filtering. +""" + +import time +from typing import Final + +import pytest + +import litellm +from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS +from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.types.router import AllowedFailsPolicy, Deployment +from tests.unit.test_router.core.retries_health_helpers import ( + BackgroundSuccessLoggingHost, + ErrorBuilder, + MakeCore, + api_connection, + api_status, + attempt, + cooled_down, + fresh, + not_found, + permission_denied, + service_unavailable, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + bad_request, + deployment, + rate_limit, + server_error, + timeout, +) + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, **params) for deployment_id in ids] + + +def unhealthy(*ids: str, age: float = 0.0) -> dict[str, dict[str, object]]: # mutable-ok: health states + return {i: {"is_healthy": False, "timestamp": time.time() - age, "reason": "probe failed"} for i in ids} + + +def failure_kwargs( + model_info: object, exception: Exception, **extra: object +) -> dict[str, object]: # kwargs-ok: logging kwargs # mutable-ok: litellm's logging kwargs + litellm_params: Final = {"metadata": {"model_group": "g"}, "model_info": model_info} + litellm_params.update(extra.pop("litellm_params", {})) # pyright: ignore[reportCallIssue, reportArgumentType] # test extras + return {"exception": exception, "litellm_params": litellm_params, **extra} + + +# --- which failures cool a deployment down ------------------------------------------------------- + + +@pytest.mark.parametrize("builder", [auth_error, not_found]) +@pytest.mark.asyncio +async def test_a_lone_deployment_cools_down_on_a_401_or_404( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=0) + host.answer("openai/a", fresh(builder)) + + first: Final = await attempt(core) + second: Final = await attempt(core) + + assert type(first) is type(builder()) + assert set(cooled_down(core, ["a"])) == {"a"} + assert isinstance(second, litellm.RouterRateLimitError) + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_404_cools_a_group_deployment_down_at_once(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0) + host.answer("*", fresh(not_found)) + + await attempt(core) + + assert list(cooled_down(core, ["a", "b"])) == host.deployments_called("acompletion") + + +@pytest.mark.parametrize("builder", [permission_denied, bad_request, api_status(422), api_status(409)]) +@pytest.mark.asyncio +async def test_other_4xx_errors_never_cool_down(make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + host.answer("*", fresh(builder)) + + for _ in range(4): + await attempt(core) + + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.parametrize("builder", [server_error, timeout, service_unavailable]) +@pytest.mark.asyncio +async def test_by_default_a_group_deployment_cools_down_on_its_fifth_failure_this_minute( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0) + host.answer("*", fresh(builder)) + + while not cooled_down(core, ["a", "b"]): + await attempt(core) + + called: Final = host.deployments_called("acompletion") + (cooled,) = cooled_down(core, ["a", "b"]) + assert called[-1] == cooled + assert called.count(cooled) == DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS + + +@pytest.mark.asyncio +async def test_successes_this_minute_keep_the_failure_rate_under_the_default_threshold(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b")) + for _ in range(5): + core.sync_deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}}, None, 0.0, 0.0 + ) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(6)] + + assert outcomes == [False] * 5 + [True] + + +@pytest.mark.parametrize("allowed_fails", [3, 1]) +@pytest.mark.asyncio +async def test_an_explicit_allowed_fails_counts_failures_even_at_the_default_value( + make_core: MakeCore, allowed_fails: int +) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=allowed_fails) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(allowed_fails + 1)] + + assert outcomes == [False] * allowed_fails + [True] + + +@pytest.mark.asyncio +async def test_a_connection_error_never_cools_down(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + host.answer("*", fresh(api_connection)) + + for _ in range(3): + await attempt(core) + + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.parametrize(("ids", "cools"), [(("a", "b"), True), (("a",), False)], ids=["group", "lone"]) +@pytest.mark.asyncio +async def test_a_402_cools_a_group_deployment_but_not_a_lone_one( + make_core: MakeCore, host: ScriptedHost, ids: tuple[str, ...], cools: bool +) -> None: + core: Final = make_core(model_list=group(*ids), num_retries=0) + host.answer("*", fresh(api_status(402))) + + await attempt(core) + + assert bool(cooled_down(core, list(ids))) is cools + + +@pytest.mark.asyncio +async def test_a_lone_deployment_with_a_fallback_group_cools_down_under_explicit_allowed_fails( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("f", "z")], + num_retries=0, + allowed_fails=0, + fallbacks=[{"g": ["f"]}], + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert set(cooled_down(core, ["a", "z"])) == {"a"} + assert host.deployments_called("acompletion") == ["a", "z"] + + +@pytest.mark.asyncio +async def test_a_lone_deployment_with_a_fallback_group_is_not_cooled_by_a_429_under_defaults( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("f", "z")], num_retries=0, fallbacks=[{"g": ["f"]}] + ) + host.answer("openai/a", fresh(rate_limit)) + + await attempt(core) + + assert cooled_down(core, ["a", "z"]) == {} + + +@pytest.mark.asyncio +async def test_record_failure_never_cools_down_when_cooldowns_are_disabled(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0, disable_cooldowns=True) + + cooled: Final = core.health().record_failure("a", rate_limit(), 429) + + assert cooled is False + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") == 1 + + +@pytest.mark.parametrize( + ("deployment_id", "status"), [("missing", 429), ("a", ""), ("a", 400)], ids=["unknown", "empty", "400"] +) +@pytest.mark.asyncio +async def test_record_failure_without_a_cooldown(make_core: MakeCore, deployment_id: str, status: object) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + assert core.health().record_failure(deployment_id, rate_limit(), status) is False # pyright: ignore[reportArgumentType] # status as given + assert core.health().failures_this_minute(deployment_id) == 1 + + +@pytest.mark.asyncio +async def test_a_router_cooldown_time_of_zero_falls_back_to_the_default( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, cooldown_time=0) + host.answer("*", fresh(rate_limit)) + + await attempt(core) + + assert core.settings().cooldown_time == DEFAULT_COOLDOWN_TIME_SECONDS + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == DEFAULT_COOLDOWN_TIME_SECONDS + + +# --- allowed_fails_policy ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_router_allowed_fails_policy_counts_every_error_type_in_one_counter(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group("a", "b"), + allowed_fails=5, + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=1), + ) + + outcomes: Final = [ + core.health().record_failure("a", server_error(), 500), + core.health().record_failure("a", server_error(), 500), + core.health().record_failure("a", rate_limit(), 429), + ] + + assert outcomes == [False, False, True] + + +@pytest.mark.parametrize( + ("field", "error", "status"), + [ + ("AuthenticationErrorAllowedFails", auth_error, 401), + ("TimeoutErrorAllowedFails", timeout, 408), + ("RateLimitErrorAllowedFails", rate_limit, 429), + ("InternalServerErrorAllowedFails", server_error, 500), + ("ServiceUnavailableErrorAllowedFails", service_unavailable, 503), + ("NotFoundErrorAllowedFails", not_found, 404), + ], +) +@pytest.mark.asyncio +async def test_router_allowed_fails_policy_threshold_per_error_type( + make_core: MakeCore, field: str, error: ErrorBuilder, status: int +) -> None: + core: Final = make_core( + model_list=group("a", "b"), allowed_fails=10, allowed_fails_policy=AllowedFailsPolicy(**{field: 2}) + ) + + outcomes: Final = [core.health().record_failure("a", error(), status) for _ in range(3)] + + assert outcomes == [False, False, True] + + +@pytest.mark.asyncio +async def test_router_allowed_fails_policy_for_bad_requests_does_not_cool_down(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group("a", "b"), allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=0) + ) + + outcomes: Final = [core.health().record_failure("a", bad_request(), 400) for _ in range(3)] + + assert outcomes == [False, False, False] + + +@pytest.mark.asyncio +async def test_deployment_allowed_fails_policy_counts_each_error_type_separately(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group( + "a", + "b", + model_info={ + "allowed_fails_policy": {"RateLimitErrorAllowedFails": 1, "InternalServerErrorAllowedFails": 1} + }, + ) + ) + + outcomes: Final = [ + core.health().record_failure("a", rate_limit(), 429), + core.health().record_failure("a", server_error(), 500), + core.health().record_failure("a", rate_limit(), 429), + ] + + assert outcomes == [False, False, True] + + +@pytest.mark.asyncio +async def test_deployment_allowed_fails_policy_wins_over_the_router_policy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model_info={"allowed_fails_policy": {"RateLimitErrorAllowedFails": 2}}), + deployment("g", "b"), + ], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + ) + + a_outcomes: Final = [core.health().record_failure("a", rate_limit(), 429) for _ in range(3)] + b_outcome: Final = core.health().record_failure("b", rate_limit(), 429) + + assert a_outcomes == [False, False, True] + assert b_outcome is True + + +@pytest.mark.asyncio +async def test_deployment_policy_not_naming_the_error_falls_back_to_router_allowed_fails(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group("a", "b", model_info={"allowed_fails_policy": {"RateLimitErrorAllowedFails": 9}}) + ) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(4)] + + assert outcomes == [False, False, False, True] + + +@pytest.mark.parametrize(("ids", "cools"), [(("a", "b"), True), (("a",), False)], ids=["group", "lone"]) +@pytest.mark.asyncio +async def test_deployment_allowed_fails_in_model_info(make_core: MakeCore, ids: tuple[str, ...], cools: bool) -> None: + core: Final = make_core(model_list=group(*ids, model_info={"allowed_fails": 1})) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(2)] + + assert outcomes == [False, cools] + + +@pytest.mark.asyncio +async def test_deployment_allowed_fails_policy_cools_down_a_lone_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", model_info={"allowed_fails_policy": {"TimeoutErrorAllowedFails": 0}})) + + assert core.health().record_failure("a", timeout(), 408) is True + + +@pytest.mark.asyncio +async def test_deployment_cooldown_time_in_model_info_sets_the_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=group( + "a", "b", model_info={"cooldown_time": 17, "allowed_fails_policy": {"RateLimitErrorAllowedFails": 0}} + ), + num_retries=0, + ) + host.answer("*", fresh(rate_limit)) + + await attempt(core) + + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == 17 + + +# --- deployment_callback_on_failure ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_failure_callback_counts_the_failure_and_cools_down(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs({"id": "a"}, server_error()), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is True + assert core.health().failures_this_minute("a") == 1 + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +@pytest.mark.parametrize("model_info", [None, "a", {}, {"name": "no id"}], ids=["none", "not-a-dict", "empty", "no-id"]) +@pytest.mark.asyncio +async def test_failure_callback_without_a_deployment_id_does_nothing(make_core: MakeCore, model_info: object) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs(model_info, server_error()), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is False + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.asyncio +async def test_failure_callback_ignores_an_advisor_orchestration_failure(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + error: Final = server_error() + mark_advisor_orchestration_failure(error) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs({"id": "a"}, error), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is False + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.parametrize( + "extra", + [ + { + "litellm_params": { + "metadata": {"model_group": "g", "internal_call_origin": "background_response_status_poll"} + } + }, + {"call_type": "avideo_status"}, + ], + ids=["background-poll", "video-lookup"], +) +@pytest.mark.asyncio +async def test_failure_callback_ignores_a_404_that_is_not_about_the_deployment( + make_core: MakeCore, extra: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + failure_kwargs({"id": "a"}, not_found(), **extra), None, 0.0, 0.0 + ) + + assert cooled is False + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.asyncio +async def test_failure_callback_ignores_a_408_from_the_callers_own_timeout(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + started: Final = time.time() + from datetime import datetime, timedelta + + start: Final = datetime.fromtimestamp(started) + kwargs: Final = failure_kwargs( + {"id": "a"}, + timeout(), + litellm_params={"client_side_timeout": True, "timeout": 1.0}, + start_time=start, + end_time=start + timedelta(seconds=2), + ) + + cooled: Final = core.deployment_callback_on_failure(kwargs, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is False + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.parametrize( + ("model_info", "litellm_params", "expected"), + [ + ({"id": "a", "cooldown_time": 3}, {"cooldown_time": 4}, 3), + ({"id": "a"}, {"cooldown_time": 4}, 4), + ({"id": "a", "cooldown_time": -1}, {}, 60), + ], + ids=["model-info-first", "litellm-params", "negative-ignored"], +) +@pytest.mark.asyncio +async def test_failure_callback_cooldown_time_resolution( + make_core: MakeCore, model_info: dict[str, object], litellm_params: dict[str, object], expected: float +) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0, cooldown_time=60) + + core.deployment_callback_on_failure( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + failure_kwargs(model_info, server_error(), litellm_params=litellm_params), None, 0.0, 0.0 + ) + + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == expected + + +# --- success callbacks and usage counters -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_sync_success_callback_counts_a_success_and_returns_its_key(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b")) + kwargs: Final = {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}} + + key: Final = core.sync_deployment_callback_on_success(kwargs, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + core.sync_deployment_callback_on_success(kwargs, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert key == "a:successes" + assert core.health().store.get_cache("a:successes", local_only=True) == 2 + + +@pytest.mark.asyncio +async def test_sync_success_callback_turns_an_integer_id_into_a_string_key(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a")) + + key: Final = core.sync_deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": 7}}}, None, 0.0, 0.0 + ) + + assert key == "7:successes" + + +@pytest.mark.parametrize( + "kwargs", + [ + {"litellm_params": {"metadata": None, "model_info": {"id": "a"}}}, + {"litellm_params": {"metadata": {}, "model_info": {"id": "a"}}}, + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": None}}, + { + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + }, + ], + ids=["no-metadata", "no-group", "no-model-info", "batch-retrieve"], +) +@pytest.mark.asyncio +async def test_sync_success_callback_counts_nothing_without_a_group_and_id( + make_core: MakeCore, kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a")) + + assert core.sync_deployment_callback_on_success(kwargs, None, 0.0, 0.0) is None # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + assert core.health().store.get_cache("a:successes", local_only=True) is None + + +@pytest.mark.parametrize( + "kwargs", + [ + {"call_type": "_arealtime"}, + {"call_type": "aretrieve_batch"}, + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}}, + ], + ids=["realtime", "batch-retrieve", "no-standard-logging-object"], +) +@pytest.mark.asyncio +async def test_async_success_callback_skips_calls_it_does_not_track( + make_core: MakeCore, kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + full: Final = { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + "standard_logging_object": {"total_tokens": 10}, + **kwargs, + } + if "litellm_params" in kwargs: + del full["standard_logging_object"] + + result: Final = await core.deployment_callback_on_success(full, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert result is None + assert core.health().store.get_cache("a:successes", local_only=True) is None + assert await core.health().usage("g") == (None, None) + + +@pytest.mark.asyncio +async def test_async_success_callback_for_an_unknown_deployment_counts_nothing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "zz"}}, + "standard_logging_object": {"total_tokens": 10}, + }, + None, + 0.0, + 0.0, + ) + + assert core.health().store.get_cache("zz:successes", local_only=True) is None + + +@pytest.mark.asyncio +async def test_async_success_callback_writes_tokens_and_a_request_for_a_tracked_deployment( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": { + "metadata": {"model_group": "g", "deployment": "openai/a"}, + "model_info": {"id": "a"}, + }, + "standard_logging_object": {"total_tokens": 12}, + }, + None, + 0.0, + 0.0, + ) + + assert core.health().store.get_cache("a:successes", local_only=True) == 1 + assert await core.health().usage("g") == (12, 1) + + +@pytest.mark.asyncio +async def test_async_success_callback_adds_only_tokens_not_yet_counted(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": { + "metadata": { + "model_group": "g", + "deployment": "openai/a", + "_litellm_router_usage_counted_tokens": 8, + }, + "model_info": {"id": "a"}, + }, + "standard_logging_object": {"total_tokens": 12}, + }, + None, + 0.0, + 0.0, + ) + + assert await core.health().usage("g") == (4, 0) + + +@pytest.mark.asyncio +async def test_async_success_callback_counts_successes_for_an_untracked_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a")) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + "standard_logging_object": {"total_tokens": 12}, + }, + None, + 0.0, + 0.0, + ) + + assert core.health().store.get_cache("a:successes", local_only=True) == 1 + assert await core.health().usage("g") == (None, None) + + +@pytest.mark.parametrize("limit", [{"rpm": 100}, {"tpm": 1000}], ids=["rpm", "tpm"]) +@pytest.mark.asyncio +async def test_a_successful_request_counts_its_tokens_and_one_request_once( + make_core: MakeCore, host: ScriptedHost, limit: dict[str, int] +) -> None: + core: Final = make_core(model_list=group("a", **limit)) + + await attempt(core) + if isinstance(host, BackgroundSuccessLoggingHost): + await host.drain() + + assert await core.health().usage("g") == (10, 1) + + +@pytest.fixture +def host(request: pytest.FixtureRequest) -> ScriptedHost: + """Background success logging for the usage test above, the plain host for the rest.""" + if request.node.name.startswith("test_a_successful_request_counts_its_tokens"): # pyright: ignore[reportAny] # pytest node + return BackgroundSuccessLoggingHost() + return ScriptedHost() + + +@pytest.mark.asyncio +async def test_a_failed_request_counts_one_request_and_no_tokens(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", rpm=100), num_retries=0) + host.answer("openai/a", fresh(bad_request)) + + await attempt(core) + + assert await core.health().usage("g") == (None, 1) + assert core.health().failures_this_minute("a") == 1 + + +@pytest.mark.asyncio +async def test_an_untracked_deployment_reports_no_usage(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a")) + + await attempt(core) + + assert await core.health().usage("g") == (None, None) + + +@pytest.mark.asyncio +async def test_a_successful_request_counts_one_success_this_minute(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a")) + + await attempt(core) + await attempt(core) + + assert core.health().store.get_cache("a:successes", local_only=True) == 2 + assert core.health().failures_this_minute("a") == 0 + + +# --- blocked deployments ------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_blocked_deployment_is_never_called(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True}), deployment("g", "b")]) + + for _ in range(4): + await attempt(core) + + assert host.deployments_called("acompletion") == ["b"] * 4 + assert core.health().fully_blocked_model_names() == set() + + +@pytest.mark.asyncio +async def test_a_group_with_every_deployment_blocked_raises_no_deployments_available( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b", model_info={"blocked": True})) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RouterRateLimitError) + assert "No deployments available for selected model" in str(error) + assert host.calls == [] + assert core.health().fully_blocked_model_names() == {"g"} + + +def test_is_blocked_and_all_blocked_read_model_info(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a")) + blocked: Final = Deployment(**deployment("g", "a", model_info={"blocked": True})) # pyright: ignore[reportArgumentType] # a model_list entry + open_: Final = Deployment(**deployment("g", "b")) # pyright: ignore[reportArgumentType] # a model_list entry + + assert core.health().is_blocked(blocked) is True + assert core.health().is_blocked(open_) is False + assert core.health().all_blocked([blocked.model_dump()]) is True # pyright: ignore[reportArgumentType] # deployment dicts + assert core.health().all_blocked([blocked.model_dump(), open_.model_dump()]) is False # pyright: ignore[reportArgumentType] # deployment dicts + + +# --- health-check-driven filtering --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_health_check_routing_skips_deployments_marked_unhealthy(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + + for _ in range(4): + await attempt(core) + + assert host.deployments_called("acompletion") == ["b"] * 4 + assert await core.health().fully_unhealthy_model_names() == set() + + +@pytest.mark.asyncio +async def test_health_check_routing_uses_every_deployment_when_all_are_unhealthy( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(unhealthy("a", "b")) # pyright: ignore[reportArgumentType] # health states + + for _ in range(12): + await attempt(core) + + assert set(host.deployments_called("acompletion")) == {"a", "b"} + assert await core.health().fully_unhealthy_model_names() == {"g"} + + +@pytest.mark.parametrize( + ("arguments", "age"), + [ + ({}, 0.0), + ( + { + "enable_health_check_routing": True, + "allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1), + }, + 0.0, + ), + ({"enable_health_check_routing": True, "health_check_staleness_threshold": 10}, 30.0), + ({"enable_health_check_routing": True, "background_health_check_model_groups": ["other"]}, 0.0), + ], + ids=["routing-off", "allowed-fails-policy", "stale", "group-not-listed"], +) +@pytest.mark.asyncio +async def test_health_check_state_is_ignored( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], age: float +) -> None: + core: Final = make_core(model_list=group("a", "b"), **arguments) + core.health().health_checks.set_deployment_health_states(unhealthy("a", age=age)) # pyright: ignore[reportArgumentType] # health states + + for _ in range(12): + await attempt(core) + + assert set(host.deployments_called("acompletion")) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_listed_health_check_group_is_filtered_even_with_an_allowed_fails_policy( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=group("a", "b"), + enable_health_check_routing=True, + background_health_check_model_groups=["g"], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=1), + ) + core.health().health_checks.set_deployment_health_states(unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + + for _ in range(4): + await attempt(core) + + assert host.deployments_called("acompletion") == ["b"] * 4 + + +@pytest.mark.asyncio +async def test_an_unhealthy_deployment_stays_excluded_when_the_healthy_one_is_cooling_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), enable_health_check_routing=True, num_retries=0) + core.health().health_checks.set_deployment_health_states(unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + host.answer("openai/b", fresh(rate_limit)) + + await attempt(core) + error: Final = await attempt(core) + + assert set(cooled_down(core, ["a", "b"])) == {"b"} + assert isinstance(error, litellm.RouterRateLimitError) + assert "cooldown_list=['b']" in str(error) + assert host.deployments_called("acompletion") == ["b"] + + +def test_a_cooldown_recorded_outside_an_event_loop_raises_after_writing_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + with pytest.raises(RuntimeError, match="no running event loop"): + core.health().record_failure("a", server_error(), 500) + + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +@pytest.mark.asyncio +async def test_a_request_cooldown_time_reaches_the_call_and_sets_the_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0) + host.answer("*", fresh(rate_limit)) + + await attempt(core, cooldown_time=11) + + assert host.calls[0].kwargs["cooldown_time"] == 11 + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == 11 diff --git a/tests/unit/test_router/core/test_retries_health_characterization_retries.py b/tests/unit/test_router/core/test_retries_health_characterization_retries.py new file mode 100644 index 00000000000..6695658332f --- /dev/null +++ b/tests/unit/test_router/core/test_retries_health_characterization_retries.py @@ -0,0 +1,636 @@ +""" +Characterization of a router core's retry loop: which errors retry, how many times, how long it sleeps, +what it records on the request and the final error, and the public retry helpers. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.constants import INITIAL_RETRY_DELAY, JITTER, MAX_RETRY_DELAY +from litellm.exceptions import GuardrailRaisedException +from litellm.types.router import RetryPolicy +from tests.unit.test_router.core.retries_health_helpers import ( + ErrorBuilder, + MakeCore, + MetadataRecorder, + api_connection, + api_status, + attempt, + attempt_sync, + bad_gateway, + cooled_down, + fresh, + headers_of, + http_response, + not_found, + permission_denied, + plain_error, + rate_limited, + service_unavailable, + unprocessable, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + bad_request, + content_policy, + context_window, + deployment, + rate_limit, + server_error, + timeout, +) + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, **params) for deployment_id in ids] + + +def within(slept: float, base: float) -> bool: + return base <= slept < base + JITTER + + +# --- which errors retry, by error class and status ------------------------------------------------ + + +@pytest.mark.parametrize( + ("builder", "lone_calls", "pair_calls"), + [ + (not_found, 1, 1), + (permission_denied, 1, 3), + (unprocessable, 1, 1), + (service_unavailable, 3, 3), + (bad_gateway, 3, 3), + (api_connection, 3, 3), + (api_status(409), 3, 3), + (plain_error, 3, 3), + ], + ids=["404", "403", "422", "503", "502", "connection", "409", "no-status"], +) +@pytest.mark.parametrize("size", ["lone", "pair"]) +@pytest.mark.asyncio +async def test_retry_decision_per_error_class_and_group_size( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, lone_calls: int, pair_calls: int, size: str +) -> None: + core: Final = make_core(model_list=group("a") if size == "lone" else group("a", "b"), num_retries=2) + host.answer("*", fresh(builder)) + + error: Final = await attempt(core) + + assert type(error) is type(builder()) + assert len(host.calls_to("acompletion")) == (lone_calls if size == "lone" else pair_calls) + + +@pytest.mark.asyncio +async def test_an_error_without_a_litellm_type_is_retried_but_not_stamped( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(plain_error)) + + error: Final = await attempt(core) + + assert isinstance(error, ValueError) + assert getattr(error, "num_retries", None) is None + assert getattr(error, "max_retries", None) is None + assert "LiteLLM Retried" not in str(error) + + +@pytest.mark.asyncio +async def test_a_403_retry_moves_to_a_deployment_that_has_not_refused_yet( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b", "c"), num_retries=4) + host.answer("*", fresh(permission_denied)) + + error: Final = await attempt(core) + + called: Final = host.deployments_called("acompletion") + assert isinstance(error, litellm.PermissionDeniedError) + assert len(called) == 5 + assert sorted(called[:3]) == ["a", "b", "c"] + assert cooled_down(core, ["a", "b", "c"]) == {} + + +@pytest.mark.asyncio +async def test_a_401_on_each_deployment_cools_each_down_and_stops_when_none_is_left( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b", "c"), num_retries=4) + host.answer("*", fresh(auth_error)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.AuthenticationError) + assert sorted(host.deployments_called("acompletion")) == ["a", "b", "c"] + assert set(cooled_down(core, ["a", "b", "c"])) == {"a", "b", "c"} + + +@pytest.mark.asyncio +async def test_a_non_retryable_error_mid_loop_stops_the_retries(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=4) + host.answer("openai/a", fresh(server_error), fresh(bad_request)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.BadRequestError) + assert len(host.calls_to("acompletion")) == 2 + assert getattr(error, "num_retries", None) is None + + +@pytest.mark.asyncio +async def test_a_retry_policy_keeps_retrying_through_a_later_non_retryable_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=RetryPolicy(InternalServerErrorRetries=3)) + host.answer("openai/a", fresh(server_error), fresh(bad_request)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.BadRequestError) + assert len(host.calls_to("acompletion")) == 4 + assert getattr(error, "num_retries", None) == 3 + + +@pytest.mark.asyncio +async def test_a_guardrail_block_is_raised_without_retrying(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=3) + host.answer("*", lambda **_: GuardrailRaisedException(guardrail_name="pii", message="blocked by guardrail")) + + error: Final = await attempt(core) + + assert isinstance(error, GuardrailRaisedException) + assert len(host.calls_to("acompletion")) == 1 + assert host.sleeps == [] + + +# --- retry policies, more error types -------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("policy", "builder", "expected_calls"), + [ + (RetryPolicy(ServiceUnavailableErrorRetries=2), service_unavailable, 3), + (RetryPolicy(InternalServerErrorRetries=2), service_unavailable, 1 + 2), + (RetryPolicy(DefaultRetries=0), api_status(409), 1), + (RetryPolicy(DefaultRetries=2), unprocessable, 3), + (RetryPolicy(BadRequestErrorRetries=1), content_policy, 2), + (RetryPolicy(BadRequestErrorRetries=1, ContentPolicyViolationErrorRetries=0), content_policy, 1), + ], + ids=[ + "service-unavailable", + "503-falls-back-to-internal-server-error", + "default-for-409", + "default-for-422", + "content-policy-inherits-bad-request", + "content-policy-specific-wins", + ], +) +@pytest.mark.asyncio +async def test_retry_policy_field_chosen_for_each_error( + make_core: MakeCore, host: ScriptedHost, policy: RetryPolicy, builder: ErrorBuilder, expected_calls: int +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=policy) + host.answer("openai/a", fresh(builder)) + + error: Final = await attempt(core) + + assert type(error) is type(builder()) + assert len(host.calls_to("acompletion")) == expected_calls + + +@pytest.mark.parametrize( + ("policy", "builder"), + [ + (RetryPolicy(AuthenticationErrorRetries=2), auth_error), + (RetryPolicy(NotFoundErrorRetries=2), not_found), + (RetryPolicy(DefaultRetries=1), not_found), + ], + ids=["auth", "not-found", "default-for-404"], +) +@pytest.mark.asyncio +async def test_retry_policy_retries_a_401_or_404_into_the_cooled_lone_deployment( + make_core: MakeCore, host: ScriptedHost, policy: RetryPolicy, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=policy) + host.answer("openai/a", fresh(builder)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RouterRateLimitError) + assert "All deployments for selected model are in cooldown" in str(error) + assert len(host.calls_to("acompletion")) == 1 + + +@pytest.mark.parametrize( + ("policy", "builder"), + [(RetryPolicy(AuthenticationErrorRetries=2), auth_error), (RetryPolicy(NotFoundErrorRetries=2), not_found)], + ids=["auth", "not-found"], +) +@pytest.mark.asyncio +async def test_retry_policy_retries_a_401_or_404_when_cooldowns_are_off( + make_core: MakeCore, host: ScriptedHost, policy: RetryPolicy, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=policy, disable_cooldowns=True) + host.answer("openai/a", fresh(builder)) + + error: Final = await attempt(core) + + assert type(error) is type(builder()) + assert len(host.calls_to("acompletion")) == 3 + + +@pytest.mark.asyncio +async def test_retry_policy_wins_over_the_deployment_num_retries(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=group("a", num_retries=4), retry_policy=RetryPolicy(InternalServerErrorRetries=1) + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert len(host.calls_to("acompletion")) == 2 + + +@pytest.mark.asyncio +async def test_request_model_group_retry_policy_overrides_the_router_one( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=group("a"), model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=0)} + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core, model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=3)}) + + assert len(host.calls_to("acompletion")) == 4 + + +@pytest.mark.asyncio +async def test_model_group_retry_policy_given_as_a_dict(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=group("a"), num_retries=0, model_group_retry_policy={"g": {"InternalServerErrorRetries": 2}} + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert len(host.calls_to("acompletion")) == 3 + + +@pytest.mark.parametrize( + ("error", "model_group", "expected"), + [ + (server_error(), "g", 1), + (server_error(), "h", 4), + (server_error(), None, 4), + (rate_limit(), "g", None), + (rate_limit(), "h", 2), + (bad_request(), "h", 7), + (context_window(), "h", 7), + (not_found(), "h", 7), + ], +) +def test_get_num_retries_from_retry_policy( + make_core: MakeCore, error: Exception, model_group: str | None, expected: int | None +) -> None: + core: Final = make_core( + model_list=group("a"), + retry_policy=RetryPolicy(InternalServerErrorRetries=4, RateLimitErrorRetries=2, DefaultRetries=7), + model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=1)}, + ) + + assert core.get_num_retries_from_retry_policy(error, model_group=model_group) == expected # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + +def test_get_num_retries_from_retry_policy_without_any_policy_is_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a"), num_retries=3) + + assert core.get_num_retries_from_retry_policy(server_error(), model_group="g") is None # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + +# --- should_retry_this_error --------------------------------------------------------------------- + +HEALTHY_ONE: Final = [{"model_info": {"id": "a"}}] +ALL_TWO: Final = [{"model_info": {"id": "a"}}, {"model_info": {"id": "b"}}] + + +@pytest.mark.parametrize( + ("error", "arguments", "retries"), + [ + (server_error(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": HEALTHY_ONE}, True), + (server_error(), {"healthy_deployments": [], "all_deployments": HEALTHY_ONE}, False), + (server_error(), {}, False), + (timeout(), {"healthy_deployments": HEALTHY_ONE}, True), + (plain_error(), {"healthy_deployments": HEALTHY_ONE}, True), + (plain_error(), {"healthy_deployments": None}, False), + (bad_request(), {"healthy_deployments": HEALTHY_ONE}, False), + (context_window(), {"healthy_deployments": HEALTHY_ONE, "context_window_fallbacks": []}, False), + (content_policy(), {"healthy_deployments": HEALTHY_ONE, "content_policy_fallbacks": []}, False), + (not_found(), {"healthy_deployments": ALL_TWO, "all_deployments": ALL_TWO}, False), + (unprocessable(), {"healthy_deployments": ALL_TWO, "all_deployments": ALL_TWO}, False), + (auth_error(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": HEALTHY_ONE}, False), + (auth_error(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": ALL_TWO}, True), + (auth_error(), {"healthy_deployments": [], "all_deployments": ALL_TWO}, False), + (permission_denied(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": HEALTHY_ONE}, False), + (permission_denied(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": ALL_TWO}, True), + (rate_limit(), {"healthy_deployments": HEALTHY_ONE, "regular_fallbacks": [{"g": ["h"]}]}, True), + (rate_limit(), {"healthy_deployments": [], "regular_fallbacks": [{"g": ["h"]}]}, False), + (rate_limit(), {"healthy_deployments": [], "regular_fallbacks": []}, False), + (api_status(409)(), {"healthy_deployments": HEALTHY_ONE}, True), + (service_unavailable(), {"healthy_deployments": HEALTHY_ONE}, True), + ], +) +def test_should_retry_this_error_returns_true_or_raises_the_error( + make_core: MakeCore, error: Exception, arguments: dict[str, object], retries: bool +) -> None: + core: Final = make_core(model_list=group("a", "b")) + + if retries: + assert core.should_retry_this_error(error=error, **arguments) is True # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + else: + with pytest.raises(type(error)) as raised: + core.should_retry_this_error(error=error, **arguments) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + assert raised.value is error + + +# --- mock rate limit errors ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_mock_rate_limit_error_raises_without_calling_the_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=0) + + error: Final = await attempt(core, mock_testing_rate_limit_error=True) + + assert isinstance(error, litellm.RateLimitError) + assert "This is a mock exception for model=g, to trigger a rate limit error." in str(error) + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_mock_rate_limit_error_is_retried_and_the_retry_reaches_the_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + + response: Final = await attempt(core, mock_testing_rate_limit_error=True) + + assert not isinstance(response, Exception) + assert host.deployments_called("acompletion") == ["a"] + assert "mock_testing_rate_limit_error" not in host.calls[0].kwargs + assert headers_of(response)["x-litellm-attempted-retries"] == 1 + assert headers_of(response)["x-litellm-max-retries"] == 2 + + +@pytest.mark.asyncio +async def test_mock_rate_limit_error_carries_a_lone_deployments_num_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", num_retries=1), num_retries=5) + + response: Final = await attempt(core, mock_testing_rate_limit_error=True) + + assert headers_of(response)["x-litellm-max-retries"] == 1 + + +@pytest.mark.parametrize("value", [False, None]) +@pytest.mark.asyncio +async def test_mock_rate_limit_error_only_fires_when_true( + make_core: MakeCore, host: ScriptedHost, value: object +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=0) + + response: Final = await attempt(core, mock_testing_rate_limit_error=value) + + assert not isinstance(response, Exception) + assert host.deployments_called("acompletion") == ["a"] + + +# --- sleeps -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_retries", [1, 2, 3]) +@pytest.mark.asyncio +async def test_lone_deployment_sleeps_once_more_than_it_retries( + make_core: MakeCore, host: ScriptedHost, num_retries: int +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=num_retries) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + bases: Final = [INITIAL_RETRY_DELAY * 2**n for n in range(num_retries + 1)] + assert len(host.sleeps) == num_retries + 1 + assert all(within(slept, base) for slept, base in zip(host.sleeps, bases)) + + +@pytest.mark.asyncio +async def test_backoff_is_capped_at_max_retry_delay(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=6) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + expected: Final = [min(INITIAL_RETRY_DELAY * 2**n, MAX_RETRY_DELAY) for n in range(7)] + assert all(within(slept, base) for slept, base in zip(host.sleeps, expected)) + assert within(host.sleeps[-1], MAX_RETRY_DELAY) + + +@pytest.mark.asyncio +async def test_retry_after_setting_above_max_retry_delay_is_capped(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=1, retry_after=20) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert all(within(slept, MAX_RETRY_DELAY) for slept in host.sleeps) + + +@pytest.mark.asyncio +async def test_a_retry_after_header_of_zero_falls_back_to_backoff(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=1) + host.answer("openai/a", lambda **_: rate_limited(retry_after="0")) + + await attempt(core) + + assert within(host.sleeps[0], INITIAL_RETRY_DELAY) + + +@pytest.mark.asyncio +async def test_retry_after_header_on_any_error_type_sets_the_sleep(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=1) + host.answer( + "openai/a", + lambda **_: litellm.InternalServerError( + message="boom", + llm_provider="openai", + model="gpt-4o", + response=http_response(500, {"retry-after": "4"}), + ), + ) + + await attempt(core) + + assert all(within(slept, 4) for slept in host.sleeps) + + +@pytest.mark.asyncio +async def test_litellm_response_headers_on_the_error_set_the_sleep(make_core: MakeCore, host: ScriptedHost) -> None: + def failure(**_: object) -> Exception: + error: Final = server_error() + error.litellm_response_headers = {"retry-after": "6"} # pyright: ignore[reportAttributeAccessIssue] # what provider handlers attach + return error + + core: Final = make_core(model_list=group("a"), num_retries=1) + host.answer("openai/a", failure) + + await attempt(core) + + assert all(within(slept, 6) for slept in host.sleeps) + + +@pytest.mark.asyncio +async def test_retry_after_header_is_ignored_while_siblings_are_healthy( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=2, allowed_fails=10) + host.answer("*", fresh(server_error)) + + await attempt(core) + + assert host.sleeps == [0, 0, 0] + + +# --- what the retry loop writes on the request ------------------------------------------------- + + +@pytest.mark.asyncio +async def test_each_attempt_carries_its_retry_count_and_the_failures_before_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=group("a", "b"), num_retries=2, allowed_fails=10) + host.answer("*", recorder.then(server_error), recorder.then(server_error), recorder.then(None)) + + response: Final = await attempt(core) + + assert not isinstance(response, Exception) + snapshots: Final = recorder.snapshots + assert [s["attempted_retries"] for s in snapshots] == [0, 1, 2] + assert [s["max_retries"] for s in snapshots] == [2, 2, 2] + assert [s["model_group_size"] for s in snapshots] == [2, 2, 2] + assert [s["request_retry_count"] for s in snapshots] == [None, 1, 2] + assert snapshots[0]["previous_models"] is None + trail: Final = snapshots[2]["previous_models"] + assert isinstance(trail, tuple) + assert [entry["attempted_retries"] for entry in trail] == [0, 1] + assert [entry["exception_type"] for entry in trail] == ["InternalServerError"] * 2 + assert all(entry["model_group"] == "g" for entry in trail) + called: Final = host.deployments_called("acompletion") + assert [entry["deployment_id"] for entry in trail] == called[:2] + assert [entry["model"] for entry in trail] == [f"openai/{d}" for d in called[:2]] + assert all(entry["exception_string"] == "litellm.InternalServerError: boom" for entry in trail) + + +@pytest.mark.asyncio +async def test_max_retries_on_the_request_follows_the_deployment_num_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=group("a", num_retries=1), num_retries=3) + host.answer("openai/a", recorder.then(server_error), recorder.then(None)) + + response: Final = await attempt(core) + + assert [s["max_retries"] for s in recorder.snapshots] == [3, 1] + assert headers_of(response)["x-litellm-max-retries"] == 1 + + +@pytest.mark.asyncio +async def test_exhausted_deployment_num_retries_are_stamped_on_the_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", num_retries=2), num_retries=5) + host.answer("openai/a", fresh(server_error)) + + error: Final = await attempt(core) + + assert getattr(error, "num_retries", None) == 2 + assert getattr(error, "max_retries", None) == 2 + + +@pytest.mark.parametrize("bad_value", ["many", None]) +@pytest.mark.asyncio +async def test_a_deployment_num_retries_that_is_not_a_number_falls_back_to_the_router( + make_core: MakeCore, host: ScriptedHost, bad_value: object +) -> None: + core: Final = make_core(model_list=group("a", num_retries=bad_value), num_retries=1) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert len(host.calls_to("acompletion")) == 2 + + +# --- sync and other endpoints ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_sync_completion_retries_like_the_async_one(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(server_error)) + + error: Final = attempt_sync(core) + + assert isinstance(error, litellm.InternalServerError) + assert host.deployments_called("completion") == ["a", "a", "a"] + assert len(host.sleeps) == 3 + assert getattr(error, "num_retries", None) == 2 + + +@pytest.mark.asyncio +async def test_sync_completion_success_after_a_retry_reports_retry_headers( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(server_error), None) + + response: Final = attempt_sync(core) + + assert headers_of(response)["x-litellm-attempted-retries"] == 1 + assert headers_of(response)["x-litellm-max-retries"] == 2 + + +@pytest.mark.parametrize("op", ["aembedding", "embedding"]) +@pytest.mark.asyncio +async def test_embeddings_retry_with_the_router_num_retries(make_core: MakeCore, host: ScriptedHost, op: str) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(server_error)) + + async def embed() -> object: + if op == "aembedding": + return await core.aembedding(model="g", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # every core serves aembedding + return core.embedding(model="g", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # every core serves embedding + + with pytest.raises(litellm.InternalServerError): + await embed() + + assert host.deployments_called(op) == ["a", "a", "a"] + + +@pytest.mark.asyncio +async def test_embedding_success_after_a_retry_reports_retry_headers(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(timeout), None) + + response: Final = await core.aembedding(model="g", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # every core serves aembedding + + assert headers_of(response)["x-litellm-attempted-retries"] == 1 + assert headers_of(response)["x-litellm-max-retries"] == 2 diff --git a/tests/unit/test_router/core/test_selection.py b/tests/unit/test_router/core/test_selection.py index 4b4c320c428..b33316ff380 100644 --- a/tests/unit/test_router/core/test_selection.py +++ b/tests/unit/test_router/core/test_selection.py @@ -26,7 +26,7 @@ from tests.unit.test_router.core.scripted_host import ( ) MESSAGES: Final = [{"role": "user", "content": "hello there, router"}] -PICKS: Final = 20 +PICKS: Final = 60 # a shuffle misses a deployment in 20 picks about once in a thousand runs MakeCore = Callable[..., RouterCore] @@ -622,3 +622,40 @@ async def test_cooling_down_one_group_leaves_other_groups_pickable(make_core: Ma with pytest.raises(RouterRateLimitError): await core.select("g") assert await _picked_ids(core, "h") <= {"c", "d"} + + +# --- custom strategies --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_registered_strategy_picks_every_deployment_an_endpoint_call_reaches( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.types.router import CustomRoutingStrategyBase + + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b"), deployment("g", "c")]) + seen: Final[list[Mapping[str, object]]] = [] # mutable-ok: the request each pick was asked for + + class HighestId(CustomRoutingStrategyBase): + async def async_get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + seen.append(dict(request_kwargs or {})) + deployments: Final = core.get_model_list(model_name=model) # pyright: ignore[reportAttributeAccessIssue] # both cores serve the Python core's lookup + return max(deployments, key=lambda d: d["model_info"]["id"]) + + def get_available_deployment(self, model: str) -> dict: # mutable-ok: the strategy interface's own type + raise NotImplementedError + + core.register_plugin(HighestId()) + for _ in range(PICKS): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert set(host.deployments_called("acompletion")) == {"c"} + # The strategy picks from the whole group: no other core narrowed the pick before it. + assert len(seen) == PICKS + assert not any("_retry_skipped_deployment_ids" in request for request in seen) diff --git a/tests/unit/test_router/core/test_selection_characterization_groups.py b/tests/unit/test_router/core/test_selection_characterization_groups.py new file mode 100644 index 00000000000..465118e1c65 --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_groups.py @@ -0,0 +1,747 @@ +""" +Characterization of which deployments a router core considers for a request before a strategy +picks: routing groups and priority routing groups, tag filtering, region filtering, model access +groups, team-only models, default deployments and default fallbacks, paused (blocked) +deployments, pass-through selection, the healthy listing and the request-scoped narrowing keys. +""" + +import datetime +import re +from collections.abc import Mapping +from typing import Final + +import pydantic +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_api import ChangeSet +from litellm.types.router import RouterRateLimitError +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, ids_of, picked_ids + +SPREAD_PICKS: Final = 60 + +NO_DEPLOYMENTS_MESSAGE: Final = ( + "No deployments available for selected model, Try again in 5 seconds. Passed model={model}. " + "pre-call-checks={checks}, cooldown_list=[]" +) + +NO_HEALTHY_MESSAGE: Final = ( + "litellm.BadRequestError: You passed in model={model}. There are no healthy deployments for this model" +) + +GROUP_MODELS: Final = [deployment("m1", "a"), deployment("m2", "b"), deployment("m3", "c")] + + +def _group( + name: str, models: list[str], strategy: str, **extra: object +) -> dict[str, object]: # mutable-ok: a routing_groups entry # kwargs-ok: extra RoutingGroup fields + return {"group_name": name, "models": models, "routing_strategy": strategy, **extra} + + +def _priority(name: str, priorities: dict[str, int]) -> dict[str, object]: # mutable-ok: a routing_groups entry + return _group(name, list(priorities), "priority", model_priorities=priorities) + + +# --- routing groups ------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_routing_group_name_is_callable_and_unions_its_members(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1", "m2"], "simple-shuffle")] + ) + assert await picked_ids(core, "grp", times=SPREAD_PICKS) == {"a", "b"} + assert ids_of(await core.select("grp", kind="healthy")) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_routing_group_call_uses_the_groups_strategy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1", "m2"], "least-busy")]) + assert await picked_ids(core, "grp", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_a_member_called_by_name_uses_its_groups_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m1", "b")], + routing_groups=[_group("grp", ["m1"], "least-busy")], + ) + assert await picked_ids(core, "m1", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_models_outside_every_group_keep_the_top_level_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("free", "x"), deployment("free", "y")], + routing_groups=[_group("grp", ["m1"], "least-busy")], + ) + assert await picked_ids(core, "free", times=SPREAD_PICKS, messages=MESSAGES) == {"x", "y"} + + +@pytest.mark.asyncio +async def test_an_alias_to_a_routing_group_resolves_to_the_group(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(GROUP_MODELS), + model_group_alias={"al": "grp"}, + routing_groups=[_group("grp", ["m1", "m2"], "simple-shuffle")], + ) + assert await picked_ids(core, "al", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_a_group_named_like_a_model_is_shadowed_by_the_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("m1", ["m2"], "least-busy")]) + assert await picked_ids(core, "m1") == {"a"} + + +@pytest.mark.asyncio +async def test_a_group_whose_members_have_no_deployments_has_no_healthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("empty", ["zz"], "least-busy")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("empty") + assert str(raised.value) == NO_HEALTHY_MESSAGE.format(model="empty") + + +@pytest.mark.asyncio +async def test_a_routing_group_completion_reports_the_group_as_model_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1"], "simple-shuffle")]) + response: Final = await core.acompletion(model="grp", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + assert host.calls[-1].model_group == "grp" + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "grp" + + +def test_routing_group_names_are_listed_among_the_routable_names(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1", "m2"], "least-busy")]) + # The group name is listed twice today. + assert sorted(core.registry().names(kind="groups")) == ["grp", "grp", "m1", "m2", "m3"] + + +@pytest.mark.asyncio +async def test_routing_groups_applied_at_runtime_become_callable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS)) + with pytest.raises(litellm.BadRequestError): + await core.select("grp") + core.apply(ChangeSet(settings={"routing_groups": [_group("grp", ["m1", "m2"], "least-busy")]})) + assert await picked_ids(core, "grp", messages=MESSAGES) == {"a"} + + +@pytest.mark.parametrize( + ("groups", "message"), + [ + ( + [_group("x", ["m1"], "least-busy"), _group("x", ["m2"], "least-busy")], + "routing_groups: group names must be unique, duplicate group_name 'x'.", + ), + ( + [_group("default", ["m1"], "least-busy")], + "routing_groups: 'default' is reserved for the implicit fallback group.", + ), + ([_group("", ["m1"], "least-busy")], "routing_groups: group_name must be non-empty."), + ( + [_group("x", ["m1"], "least-busy"), _group("y", ["m1"], "simple-shuffle")], + "routing_groups: model_name 'm1' appears in 'x' and 'y'. Each model may belong to at most one group.", + ), + ([_group("x", ["m1"], "nope")], "Invalid routing_strategy: 'nope'."), + ([_priority("m1", {"m2": 1})], "Priority routing group names must not shadow a model or alias"), + ], + ids=["duplicate", "reserved-default", "empty-name", "overlap", "bad-strategy", "priority-shadow"], +) +def test_invalid_routing_groups_are_rejected_at_construction( + make_core: MakeCore, groups: list[dict[str, object]], message: str +) -> None: + with pytest.raises(ValueError, match=re.escape(message)): + make_core(model_list=list(GROUP_MODELS), routing_groups=groups) + + +@pytest.mark.parametrize( + ("group", "message"), + [ + (_group("p", ["m1", "m2"], "priority", model_priorities={"m1": 1}), "model_priorities must contain exactly"), + ( + _group("p", ["m1"], "priority", model_priorities={"m1": 1}, routing_strategy_args={"x": 1}), + "Priority routing groups use model_priorities, not routing_strategy_args", + ), + (_group("p", ["m1"], "least-busy", model_priorities={"m1": 1}), "model_priorities requires routing_strategy"), + (_group("p", ["m1"], "priority", model_priorities={"m1": 0}), "greater than or equal to 1"), + ( + _group("p", ["m1", "m1"], "priority", model_priorities={"m1": 1}), + "Priority routing groups require nonempty, distinct models", + ), + ], + ids=["missing-priority", "strategy-args", "priorities-without-priority", "zero-priority", "duplicate-member"], +) +def test_invalid_priority_groups_fail_validation_at_construction( + make_core: MakeCore, group: dict[str, object], message: str +) -> None: + with pytest.raises(pydantic.ValidationError, match=re.escape(message)): + make_core(model_list=list(GROUP_MODELS), routing_groups=[group]) + + +# --- priority routing groups --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_priority_group_picks_only_its_lowest_priority_members(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*GROUP_MODELS, deployment("m2", "b2")], + routing_groups=[_priority("p", {"m1": 2, "m2": 1, "m3": 2})], + ) + assert await picked_ids(core, "p", times=SPREAD_PICKS) == {"b", "b2"} + assert sorted(ids_of(await core.select("p", kind="healthy"))) == ["b", "b2"] + + +@pytest.mark.asyncio +async def test_a_priority_group_pick_carries_the_member_priority_as_order(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_priority("p", {"m1": 3, "m2": 1})]) + chosen: Final = await core.select("p") + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + assert chosen["model_name"] == "m2" # pyright: ignore[reportIndexIssue] + assert chosen["litellm_params"]["order"] == 1 # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_priority_group_member_called_directly_ignores_the_priority(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_priority("p", {"m1": 3, "m2": 1})]) + chosen: Final = await core.select("m1") + assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + assert "order" not in chosen["litellm_params"] # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_priority_group_fails_over_to_the_next_priority_and_stays_there_while_the_first_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=list(GROUP_MODELS), + routing_groups=[_priority("p", {"m1": 1, "m2": 2})], + allowed_fails=0, + num_retries=0, + cooldown_time=60, + ) + host.answer("openai/a", rate_limit()) + await core.acompletion(model="p", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a", "b"] + assert await picked_ids(core, "p") == {"b"} + + +# --- tag filtering ----------------------------------------------------------------------------------- + +TAGGED: Final = [ + deployment("g", "free", tags=["free"]), + deployment("g", "paid", tags=["paid"]), + deployment("g", "dflt", tags=["default"]), + deployment("g", "untagged"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tags", "expected"), + [(["free"], {"free"}), (["paid"], {"paid"}), (["free", "paid"], {"free", "paid"})], +) +async def test_tag_filtering_picks_deployments_carrying_a_request_tag( + make_core: MakeCore, tags: list[str], expected: set[str] +) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + assert await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"metadata": {"tags": tags}}) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", [{}, {"metadata": {"tags": ["unknown"]}}], ids=["untagged", "unmatched"]) +async def test_tag_filtering_sends_untagged_or_unmatched_requests_to_default_tagged_deployments( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + assert await picked_ids(core, "g", request_kwargs=request_kwargs) == {"dflt"} + + +@pytest.mark.asyncio +async def test_tag_filtering_without_a_default_rejects_an_unmatched_tag(make_core: MakeCore) -> None: + core: Final = make_core(model_list=TAGGED[:2], enable_tag_filtering=True) + with pytest.raises(ValueError, match="Not allowed to access model") as raised: + await core.select("g", {"metadata": {"tags": ["unknown"]}}) + assert str(raised.value) == ( + "Not allowed to access model due to tags configuration. Passed model=g and tags=['unknown']" + ) + + +@pytest.mark.asyncio +async def test_tag_filtering_without_a_default_serves_an_untagged_request_from_every_deployment( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=TAGGED[:2], enable_tag_filtering=True) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"free", "paid"} + + +@pytest.mark.asyncio +async def test_request_tags_are_ignored_without_tag_filtering(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TAGGED)) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"metadata": {"tags": ["free"]}}) + assert picked == {"free", "paid", "dflt", "untagged"} + + +@pytest.mark.asyncio +async def test_a_request_can_enable_tag_filtering_for_itself(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TAGGED)) + request: Final = {"metadata": {"tags": ["free"]}, "enable_tag_filtering": True} + assert await picked_ids(core, "g", request_kwargs=request) == {"free"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("match_any", "expected"), [(True, {"ab", "a"}), (False, {"ab"})]) +async def test_tag_filtering_match_any_decides_whether_one_shared_tag_is_enough( + make_core: MakeCore, match_any: bool, expected: set[str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "ab", tags=["a", "b"]), deployment("g", "a", tags=["a"])], + enable_tag_filtering=True, + tag_filtering_match_any=match_any, + ) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"metadata": {"tags": ["a", "b"]}}) + assert picked == expected + + +@pytest.mark.asyncio +async def test_tags_in_litellm_metadata_are_read_too(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + assert await picked_ids(core, "g", request_kwargs={"litellm_metadata": {"tags": ["paid"]}}) == {"paid"} + + +@pytest.mark.asyncio +async def test_completion_with_tags_calls_the_tagged_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + for _ in range(5): + await core.acompletion(model="g", messages=MESSAGES, metadata={"tags": ["paid"]}) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["paid"] * 5 + + +# --- region filtering -------------------------------------------------------------------------------- + +REGIONAL: Final = [ + deployment("g", "eu", region_name="eu"), + deployment("g", "us", region_name="us"), + deployment("g", "unknown"), +] + + +@pytest.mark.asyncio +async def test_allowed_model_region_keeps_only_deployments_in_that_region(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL), enable_pre_call_checks=True) + assert await picked_ids(core, "g", request_kwargs={"allowed_model_region": "eu"}, messages=MESSAGES) == {"eu"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_request_input(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL), enable_pre_call_checks=True) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"allowed_model_region": "eu"}) + assert picked == {"eu", "us", "unknown"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_pre_call_checks(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL)) + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"allowed_model_region": "eu"}, messages=MESSAGES + ) + assert picked == {"eu", "us", "unknown"} + + +@pytest.mark.asyncio +async def test_an_allowed_model_region_no_deployment_is_in_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL), enable_pre_call_checks=True) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", {"allowed_model_region": "ap"}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=True) + + +# --- model access groups ------------------------------------------------------------------------------- + +ACCESS_GROUPED: Final = [ + deployment("g", "in_a", model_info={"access_groups": ["grpA"]}), + deployment("g", "in_b", model_info={"access_groups": ["grpB"]}), +] + + +def _auth( + bucket: str = "metadata", **fields: object +) -> dict[str, object]: # mutable-ok: request kwargs # kwargs-ok: UserAPIKeyAuth fields + return {bucket: {"user_api_key_auth": UserAPIKeyAuth(**fields)}} # pyright: ignore[reportArgumentType] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected"), + [ + (_auth(models=["grpA"]), {"in_a"}), + (_auth(models=["grpB"]), {"in_b"}), + (_auth(team_models=["grpB"]), {"in_b"}), + (_auth("litellm_metadata", models=["grpA"]), {"in_a"}), + ], + ids=["key-group-a", "key-group-b", "team-group-b", "litellm-metadata"], +) +async def test_an_access_group_only_caller_is_narrowed_to_its_groups_deployments( + make_core: MakeCore, request_kwargs: dict[str, object], expected: set[str] +) -> None: + core: Final = make_core(model_list=list(ACCESS_GROUPED)) + assert await picked_ids(core, "g", request_kwargs=request_kwargs) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "models", [["g"], ["*"], ["all-proxy-models"], [], ["unrelated"]], ids=["direct", "star", "all", "none", "other"] +) +async def test_a_caller_not_limited_to_access_groups_can_reach_every_deployment( + make_core: MakeCore, models: list[str] +) -> None: + core: Final = make_core(model_list=list(ACCESS_GROUPED)) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs=_auth(models=models)) + assert picked == {"in_a", "in_b"} + + +# --- team-only models ------------------------------------------------------------------------------------ + +TEAM_ONLY: Final = [ + deployment("internal-x", "tx", model_info={"team_id": "team-x", "team_public_model_name": "pub"}), + deployment("internal-y", "ty", model_info={"team_id": "team-y", "team_public_model_name": "pub"}), +] + + +@pytest.mark.asyncio +async def test_a_team_id_in_litellm_metadata_selects_the_team_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAM_ONLY)) + assert await picked_ids(core, "pub", request_kwargs={"litellm_metadata": {"user_api_key_team_id": "team-y"}}) == { + "ty" + } + + +@pytest.mark.asyncio +async def test_a_proxy_admin_without_a_team_is_refused_a_public_name_several_teams_share(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAM_ONLY)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("pub", _auth(user_role="proxy_admin")) + assert str(raised.value) == ( + "litellm.BadRequestError: Model name 'pub' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ) + + +@pytest.mark.asyncio +async def test_a_proxy_admin_without_a_team_finds_one_teams_public_name_but_cannot_use_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=TEAM_ONLY[:1]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("pub", _auth(user_role="proxy_admin")) + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="pub", checks=False) + + +# --- default deployment, default fallbacks ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_an_unknown_model_goes_to_a_catch_all_deployment_under_its_own_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("*", "catch_all", model="openai/*")]) + chosen: Final = await core.select("brand-new") + assert chosen["model_info"]["id"] == "catch_all" # pyright: ignore[reportIndexIssue] + assert chosen["litellm_params"]["model"] == "openai/brand-new" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_an_unknown_model_goes_to_the_first_default_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("h", "b")], default_fallbacks=["h", "g"]) + chosen: Final = await core.select("unknown") + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + await core.acompletion(model="unknown", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["b"] + + +@pytest.mark.asyncio +async def test_an_unknown_model_goes_to_the_process_default_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(default_fallbacks=["g"]) + core: Final = make_core(model_list=[deployment("g", "a")]) + chosen: Final = await core.select("unknown") + assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + + +# --- paused (blocked) deployments ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_blocked_deployment_is_never_picked_for_its_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True}), deployment("g", "b")]) + assert await picked_ids(core, "g") == {"b"} + assert ids_of(await core.select("g", kind="healthy")) == ["b"] + + +@pytest.mark.asyncio +async def test_selecting_a_blocked_deployment_by_id_raises_service_unavailable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True}), deployment("g", "b")]) + with pytest.raises(litellm.ServiceUnavailableError) as raised: + await core.select("a") + # The message names the deployment's litellm model, not the id that was asked for. + assert str(raised.value) == ( + "litellm.ServiceUnavailableError: Model 'openai/a' is currently paused and cannot accept requests." + ) + assert raised.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entry", ["select", "acompletion", "completion"]) +async def test_a_group_whose_every_deployment_is_blocked_rate_limits( + make_core: MakeCore, host: ScriptedHost, entry: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True})]) + call: Final = core.select("g") if entry == "select" else core.call(entry, {"model": "g", "messages": MESSAGES}) + with pytest.raises(RouterRateLimitError) as raised: + await call + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=False) + assert host.calls == [] + + +# --- pass-through selection ------------------------------------------------------------------------------------ + +PASS_THROUGH: Final = [deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")] + + +@pytest.mark.asyncio +async def test_pass_through_selection_picks_only_pass_through_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + assert await picked_ids(core, "g", kind="pass_through") == {"pt"} + + +@pytest.mark.asyncio +async def test_pass_through_selection_by_id_returns_a_pass_through_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + chosen: Final = await core.select("pt", kind="pass_through") + assert chosen["model_info"]["id"] == "pt" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_pass_through_selection_by_id_refuses_a_plain_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("plain", kind="pass_through") + assert str(raised.value) == ( + "litellm.BadRequestError: Deployment plain does not support pass-through endpoint (use_in_pass_through=False)" + ) + + +@pytest.mark.asyncio +async def test_pass_through_selection_without_pass_through_deployments_raises_bad_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("g", kind="pass_through") + assert str(raised.value) == ( + "litellm.BadRequestError: Model g has no deployment configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true in the deployment configuration" + ) + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_words_the_missing_deployment_error_differently( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == ( + "litellm.BadRequestError: Model g has no deployments configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true to the deployment configuration" + ) + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_an_unknown_model_has_no_healthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("nope", kind="pass_through") + assert str(raised.value) == NO_HEALTHY_MESSAGE.format(model="nope") + + +@pytest.mark.asyncio +async def test_pass_through_selection_skips_a_cooling_pass_through_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "pt1", use_in_pass_through=True), + deployment("g", "pt2", use_in_pass_through=True), + ], + allowed_fails=0, + num_retries=0, + cooldown_time=60, + ) + host.answer("openai/pt1", rate_limit()) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="pt1", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert await picked_ids(core, "g", kind="pass_through") == {"pt2"} + + +def test_deprecated_pass_through_pick_matches_select(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + assert core.get_available_deployment_for_pass_through("g")["model_info"]["id"] == "pt" # pyright: ignore[reportAttributeAccessIssue] + + +# --- the healthy listing -------------------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_healthy_listing_keeps_only_the_lowest_order(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", order=1), deployment("g", "b", order=2)]) + assert ids_of(await core.select("g", kind="healthy")) == ["a"] + + +@pytest.mark.asyncio +async def test_the_healthy_listing_for_a_deployment_id_is_that_single_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + healthy: Final = await core.select("b", kind="healthy") + assert isinstance(healthy, Mapping) + assert healthy["model_info"]["id"] == "b" + + +@pytest.mark.asyncio +async def test_the_healthy_listing_of_an_unknown_model_raises_bad_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("nope", kind="healthy") + assert str(raised.value) == NO_HEALTHY_MESSAGE.format(model="nope") + + +@pytest.mark.asyncio +async def test_the_healthy_listing_leaves_out_cooling_deployments(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60 + ) + host.answer("openai/a", rate_limit()) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert ids_of(await core.select("g", kind="healthy")) == ["b"] + + +# --- request-scoped narrowing keys ----------------------------------------------------------------------------------- + +THREE: Final = [deployment("g", "a"), deployment("g", "b"), deployment("g", "c")] + + +@pytest.mark.asyncio +async def test_target_order_picks_that_order_instead_of_the_lowest(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", order=1), deployment("g", "b", order=2)]) + assert await picked_ids(core, "g", request_kwargs={"_target_order": 2}) == {"b"} + assert ids_of(await core.select("g", {"_target_order": 2}, kind="healthy")) == ["b"] + + +@pytest.mark.asyncio +async def test_a_target_order_no_deployment_has_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", order=1)]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", {"_target_order": 3}) + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=False) + + +@pytest.mark.asyncio +async def test_excluded_deployment_ids_are_never_picked(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + assert await picked_ids(core, "g", request_kwargs={"_excluded_deployment_ids": ["a", "b"]}) == {"c"} + + +@pytest.mark.asyncio +async def test_excluding_every_deployment_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", {"_excluded_deployment_ids": ["a", "b", "c"]}) + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=False) + + +@pytest.mark.asyncio +async def test_retry_skipped_deployment_ids_are_not_picked_while_others_remain(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + assert await picked_ids(core, "g", request_kwargs={"_retry_skipped_deployment_ids": ("a", "b")}) == {"c"} + + +@pytest.mark.asyncio +async def test_skipping_every_deployment_on_retry_keeps_them_all(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + healthy: Final = await core.select("g", {"_retry_skipped_deployment_ids": ("a", "b", "c")}, kind="healthy") + assert ids_of(healthy) == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_retry_skipped_deployment_ids_given_as_a_list_are_ignored(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + healthy: Final = await core.select("g", {"_retry_skipped_deployment_ids": ["a", "b"]}, kind="healthy") + assert ids_of(healthy) == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_select_pops_the_narrowing_keys_from_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + request: Final[dict[str, object]] = { + "_target_order": None, + "_excluded_deployment_ids": ["a"], + "_retry_skipped_deployment_ids": ("b",), + } + assert ids_of(await core.select("g", request, kind="healthy")) == ["c"] + assert request == {} + + +# --- access windows (reserved deployments) ---------------------------------------------------------------------------- + + +def _window_for(team_ids: list[str]) -> dict[str, object]: # mutable-ok: an access_windows entry + now: Final = datetime.datetime.now(datetime.timezone.utc) + return { + "start": (now - datetime.timedelta(hours=1)).strftime("%H:%M"), + "end": (now + datetime.timedelta(hours=1)).strftime("%H:%M"), + "timezone": "UTC", + "team_ids": team_ids, + } + + +def _reserved(window: dict[str, object]) -> dict[str, object]: # mutable-ok: a model_list entry + return deployment("g", "reserved", model_info={"access_windows": [window]}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", [{"metadata": {"user_api_key_team_id": "team-b"}}, {}]) +async def test_a_deployment_reserved_for_another_team_is_hidden_during_its_window( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=[_reserved(_window_for(["team-a"])), deployment("g", "open")]) + assert await picked_ids(core, "g", request_kwargs=request_kwargs) == {"open"} + + +@pytest.mark.asyncio +async def test_the_reserving_team_reaches_its_reserved_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[_reserved(_window_for(["team-a"]))]) + assert await picked_ids(core, "g", request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}) == {"reserved"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "named"), [("g", "g"), ("reserved", "openai/reserved")], ids=["group", "id"]) +async def test_only_reserved_deployments_left_raises_bad_request_naming_the_window_end( + make_core: MakeCore, model: str, named: str +) -> None: + window: Final = _window_for(["team-a"]) + core: Final = make_core(model_list=[_reserved(window)]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select(model, {"metadata": {"user_api_key_team_id": "team-b"}}) + assert str(raised.value) == ( + f"litellm.BadRequestError: Deployment {named} is reserved for another team until {window['end']} UTC" + ) + + +# --- specific_deployment ------------------------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["openai/unknown", "g"], ids=["unknown-litellm-model", "group-name"]) +async def test_specific_deployment_for_a_name_no_litellm_model_matches_rate_limits( + make_core: MakeCore, model: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select(model, specific_deployment=True) # pyright: ignore[reportCallIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model=model, checks=False) diff --git a/tests/unit/test_router/core/test_selection_characterization_prerouting.py b/tests/unit/test_router/core/test_selection_characterization_prerouting.py new file mode 100644 index 00000000000..317fbf4be46 --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_prerouting.py @@ -0,0 +1,505 @@ +""" +Characterization of what runs before a router core's strategy picks: pre-routing strategies (the +complexity router, chosen among tag-scoped markers), the strategy marker deployments they are +registered from, tier ``litellm_params`` pinned onto the request, alias-marker params forwarded to +the routed call, the Claude Code session router binding, and the ``Router(plugins=[...])`` +routing-plugin pipeline. +""" + +import asyncio +import re +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RoutingContext +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, picked_ids + +HARD_MESSAGES: Final = [ + { + "role": "user", + "content": ( + "Design a distributed consensus algorithm, prove its correctness step by step, analyze the time " + "complexity, and implement it in Rust with a kubernetes deployment and database sharding strategy. " + "Think step by step and reason about edge cases." + ), + } +] + +TIERS: Final = {"SIMPLE": "small", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} +ALL_BIG: Final = {"SIMPLE": "big", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} + +ONLY_MARKERS_MESSAGE: Final = ( + "litellm.BadRequestError: You passed in model={model}. Every deployment for it is a strategy router marker " + "(auto_router/...), which is not a callable model, and no pre-routing strategy selected a deployment for " + "this request" +) + +SESSION_ID: Final = "0f8fad5b-d9cb-469f-a165-70867728950e" + + +def _complexity_router( + name: str = "smart", + tiers: Mapping[str, object] = TIERS, + deployment_id: str = "marker", + **params: object, # kwargs-ok: extra marker litellm_params +) -> dict[str, object]: # mutable-ok: a model_list entry + return { + "model_name": name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": dict(tiers)}, + "complexity_router_default_model": "big", + **params, + }, + "model_info": {"id": deployment_id}, + } + + +def _tiered( + **marker_params: object, +) -> list[dict[str, object]]: # mutable-ok: a model_list # kwargs-ok: marker litellm_params + return [deployment("small", "s"), deployment("big", "b"), _complexity_router(**marker_params)] + + +def _sync_entry(core: RouterCore, entry: str, model: str) -> object: + if entry == "completion": + return core.completion(model=model, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + return core.get_available_deployment(model, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + + +def _metadata(request: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = request.get("metadata") + assert isinstance(metadata, Mapping) + return metadata + + +# --- the complexity router ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("messages", "expected"), [(MESSAGES, "s"), (HARD_MESSAGES, "b")], ids=["simple", "hard"]) +async def test_a_complexity_router_picks_the_tier_deployment_for_the_prompt( + make_core: MakeCore, messages: list[dict[str, str]], expected: str +) -> None: + core: Final = make_core(model_list=_tiered()) + chosen: Final = await core.select("smart", {}, messages=messages) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == expected # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_complexity_router_pick_stamps_the_routing_decision_on_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + metadata: Final = _metadata(request) + decision: Final = metadata["routing_decision"] + assert isinstance(decision, Mapping) + assert decision["router_model_name"] == "smart" + assert decision["router_type"] == "complexity" + assert decision["router_config_id"] == "marker" + assert decision["routed_model"] == "small" + assert decision["tier"] == "SIMPLE" + assert metadata["pre_routing_selected_model"] == "small" + + +@pytest.mark.asyncio +async def test_completion_through_a_complexity_router_reports_the_router_name_as_model_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_tiered()) + await core.acompletion(model="smart", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["s"] + assert host.calls[-1].model == "openai/s" + assert host.calls[-1].model_group == "smart" + + +@pytest.mark.asyncio +async def test_an_alias_to_a_complexity_router_routes_through_it(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_tiered(), model_group_alias={"al": "smart"}) + request: Final[dict[str, object]] = {} + chosen: Final = await core.select("al", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + assert _metadata(request)["routing_decision"]["router_model_name"] == "smart" # pyright: ignore[reportIndexIssue] + await core.acompletion(model="al", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].model_group == "al" + + +@pytest.mark.asyncio +async def test_a_complexity_router_takes_precedence_over_plain_deployments_of_its_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*_tiered(), deployment("smart", "plain")]) + assert await picked_ids(core, "smart", times=10, messages=MESSAGES) == {"s"} + + +@pytest.mark.asyncio +async def test_a_complexity_router_tier_naming_an_unknown_model_falls_back_to_the_default_model( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("big", "b"), _complexity_router(tiers={**TIERS, "SIMPLE": "missing"})] + ) + request: Final[dict[str, object]] = {} + chosen: Final = await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + decision: Final = _metadata(request)["routing_decision"] + assert isinstance(decision, Mapping) + assert decision["cause"] == "health_default_fallback" + assert decision["routed_model"] == "big" + + +def test_a_complexity_router_without_tiers_or_a_default_model_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="complexity_router_default_model is required for complexity-router"): + make_core(model_list=[deployment("smart", "marker", model="auto_router/complexity_router")]) + + +# --- strategy markers are never callable ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_healthy_listing_of_a_router_name_refuses_the_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("smart", {}, kind="healthy") + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="smart") + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_router_name_refuses_the_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("smart", {}, kind="pass_through", messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="smart") + + +@pytest.mark.asyncio +async def test_selecting_a_marker_by_its_id_refuses_it_naming_its_litellm_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("marker", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="auto_router/complexity_router") + + +@pytest.mark.parametrize("entry", ["get_available_deployment", "completion"]) +def test_the_sync_path_never_runs_pre_routing_and_refuses_the_marker( + make_core: MakeCore, host: ScriptedHost, entry: str +) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + _sync_entry(core, entry, "smart") + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="smart") + assert host.calls == [] + + +# --- choosing among tag-scoped markers ----------------------------------------------------------------- + +TAG_SCOPED: Final = [ + deployment("small", "s"), + deployment("big", "b"), + _complexity_router(deployment_id="cheap_marker", tags=["cheap"]), + _complexity_router(tiers=ALL_BIG, deployment_id="premium_marker", tags=["premium"]), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected"), + [({"metadata": {"tags": ["cheap"]}}, "s"), ({"metadata": {"tags": ["premium"]}}, "b"), ({}, "s")], + ids=["cheap-tag", "premium-tag", "untagged-takes-the-first"], +) +async def test_request_tags_pick_among_markers_of_one_name( + make_core: MakeCore, request_kwargs: dict[str, object], expected: str +) -> None: + core: Final = make_core(model_list=list(TAG_SCOPED)) + chosen: Final = await core.select("smart", request_kwargs, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == expected # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_an_untagged_request_takes_the_default_tagged_marker(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s"), + deployment("big", "b"), + _complexity_router(deployment_id="cheap_marker", tags=["cheap"]), + _complexity_router(tiers=ALL_BIG, deployment_id="default_marker", tags=["default"]), + ] + ) + chosen: Final = await core.select("smart", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +SCOPED_WITH_PLAIN: Final = [ + deployment("small", "s"), + deployment("big", "b"), + _complexity_router(deployment_id="cheap_marker", tags=["cheap"]), + deployment("smart", "plain"), +] + + +@pytest.mark.asyncio +async def test_with_tag_filtering_an_unmatched_tag_skips_the_markers_for_plain_tag_selection( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=list(SCOPED_WITH_PLAIN), enable_tag_filtering=True) + with pytest.raises( + ValueError, match=re.escape("Not allowed to access model due to tags configuration. Passed model=smart") + ): + await core.select("smart", {"metadata": {"tags": ["other"]}}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + chosen: Final = await core.select("smart", {"metadata": {"tags": ["cheap"]}}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_without_tag_filtering_an_unmatched_tag_still_takes_the_first_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(SCOPED_WITH_PLAIN)) + chosen: Final = await core.select("smart", {"metadata": {"tags": ["other"]}}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + + +# --- tier litellm_params and alias-marker params ------------------------------------------------------------- + +PINNED_TIERS: Final = { + **TIERS, + "SIMPLE": {"model_name": "small", "litellm_params": {"temperature": 0.1, "max_tokens": 7}}, +} + + +@pytest.mark.asyncio +async def test_tier_litellm_params_override_the_callers_params_on_the_call( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_tiered(tiers=PINNED_TIERS)) + await core.acompletion(model="smart", messages=MESSAGES, temperature=0.9, max_completion_tokens=100) # pyright: ignore[reportAttributeAccessIssue] + sent: Final = host.calls[-1].kwargs + assert host.deployments_called() == ["s"] + assert sent["temperature"] == 0.1 + assert sent["max_tokens"] == 7 + assert "max_completion_tokens" not in sent + + +@pytest.mark.asyncio +async def test_tier_litellm_params_are_pinned_onto_the_request_by_select(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=PINNED_TIERS)) + request: Final[dict[str, object]] = {"temperature": 0.9, "max_completion_tokens": 100} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["temperature"] == 0.1 + assert request["max_tokens"] == 7 + assert "max_completion_tokens" not in request + + +@pytest.mark.asyncio +async def test_tier_litellm_params_apply_only_to_their_tier(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_tiered(tiers=PINNED_TIERS)) + await core.acompletion(model="smart", messages=HARD_MESSAGES, temperature=0.9) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["b"] + assert host.calls[-1].kwargs["temperature"] == 0.9 + assert "max_tokens" not in host.calls[-1].kwargs + + +@pytest.mark.asyncio +async def test_marker_params_fill_gaps_on_the_routed_call_but_never_override_the_caller( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_tiered(temperature=0.3)) + await core.acompletion(model="smart", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["temperature"] == 0.3 + await core.acompletion(model="smart", messages=MESSAGES, temperature=0.9) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["temperature"] == 0.9 + + +@pytest.mark.asyncio +async def test_marker_credentials_and_pricing_are_never_forwarded(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_tiered(api_base="http://marker.example", input_cost_per_token=0.5)) + await core.acompletion(model="smart", messages=HARD_MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + sent: Final = host.calls[-1].kwargs + assert host.deployments_called() == ["b"] + assert sent.get("api_base") is None + assert "input_cost_per_token" not in sent + + +@pytest.mark.asyncio +async def test_a_deployments_own_param_beats_the_forwarded_marker_param( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s", aws_region_name="deployment-region"), + deployment("big", "b"), + _complexity_router(aws_region_name="marker-region"), + ] + ) + await core.acompletion(model="smart", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["aws_region_name"] == "deployment-region" + await core.acompletion(model="smart", messages=HARD_MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["aws_region_name"] == "marker-region" + + +# --- the Claude Code session router binding -------------------------------------------------------------------- + + +def _session_request( + agent_id: str | None = None, caller: str = "caller-hash" +) -> dict[str, object]: # mutable-ok: request kwargs + headers: Final[dict[str, str]] = {"x-claude-code-session-id": SESSION_ID} + if agent_id is not None: + headers["x-claude-code-agent-id"] = agent_id + return {"proxy_server_request": {"headers": headers}, "metadata": {"user_api_key_hash": caller}} + + +@pytest.mark.asyncio +async def test_a_subagent_request_follows_the_router_its_session_started_on(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + await core.select("smart", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + subagent: Final = _session_request(agent_id="agent-1") + chosen: Final = await core.select("big", subagent, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + assert _metadata(subagent)["model_group"] == "smart" + + +@pytest.mark.asyncio +async def test_a_session_binding_is_scoped_to_the_caller(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + await core.select("smart", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + other_caller: Final = _session_request(agent_id="agent-1", caller="other-hash") + chosen: Final = await core.select("big", other_caller, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_main_agent_request_is_not_redirected_by_its_session_binding(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + await core.select("smart", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + chosen: Final = await core.select("big", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_subagent_request_without_a_session_binding_keeps_its_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + chosen: Final = await core.select("big", _session_request(agent_id="agent-1"), messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_id", ["short", "has spaces in it"]) +async def test_an_invalid_session_id_never_binds(make_core: MakeCore, session_id: str) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["proxy_server_request"] = {"headers": {"x-claude-code-session-id": session_id}} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + subagent: Final = _session_request(agent_id="agent-1") + subagent["proxy_server_request"] = { + "headers": {"x-claude-code-session-id": session_id, "x-claude-code-agent-id": "agent-1"} + } + chosen: Final = await core.select("big", subagent, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +# --- the routing-plugin pipeline ------------------------------------------------------------------------------ + + +class Narrowing: + """Keeps only ``keep`` among the candidate models and records what it was handed.""" + + def __init__(self, keep: frozenset[str]) -> None: + self.keep: Final = keep + self.seen: list[list[str]] = [] # mutable-ok: the candidate lists each run saw + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.candidate_models)) + context.candidate_models = [model for model in context.candidate_models if model in self.keep] + context.signals["narrowing"] = {"kept": sorted(self.keep)} + return context + + +TWO: Final = [deployment("g", "a"), deployment("g", "b")] + + +@pytest.mark.asyncio +async def test_a_routing_plugin_narrows_the_candidates_to_the_models_it_keeps(make_core: MakeCore) -> None: + plugin: Final = Narrowing(frozenset({"openai/b"})) + core: Final = make_core(model_list=list(TWO), plugins=[plugin]) + assert await picked_ids(core, "g", messages=MESSAGES) == {"b"} + assert plugin.seen[0] == ["openai/a", "openai/b"] + + +@pytest.mark.asyncio +async def test_routing_plugin_signals_and_narrowed_candidates_are_stamped_on_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/b"}))]) + request: Final[dict[str, object]] = {} + await core.select("g", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request == { + "metadata": { + "routing_plugin_signals": {"narrowing": {"kept": ["openai/b"]}}, + "_routing_plugin_candidate_models": ["openai/b"], + } + } + + +@pytest.mark.asyncio +async def test_a_routing_plugin_that_keeps_every_candidate_stamps_only_its_signals(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a", "openai/b"}))]) + request: Final[dict[str, object]] = {} + await core.select("g", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request == {"metadata": {"routing_plugin_signals": {"narrowing": {"kept": ["openai/a", "openai/b"]}}}} + + +@pytest.mark.asyncio +async def test_routing_plugins_run_in_order_each_seeing_the_previous_narrowing(make_core: MakeCore) -> None: + first: Final = Narrowing(frozenset({"openai/b", "openai/c"})) + second: Final = Narrowing(frozenset({"openai/c"})) + core: Final = make_core(model_list=[*TWO, deployment("g", "c")], plugins=[first, second]) + assert await picked_ids(core, "g", times=1, messages=MESSAGES) == {"c"} + assert second.seen == [["openai/b", "openai/c"]] + + +@pytest.mark.asyncio +async def test_a_routing_plugin_that_narrows_to_nothing_fails_the_pick(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset())]) + with pytest.raises(ValueError, match="No deployments left") as raised: + await core.select("g", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == "No deployments left after routing-plugin filtering. candidate_models=[]" + + +SYNC_PLUGIN_MESSAGE_START: Final = ( + "Router(plugins=[...]) is configured, but this call resolved to the synchronous deployment-selection path" +) + + +@pytest.mark.parametrize("entry", ["get_available_deployment", "completion"]) +def test_routing_plugins_refuse_the_sync_selection_path(make_core: MakeCore, host: ScriptedHost, entry: str) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a"}))]) + with pytest.raises(ValueError, match=re.escape(SYNC_PLUGIN_MESSAGE_START)) as raised: + _sync_entry(core, entry, "g") + assert str(raised.value).startswith(SYNC_PLUGIN_MESSAGE_START) + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_routing_plugins_refuse_usage_based_routing_v1_even_from_async_callers(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a"}))], routing_strategy="usage-based-routing" + ) + with pytest.raises(ValueError, match=re.escape(SYNC_PLUGIN_MESSAGE_START)) as raised: + await core.select("g", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value).startswith(SYNC_PLUGIN_MESSAGE_START) + + +@pytest.mark.asyncio +async def test_a_routing_plugin_does_not_narrow_a_pick_by_deployment_id(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/b"}))]) + chosen: Final = await core.select("a", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_completion_with_a_routing_plugin_calls_a_kept_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a"}))]) + await asyncio.gather(*(core.acompletion(model="g", messages=MESSAGES) for _ in range(5))) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] * 5 diff --git a/tests/unit/test_router/core/test_selection_characterization_strategies.py b/tests/unit/test_router/core/test_selection_characterization_strategies.py new file mode 100644 index 00000000000..3999c24857f --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_strategies.py @@ -0,0 +1,394 @@ +""" +Characterization of how a router core's routing strategies pick a deployment: strategy validation, +simple-shuffle weighting, the stateful strategies without traffic, per-request strategy overrides, +runtime strategy changes and custom strategies registered with ``register_plugin``. +""" + +import re +from typing import Final + +import pydantic +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.types.router import CustomRoutingStrategyBase, RouterRateLimitError, RoutingStrategy +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, picked_ids + +SPREAD_PICKS: Final = 60 +"""Enough uniform picks between two deployments that missing one is a 2**-59 event.""" + +INVALID_STRATEGY_MESSAGE: Final = ( + "Invalid routing_strategy: '{name}'. Valid options: ['simple-shuffle', 'lar1', 'least-busy', " + "'latency-based-routing', 'cost-based-routing', 'usage-based-routing-v2', 'usage-based-routing', " + "'provider-budget-routing']. Check 'router_settings.routing_strategy' in your config.yaml or the " + "'routing_strategy' parameter if using the Router SDK directly." +) + +NO_DEPLOYMENTS_MESSAGE: Final = ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=g. " + "pre-call-checks=False, cooldown_list=[]" +) + + +def _two(**params: object) -> list[dict[str, object]]: # mutable-ok: a model_list # kwargs-ok: shared litellm_params + return [deployment("g", "a", **params), deployment("g", "b", **params)] + + +# --- strategy validation --------------------------------------------------------------------- + + +@pytest.mark.parametrize("routing_strategy", ["bogus", "priority", "Simple-Shuffle"]) +def test_an_unknown_routing_strategy_is_rejected_at_construction(make_core: MakeCore, routing_strategy: str) -> None: + with pytest.raises(ValueError, match="Invalid routing_strategy") as raised: + make_core(model_list=_two(), routing_strategy=routing_strategy) + assert str(raised.value) == INVALID_STRATEGY_MESSAGE.format(name=routing_strategy) + + +@pytest.mark.asyncio +async def test_lar1_is_an_accepted_strategy_that_picks_within_the_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*_two(), deployment("other", "c")], routing_strategy="lar1") + assert await picked_ids(core, "g", messages=MESSAGES) <= {"a", "b"} + + +@pytest.mark.asyncio +async def test_provider_budget_routing_without_budgets_has_no_selector_and_rate_limits_every_pick( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="provider-budget-routing") + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE + with pytest.raises(RouterRateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ids", [("a", "b"), ("b", "a")]) +async def test_a_routing_strategy_enum_member_is_accepted_like_its_string( + make_core: MakeCore, ids: tuple[str, str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", deployment_id) for deployment_id in ids], + routing_strategy=RoutingStrategy.LEAST_BUSY, + ) + assert await picked_ids(core, "g", messages=MESSAGES) == {ids[0]} + + +@pytest.mark.asyncio +async def test_select_with_an_unknown_kind_raises_value_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + with pytest.raises(ValueError, match=re.escape("unknown select kind 'nope'")): + await core.select("g", kind="nope") # pyright: ignore[reportCallIssue, reportArgumentType] + + +# --- simple-shuffle -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_simple_shuffle_spreads_picks_over_unweighted_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_spreads_picks_when_every_weight_is_zero(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(weight=0, rpm=0, tpm=0)) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_falls_through_all_zero_weights_to_rpm(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight=0, rpm=5), deployment("g", "b", weight=0, rpm=0)]) + assert await picked_ids(core, "g") == {"a"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_falls_through_unset_weight_and_rpm_to_tpm(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", tpm=0), deployment("g", "b", tpm=100)]) + assert await picked_ids(core, "g") == {"b"} + + +@pytest.mark.asyncio +async def test_rpm_takes_precedence_over_tpm(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=1, tpm=0), deployment("g", "b", rpm=0, tpm=1000)]) + assert await picked_ids(core, "g") == {"a"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_with_a_non_numeric_weight_raises_type_error_at_pick(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight="heavy"), deployment("g", "b")]) + with pytest.raises(TypeError, match="Deployment weight must be numeric"): + await core.select("g") + + +@pytest.mark.asyncio +async def test_request_router_weights_are_keyed_by_the_alias_target_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), model_group_alias={"alias": "g"}) + assert await picked_ids(core, "alias", request_kwargs={"_router_weights": {"g": {"b": 1.0}}}) == {"b"} + + +@pytest.mark.asyncio +async def test_request_router_weights_for_another_group_are_ignored(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight=1), deployment("g", "b", weight=0)]) + assert await picked_ids(core, "g", request_kwargs={"_router_weights": {"other": {"b": 1.0}}}) == {"a"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("weights", [{"g": {"a": 0.0, "b": 0.0}}, {"g": {"a": -1.0}}, {"g": {"a": "1"}}]) +async def test_invalid_request_router_weights_fail_the_pick_with_a_validation_error( + make_core: MakeCore, weights: dict[str, object] +) -> None: + core: Final = make_core(model_list=_two()) + with pytest.raises(pydantic.ValidationError): + await core.select("g", {"_router_weights": weights}) + + +@pytest.mark.asyncio +async def test_an_empty_request_router_weights_group_falls_back_to_deployment_weights(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight=0), deployment("g", "b", weight=2)]) + assert await picked_ids(core, "g", request_kwargs={"_router_weights": {"g": {}}}) == {"b"} + + +@pytest.mark.asyncio +async def test_request_router_weights_give_unlisted_deployments_zero_weight(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*_two(), deployment("g", "c")]) + assert await picked_ids(core, "g", request_kwargs={"_router_weights": {"g": {"c": 0.5}}}) == {"c"} + + +# --- stateful strategies without traffic --------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["latency-based-routing", "usage-based-routing-v2"]) +async def test_strategy_without_traffic_spreads_picks(make_core: MakeCore, routing_strategy: str) -> None: + core: Final = make_core(model_list=_two(), routing_strategy=routing_strategy) + assert await picked_ids(core, "g", times=SPREAD_PICKS, messages=MESSAGES) == {"a", "b"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ids", [("a", "b"), ("b", "a")]) +async def test_cost_based_routing_without_known_costs_picks_the_first_configured( + make_core: MakeCore, ids: tuple[str, str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", deployment_id) for deployment_id in ids], routing_strategy="cost-based-routing" + ) + assert await picked_ids(core, "g", messages=MESSAGES) == {ids[0]} + + +def test_sync_pick_under_usage_based_routing_v1_picks_the_first_configured(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="usage-based-routing") + for _ in range(5): + assert core.get_available_deployment("g", messages=MESSAGES)["model_info"]["id"] == "a" # pyright: ignore[reportAttributeAccessIssue] + + +def test_sync_completion_under_least_busy_calls_the_first_configured(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="least-busy") + for _ in range(3): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("completion") == ["a", "a", "a"] + + +# --- per-request routing strategy override -------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("override", ["least-busy", RoutingStrategy.LEAST_BUSY], ids=["string", "enum"]) +async def test_a_request_routing_strategy_overrides_simple_shuffle(make_core: MakeCore, override: object) -> None: + core: Final = make_core(model_list=_two()) + assert await picked_ids(core, "g", request_kwargs={"routing_strategy": override}, messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_a_simple_shuffle_request_override_spreads_picks_on_a_least_busy_router(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="least-busy") + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"routing_strategy": "simple-shuffle"}, messages=MESSAGES + ) + assert picked == {"a", "b"} + + +@pytest.mark.asyncio +async def test_a_cost_based_request_override_picks_the_cheapest(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "pricey", input_cost_per_token=0.00001, output_cost_per_token=0.00002), + deployment("g", "cheap", input_cost_per_token=0.000001, output_cost_per_token=0.000002), + ] + ) + picked: Final = await picked_ids( + core, "g", request_kwargs={"routing_strategy": "cost-based-routing"}, messages=MESSAGES + ) + assert picked == {"cheap"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("override", ["bogus", "lar1", "provider-budget-routing", 5]) +async def test_an_unsupported_request_routing_strategy_is_ignored(make_core: MakeCore, override: object) -> None: + core: Final = make_core(model_list=_two()) + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"routing_strategy": override}, messages=MESSAGES + ) + assert picked == {"a", "b"} + + +@pytest.mark.asyncio +async def test_a_request_routing_strategy_overrides_the_routing_group_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=_two(), routing_groups=[{"group_name": "grp", "models": ["g"], "routing_strategy": "least-busy"}] + ) + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"routing_strategy": "simple-shuffle"}, messages=MESSAGES + ) + assert picked == {"a", "b"} + + +@pytest.mark.asyncio +async def test_completion_with_a_request_routing_strategy_calls_the_overridden_pick( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_two()) + for _ in range(5): + await core.acompletion(model="g", messages=MESSAGES, routing_strategy="least-busy") # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("acompletion") == ["a"] * 5 + + +# --- runtime strategy changes ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_applying_a_routing_strategy_switches_how_picks_are_made(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.apply(ChangeSet(settings={"routing_strategy": "least-busy"})) + assert core.settings().routing_strategy == "least-busy" + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_applying_an_invalid_routing_strategy_raises_and_keeps_the_previous_one(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="least-busy") + with pytest.raises(ValueError, match="Invalid routing_strategy") as raised: + core.apply(ChangeSet(settings={"routing_strategy": "bogus"})) + assert str(raised.value) == INVALID_STRATEGY_MESSAGE.format(name="bogus") + assert core.settings().routing_strategy == "least-busy" + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + + +# --- custom routing strategies ----------------------------------------------------------------- + + +def _fixed(deployment_id: str) -> dict[str, object]: # mutable-ok: the strategy interface's own type + return { + "model_name": "x", + "litellm_params": {"model": f"openai/{deployment_id}"}, + "model_info": {"id": deployment_id}, + } + + +class FixedStrategy(CustomRoutingStrategyBase): + """Always answers ``custom`` from the async pick and ``sync-custom`` from the sync pick.""" + + def __init__(self) -> None: + self.asked: list[str] = [] # mutable-ok: the models each pick was asked for + + async def async_get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + self.asked.append(model) + return _fixed("custom") + + def get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + self.asked.append(model) + return _fixed("sync-custom") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["g", "unknown-model"]) +async def test_a_custom_strategy_answer_is_returned_as_is_for_any_model(make_core: MakeCore, model: str) -> None: + core: Final = make_core(model_list=_two()) + strategy: Final = FixedStrategy() + core.register_plugin(strategy) + chosen: Final = await core.select(model) + assert chosen == _fixed("custom") + assert strategy.asked == [model] + + +@pytest.mark.asyncio +async def test_a_custom_strategy_does_not_change_the_healthy_listing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + healthy: Final = await core.select("g", kind="healthy") + assert sorted(d["model_info"]["id"] for d in healthy) == ["a", "b"] # pyright: ignore[reportIndexIssue, reportArgumentType] + + +@pytest.mark.asyncio +async def test_completion_calls_the_deployment_a_custom_strategy_returns( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("acompletion") == ["custom"] + assert host.calls[-1].model == "openai/custom" + + +def test_sync_completion_uses_the_custom_strategys_sync_pick(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("completion") == ["sync-custom"] + assert core.get_available_deployment("g") == _fixed("sync-custom") # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.asyncio +async def test_set_custom_routing_strategy_registers_like_register_plugin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.set_custom_routing_strategy(FixedStrategy()) # pyright: ignore[reportAttributeAccessIssue] + assert await core.select("g") == _fixed("custom") + + +@pytest.mark.asyncio +async def test_applying_a_different_routing_strategy_drops_the_custom_strategy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + core.apply(ChangeSet(settings={"routing_strategy": "least-busy"})) + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_applying_the_current_routing_strategy_keeps_the_custom_strategy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + core.apply(ChangeSet(settings={"routing_strategy": "simple-shuffle"})) + assert await core.select("g") == _fixed("custom") + + +@pytest.mark.asyncio +async def test_a_custom_strategy_error_reaches_the_caller(make_core: MakeCore, host: ScriptedHost) -> None: + class Refusing(FixedStrategy): + async def async_get_available_deployment( + self, model: str, *args: object, **kwargs: object + ) -> dict: # mutable-ok: the strategy interface's own type # kwargs-ok: the strategy interface + raise litellm.BadRequestError(message="refused by strategy", model=model, llm_provider="") + + core: Final = make_core(model_list=_two()) + core.register_plugin(Refusing()) + with pytest.raises(litellm.BadRequestError, match="refused by strategy"): + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] diff --git a/tests/unit/test_router/core/test_selection_characterization_traffic.py b/tests/unit/test_router/core/test_selection_characterization_traffic.py new file mode 100644 index 00000000000..7793b14ef6e --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_traffic.py @@ -0,0 +1,145 @@ +""" +Characterization of the stateful routing strategies once traffic has gone through the core: +least-busy follows in-flight requests, latency-based follows measured latency, and usage-based +routing follows the tokens and requests each deployment served. The host here is a +``TrafficHost``, which feeds the strategy loggers the way litellm's logging does around a call. +""" + +import asyncio +import re +from typing import Final + +import pytest + +import litellm +from tests.unit.test_router.core.scripted_host import bad_request, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, TrafficHost, picked_ids + + +@pytest.fixture +def host() -> TrafficHost: + return TrafficHost() + + +async def _until_called(host: TrafficHost, count: int) -> None: + for _ in range(200): + if len(host.started) >= count: + return + await asyncio.sleep(0.005) + raise AssertionError(f"expected {count} started calls, saw {len(host.started)}") + + +# --- least-busy ------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_least_busy_avoids_a_deployment_with_a_request_in_flight(make_core: MakeCore, host: TrafficHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="least-busy") + host.gates["openai/a"] = asyncio.Event() + in_flight: Final = asyncio.create_task(core.acompletion(model="g", messages=MESSAGES)) # pyright: ignore[reportAttributeAccessIssue] + await _until_called(host, 1) + assert host.started == ["openai/a"] + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"b"} + host.gates["openai/a"].set() + await in_flight + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_least_busy_releases_a_deployment_whose_request_failed(make_core: MakeCore, host: TrafficHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="least-busy") + host.gates["openai/a"] = asyncio.Event() + host.answer("openai/a", bad_request()) + in_flight: Final = asyncio.create_task(core.acompletion(model="g", messages=MESSAGES, num_retries=0)) # pyright: ignore[reportAttributeAccessIssue] + await _until_called(host, 1) + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"b"} + host.gates["openai/a"].set() + with pytest.raises(litellm.BadRequestError): + await in_flight + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_least_busy_spreads_concurrent_requests_over_idle_deployments( + make_core: MakeCore, host: TrafficHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("g", "c")], routing_strategy="least-busy" + ) + gate: Final = asyncio.Event() + for model in ("openai/a", "openai/b", "openai/c"): + host.gates[model] = gate + calls: Final = [] + for count in range(1, 4): + calls.append(asyncio.create_task(core.acompletion(model="g", messages=MESSAGES))) # pyright: ignore[reportAttributeAccessIssue] + await _until_called(host, count) + assert host.started == ["openai/a", "openai/b", "openai/c"] + gate.set() + await asyncio.gather(*calls) + + +# --- latency-based ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_latency_based_routing_moves_to_the_faster_deployment(make_core: MakeCore, host: TrafficHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "slow"), deployment("g", "fast")], routing_strategy="latency-based-routing" + ) + host.delays["openai/slow"] = 0.15 + for _ in range(6): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert "fast" in host.deployments_called() + assert await picked_ids(core, "g", times=10, messages=MESSAGES) == {"fast"} + + +# --- usage-based --------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["usage-based-routing", "usage-based-routing-v2"]) +async def test_usage_based_routing_alternates_between_deployments_by_tokens_served( + make_core: MakeCore, host: TrafficHost, routing_strategy: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy=routing_strategy) + for _ in range(4): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + called: Final = host.deployments_called() + assert sorted(called[:2]) == ["a", "b"] + assert sorted(called[2:]) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_usage_based_v2_stops_picking_a_deployment_once_its_tpm_is_used( + make_core: MakeCore, host: TrafficHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", tpm=15), deployment("g", "b", tpm=1000)], + routing_strategy="usage-based-routing-v2", + ) + for _ in range(4): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + called: Final = host.deployments_called() + assert called.count("a") <= 1 + assert called[-2:] == ["b", "b"] + + +@pytest.mark.asyncio +async def test_usage_based_v2_rate_limits_once_every_deployment_used_its_rpm( + make_core: MakeCore, host: TrafficHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", rpm=1), deployment("g", "b", rpm=1)], + routing_strategy="usage-based-routing-v2", + ) + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + assert sorted(host.deployments_called()) == ["a", "b"] + with pytest.raises(litellm.RateLimitError) as raised: + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + assert re.search( + r"No deployments available for selected model\. Passed model=g\. Deployments=\{'[ab]': \{'current_tpm': 10, " + r"'tpm_limit': inf, 'current_rpm': 1, 'rpm_limit': 1\}", + str(raised.value), + ) + assert len(host.calls) == 2 diff --git a/tests/unit/test_router/core/test_streaming_characterization_anthropic.py b/tests/unit/test_router/core/test_streaming_characterization_anthropic.py new file mode 100644 index 00000000000..3330b350155 --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_anthropic.py @@ -0,0 +1,485 @@ +""" +Characterization of /v1/messages streaming through a router core: lifecycle frames held back while a +retry or fallback can still take over, provider ``event: error`` frames, raised stream errors, safeguard +refusals, same-group stream retries and their budget, and the hidden params the returned stream carries. + +Scripted streams are async iterators of raw SSE byte frames (``streaming_helpers.anthropic_stream``). +""" + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RetryPolicy +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + bad_request, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + BLOCK_START, + MESSAGE_START, + MESSAGE_STOP, + MESSAGES, + PING, + ByteStream, + anthropic_stream, + collect_bytes, + error_frame, + install_orjson_stub_if_missing, + refusal_delta, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] +REFUSAL_DETAILS: Final = {"type": "refusal", "category": "cyber"} + + +def event_names(frames: list[bytes]) -> list[bytes]: + return [frame.split(b"\n", 1)[0].removeprefix(b"event: ") for frame in frames] + + +def fallback_ok(opened: list[ByteStream] | None = None) -> Callable[..., ByteStream]: + return anthropic_stream( + MESSAGE_START, + text_delta("B"), + MESSAGE_STOP, + hidden_params={"model_id": "b-stream", "additional_headers": {"x-provider-request-id": "req-b"}}, + opened=opened, + ) + + +def answered_ok(opened: list[ByteStream] | None = None) -> Callable[..., ByteStream]: + return anthropic_stream(MESSAGE_START, text_delta("A"), MESSAGE_STOP, opened=opened) + + +async def open_stream(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=5, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def hidden(stream: object) -> Mapping[str, object]: + return getattr(stream, "_hidden_params") + + +def headers(stream: object) -> Mapping[str, object]: + return hidden(stream)["additional_headers"] # pyright: ignore[reportReturnType] # a dict of headers + + +# --- provider error frames before content ------------------------------------------------------ + + +async def test_an_error_frame_before_content_falls_back_and_drops_the_primarys_held_frames( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, BLOCK_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_an_error_frame_before_content_with_nothing_to_take_over_is_forwarded_verbatim( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + core = make_core(model_list=GROUPS_AB, num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a"] + + +async def test_an_error_frame_split_across_byte_chunks_is_still_recognised( + make_core: MakeCore, host: ScriptedHost +) -> None: + frame = error_frame() + host.answer("openai/a", anthropic_stream(MESSAGE_START, frame[:25], frame[25:])) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_an_error_frame_after_content_is_forwarded_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, text_delta("x"), error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("x"), error_frame()] + assert host.deployments_called() == ["a"] + + +async def test_a_non_retriable_error_frame_is_forwarded_even_with_a_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + frame = error_frame("invalid_request_error", "bad request") + host.answer("openai/a", anthropic_stream(MESSAGE_START, frame)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, frame] + assert host.deployments_called() == ["a"] + + +async def test_disable_fallbacks_forwards_a_retriable_error_frame(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core, disable_fallbacks=True)) + assert frames == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a"] + + +async def test_when_the_fallback_after_an_error_frame_fails_the_frames_error_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame("overloaded_error", "primary overloaded"))) + host.answer("openai/b", server_error()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.InternalServerError, match="primary overloaded"): + await collect_bytes(stream) + assert host.deployments_called() == ["a", "b"] + + +# --- held lifecycle frames -------------------------------------------------------------------- + + +async def test_a_ping_is_forwarded_live_while_lifecycle_frames_are_held( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, PING, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [PING, MESSAGE_START, text_delta("B"), MESSAGE_STOP] + + +async def test_with_nothing_to_take_over_every_frame_is_forwarded_live_in_order( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, PING, BLOCK_START, error_frame())) + core = make_core(model_list=GROUPS_AB, num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, PING, BLOCK_START, error_frame()] + + +async def test_held_frames_are_flushed_in_order_when_content_arrives(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, BLOCK_START, text_delta("A"), MESSAGE_STOP)) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, BLOCK_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +async def test_held_frames_are_flushed_when_the_stream_ends_without_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, BLOCK_START, MESSAGE_STOP)) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, BLOCK_START, MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +async def test_two_hundred_held_frames_commit_the_stream_so_a_later_error_frame_is_forwarded( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(*([BLOCK_START] * 201), error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [*([BLOCK_START] * 201), error_frame()] + assert host.deployments_called() == ["a"] + + +# --- raised stream errors ---------------------------------------------------------------------- + + +async def test_a_raised_rate_limit_before_content_falls_back(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=rate_limit())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_raised_bad_request_before_content_is_raised_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=bad_request(message="raised 400"))) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.BadRequestError, match="raised 400"): + await collect_bytes(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_raised_error_after_content_is_raised_without_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, text_delta("x"), error=rate_limit(message="after"))) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="after"): + await collect_bytes(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_raised_server_error_before_content_is_retried_within_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", anthropic_stream(MESSAGE_START, error=server_error()), answered_ok()) + core = make_core(model_list=[deployment("a", "a1"), deployment("a", "a2")], num_retries=1) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert len(host.deployments_called()) == 2 + assert len(host.sleeps) == 1 + + +# --- same-group stream retries ----------------------------------------------------------------- + + +async def test_an_error_frame_before_content_is_retried_within_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame()), answered_ok()) + core = make_core(model_list=[deployment("a", "a")], num_retries=2) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a"] + assert len(host.sleeps) == 1 + assert headers(stream)["x-litellm-attempted-retries"] == 1 + assert headers(stream)["x-litellm-max-retries"] == 2 + + +async def test_a_retry_policy_grants_stream_retries_when_num_retries_is_zero( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame("rate_limit_error", "slow")), answered_ok()) + core = make_core( + model_list=[deployment("a", "a")], num_retries=0, retry_policy=RetryPolicy(RateLimitErrorRetries=2) + ) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a"] + assert headers(stream)["x-litellm-max-retries"] == 2 + + +async def test_request_num_retries_zero_opts_out_of_the_retry_policy_for_a_stream( + make_core: MakeCore, host: ScriptedHost +) -> None: + frame = error_frame("rate_limit_error", "slow") + host.answer("openai/a", anthropic_stream(MESSAGE_START, frame), answered_ok()) + core = make_core( + model_list=[deployment("a", "a")], num_retries=0, retry_policy=RetryPolicy(RateLimitErrorRetries=2) + ) + frames = await collect_bytes(await open_stream(core, num_retries=0)) + assert frames == [MESSAGE_START, frame] + assert host.deployments_called() == ["a"] + + +async def test_the_deployment_num_retries_is_the_stream_retry_budget(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame()), answered_ok()) + core = make_core(model_list=[deployment("a", "a", num_retries=1)], num_retries=0) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert headers(stream)["x-litellm-max-retries"] == 1 + + +async def test_stream_retries_stop_at_the_budget_and_forward_the_last_error_frame( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + core = make_core(model_list=[deployment("a", "a")], num_retries=2) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a", "a", "a"] + assert len(host.sleeps) == 2 + + +async def test_spent_stream_retries_hand_over_to_the_fallback_chain(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=1) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a", "b"] + assert len(host.sleeps) == 1 + + +async def test_a_superseded_attempt_and_every_later_stream_are_closed(make_core: MakeCore, host: ScriptedHost) -> None: + opened: list[ByteStream] = [] + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame(), opened=opened)) + host.answer("openai/b", fallback_ok(opened)) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=1) + await collect_bytes(await open_stream(core)) + assert len(opened) == 3 + assert all(stream.closed for stream in opened) + + +# --- safeguard refusals ------------------------------------------------------------------------ + + +async def test_a_refusal_with_stop_details_before_content_falls_back(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_refusal_falls_back_through_the_content_policy_chain(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, content_policy_fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_refusal_without_stop_details_is_forwarded(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, refusal_delta(), MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +@pytest.mark.parametrize( + ("router_arguments", "request_kwargs"), + [({"fallbacks": [{"a": ["b"]}]}, {"disable_fallbacks": True}), ({}, {})], + ids=["fallbacks-disabled", "no-fallbacks"], +) +async def test_a_refusal_with_no_reachable_fallback_is_forwarded( + make_core: MakeCore, host: ScriptedHost, router_arguments: dict[str, object], request_kwargs: dict[str, object] +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, num_retries=0, **router_arguments) + frames = await collect_bytes(await open_stream(core, **request_kwargs)) + assert frames == [MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +async def test_a_refusal_after_content_is_forwarded(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, text_delta("x"), refusal_delta(REFUSAL_DETAILS))) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("x"), refusal_delta(REFUSAL_DETAILS)] + assert host.deployments_called() == ["a"] + + +# --- what the fallback hands back -------------------------------------------------------------- + + +async def test_a_fallback_that_answers_without_streaming_is_replayed_as_an_sse_lifecycle( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer( + "openai/b", + { + "id": "msg_b", + "type": "message", + "role": "assistant", + "model": "b", + "content": [{"type": "text", "text": "B"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert event_names(frames) == [ + b"message_start", + b"content_block_start", + b"content_block_delta", + b"content_block_stop", + b"message_delta", + b"message_stop", + ] + assert b'"B"' in frames[2] + + +# --- hidden params on the returned stream ------------------------------------------------------ + + +async def test_a_stream_without_recovery_reports_the_serving_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", answered_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + await collect_bytes(stream) + assert hidden(stream)["model_id"] == "a" + assert headers(stream) == { + "x-litellm-model-group": "a", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 0, + } + + +async def test_after_a_stream_fallback_the_stream_carries_the_fallbacks_hidden_params_and_headers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + await collect_bytes(stream) + assert hidden(stream)["model_id"] == "b-stream" + assert headers(stream) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + "x-provider-request-id": "req-b", + } + + +async def test_after_retries_then_a_fallback_the_retry_budget_header_lingers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", answered_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=1) + stream = await open_stream(core) + await collect_bytes(stream) + assert headers(stream) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + "x-litellm-max-retries": 1, + } + + +async def test_a_stream_retry_keeps_reporting_the_first_deployment_when_the_retry_has_no_hidden_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", anthropic_stream(MESSAGE_START, error=server_error()), answered_ok()) + core = make_core(model_list=[deployment("a", "a1"), deployment("a", "a2")], num_retries=1) + stream = await open_stream(core) + await collect_bytes(stream) + first, _retried = host.deployments_called() + assert hidden(stream)["model_id"] == first + assert headers(stream)["x-litellm-attempted-retries"] == 1 diff --git a/tests/unit/test_router/core/test_streaming_characterization_chat.py b/tests/unit/test_router/core/test_streaming_characterization_chat.py new file mode 100644 index 00000000000..992404f80d7 --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_chat.py @@ -0,0 +1,669 @@ +""" +Characterization of chat-completion streaming through a router core: the first chunk read inside the +attempt, mid-stream errors and their fallback, mid-stream continuation, the sync twin, and the stream +timeout sent to the host. + +Each scripted stream is a ``CustomStreamWrapper`` built by ``streaming_helpers.chat_stream``; a stream +that raises before its first chunk fails inside the attempt (retries and fallbacks run as for any error), +while one that raises later fails while the caller iterates. +""" + +import asyncio +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + bad_request, + content_policy, + context_window, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGES, + chat_stream, + collect_texts, + collect_texts_sync, + drain_into, +) + +MakeCore = Callable[..., RouterCore] +PREFILL_MODEL: Final = "anthropic/claude-sonnet-4-5" +"""A model whose deployments accept an assistant prefill, so a continuation may land on them.""" +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_AB, "fallbacks": [{"a": ["b"]}], "num_retries": 0, **arguments}) + + +async def open_stream(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.acompletion(model="a", messages=MESSAGES, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def headers(response: object) -> Mapping[str, object]: + return getattr(response, "_hidden_params")["additional_headers"] + + +def model_id(response: object) -> object: + return getattr(response, "_hidden_params").get("model_id") + + +# --- the first chunk is read inside the attempt ----------------------------------------------- + + +async def test_a_stream_that_fails_on_its_first_chunk_is_retried_inside_the_attempt( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0), chat_stream("ok")) + core = make_core(model_list=[deployment("a", "a")], num_retries=2) + stream = await open_stream(core) + assert await collect_texts(stream) == "ok" + assert host.deployments_called() == ["a", "a"] + assert len(host.sleeps) == 1 + assert headers(stream)["x-litellm-attempted-retries"] == 1 + + +async def test_a_stream_that_fails_on_its_first_chunk_falls_back_after_its_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, num_retries=2) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "a", "a", "b"] + assert len(host.sleeps) == 3 + assert model_id(stream) == "b" + assert headers(stream)["x-litellm-attempted-fallbacks"] == 1 + + +async def test_a_first_chunk_failure_with_nowhere_to_go_raises_the_provider_error_from_the_call( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(message="first chunk"), fail_after=0)) + core = make_core(model_list=[deployment("a", "a")], num_retries=0) + with pytest.raises(litellm.RateLimitError, match="first chunk"): + await open_stream(core) + + +async def test_a_first_chunk_bad_request_is_raised_from_the_call_without_retry( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=bad_request(message="nope"), fail_after=0)) + core = make_core(model_list=GROUPS_AB, num_retries=2) + with pytest.raises(litellm.BadRequestError, match="nope"): + await open_stream(core) + assert host.deployments_called() == ["a"] + + +async def test_the_first_chunk_already_read_is_still_delivered_to_the_caller( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("first", "second", "third")) + core = make_core(model_list=[deployment("a", "a")]) + stream = await open_stream(core) + assert await collect_texts(stream) == "firstsecondthird" + + +async def test_a_stream_that_ends_without_any_chunk_or_finish_reason_ends_quietly_without_text( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(finish=False)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "" + assert host.deployments_called() == ["a"] + + +async def test_a_stream_that_ends_after_content_without_a_finish_reason_raises_while_iterating( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", finish=False)) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.APIConnectionError, match="Stream ended before a finish_reason was received"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +# --- mid-stream errors ------------------------------------------------------------------------- + + +async def test_a_mid_stream_error_after_content_raises_the_provider_error_not_mid_stream_fallback_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="cut after content"))) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="cut after content") as raised: + await collect_texts(stream) + assert not isinstance(raised.value, MidStreamFallbackError) + assert host.deployments_called() == ["a"] + + +async def test_a_mid_stream_error_after_content_delivers_the_text_before_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + core = build(make_core) + stream = await open_stream(core) + received: list[object] = [] + with pytest.raises(litellm.RateLimitError): + await drain_into(stream, received) + assert "".join(getattr(chunk, "choices")[0].delta.content or "" for chunk in received) == "hello" + + +async def test_a_mid_stream_error_before_content_falls_back_to_the_next_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "b"] + + +async def test_a_mid_stream_error_before_content_does_not_retry_the_same_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, num_retries=2) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "b"] + assert host.sleeps == [] + + +async def test_a_mid_stream_fallback_sends_the_original_messages_to_the_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + await collect_texts(stream) + fallback_call = host.calls_to("acompletion")[1] + assert fallback_call.kwargs["messages"] == MESSAGES + assert fallback_call.model_group == "b" + assert fallback_call.kwargs["stream"] is True + + +async def test_a_mid_stream_error_before_content_without_fallbacks_raises_the_provider_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="no fallback"), leading_role_chunk=True)) + core = make_core(model_list=GROUPS_AB, num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="no fallback") as raised: + await collect_texts(stream) + assert not isinstance(raised.value, MidStreamFallbackError) + assert host.deployments_called() == ["a"] + + +async def test_disable_fallbacks_on_the_request_raises_a_mid_stream_error_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="disabled"), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core, disable_fallbacks=True) + with pytest.raises(litellm.RateLimitError, match="disabled"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_per_request_fallback_list_is_used_for_a_mid_stream_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/c", chat_stream("C")) + core = build(make_core, model_list=[*GROUPS_AB, deployment("c", "c")]) + stream = await open_stream(core, fallbacks=["c"]) + assert await collect_texts(stream) == "C" + assert host.deployments_called() == ["a", "c"] + + +async def test_a_mid_stream_fallback_walks_the_chain_until_a_group_answers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", server_error()) + host.answer("openai/c", chat_stream("C")) + core = build(make_core, model_list=[*GROUPS_AB, deployment("c", "c")], fallbacks=[{"a": ["b", "c"]}]) + stream = await open_stream(core) + assert await collect_texts(stream) == "C" + assert host.deployments_called() == ["a", "b", "c"] + + +async def test_when_the_mid_stream_fallback_also_fails_the_primary_provider_error_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="primary"), leading_role_chunk=True)) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="primary"): + await collect_texts(stream) + assert host.deployments_called() == ["a", "b"] + + +async def test_a_mid_stream_server_error_before_content_falls_through_order_tiers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a1", chat_stream(error=server_error(), leading_role_chunk=True)) + host.answer("openai/a2", chat_stream("A2")) + core = make_core(model_list=[deployment("a", "a1", order=1), deployment("a", "a2", order=2)], num_retries=0) + stream = await open_stream(core) + assert await collect_texts(stream) == "A2" + assert host.deployments_called() == ["a1", "a2"] + assert model_id(stream) == "a2" + assert headers(stream)["x-litellm-model-group"] == "a" + + +async def test_a_mid_stream_bad_request_before_content_is_raised_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=bad_request(message="client error"), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.BadRequestError, match="client error"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_mid_stream_content_policy_error_before_content_is_raised_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=content_policy(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, fallbacks=None, content_policy_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContentPolicyViolationError): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_mid_stream_context_window_error_before_content_uses_the_context_window_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=context_window(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, fallbacks=None, context_window_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "b"] + + +async def test_a_mid_stream_context_window_error_after_content_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=context_window())) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, context_window_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContextWindowExceededError): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +# --- hidden params on a streamed response ----------------------------------------------------- + + +async def test_a_stream_without_fallback_reports_the_serving_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x")) + core = build(make_core) + stream = await open_stream(core) + await collect_texts(stream) + assert model_id(stream) == "a" + assert headers(stream)["x-litellm-model-group"] == "a" + assert headers(stream)["x-litellm-attempted-fallbacks"] == 0 + + +async def test_after_a_mid_stream_fallback_the_stream_reports_the_fallback_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + await collect_texts(stream) + assert model_id(stream) == "b" + assert headers(stream) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +async def test_after_a_mid_stream_fallback_each_fallback_chunk_carries_the_fallback_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + chunks = [chunk async for chunk in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + fallback_chunks = chunks[1:] + assert fallback_chunks + assert all(chunk._hidden_params.get("model_id") == "b" for chunk in fallback_chunks) + assert fallback_chunks[-1]._hidden_params["additional_headers"]["x-litellm-model-group"] == "b" + + +async def test_when_the_mid_stream_fallback_fails_the_stream_still_reports_the_primary( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError): + await collect_texts(stream) + assert model_id(stream) == "a" + assert headers(stream)["x-litellm-model-group"] == "a" + + +# --- mid-stream continuation ------------------------------------------------------------------- + + +def continuation_core(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return build( + make_core, + model_list=[deployment("a", "a"), deployment("b", "b", model=PREFILL_MODEL)], + enable_mid_stream_fallback_continuation=True, + **arguments, + ) + + +async def test_continuation_off_by_default_raises_after_content(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="cut"))) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = build(make_core, model_list=[deployment("a", "a"), deployment("b", "b", model=PREFILL_MODEL)]) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="cut"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_continues_after_content_on_a_prefill_capable_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "hello world" + assert host.deployments_called() == ["a", "b"] + + +async def test_continuation_sends_the_partial_text_as_an_assistant_prefill( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core) + await collect_texts(stream) + assert host.calls[1].kwargs["messages"] == [*MESSAGES, {"role": "assistant", "content": "hello", "prefix": True}] + + +async def test_continuation_extends_a_trailing_prefill_turn_in_place(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("Once", " upon", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" a time")) + core = continuation_core(make_core) + messages = [*MESSAGES, {"role": "assistant", "content": "Once", "prefix": True}] + stream = await core.acompletion(model="a", messages=messages, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await collect_texts(stream) + sent = host.calls[1].kwargs["messages"] + assert sent[-1] == {"role": "assistant", "content": "Once upon", "prefix": True} + assert len(sent) == 2 + + +async def test_continuation_declines_when_the_request_ends_in_a_plain_assistant_turn( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(message="cut"))) + host.answer(PREFILL_MODEL, chat_stream("y")) + core = continuation_core(make_core) + messages = [*MESSAGES, {"role": "assistant", "content": "earlier"}] + stream = await core.acompletion(model="a", messages=messages, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="cut"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_reduces_the_output_ceiling_by_the_tokens_already_emitted( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("hello", " there", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" friend")) + core = continuation_core(make_core) + stream = await open_stream(core, max_tokens=50) + await collect_texts(stream) + assert [call.kwargs["max_tokens"] for call in host.calls] == [50, 48] + + +async def test_continuation_raises_the_provider_error_when_the_output_ceiling_is_spent( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("hello", " there", error=rate_limit(message="spent"))) + host.answer(PREFILL_MODEL, chat_stream(" friend")) + core = continuation_core(make_core) + stream = await open_stream(core, max_tokens=2) + with pytest.raises(litellm.RateLimitError, match="spent"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_never_lands_on_a_fallback_that_cannot_take_a_prefill( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="no prefill target"))) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, enable_mid_stream_fallback_continuation=True) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="no prefill target"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +@pytest.mark.parametrize( + "request_kwargs", + [ + {"response_format": {"type": "json_object"}}, + {"tool_choice": "required", "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}]}, + ], + ids=["json-response-format", "forced-tool-choice"], +) +async def test_continuation_declines_a_constrained_request( + make_core: MakeCore, host: ScriptedHost, request_kwargs: dict[str, object] +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="constrained"))) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core, **request_kwargs) + with pytest.raises(litellm.RateLimitError, match="constrained"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_declines_a_fallback_target_that_overrides_request_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="override"))) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core, fallbacks=[{"a": [{"model": "b", "temperature": 0.1}]}]) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="override"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_runs_a_plain_fallback_for_an_error_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer(PREFILL_MODEL, chat_stream("B")) + core = continuation_core(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.calls[1].kwargs["messages"] == MESSAGES + + +# --- the deployment slot held by an open stream ----------------------------------------------- + + +async def test_an_open_stream_holds_its_max_parallel_requests_slot_until_consumed( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x")) + core = make_core(model_list=[deployment("a", "a", max_parallel_requests=1)], num_retries=0) + first = await open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(open_stream(core), 5) + await collect_texts(first) + second = await asyncio.wait_for(open_stream(core), 5) + assert await collect_texts(second) == "x" + + +# --- the sync twin ----------------------------------------------------------------------------- + + +def test_sync_stream_falls_back_on_a_mid_stream_error_before_content(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True, sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert collect_texts_sync(stream) == "B" + assert [(call.op, call.deployment_id) for call in host.calls] == [("completion", "a"), ("completion", "b")] + assert model_id(stream) == "b" + + +def test_sync_stream_raises_the_provider_error_after_content(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(message="sync cut"), sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync cut"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a"] + + +def test_sync_stream_never_continues_after_content_even_with_continuation_on( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="sync no continuation"), sync=True)) + host.answer(PREFILL_MODEL, chat_stream(" world", sync=True)) + core = continuation_core(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync no continuation"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a"] + + +def test_sync_stream_disable_fallbacks_raises_a_mid_stream_error_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="sync disabled"), leading_role_chunk=True, sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True, disable_fallbacks=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync disabled"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a"] + + +def test_sync_stream_first_chunk_failure_falls_back_inside_the_call(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0, sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert collect_texts_sync(stream) == "B" + assert host.deployments_called() == ["a", "b"] + + +def test_sync_stream_when_the_fallback_fails_raises_the_primary_error(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="sync primary"), leading_role_chunk=True, sync=True)) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync primary"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a", "b"] + + +# --- the stream timeout sent to the host ------------------------------------------------------- + + +@pytest.mark.parametrize( + ("router_arguments", "deployment_params", "request_kwargs", "expected"), + [ + ({"stream_timeout": 4}, {}, {}, 4), + ({"stream_timeout": 4}, {"stream_timeout": 3}, {}, 3.0), + ({"stream_timeout": 4}, {"stream_timeout": 3}, {"stream_timeout": 2}, 2), + ({}, {"stream_timeout_by_service_tier": {"flex": 9}, "stream_timeout": 3}, {"service_tier": "flex"}, 9.0), + ({}, {"stream_timeout_by_service_tier": {"flex": 9}, "stream_timeout": 3}, {"service_tier": "auto"}, 3.0), + ( + {}, + {"stream_timeout_by_service_tier": {"flex": 9}}, + {"service_tier": "flex", "stream_timeout": 2}, + 2, + ), + ({"default_litellm_params": {"stream_timeout": 8}}, {}, {}, 8), + ], + ids=[ + "router", + "deployment-over-router", + "request-over-deployment", + "service-tier-over-deployment", + "unmatched-service-tier-uses-deployment", + "request-over-service-tier", + "default-litellm-params", + ], +) +async def test_the_stream_timeout_comes_from_request_then_tier_then_deployment_then_router( + make_core: MakeCore, + host: ScriptedHost, + router_arguments: dict[str, object], + deployment_params: dict[str, object], + request_kwargs: dict[str, object], + expected: float, +) -> None: + host.answer("openai/a", chat_stream("x")) + core = make_core(model_list=[deployment("a", "a", **deployment_params)], **router_arguments) + stream = await open_stream(core, **request_kwargs) + await collect_texts(stream) + assert host.calls[0].kwargs["stream_timeout"] == expected + + +@pytest.mark.parametrize("router_arguments", [{}, {"timeout": 42}], ids=["no-router-timeout", "router-timeout"]) +async def test_without_any_stream_timeout_none_is_sent_and_the_plain_timeout_applies( + make_core: MakeCore, host: ScriptedHost, router_arguments: dict[str, object] +) -> None: + host.answer("openai/a", chat_stream("x")) + core = make_core(model_list=[deployment("a", "a")], **router_arguments) + stream = await open_stream(core) + await collect_texts(stream) + sent = host.calls[0].kwargs + assert "stream_timeout" not in sent + assert sent["timeout"] == router_arguments.get("timeout", 6000.0) + + +async def test_a_non_streaming_request_sends_no_stream_timeout(make_core: MakeCore, host: ScriptedHost) -> None: + core = make_core(model_list=[deployment("a", "a")], stream_timeout=4) + await core.acompletion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert "stream_timeout" not in host.calls[0].kwargs diff --git a/tests/unit/test_router/core/test_streaming_characterization_outcomes.py b/tests/unit/test_router/core/test_streaming_characterization_outcomes.py new file mode 100644 index 00000000000..99b3add153b --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_outcomes.py @@ -0,0 +1,382 @@ +""" +Characterization of what a streamed attempt records against its deployment: the success, tokens and +request a stream counts once it is fully consumed, and the failure and cooldown a stream that breaks +after its first chunk counts, for chat, Responses and /v1/messages streams. + +Every scripted stream here logs through the host when it ends (``logged=True``), as litellm's stream +logging does, so a core that learns outcomes from logging callbacks sees the stream's outcome only +then, with the assembled response's usage. +""" + +import asyncio +import time +from collections.abc import Callable +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGE_STOP, + MESSAGES, + Event, + anthropic_stream, + chat_stream, + collect_bytes, + collect_event_types, + collect_texts, + collect_texts_sync, + install_orjson_stub_if_missing, + responses_stream, + sse, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +USAGE: Final = {"prompt_tokens": 4, "completion_tokens": 8, "total_tokens": 12} + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, None, **params) for deployment_id in ids] + + +async def open_chat(core: RouterCore, model: str = "g") -> object: + stream: Final[object] = await core.acompletion(model=model, messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +def open_chat_sync(core: RouterCore) -> object: + stream: Final[object] = core.completion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +async def open_responses(core: RouterCore) -> object: + stream: Final[object] = await core.aresponses(model="g", input="hi", stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +async def open_messages(core: RouterCore) -> object: + stream: Final[object] = await core.aanthropic_messages(model="g", messages=MESSAGES, max_tokens=5, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +async def settle() -> None: + """Let the logging and recording tasks a finished stream scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +def eventually(check: Callable[[], bool]) -> None: + """Wait for ``check``: litellm logs a sync stream's outcome from a worker thread.""" + deadline: Final = time.monotonic() + 5.0 + while not check() and time.monotonic() < deadline: + time.sleep(0.01) + + +def successes(core: RouterCore, deployment_id: str) -> object: + count: Final[object] = core.health().store.get_cache(f"{deployment_id}:successes", local_only=True) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return count # pyright: ignore[reportUnknownVariableType] # see above + + +def cooled(core: RouterCore, *ids: str) -> list[str]: # mutable-ok: a fresh list + return sorted(dict(core.health().cooldowns.get_active_cooldowns(list(ids), None))) + + +# --- a stream consumed to its end ---------------------------------------------------------------- + + +async def test_a_consumed_stream_counts_its_final_tokens_and_one_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("hello", " world", usage=USAGE, logged=True)) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = await open_chat(core) + assert await collect_texts(stream) == "hello world" + await settle() + + assert await core.health().usage("g") == (12, 1) + assert successes(core, "a") == 1 + + +async def test_a_stream_counts_only_its_request_before_it_is_consumed(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("hello", usage=USAGE, logged=True)) + core: Final = make_core(model_list=group("a", rpm=100)) + + await open_chat(core) + await settle() + + # The request counts when the stream opens, with the tokens it has read so far (none). + assert await core.health().usage("g") == (0, 1) + assert successes(core, "a") is None + + +def test_a_sync_stream_counts_its_success_only_once_consumed(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("hello", usage=USAGE, logged=True, sync=True)) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = open_chat_sync(core) + assert successes(core, "a") is None + assert collect_texts_sync(stream) == "hello" + eventually(lambda: successes(core, "a") is not None) + + assert successes(core, "a") == 1 + # A sync call's success callback counts the success alone: the request counted when the stream + # opened, and no tokens. + assert asyncio.run(core.health().usage("g")) == (0, 1) + + +def test_a_sync_stream_that_breaks_after_content_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True, sync=True)) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = open_chat_sync(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + collect_texts_sync(stream) + eventually(lambda: bool(cooled(core, "a", "b"))) + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + assert asyncio.run(core.health().usage("g")) == (0, 1) + + +# --- a stream that breaks after its first chunk -------------------------------------------------- + + +async def test_a_stream_that_breaks_after_content_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await open_chat(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_texts(stream) + await settle() + + assert host.deployments_called() == ["a"] + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + # One request when the stream opened, one more for the failure. + assert await core.health().usage("g") == (0, 2) + + +async def test_a_stream_that_breaks_after_content_within_the_allowed_fails_does_not_cool_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = make_core(model_list=group("a"), num_retries=0, allowed_fails=3) + + stream: Final = await open_chat(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_texts(stream) + await settle() + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a") == [] + + +async def test_a_stream_that_breaks_before_content_and_falls_back_counts_one_failure_and_one_success( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + chat_stream("x", error=rate_limit(message="cut"), fail_after=1, leading_role_chunk=True, logged=True), + ) + host.answer("openai/b", chat_stream("B", usage=USAGE, logged=True)) + core: Final = make_core( + model_list=[deployment("a", "a", rpm=100), deployment("b", "b", rpm=100)], + fallbacks=[{"a": ["b"]}], + num_retries=0, + allowed_fails=3, + ) + + stream: Final = await open_chat(core, model="a") + assert await collect_texts(stream) == "B" + await settle() + + assert host.deployments_called() == ["a", "b"] + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert await core.health().usage("a") == (0, 2) + assert successes(core, "b") == 1 + assert await core.health().usage("b") == (12, 1) + + +async def test_a_stream_that_fails_on_its_first_chunk_counts_its_failure_once( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + chat_stream("x", error=rate_limit(), fail_after=0, logged=True), + chat_stream("ok", usage=USAGE, logged=True), + ) + core: Final = make_core(model_list=group("a", rpm=100), num_retries=1, allowed_fails=3) + + stream: Final = await open_chat(core) + assert await collect_texts(stream) == "ok" + await settle() + + assert host.deployments_called() == ["a", "a"] + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") == 1 + assert await core.health().usage("g") == (12, 2) + + +# --- Responses streams --------------------------------------------------------------------------- + + +def completed_event(input_tokens: int, output_tokens: int) -> ResponseCompletedEvent: + response: Final = ResponsesAPIResponse.model_construct( + id="resp", + background=None, + usage=ResponseAPIUsage( + input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens + ), + ) + return ResponseCompletedEvent(type="response.completed", response=response) # pyright: ignore[reportArgumentType] # the enum accepts its value + + +async def test_a_consumed_responses_stream_counts_its_final_tokens_and_one_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + responses_stream( + Event("response.created"), + Event("response.output_text.delta"), + completed_event(4, 8), + logged=True, + total_tokens=12, + ), + ) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = await open_responses(core) + assert (await collect_event_types(stream))[-1] == "response.completed" + await settle() + + assert await core.health().usage("g") == (12, 1) + assert successes(core, "a") == 1 + + +async def test_a_responses_stream_that_breaks_after_output_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + responses_stream( + Event("response.created"), + Event("response.output_text.delta", delta="partial"), + error=rate_limit(message="mid stream"), + logged=True, + ), + ) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await open_responses(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_event_types(stream) + await settle() + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + assert await core.health().usage("g") == (0, 2) + + +# --- /v1/messages streams ------------------------------------------------------------------------ + +MESSAGE_START_WITH_USAGE: Final = sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 3, "cache_read_input_tokens": 1, "output_tokens": 1}, + }, + }, +) +MESSAGE_DELTA_WITH_USAGE: Final = sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 8}}, +) + + +async def test_a_consumed_messages_stream_counts_its_final_tokens_and_one_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + anthropic_stream( + MESSAGE_START_WITH_USAGE, + text_delta("hello"), + # the delta frame split across two byte chunks + MESSAGE_DELTA_WITH_USAGE[:30], + MESSAGE_DELTA_WITH_USAGE[30:], + MESSAGE_STOP, + logged=True, + total_tokens=12, + ), + ) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = await open_messages(core) + assert len(await collect_bytes(stream)) == 5 + await settle() + + assert await core.health().usage("g") == (12, 1) + assert successes(core, "a") == 1 + + +async def test_a_messages_stream_that_breaks_after_content_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + anthropic_stream( + MESSAGE_START_WITH_USAGE, text_delta("x"), error=rate_limit(message="mid stream"), logged=True + ), + ) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await open_messages(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_bytes(stream) + await settle() + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + assert await core.health().usage("g") == (0, 2) diff --git a/tests/unit/test_router/core/test_streaming_characterization_responses.py b/tests/unit/test_router/core/test_streaming_characterization_responses.py new file mode 100644 index 00000000000..788e79d368e --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_responses.py @@ -0,0 +1,403 @@ +""" +Characterization of Responses API streaming through a router core: pre-output lifecycle events held +back while a fallback can still take over, mid-stream errors and their fallback, the continuation a +fallback receives after text was streamed, usage combined across a stream fallback, and the hidden +params the returned stream and its events carry. + +Scripted streams are ``BaseResponsesAPIStreamingIterator`` subclasses +(``streaming_helpers.responses_stream``) that raise a ``MidStreamFallbackError`` the way the +completion-bridge iterator does. +""" + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from litellm.router_core import RouterCore +from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + content_policy, + context_window, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + Event, + ScriptedResponsesStream, + collect_event_types, + drain_into, + responses_stream, +) + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] +CREATED: Final = "response.created" +IN_PROGRESS: Final = "response.in_progress" +TEXT_DELTA: Final = "response.output_text.delta" +COMPLETED: Final = "response.completed" + + +def cut(generated_content: str = "", original: Exception | None = None) -> MidStreamFallbackError: + """The error the bridge iterator raises when its stream breaks (before output when nothing was generated).""" + return MidStreamFallbackError( + message="stream cut", + model="d1", + llm_provider="anthropic", + original_exception=original if original is not None else rate_limit(message="stream cut"), + generated_content=generated_content, + is_pre_first_chunk=not generated_content, + ) + + +def fallback_ok( + *events: object, opened: list[ScriptedResponsesStream] | None = None +) -> Callable[..., ScriptedResponsesStream]: + return responses_stream( + *(events or (Event(CREATED), Event(TEXT_DELTA), Event(COMPLETED))), + hidden_params={"model_id": "b-stream", "additional_headers": {"x-provider-request-id": "req-b"}}, + opened=opened, + ) + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_AB, "fallbacks": [{"a": ["b"]}], "num_retries": 0, **arguments}) + + +async def open_stream(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aresponses(model="a", input="hi", stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def hidden(stream: object) -> Mapping[str, object]: + return getattr(stream, "_hidden_params") + + +def completed_event(input_tokens: int, output_tokens: int) -> ResponseCompletedEvent: + response = ResponsesAPIResponse.model_construct( + id="resp", + background=None, + usage=ResponseAPIUsage( + input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens + ), + ) + return ResponseCompletedEvent(type=COMPLETED, response=response) # pyright: ignore[reportArgumentType] # the enum accepts its value + + +# --- the returned stream ------------------------------------------------------------------------ + + +async def test_a_streaming_response_is_returned_as_a_responses_streaming_iterator( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), Event(COMPLETED))) + core = build(make_core) + stream = await open_stream(core) + assert isinstance(stream, BaseResponsesAPIStreamingIterator) + assert await collect_event_types(stream) == [CREATED, TEXT_DELTA, COMPLETED] + assert hidden(stream)["model_id"] == "a" + + +async def test_the_returned_stream_latches_the_terminal_event_it_forwarded( + make_core: MakeCore, host: ScriptedHost +) -> None: + completed = Event(COMPLETED) + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), completed)) + core = build(make_core) + stream = await open_stream(core) + await collect_event_types(stream) + assert getattr(stream, "completed_response") is completed + + +async def test_lifecycle_events_are_held_until_output_and_then_forwarded_in_order( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS), Event(TEXT_DELTA), Event(COMPLETED))) + core = build(make_core) + assert await collect_event_types(await open_stream(core)) == [CREATED, IN_PROGRESS, TEXT_DELTA, COMPLETED] + + +async def test_held_lifecycle_events_are_flushed_when_the_stream_ends_without_output( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS))) + core = build(make_core) + assert await collect_event_types(await open_stream(core)) == [CREATED, IN_PROGRESS] + + +async def test_the_source_stream_is_closed_when_the_stream_ends(make_core: MakeCore, host: ScriptedHost) -> None: + opened: list[ScriptedResponsesStream] = [] + host.answer("openai/a", responses_stream(Event(CREATED), Event(COMPLETED), opened=opened)) + core = build(make_core) + await collect_event_types(await open_stream(core)) + assert [stream.closed for stream in opened] == [True] + + +# --- mid-stream errors before output ----------------------------------------------------------- + + +async def test_a_mid_stream_error_before_output_falls_back_and_drops_the_primarys_held_events( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + assert await collect_event_types(stream) == [CREATED, TEXT_DELTA, COMPLETED] + assert host.deployments_called() == ["a", "b"] + assert host.calls[1].kwargs["input"] == "hi" + assert host.calls[1].model_group == "b" + + +async def test_a_mid_stream_error_before_output_without_fallbacks_flushes_held_events_then_raises( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS), error=cut())) + core = make_core(model_list=GROUPS_AB, num_retries=0) + stream = await open_stream(core) + seen: list[object] = [] + with pytest.raises(litellm.RateLimitError, match="stream cut") as raised: + await drain_into(stream, seen) + assert not isinstance(raised.value, MidStreamFallbackError) + assert [getattr(event, "type") for event in seen] == [CREATED, IN_PROGRESS] + assert host.deployments_called() == ["a"] + + +async def test_disable_fallbacks_raises_a_mid_stream_error_before_output( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core, disable_fallbacks=True) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_when_the_fallback_fails_the_held_events_are_replayed_before_the_primary_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + seen: list[object] = [] + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await drain_into(stream, seen) + assert [getattr(event, "type") for event in seen] == [CREATED] + assert host.deployments_called() == ["a", "b"] + + +async def test_more_than_two_hundred_lifecycle_events_are_forwarded_before_a_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(*(Event(IN_PROGRESS) for _ in range(201)), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + types = await collect_event_types(await open_stream(core)) + assert types == [*([IN_PROGRESS] * 201), CREATED, TEXT_DELTA, COMPLETED] + + +async def test_a_context_window_error_before_output_uses_the_context_window_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=context_window())) + host.answer("openai/b", fallback_ok()) + core = build(make_core, fallbacks=None, context_window_fallbacks=[{"a": ["b"]}]) + assert await collect_event_types(await open_stream(core)) == [CREATED, TEXT_DELTA, COMPLETED] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_context_window_error_after_output_is_raised(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=context_window())) + host.answer("openai/b", fallback_ok()) + core = build(make_core, context_window_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContextWindowExceededError): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_content_policy_cut_before_output_uses_the_content_policy_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut(original=content_policy()))) + host.answer("openai/b", fallback_ok()) + core = build(make_core, fallbacks=None, content_policy_fallbacks=[{"a": ["b"]}]) + assert await collect_event_types(await open_stream(core)) == [CREATED, TEXT_DELTA, COMPLETED] + assert host.deployments_called() == ["a", "b"] + + +# --- continuation after output ----------------------------------------------------------------- + + +async def test_after_text_the_fallback_receives_the_partial_output_and_a_request_to_continue( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + await collect_event_types(await open_stream(core)) + sent = host.calls[1].kwargs["input"] + assert isinstance(sent, list) + assert sent[0] == {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + assert sent[1] == {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "hel"}]} + assert sent[2]["role"] == "user" + assert 'It ended with "hel"' in sent[2]["content"][0]["text"] + + +async def test_after_text_the_output_ceiling_is_reduced_by_the_emitted_tokens( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + await collect_event_types(await open_stream(core, max_output_tokens=100)) + assert [call.kwargs["max_output_tokens"] for call in host.calls] == [100, 99] + + +async def test_after_text_a_spent_output_ceiling_raises_the_primary_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hello there friend"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core, max_output_tokens=1) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_text_a_content_policy_cut_raises_instead_of_continuing( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel", content_policy()))) + host.answer("openai/b", fallback_ok()) + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContentPolicyViolationError): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_a_tool_call_began_the_primary_error_is_raised(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer( + "openai/a", + responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"), started_tool_call=True), + ) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_text_a_json_schema_output_format_raises_the_primary_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core, text={"format": {"type": "json_schema", "name": "x", "schema": {}}}) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_text_the_primarys_events_stay_and_the_fallbacks_follow( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + types = await collect_event_types(await open_stream(core)) + assert types[:2] == [CREATED, TEXT_DELTA] + assert types[-1] == COMPLETED + assert host.deployments_called() == ["a", "b"] + + +# --- usage across a stream fallback -------------------------------------------------------------- + + +async def test_the_fallbacks_completed_usage_adds_the_usage_the_primary_latched( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut(), completed_response=completed_event(3, 4))) + host.answer("openai/b", fallback_ok(Event(CREATED), completed_event(10, 20))) + core = build(make_core) + stream = await open_stream(core) + events = [event async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + usage = events[-1].response.usage + assert (usage.input_tokens, usage.output_tokens, usage.total_tokens) == (13, 24, 37) + + +async def test_without_primary_usage_the_fallbacks_completed_usage_is_unchanged( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok(Event(CREATED), completed_event(10, 20))) + core = build(make_core) + stream = await open_stream(core) + events = [event async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + usage = events[-1].response.usage + assert (usage.input_tokens, usage.output_tokens, usage.total_tokens) == (10, 20, 30) + + +# --- hidden params after a stream fallback ------------------------------------------------------- + + +async def test_after_a_stream_fallback_the_stream_carries_the_fallbacks_hidden_params_and_headers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + await collect_event_types(stream) + assert hidden(stream)["model_id"] == "b-stream" + assert hidden(stream)["additional_headers"] == { + "x-provider-request-id": "req-b", + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +async def test_after_a_stream_fallback_each_fallback_event_carries_the_fallbacks_hidden_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + events = [event async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + assert all(event._hidden_params["model_id"] == "b-stream" for event in events) + assert all(event._hidden_params["additional_headers"]["x-litellm-model-group"] == "b" for event in events) + + +async def test_when_the_stream_fallback_fails_the_stream_still_reports_the_primary( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError): + await collect_event_types(stream) + assert hidden(stream)["model_id"] == "a" + assert hidden(stream)["additional_headers"]["x-litellm-model-group"] == "a" + + +async def test_the_primary_and_fallback_streams_are_closed(make_core: MakeCore, host: ScriptedHost) -> None: + opened: list[ScriptedResponsesStream] = [] + host.answer("openai/a", responses_stream(Event(CREATED), error=cut(), opened=opened)) + host.answer("openai/b", fallback_ok(opened=opened)) + core = build(make_core) + await collect_event_types(await open_stream(core)) + assert [stream.closed for stream in opened] == [True, True]