diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6cfe07d6aaf..ed658e074ee 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4343,12 +4343,14 @@ dependencies = [ "jsonschema", "litellm-cost", "litellm-error", + "litellm-tracing", "macro_rules_attribute", "rstest", "schemars 1.2.2", "serde", "serde_json", "serde_with", + "sha1 0.10.7", "strum", "thiserror 2.0.19", "time", @@ -4403,6 +4405,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-model-catalog-refresh", "litellm-python-compat", "litellm-redis-client", "litellm-redis-connect", diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 3b22f74d95e..a9ebae427d9 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -36,7 +36,7 @@ pub use marshal::{ pub use owned::PythonOwned; pub use runtime::{ ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value, - ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync, + ready_future, reserve_process_for_forking, run_async, run_async_value, run_isolated, run_sync, run_sync_value, runtime_handle, runtime_started, }; pub use services::PythonHostCalls; diff --git a/litellm-rust/crates/host-python/src/runtime.rs b/litellm-rust/crates/host-python/src/runtime.rs index 4f91e2e085b..2a8fc598fc4 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -68,6 +68,18 @@ fn runtime() -> PyResult<&'static Runtime> { Ok(pyo3_async_runtimes::tokio::get_runtime()) } +/// Runs `future` on a current-thread runtime built and shut down inside the call, its blocking pool joined before +/// it returns. No thread outlives the call, so unlike every other door this one never claims the process in the fork +/// gate: a process that only used it can still fork. The future must not hold resources tied to another runtime +pub fn run_isolated(future: F) -> std::io::Result { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let output = runtime.block_on(future); + runtime.shutdown_timeout(Duration::from_secs(5)); + Ok(output) +} + /// The native runtime, for work that must outlive the call that schedules it. pub fn runtime_handle() -> PyResult { Ok(runtime()?.handle().clone()) @@ -818,4 +830,21 @@ asyncio.run(exercise()) .expect("async route contract should hold"); }); } + + #[rstest] + fn an_isolated_run_joins_blocking_work_it_left_running_before_returning() { + let finished = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&finished); + let output = run_isolated(async move { + drop(tokio::task::spawn_blocking(move || { + thread::sleep(Duration::from_millis(50)); + flag.store(true, Ordering::SeqCst); + })); + 7 + }) + .unwrap(); + + assert_eq!(output, 7); + assert!(finished.load(Ordering::SeqCst)); + } } diff --git a/litellm-rust/crates/host-python/tests/isolated_run.rs b/litellm-rust/crates/host-python/tests/isolated_run.rs new file mode 100644 index 00000000000..7b8603229a7 --- /dev/null +++ b/litellm-rust/crates/host-python/tests/isolated_run.rs @@ -0,0 +1,9 @@ +use litellm_host_python::{reserve_process_for_forking, run_isolated, runtime_started}; + +#[test] +fn an_isolated_run_leaves_the_process_free_to_reserve_for_forking() { + run_isolated(async { tokio::task::spawn_blocking(|| ()).await.unwrap() }).unwrap(); + + assert!(!runtime_started()); + assert!(reserve_process_for_forking().is_ok()); +} diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 3bdc63453ef..3b20bbda607 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -55,36 +55,13 @@ impl HttpClientPool { config: &HttpClientConfig, variant: ClientVariant, ) -> Result { - let effective = match variant { - ClientVariant::Media => HttpClientConfig { - client_certificate: None, - proxies: EnvironmentProxies::default(), - ..config.clone() - }, - ClientVariant::UnpinnedMedia => HttpClientConfig { - client_certificate: None, - ..config.clone() - }, - ClientVariant::WebSocket => HttpClientConfig { - http2: false, - verify: match &config.verify { - Verify::BuiltInRoots => Verify::NativeRoots, - verify => verify.clone(), - }, - ..config.clone() - }, - ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), - }; - let key = (effective, variant); + let key = (effective(config, variant), variant); if let Some(pooled) = self.lock().get(&key) && pooled.built_at.elapsed() < self.ttl { return Ok(Client::new(pooled.client.clone())); } - let client = self - .apply(variant, reqwest::ClientBuilder::try_from(&key.0)?) - .build() - .map_err(Repr::Client)?; + let client = self.build(&key.0, variant)?; self.lock().insert( key, PooledClient { @@ -95,6 +72,26 @@ impl HttpClientPool { Ok(Client::new(client)) } + /// A client nobody else holds, for a caller that drives it on a runtime of its own and drops both together + pub fn unpooled( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.build(&effective(config, variant), variant) + .map(Client::new) + } + + fn build( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.apply(variant, reqwest::ClientBuilder::try_from(config)?) + .build() + .map_err(|error| Error::from(Repr::Client(error))) + } + fn lock(&self) -> MutexGuard<'_, Clients> { self.clients.lock().unwrap_or_else(PoisonError::into_inner) } @@ -116,6 +113,29 @@ impl HttpClientPool { } } +fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientConfig { + match variant { + ClientVariant::Media => HttpClientConfig { + client_certificate: None, + proxies: EnvironmentProxies::default(), + ..config.clone() + }, + ClientVariant::UnpinnedMedia => HttpClientConfig { + client_certificate: None, + ..config.clone() + }, + ClientVariant::WebSocket => HttpClientConfig { + http2: false, + verify: match &config.verify { + Verify::BuiltInRoots => Verify::NativeRoots, + verify => verify.clone(), + }, + ..config.clone() + }, + ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), + } +} + #[cfg(test)] mod tests { use std::{ diff --git a/litellm-rust/crates/model-catalog-refresh/AGENTS.md b/litellm-rust/crates/model-catalog-refresh/AGENTS.md new file mode 100644 index 00000000000..96228716abf --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/AGENTS.md @@ -0,0 +1,9 @@ +# Model catalog refresh + +- `fetch` fetches one document over http(s) or `s3://` and never parses it +- Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final +- `Failed.retryable` means only the attempt limit ended the fetch +- A server's Retry-After overrides the backoff, capped at `max_delay` +- `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off +- Error reasons use the URL with userinfo, query and fragment stripped, since presigned URLs carry credentials +- Bodies over `MAX_BODY_BYTES` are rejected diff --git a/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs b/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs index be506130554..2a4ec996bd4 100644 --- a/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs +++ b/litellm-rust/crates/model-catalog-refresh/tests/fetch.rs @@ -148,10 +148,6 @@ async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result` is generic over the host's value type, because Python registrations can hold `datetime` and other values JSON cannot. The host keeps its own objects (`Py` in python-bridge) and does the merging + - Entry order is latest-registration order. `landed_on` records which registered keys landed on each row key +- `write_pass` stops at the first `place` that returns `None` and reports `Pass::Swapped`. Retrying is the host's job + +`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 63b9ad14768..8cd8ee5760b 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -11,6 +11,7 @@ schema = ["dep:schemars", "litellm-cost/schema"] [dependencies] litellm-cost = { workspace = true, features = ["serde"] } litellm-error.workspace = true +litellm-tracing.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] } @@ -19,6 +20,7 @@ schemars = { workspace = true, optional = true } serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } serde_with.workspace = true +sha1.workspace = true strum.workspace = true thiserror.workspace = true time.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 55a95416da3..6c11a39b0eb 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -12,7 +12,8 @@ use crate::store::{Rows, Store}; use crate::wire::{EntryWire, Fields}; const SAMPLE_SPEC: &str = "sample_spec"; -const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; +/// The top-level key of the rules block, which is never a model row +pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; #[derive(Clone, Copy, Debug, PartialEq)] pub struct IntegrityLimits { @@ -111,6 +112,11 @@ impl Rows for CatalogRows { } impl IntegrityLimits { + pub const UNCHECKED: Self = Self { + min_model_count: 0, + min_reference_ratio: 0.0, + }; + pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option { (0.0..=1.0).contains(&min_reference_ratio).then_some(Self { min_model_count, @@ -418,6 +424,20 @@ impl Catalog { }) } + /// The row under an exact model or alias key, an alias reading its owner's row + pub fn get(&self, key: &str) -> Option<&Fields> { + self.rows().fields(key) + } + + /// A top-level value that is not a model row, such as `sample_spec` + pub fn section(&self, key: &str) -> Option<&Value> { + self.sections + .iter() + .filter(|section| is_section(§ion.key)) + .find(|section| section.key == key) + .map(|section| §ion.value) + } + pub fn model_count(&self) -> usize { self.rows().entries.len() } diff --git a/litellm-rust/crates/model-catalog/src/cost_map.rs b/litellm-rust/crates/model-catalog/src/cost_map.rs new file mode 100644 index 00000000000..bec417acae7 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/cost_map.rs @@ -0,0 +1,130 @@ +use std::sync::{Arc, LazyLock}; + +use indexmap::IndexMap; +use serde_json::Value; +use sha1::{Digest, Sha1}; + +use crate::bundled::BUNDLED_REGISTRY; +use crate::catalog::{Catalog, FALLBACK_GENERALIZATIONS, IntegrityLimits}; +use crate::error::Error; + +/// Which document is installed: the `git_blob_id` of its bytes and the ETag it was served with +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Provenance { + pub revision: String, + pub etag: Option, +} + +/// An installed cost map document. A host builds its live table from it +#[derive(Clone, Debug)] +pub struct CostMap { + base: Arc, + provenance: Provenance, +} + +#[derive(Debug)] +struct Base { + catalog: Catalog, + /// Each owner's names, the owner first, for owners with at least one alias + groups: IndexMap>, +} + +impl Base { + fn new(catalog: Catalog) -> Self { + let groups = catalog.aliases().iter().fold( + IndexMap::>::new(), + |mut groups, (alias, owner)| { + groups + .entry(owner.clone()) + .or_insert_with(|| vec![owner.clone()]) + .push(alias.clone()); + groups + }, + ); + Self { catalog, groups } + } +} + +static BUNDLED: LazyLock = LazyLock::new(|| CostMap { + base: Arc::new(Base::new( + Catalog::parse(BUNDLED_REGISTRY) + .expect("the bundled registry parses") + .value, + )), + provenance: Provenance { + revision: git_blob_id(BUNDLED_REGISTRY), + etag: None, + }, +}); + +impl CostMap { + pub fn bundled() -> Self { + BUNDLED.clone() + } + + /// Reads a fetched document. One that does not parse, or that has fewer models than `limits` allow against the + /// bundled one, is refused. Every alias problem it has is logged + pub fn read(body: &[u8], etag: Option, limits: IntegrityLimits) -> Result { + let loaded = Catalog::parse(body)?; + loaded + .value + .check_size(BUNDLED.base.catalog.model_count(), limits)?; + for problem in loaded.errors.iter().filter(|error| error.is_alias()) { + litellm_tracing::warn!( + "LiteLLM: model cost map alias problem: {}", + litellm_error::Report::new(problem) + ); + } + Ok(Self { + base: Arc::new(Base::new(loaded.value)), + provenance: Provenance { + revision: git_blob_id(body), + etag, + }, + }) + } + + pub fn catalog(&self) -> &Catalog { + &self.base.catalog + } + + pub fn provenance(&self) -> &Provenance { + &self.provenance + } + + /// Every name the row under `key` is served under, the owner first. Empty for a row without aliases + pub fn names(&self, key: &str) -> &[String] { + let owner = self + .catalog() + .aliases() + .get(key) + .map_or(key, String::as_str); + self.base + .groups + .get(owner) + .map(Vec::as_slice) + .unwrap_or_default() + } + + /// Each owner with its names, the owner first, in the order the document declared the aliases + pub fn alias_groups(&self) -> impl Iterator { + self.base + .groups + .iter() + .map(|(owner, names)| (owner.as_str(), names.as_slice())) + } + + /// The `fallback_generalizations` block as the document wrote it + pub fn fallback_generalizations(&self) -> Option<&Value> { + self.catalog().section(FALLBACK_GENERALIZATIONS) + } +} + +/// `git hash-object` of the bytes, which main reports as `source_revision` +pub fn git_blob_id(body: &[u8]) -> String { + let digest = Sha1::new() + .chain_update(format!("blob {}\0", body.len())) + .chain_update(body) + .finalize(); + format!("{digest:x}") +} diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 846398f7155..a4b54b5a551 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -167,6 +167,74 @@ impl Error { } } +impl Error { + /// The document parsed but holds fewer models than the integrity limits trust + pub fn is_integrity(&self) -> bool { + match self.repr() { + Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => true, + Repr::Json(_) + | Repr::NotAnObject + | Repr::Invalid(_) + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } + | Repr::MissingProvider + | Repr::UnreadField + | Repr::NoFallbackRules => false, + } + } + + /// A problem with a row's aliases, which the served map skips + pub fn is_alias(&self) -> bool { + match self.repr() { + Repr::AliasShadowsModel { .. } + | Repr::AliasTaken { .. } + | Repr::AliasTargetMissing { .. } + | Repr::AliasRowField => true, + Repr::Invalid(_) => matches!(self.field().as_deref(), Some("aliases" | "alias_of")), + Repr::Json(_) + | Repr::Empty + | Repr::BelowMinimum { .. } + | Repr::Shrunk { .. } + | Repr::NotAnObject + | Repr::UnknownBlockKey + | Repr::UnknownTimezone { .. } + | Repr::DuplicatePrice { .. } + | Repr::SuffixInsideBlock { .. } + | Repr::TierRange + | Repr::TierReversed + | Repr::TierGap { .. } + | Repr::MixedTierBasis + | Repr::NotPricedByResults + | Repr::BandsFromTwoSources + | Repr::NotAnOffPeakRate + | Repr::NoOffPeakPrices + | Repr::NoOffPeakHours + | Repr::EmptyList + | Repr::InvalidDate { .. } + | Repr::MissingProvider + | Repr::UnreadField + | Repr::NoFallbackRules => false, + } + } +} + #[derive(Debug)] pub struct Loaded { pub value: T, diff --git a/litellm-rust/crates/model-catalog/src/layers.rs b/litellm-rust/crates/model-catalog/src/layers.rs deleted file mode 100644 index 881f0f6af9b..00000000000 --- a/litellm-rust/crates/model-catalog/src/layers.rs +++ /dev/null @@ -1,32 +0,0 @@ -use std::sync::Arc; - -use indexmap::IndexMap; -use serde_json::Value; - -use crate::catalog::Catalog; -use crate::register::Registrations; -use crate::wire::Fields; - -/// One field of a fetched row over the layer below, RFC 7386. An absent field inherits, JSON `null` is `Remove` -#[derive(Clone, Debug, PartialEq)] -pub enum Patch { - Set(Value), - Remove, -} - -/// The four layers served rows are composed from, lowest first -#[derive(Clone, Debug)] -#[expect( - dead_code, - reason = "02 and 07 add the functions that read these layers" -)] -pub struct Layers { - /// 0, the wheel's embedded map - bundled: Arc, - /// 1, the last installed fetch, rejected rows left out - fetched: Option>, - /// 2, `register(persist=False)`, cleared by every install and refilled by the router replay - transient: IndexMap, - /// 3, `register(persist=True)`, applied over every install - persisted: Registrations, -} diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index fe419dde054..9b869a078db 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -17,6 +17,9 @@ //! failing the document. //! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known //! families. They are data in the file, not entries. +//! - `CostMap`: an installed document with its provenance, which a host builds its live table from. +//! `Registrations` is the ordered `register_model` overlay over the host's own values, and +//! `write_pass` writes registrations into a host's table through `RegistrationTable`. //! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy //! that lost most of its models is refused), and `Catalog::row_errors` validates one row. //! @@ -37,9 +40,9 @@ macro_rules_attribute::attribute_alias! { mod bundled; mod catalog; +mod cost_map; mod error; mod fallback; -mod layers; mod lower; mod model; mod model_info; @@ -51,9 +54,9 @@ pub mod wire; pub use bundled::*; pub use catalog::*; +pub use cost_map::*; pub use error::*; pub use fallback::*; -pub use layers::*; pub use model::*; pub use model_info::*; pub use register::*; diff --git a/litellm-rust/crates/model-catalog/src/register.rs b/litellm-rust/crates/model-catalog/src/register.rs index ed34082ae57..7aaf2ed3978 100644 --- a/litellm-rust/crates/model-catalog/src/register.rs +++ b/litellm-rust/crates/model-catalog/src/register.rs @@ -1,26 +1,252 @@ -use std::collections::{BTreeSet, HashMap}; +use std::collections::BTreeMap; use indexmap::IndexMap; -use crate::wire::Fields; - -/// Where one `register` call writes and which row it inherits from -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RegistrationTarget { - /// Always the caller's key, never a case variant of it - pub write: String, - /// The first candidate that is an exact key whose row passes `provider_matches` - pub inherit_from: Option, +/// `register_model(persist_across_reloads=True)` calls, replayed above every installed document +#[derive(Clone, Debug, PartialEq)] +pub struct Registrations { + /// Each registered key's merged value, ordered by its latest registration + entries: IndexMap, + /// The registered keys whose write landed on each row key, in landing order + landed_on: IndexMap>, } -/// Python's inheritance candidates in main's order, `get_model_info`'s then `_resolve_builtin_model_cost_entry`'s -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RegistrationCandidates(pub Vec); - -/// Layer 3, and the fields `register` last wrote, which litellm owns again -#[derive(Clone, Debug, Default)] -#[expect(dead_code, reason = "07 adds the functions that read registrations")] -pub struct Registrations { - persisted: IndexMap, - registered_fields: HashMap>, +impl Default for Registrations { + fn default() -> Self { + Self { + entries: IndexMap::new(), + landed_on: IndexMap::new(), + } + } +} + +/// One registration written to a host's table: the key it was registered under and the row key it landed on +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Landed { + pub registered: String, + pub written: String, +} + +impl Registrations { + /// Keeps `value` as the key's registration, already merged with `get(key)` by the host, and makes it the latest + pub fn register(&mut self, key: &str, value: V) { + self.entries.shift_remove(key); + self.entries.insert(key.to_owned(), value); + } + + pub fn get(&self, key: &str) -> Option<&V> { + self.entries.get(key) + } + + pub fn record(&mut self, landed: &[Landed]) { + for Landed { + registered, + written, + } in landed + { + let names = self.landed_on.entry(written.clone()).or_default(); + if !names.contains(registered) { + names.push(registered.clone()); + } + } + } + + /// Every registration, oldest first + pub fn entries(&self) -> impl Iterator { + self.entries + .iter() + .map(|(key, value)| (key.as_str(), value)) + } + + /// The registrations that apply to the row under `key`, oldest first: `key`'s own and every one that landed on it + pub fn over<'a>(&'a self, key: &str) -> Vec<&'a V> { + let landed = self.landed_on.get(key).into_iter().flatten(); + let by_order: BTreeMap = landed + .map(String::as_str) + .chain([key]) + .filter_map(|name| self.entries.get_full(name)) + .map(|(order, _, value)| (order, value)) + .collect(); + by_order.into_values().collect() + } + + pub fn map(&self, mut convert: impl FnMut(&V) -> W) -> Registrations { + Registrations { + entries: self + .entries + .iter() + .map(|(key, value)| (key.clone(), convert(value))) + .collect(), + landed_on: self.landed_on.clone(), + } + } +} + +/// A host's live table that registrations are written into +pub trait RegistrationTable { + type Registration; + type Error; + + /// Writes `registration` where it lands, `None` when this table is no longer the one the host publishes + fn place( + &mut self, + key: &str, + registration: &Self::Registration, + ) -> Result, Self::Error>; +} + +/// Where a registration was written and whether a built-in row matched it +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Placed { + pub key: String, + pub matched_builtin: bool, +} + +#[derive(Debug, PartialEq)] +pub enum Pass<'a, R> { + /// Every entry was written. `unmatched` are those no built-in row matched + Finished { + landed: Vec, + unmatched: Vec<(&'a str, &'a R)>, + }, + /// The table stopped being the published one before every entry was written + Swapped { landed: Vec }, +} + +impl Pass<'_, R> { + pub fn landed(&self) -> &[Landed] { + match self { + Pass::Finished { landed, .. } | Pass::Swapped { landed } => landed, + } + } +} + +/// Writes `entries` into `table` in order +pub fn write_pass<'a, T: RegistrationTable>( + table: &mut T, + entries: &'a [(String, T::Registration)], +) -> Result, T::Error> { + let mut landed = Vec::with_capacity(entries.len()); + let mut unmatched = Vec::new(); + for (key, registration) in entries { + let Some(written) = table.place(key, registration)? else { + return Ok(Pass::Swapped { landed }); + }; + if !written.matched_builtin { + unmatched.push((key.as_str(), registration)); + } + landed.push(Landed { + registered: key.clone(), + written: written.key, + }); + } + Ok(Pass::Finished { landed, unmatched }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn landed(registered: &str, written: &str) -> Landed { + Landed { + registered: registered.to_owned(), + written: written.to_owned(), + } + } + + #[rstest] + fn reregistering_a_key_makes_it_the_latest() { + let mut registrations = Registrations::default(); + registrations.register("a", 1); + registrations.register("b", 2); + registrations.register("a", 3); + assert_eq!( + registrations.entries().collect::>(), + [("b", &2), ("a", &3)] + ); + } + + #[rstest] + fn over_applies_names_landed_on_the_row_in_registration_order() { + let mut registrations = Registrations::default(); + registrations.register("openai/o3-pro", 1); + registrations.register("o3-pro", 2); + registrations.register("other", 3); + registrations.record(&[ + landed("openai/o3-pro", "o3-pro"), + landed("o3-pro", "o3-pro"), + ]); + registrations.register("openai/o3-pro", 4); + assert_eq!(registrations.over("o3-pro"), [&2, &4]); + } + + #[rstest] + fn over_applies_a_keys_own_registration_that_landed_on_another_row() { + let mut registrations = Registrations::default(); + registrations.register("openai/o3-pro", 1); + registrations.record(&[landed("openai/o3-pro", "o3-pro")]); + assert_eq!( + ( + registrations.over("openai/o3-pro"), + registrations.over("o3-pro") + ), + (vec![&1], vec![&1]) + ); + } + + /// A table published until `swaps_after` writes, whose rows match a built-in only when the key starts with `b` + struct Table { + writes: Vec<(String, i32)>, + swaps_after: usize, + } + + impl RegistrationTable for Table { + type Registration = i32; + type Error = (); + + fn place(&mut self, key: &str, registration: &i32) -> Result, ()> { + if self.writes.len() >= self.swaps_after { + return Ok(None); + } + let written = key.to_uppercase(); + self.writes.push((written.clone(), *registration)); + Ok(Some(Placed { + matched_builtin: written.starts_with('B'), + key: written, + })) + } + } + + #[rstest] + #[case::published_throughout( + usize::MAX, + Pass::Finished { landed: vec![landed("a", "A"), landed("b", "B")], unmatched: vec![("a", &1)] }, + &[("A", 1), ("B", 2)], + )] + #[case::swapped_after_the_first_write( + 1, + Pass::Swapped { landed: vec![landed("a", "A")] }, + &[("A", 1)], + )] + fn a_pass_writes_entries_in_order_until_the_table_is_swapped( + #[case] swaps_after: usize, + #[case] expected: Pass<'static, i32>, + #[case] writes: &[(&str, i32)], + ) { + let entries = [("a".to_owned(), 1), ("b".to_owned(), 2)]; + let mut table = Table { + writes: Vec::new(), + swaps_after, + }; + assert_eq!(write_pass(&mut table, &entries), Ok(expected)); + assert_eq!( + table.writes, + writes + .iter() + .map(|(key, value)| ((*key).to_owned(), *value)) + .collect::>() + ); + } } diff --git a/litellm-rust/crates/model-catalog/src/store/AGENTS.md b/litellm-rust/crates/model-catalog/src/store/AGENTS.md index f1dacd2e217..580798da895 100644 --- a/litellm-rust/crates/model-catalog/src/store/AGENTS.md +++ b/litellm-rust/crates/model-catalog/src/store/AGENTS.md @@ -24,7 +24,7 @@ | Source | Stamp | Blind spot | |---|---|---| | `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none | -| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key | +| python-bridge's bound `litellm.model_cost` | the dict object, its `len` and its last key | deletes plus adds between two reads that put the old last key back last, or a `litellm_provider` rewrite on an existing key | A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned diff --git a/litellm-rust/crates/model-catalog/tests/cost_map.rs b/litellm-rust/crates/model-catalog/tests/cost_map.rs new file mode 100644 index 00000000000..04a04170df6 --- /dev/null +++ b/litellm-rust/crates/model-catalog/tests/cost_map.rs @@ -0,0 +1,53 @@ +use litellm_model_catalog::{ + BUNDLED_REGISTRY, CostMap, ErrorKind, IntegrityLimits, Provenance, git_blob_id, +}; +use rstest::rstest; +use serde_json::json; + +fn any_size() -> IntegrityLimits { + IntegrityLimits::new(0, 0.0).unwrap() +} + +#[rstest] +#[case::not_json(b"{".as_slice(), false, ErrorKind::Json)] +#[case::empty(b"{}".as_slice(), true, ErrorKind::Empty)] +fn a_refused_document_says_whether_it_failed_the_integrity_limits( + #[case] body: &[u8], + #[case] integrity: bool, + #[case] kind: ErrorKind, +) { + let error = CostMap::read(body, None, any_size()).unwrap_err(); + assert_eq!((error.is_integrity(), error.kind()), (integrity, kind)); +} + +#[rstest] +fn provenance_is_the_read_bytes_and_etag() { + let body = br#"{"m":{"litellm_provider":"p"}}"#; + let cost_map = CostMap::read(body, Some("\"v1\"".into()), any_size()).unwrap(); + assert_eq!( + cost_map.provenance(), + &Provenance { + revision: "eaa747f210a1527d5ea7138be0f768c47578ca68".to_owned(), + etag: Some("\"v1\"".into()), + } + ); + assert_eq!( + CostMap::bundled().provenance().revision, + git_blob_id(BUNDLED_REGISTRY) + ); +} + +#[rstest] +fn every_name_of_an_aliased_row_lists_its_owner_first( + #[values("owner", "short", "owner-latest")] key: &str, +) { + let body = json!({ + "owner": {"litellm_provider": "p", "aliases": ["short"]}, + "owner-latest": {"alias_of": "owner"}, + "plain": {"litellm_provider": "p"}, + }) + .to_string(); + let cost_map = CostMap::read(body.as_bytes(), None, any_size()).unwrap(); + assert_eq!(cost_map.names(key), ["owner", "short", "owner-latest"]); + assert!(cost_map.names("plain").is_empty()); +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 79773ca5559..48206d242e4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -60,6 +60,7 @@ litellm-router-types.workspace = true litellm-router-core.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true +litellm-model-catalog-refresh.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } indexmap = "2.14.0" pyo3 = { workspace = true, features = ["indexmap"] } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index b4e20a8f750..c6a8b98747c 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -108,13 +108,35 @@ pub(crate) fn call_config( Some(snapshot) => configured(&snapshot)?, None => HttpSettingsLayer::default(), }, - ]) - .without_missing_files(&|path: &Path| path.exists()); - let resolution = Resolution::from(&settings); + ]); + Ok(resolved(py, settings)) +} + +/// The client config for a value of `litellm.rust_bridge.host.settings.http_settings()` the caller handed in, under +/// the process environment. `None` reads the environment alone +pub(crate) fn settings_config( + py: Python<'_>, + http_settings: Option<&Bound<'_, PyAny>>, +) -> PyResult { + let configured = match http_settings { + Some(value) => configured(&PythonSettings::Http.snapshot(value.clone()))?, + None => HttpSettingsLayer::default(), + }; + let settings = HttpSettings::from_layers([ + HttpSettingsLayer::from_environment(&ProcessEnvironment), + configured, + ]); + Ok(resolved(py, settings)) +} + +/// The client config `settings` resolve to, warning once per process about each setting the client cannot apply +fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig { + let resolution = + Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists())); for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) { crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}")); } - Ok(resolution.config) + resolution.config } pub(crate) fn provider_client( diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 7513ff6b87b..7edf13035c1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -9,6 +9,7 @@ mod execution; mod http; mod logger; mod marshal; +mod model_cost; mod model_registry; mod preflight; mod python_settings; @@ -36,6 +37,12 @@ mod _native { #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] + use crate::model_cost::{ + NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched, NativeInstalled, + NativeModelCostMap, NativeModelCostMapState, NativeProvenance, NativeRegistered, + NativeRejected, NativeTableSwapped, + }; + #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states}; @@ -166,6 +173,17 @@ mod tests { "NativeDatabase", "trace_encode_error", "bundled_model_cost_map", + "NativeModelCostMap", + "NativeModelCostMapState", + "NativeInstalled", + "NativeRejected", + "NativeRegistered", + "NativeTableSwapped", + "NativeProvenance", + "NativeCostMapSource", + "NativeFetchPolicy", + "NativeFetched", + "NativeFetchFailed", "TokenCounter", "Tokenizer", "RouterCoreHandle", diff --git a/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md b/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md new file mode 100644 index 00000000000..4ba23aac1ce --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md @@ -0,0 +1,16 @@ +# Model cost bridge + +- `NativeModelCostMap` holds inputs only, the installed `CostMap`, kept `Registrations` and the name index. Python's `litellm.model_cost` is the only copy of the live rows +- Every call that reads or writes rows takes the dict as an argument. Never store rows or build a second table here +- `LiveDict` holds a reference to that same dict and reads it under the GIL on every query, so a Python write is seen on the next read + - It needs the GIL, so Rust code running without the GIL cannot use it + - Its stamp is the dict's address, length and last key. [model-catalog/src/store/AGENTS.md](../../../model-catalog/src/store/AGENTS.md) says what that misses +- `host.rs` is the only code that calls Python, by name through `MODULE`. Calls happen with the GIL held and no `NativeModelCostMap` lock held + - `resolve` reaches `find_key` on this same object, so holding a lock across a host call deadlocks + - Test `a_host_lookup_that_reaches_find_key_on_the_same_map_finishes` guards this +- Rows and registrations cross as Python objects, never JSON, because router `model_info` holds `datetime` +- `table::build` serves the document as main's `_expand_model_aliases` did, so all names of a model share one row object made by `host.entry` +- `install` touches no state until `CostMap::read` succeeds. `NativeRejected` means nothing changed +- `register` keeps a registration only after its pass finished, so an exception in a host call keeps nothing +- `rows=None` writes each entry into whichever table is published at that moment and never reports a swap +- `fetch_blocking` runs on a throwaway runtime with an unpooled client so an importing process can still fork. `fetch` is for workers after the fork diff --git a/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs new file mode 100644 index 00000000000..77ec88cc1b9 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/fetch.rs @@ -0,0 +1,239 @@ +use std::time::Duration; + +use litellm_http::{ClientVariant, HttpClientConfig}; +use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +/// A cost map URL and how to connect to it. An `s3://` URL reaches S3 as an AWS SDK finds it in the environment, +/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds +#[pyclass( + frozen, + name = "NativeCostMapSource", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeCostMapSource { + url: String, + http: HttpClientConfig, +} + +#[pymethods] +impl NativeCostMapSource { + /// `http_settings` is `litellm.rust_bridge.host.settings.http_settings()`, or `None` for the environment alone + #[new] + #[pyo3(signature = (url, *, http_settings))] + fn new(py: Python<'_>, url: String, http_settings: Option>) -> PyResult { + Ok(Self { + url, + http: crate::http::settings_config(py, http_settings.as_ref())?, + }) + } + + /// Fetches on the shared runtime with a pooled client + fn fetch<'py>( + &self, + py: Python<'py>, + policy: PyRef<'_, NativeFetchPolicy>, + ) -> PyResult> { + let client = crate::http::pool() + .client(&self.http, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + let url = self.url.clone(); + let policy = policy.0; + crate::execution::run_async_value(py, async move { + Ok(Outcome( + litellm_model_catalog_refresh::fetch(&client, &s3(), &url, policy, Start::Fresh) + .await, + )) + }) + } + + /// Fetches on a throwaway runtime with a client of its own. `resume_from` continues a fetch whose first attempt + /// failed retryably + #[pyo3(signature = (policy, resume_from = None))] + fn fetch_blocking( + &self, + py: Python<'_>, + policy: PyRef<'_, NativeFetchPolicy>, + resume_from: Option>, + ) -> PyResult> { + let start = resume_from.map_or(Start::Fresh, |failed| Start::After(failed.0.clone())); + let (config, url, policy) = (self.http.clone(), self.url.clone(), policy.0); + let logger = crate::logger::capture(py); + let fetched = litellm_host_python::release_gil(py, move || { + logger.scope(|| { + let client = crate::http::pool() + .unpooled(&config, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + litellm_host_python::run_isolated(litellm_model_catalog_refresh::fetch( + &client, + &s3(), + &url, + policy, + start, + )) + .map_err(PyErr::from) + }) + })?; + Outcome(fetched).into_pyobject(py).map(Bound::unbind) + } +} + +fn s3() -> S3Access { + S3Access::Environment { + region: std::env::var("AWS_REGION_NAME").ok(), + } +} + +#[pyclass( + frozen, + name = "NativeFetchPolicy", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeFetchPolicy(RetryPolicy); + +#[pymethods] +impl NativeFetchPolicy { + #[new] + /// Retries wait an exponential backoff with jitter from `min_delay_seconds` up to `max_wait_seconds`, or the + /// server's `Retry-After` up to that cap + #[pyo3(signature = (*, max_attempts, max_wait_seconds, timeout_seconds, min_delay_seconds))] + fn new( + max_attempts: u32, + max_wait_seconds: f64, + timeout_seconds: f64, + min_delay_seconds: f64, + ) -> PyResult { + if max_attempts == 0 { + return Err(PyValueError::new_err("max_attempts must be at least 1")); + } + let seconds = |name: &str, value: f64| { + Duration::try_from_secs_f64(value).map_err(|_| { + PyValueError::new_err(format!("{name} must be a non-negative number of seconds")) + }) + }; + let min_delay = seconds("min_delay_seconds", min_delay_seconds)?; + let max_delay = seconds("max_wait_seconds", max_wait_seconds)?; + if min_delay > max_delay { + return Err(PyValueError::new_err( + "min_delay_seconds must not exceed max_wait_seconds", + )); + } + Ok(Self(RetryPolicy { + max_attempts, + timeout: seconds("timeout_seconds", timeout_seconds)?, + min_delay, + max_delay, + })) + } +} + +#[pyclass(frozen, name = "NativeFetched", module = "litellm.rust_bridge._native")] +pub(crate) struct NativeFetched { + #[pyo3(get)] + body: Py, + #[pyo3(get)] + etag: Option, +} + +#[pyclass( + frozen, + skip_from_py_object, + name = "NativeFetchFailed", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeFetchFailed(Failed); + +#[pymethods] +impl NativeFetchFailed { + #[getter] + fn reason(&self) -> &str { + &self.0.reason + } + + /// Another attempt could succeed. Only the attempt limit ended the fetch + #[getter] + fn retryable(&self) -> bool { + self.0.retryable + } + + #[getter] + fn retry_after(&self) -> Option { + self.0.retry_after.map(|wait| wait.as_secs_f64()) + } +} + +struct Outcome(Result); + +impl<'py> IntoPyObject<'py> for Outcome { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + match self.0 { + Ok(fetched) => Bound::new( + py, + NativeFetched { + body: PyBytes::new(py, &fetched.body).unbind(), + etag: fetched.etag, + }, + ) + .map(Bound::into_any), + Err(failed) => Bound::new(py, NativeFetchFailed(failed)).map(Bound::into_any), + } + } +} + +#[cfg(test)] +mod tests { + use std::io::{BufRead, BufReader, Write}; + use std::net::TcpListener; + + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::equal_bounds(2.0, 2.0, true)] + #[case::min_below_max(2.0, 30.0, true)] + #[case::min_above_max(31.0, 30.0, false)] + fn a_policy_accepts_a_min_delay_only_up_to_the_max_wait( + #[case] min_delay_seconds: f64, + #[case] max_wait_seconds: f64, + #[case] accepted: bool, + ) { + Python::initialize(); + let policy = NativeFetchPolicy::new(3, max_wait_seconds, 5.0, min_delay_seconds); + assert_eq!(policy.is_ok(), accepted); + } + + #[rstest] + fn a_blocking_fetch_hands_python_the_document_without_starting_the_shared_runtime() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/cost.json", listener.local_addr().unwrap()); + std::thread::spawn(move || { + let mut stream = listener.incoming().next().unwrap().unwrap(); + let mut reader = BufReader::new(stream.try_clone().unwrap()); + let mut line = String::new(); + while reader.read_line(&mut line).unwrap() > 2 { + line.clear(); + } + stream + .write_all(b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}") + .unwrap(); + }); + Python::initialize(); + Python::attach(|py| { + let source = NativeCostMapSource::new(py, url, None).unwrap(); + let policy = + Bound::new(py, NativeFetchPolicy::new(3, 30.0, 5.0, 2.0).unwrap()).unwrap(); + let outcome = source.fetch_blocking(py, policy.borrow(), None).unwrap(); + let fetched = outcome.bind(py).cast::().unwrap().get(); + assert_eq!(fetched.body.bind(py).as_bytes(), b"{}"); + assert_eq!(fetched.etag.as_deref(), Some("\"v1\"")); + }); + assert!(!litellm_host_python::runtime_started()); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/host.rs b/litellm-rust/crates/python-bridge/src/model_cost/host.rs new file mode 100644 index 00000000000..25fbbefbb73 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/host.rs @@ -0,0 +1,81 @@ +//! Calls into `litellm.rust_bridge.host.model_cost` + +use pyo3::call::PyCallArgs; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +pub(crate) const MODULE: &str = "litellm.rust_bridge.host.model_cost"; + +/// Where a registration lands in the published table and the row it starts from +#[derive(FromPyObject)] +pub(crate) struct Landing<'py> { + #[pyo3(attribute)] + pub(crate) key: String, + #[pyo3(attribute)] + pub(crate) base: Bound<'py, PyAny>, + #[pyo3(attribute)] + pub(crate) matched_builtin: bool, +} + +fn call<'py>( + py: Python<'py>, + name: &str, + args: impl PyCallArgs<'py>, +) -> PyResult> { + py.import(MODULE)?.getattr(name)?.call1(args) +} + +/// The row object every name in `names` shares, the owner first +pub(crate) fn entry<'py>( + fields: Bound<'py, PyAny>, + names: &[String], +) -> PyResult> { + let py = fields.py(); + call(py, "entry", (fields, PyTuple::new(py, names)?)) +} + +/// Where `registration` lands in `rows`, `None` when `rows` is no longer the published table. `rows` of `None` is the +/// published table +pub(crate) fn resolve<'py>( + rows: Option<&Bound<'py, PyDict>>, + key: &str, + registration: &Bound<'py, PyAny>, + custom_llm_provider: Option<&str>, +) -> PyResult>> { + call( + registration.py(), + "resolve", + (rows, key, registration, custom_llm_provider), + )? + .extract() +} + +/// The row `registration` writes over `current` +pub(crate) fn merged<'py>( + current: Option<&Bound<'py, PyAny>>, + landing: &Landing<'py>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + call( + registration.py(), + "merged", + (current, &landing.base, registration), + ) +} + +/// The registration kept for a key: `registration` merged into the one kept before it +pub(crate) fn kept<'py>( + previous: Option<&Bound<'py, PyDict>>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + Ok(call(registration.py(), "kept", (previous, registration))?.cast_into::()?) +} + +pub(crate) fn published(py: Python<'_>) -> PyResult> { + Ok(call(py, "published", ())?.cast_into::()?) +} + +/// Writes `written` into `rows`, the published table, through the host, which tracks the table's version +pub(crate) fn write(rows: &Bound<'_, PyDict>, written: &Bound<'_, PyDict>) -> PyResult<()> { + call(rows.py(), "write", (rows, written)).map(drop) +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/live.rs b/litellm-rust/crates/python-bridge/src/model_cost/live.rs new file mode 100644 index 00000000000..0b7f071127d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/live.rs @@ -0,0 +1,176 @@ +use std::borrow::Cow; + +use litellm_model_catalog::Rows; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +/// The dict Python last handed in as `litellm.model_cost`, read under the GIL on every query +pub(crate) struct LiveDict(Py); + +/// One row exactly as the dict holds it +pub(crate) struct LiveRow(Py); + +impl Clone for LiveRow { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl LiveDict { + pub(crate) fn new(dict: &Bound<'_, PyDict>) -> Self { + Self(dict.clone().unbind()) + } + + pub(crate) fn is(&self, dict: &Bound<'_, PyDict>) -> bool { + self.0.bind(dict.py()).is(dict) + } +} + +impl Rows for LiveDict { + type Row = LiveRow; + /// The dict's address, length and last key + type Stamp = (usize, usize, Option); + + fn fold(name: &str) -> String { + name.to_lowercase() + } + + fn stamp(&self) -> Self::Stamp { + Python::attach(|py| { + let rows = self.0.bind(py); + let last = rows + .call_method0(intern!(py, "__reversed__")) + .and_then(|keys| keys.try_iter()) + .ok() + .and_then(|mut keys| keys.next()) + .and_then(Result::ok) + .and_then(|key| key.extract::().ok()); + (self.0.as_ptr() as usize, rows.len(), last) + }) + } + + fn contains(&self, key: &str) -> bool { + Python::attach(|py| self.0.bind(py).contains(key).unwrap_or(false)) + } + + fn row(&self, key: &str) -> Option> { + Python::attach(|py| { + let row = self.0.bind(py).get_item(key).ok().flatten()?; + Some(Cow::Owned(LiveRow(row.unbind()))) + }) + } + + /// Groups each key by its row's `litellm_provider` when that is a string + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + Python::attach(|py| { + for (key, row) in self.0.bind(py).iter() { + let Ok(key) = key.extract::() else { + continue; + }; + let provider = row + .cast::() + .ok() + .and_then(|row| row.get_item("litellm_provider").ok().flatten()) + .and_then(|provider| provider.extract::().ok()); + visit(&key, provider.as_deref()); + } + }); + } +} + +#[cfg(test)] +mod tests { + use litellm_model_catalog::Store; + use pyo3::ffi::c_str; + use rstest::rstest; + + use super::*; + + fn dict<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + py.eval(source, None, None) + .unwrap() + .cast_into::() + .unwrap() + } + + #[rstest] + fn a_key_added_to_the_dict_is_found_on_the_next_query() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}")); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.find_key("GPT-4O"), Some("gpt-4o".to_owned())); + rows.set_item("Claude-X", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("claude-x"), Some("Claude-X".to_owned())); + }); + } + + #[rstest] + fn a_removed_key_is_never_returned() { + Python::initialize(); + Python::attach(|py| { + let rows = dict( + py, + c_str!("{'gpt-4o': {'litellm_provider': 'openai'}, 'o3': {}}"), + ); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.keys_in("openai"), ["gpt-4o"]); + rows.del_item("gpt-4o").unwrap(); + rows.set_item("o4", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("GPT-4O"), None); + assert!(store.keys_in("openai").is_empty()); + }); + } + + #[rstest] + fn a_key_swapped_in_for_a_removed_one_is_found() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {}, 'Lit-Gone': {}}")); + let store = Store::new(LiveDict::new(&rows)); + assert_eq!(store.find_key("lit-gone"), Some("Lit-Gone".to_owned())); + rows.del_item("Lit-Gone").unwrap(); + rows.set_item("Lit-New", PyDict::new(py)).unwrap(); + assert_eq!(store.find_key("lit-new"), Some("Lit-New".to_owned())); + assert_eq!(store.find_key("lit-gone"), None); + }); + } + + #[rstest] + fn two_keys_folding_alike_resolve_to_the_later_one() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'Model': {}, 'MODEL': {}}")); + assert_eq!( + Store::new(LiveDict::new(&rows)).find_key("model"), + Some("MODEL".to_owned()) + ); + }); + } + + #[rstest] + fn rows_are_the_objects_the_dict_holds() { + Python::initialize(); + Python::attach(|py| { + let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}")); + let store = Store::new(LiveDict::new(&rows)); + let row = store.get("gpt-4o").unwrap().into_owned(); + assert!(row.0.bind(py).is(rows.get_item("gpt-4o").unwrap().unwrap())); + }); + } + + #[rstest] + fn providers_group_only_string_providers() { + Python::initialize(); + Python::attach(|py| { + let rows = dict( + py, + c_str!( + "{'a': {'litellm_provider': 'p'}, 'b': {'litellm_provider': None}, 'c': 1, 1: {'litellm_provider': 'p'}}" + ), + ); + assert_eq!(Store::new(LiveDict::new(&rows)).keys_in("p"), ["a"]); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/mod.rs b/litellm-rust/crates/python-bridge/src/model_cost/mod.rs new file mode 100644 index 00000000000..99451033122 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/mod.rs @@ -0,0 +1,497 @@ +mod fetch; +mod host; +mod live; +mod table; + +use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard}; + +use litellm_host_python::to_py; +use litellm_model_catalog::{ + CostMap, IntegrityLimits, Landed, Pass, Placed, Provenance, RegistrationTable, Registrations, + Store, write_pass, +}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +pub(crate) use fetch::{NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched}; +use live::LiveDict; + +/// The installed document and the `register_model` calls kept above it, each as the dict Python merged +struct Inputs { + cost_map: CostMap, + registrations: Registrations>, +} + +impl Inputs { + fn clone_ref(&self, py: Python<'_>) -> Self { + Self { + cost_map: self.cost_map.clone(), + registrations: self.registrations.map(|kept| kept.clone_ref(py)), + } + } +} + +/// The cost map's inputs and the name index over `litellm.model_cost` +#[pyclass( + frozen, + name = "NativeModelCostMap", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeModelCostMap { + inputs: RwLock, + limits: IntegrityLimits, + index: RwLock>>>, +} + +#[pyclass( + frozen, + name = "NativeInstalled", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeInstalled { + /// The new live table, the document without registrations + #[pyo3(get)] + rows: Py, + /// How many keys the document itself serves, aliases included and registrations left out + #[pyo3(get)] + model_count: usize, + /// The document's `fallback_generalizations` block as written, `None` when it has none + #[pyo3(get)] + fallback_generalizations: Py, +} + +#[pyclass( + frozen, + name = "NativeRejected", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeRejected { + #[pyo3(get)] + reason: String, + /// The document parsed but has fewer models than the integrity limits allow + #[pyo3(get)] + integrity: bool, +} + +/// A finished registration pass, with the registrations no built-in row matched, by registered key +#[pyclass( + frozen, + name = "NativeRegistered", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeRegistered { + #[pyo3(get)] + unmatched: Py, +} + +/// A registration pass that stopped because the table it was handed is no longer the published one +#[pyclass( + frozen, + name = "NativeTableSwapped", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeTableSwapped; + +/// The inputs at one moment, for a caller that has to put them back, such as test isolation +#[pyclass( + frozen, + name = "NativeModelCostMapState", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeModelCostMapState(Inputs); + +#[pyclass( + frozen, + name = "NativeProvenance", + module = "litellm.rust_bridge._native" +)] +pub(crate) struct NativeProvenance { + #[pyo3(get)] + revision: String, + #[pyo3(get)] + etag: Option, +} + +#[pymethods] +impl NativeModelCostMap { + /// Starts from the bundled document. A fetched document with fewer than `min_model_count` models, or fewer than + /// `max_shrink_ratio` times the bundled count, is rejected + #[staticmethod] + fn bundled(min_model_count: i64, max_shrink_ratio: f64) -> PyResult { + let min_model_count = usize::try_from(min_model_count).map_err(|_| { + PyValueError::new_err(format!( + "min_model_count must be 0 or more, got {min_model_count}" + )) + })?; + let limits = IntegrityLimits::new(min_model_count, max_shrink_ratio).ok_or_else(|| { + PyValueError::new_err(format!( + "max_shrink_ratio must be between 0 and 1, got {max_shrink_ratio}" + )) + })?; + Ok(Self { + inputs: RwLock::new(Inputs { + cost_map: CostMap::bundled(), + registrations: Registrations::default(), + }), + limits, + index: RwLock::new(None), + }) + } + + /// The live table built from a fetched document, or why it was rejected, in which case nothing changed + #[pyo3(signature = (body, etag = None))] + fn install(&self, py: Python<'_>, body: &[u8], etag: Option) -> PyResult> { + match crate::logger::capture(py).scope(|| CostMap::read(body, etag, self.limits)) { + Ok(cost_map) => self.install_document(py, cost_map), + Err(error) => rejected(py, &error), + } + } + + /// A caller's document as a table, read without size limits and installed nowhere + #[staticmethod] + fn finalize(py: Python<'_>, body: &[u8]) -> PyResult> { + match crate::logger::capture(py) + .scope(|| CostMap::read(body, None, IntegrityLimits::UNCHECKED)) + { + Ok(cost_map) => Ok(Py::new(py, built(py, &cost_map)?)?.into_any()), + Err(error) => rejected(py, &error), + } + } + + /// The live table built from the bundled document, which is never size checked + fn install_bundled(&self, py: Python<'_>) -> PyResult> { + self.install_document(py, CostMap::bundled()) + } + + /// Writes each entry into `rows`, or into the published table when `rows` is `None` + #[pyo3(signature = (rows, entries, persist, custom_llm_provider = None))] + fn register( + &self, + py: Python<'_>, + rows: Option<&Bound<'_, PyDict>>, + entries: &Bound<'_, PyDict>, + persist: bool, + custom_llm_provider: Option<&str>, + ) -> PyResult> { + let entries: Vec<(String, Bound<'_, PyAny>)> = entries + .iter() + .map(|(key, registration)| Ok((key.str()?.to_string(), registration))) + .collect::>()?; + let kept = match persist { + true => self.merged_with_kept(&entries)?, + false => Vec::new(), + }; + let pass = self.write_pass(rows, &entries, custom_llm_provider)?; + if persist { + self.keep(kept, pass.landed()); + } + into_py(py, pass) + } + + /// Writes every kept registration into `rows` again, oldest first + fn reapply(&self, rows: &Bound<'_, PyDict>) -> PyResult> { + let py = rows.py(); + let entries: Vec<(String, Bound<'_, PyAny>)> = self + .read_inputs() + .registrations + .entries() + .map(|(key, kept)| (key.to_owned(), kept.bind(py).clone().into_any())) + .collect(); + let pass = self.write_pass(Some(rows), &entries, None)?; + self.write_inputs().registrations.record(pass.landed()); + into_py(py, pass) + } + + /// The installed row under `key` with the kept registrations that apply to it, without anything the router wrote + fn registered_row<'py>( + &self, + py: Python<'py>, + key: &str, + ) -> PyResult>> { + let (cost_map, applied) = { + let inputs = self.read_inputs(); + let applied: Vec> = inputs + .registrations + .over(key) + .into_iter() + .map(|kept| kept.bind(py).clone()) + .collect(); + (inputs.cost_map.clone(), applied) + }; + table::registered_row(py, &cost_map, key, &applied) + } + + fn registered_keys(&self) -> Vec { + let inputs = self.read_inputs(); + table::registered_keys( + &inputs.cost_map, + inputs.registrations.entries().map(|(key, _)| key), + ) + } + + /// The installed document as a table, without any registration + fn loaded<'py>(&self, py: Python<'py>) -> PyResult> { + let cost_map = self.read_inputs().cost_map.clone(); + table::build(py, &cost_map) + } + + /// The key itself when `rows` has it, otherwise the key that differs from `name` only in case + fn find_key(&self, rows: &Bound<'_, PyDict>, name: &str) -> Option { + self.index_over(rows).find_key(name) + } + + /// Keys whose row's `litellm_provider` is exactly `provider`, in row order + fn keys_in(&self, rows: &Bound<'_, PyDict>, provider: &str) -> Vec { + self.index_over(rows).keys_in(provider) + } + + /// Drops the name index, for a caller that changed keys in a way the dict's length does not show + fn invalidate(&self) { + if let Some(index) = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + { + index.invalidate(); + } + } + + fn snapshot(&self, py: Python<'_>) -> NativeModelCostMapState { + NativeModelCostMapState(self.read_inputs().clone_ref(py)) + } + + /// Puts back the document, registrations and provenance a `snapshot` saw + fn restore(&self, py: Python<'_>, state: PyRef<'_, NativeModelCostMapState>) { + *self.write_inputs() = state.0.clone_ref(py); + self.invalidate(); + } + + #[getter] + fn provenance(&self) -> NativeProvenance { + let Provenance { revision, etag } = self.read_inputs().cost_map.provenance().clone(); + NativeProvenance { revision, etag } + } +} + +impl NativeModelCostMap { + fn read_inputs(&self) -> RwLockReadGuard<'_, Inputs> { + self.inputs.read().unwrap_or_else(PoisonError::into_inner) + } + + fn write_inputs(&self) -> RwLockWriteGuard<'_, Inputs> { + self.inputs.write().unwrap_or_else(PoisonError::into_inner) + } + + fn install_document(&self, py: Python<'_>, cost_map: CostMap) -> PyResult> { + let installed = built(py, &cost_map)?; + self.write_inputs().cost_map = cost_map; + *self.index.write().unwrap_or_else(PoisonError::into_inner) = + Some(Arc::new(Store::new(LiveDict::new(installed.rows.bind(py))))); + Ok(Py::new(py, installed)?.into_any()) + } + + /// Each entry merged into the registration kept for its key + fn merged_with_kept<'py>( + &self, + entries: &[(String, Bound<'py, PyAny>)], + ) -> PyResult)>> { + entries + .iter() + .map(|(key, registration)| { + let previous = self + .read_inputs() + .registrations + .get(key) + .map(|kept| kept.bind(registration.py()).clone()); + Ok((key.clone(), host::kept(previous.as_ref(), registration)?)) + }) + .collect() + } + + fn keep(&self, kept: Vec<(String, Bound<'_, PyDict>)>, landed: &[Landed]) { + let mut inputs = self.write_inputs(); + for (key, registration) in kept { + inputs.registrations.register(&key, registration.unbind()); + } + inputs.registrations.record(landed); + } + + fn write_pass<'a, 'py>( + &self, + rows: Option<&Bound<'py, PyDict>>, + entries: &'a [(String, Bound<'py, PyAny>)], + custom_llm_provider: Option<&str>, + ) -> PyResult>> { + let mut table = PublishedTable { + map: self, + rows, + custom_llm_provider, + }; + write_pass(&mut table, entries) + } + + /// The index over `rows`, rebuilt when `rows` is not the dict it was built over + fn index_over(&self, rows: &Bound<'_, PyDict>) -> Arc> { + let current = self + .index + .read() + .unwrap_or_else(PoisonError::into_inner) + .clone(); + if let Some(index) = current.filter(|index| index.rows().is(rows)) { + return index; + } + let index = Arc::new(Store::new(LiveDict::new(rows))); + *self.index.write().unwrap_or_else(PoisonError::into_inner) = Some(Arc::clone(&index)); + index + } +} + +/// The dict Python publishes as `litellm.model_cost`, written through the host +struct PublishedTable<'a, 'py> { + map: &'a NativeModelCostMap, + rows: Option<&'a Bound<'py, PyDict>>, + custom_llm_provider: Option<&'a str>, +} + +impl<'py> RegistrationTable for PublishedTable<'_, 'py> { + type Registration = Bound<'py, PyAny>; + type Error = PyErr; + + fn place(&mut self, key: &str, registration: &Self::Registration) -> PyResult> { + let Some(landing) = host::resolve(self.rows, key, registration, self.custom_llm_provider)? + else { + return Ok(None); + }; + let rows = match self.rows { + Some(rows) => rows.clone(), + None => host::published(registration.py())?, + }; + table::write(&rows, &landing, registration)?; + self.map.invalidate(); + Ok(Some(Placed { + key: landing.key, + matched_builtin: landing.matched_builtin, + })) + } +} + +fn built(py: Python<'_>, cost_map: &CostMap) -> PyResult { + let rows = table::build(py, cost_map)?; + let fallback_generalizations = match cost_map.fallback_generalizations() { + Some(block) => to_py(py, block)?, + None => py.None(), + }; + Ok(NativeInstalled { + model_count: rows.len(), + rows: rows.unbind(), + fallback_generalizations, + }) +} + +fn rejected(py: Python<'_>, error: &litellm_model_catalog::Error) -> PyResult> { + let rejected = NativeRejected { + reason: litellm_error::Report::new(error).to_string(), + integrity: error.is_integrity(), + }; + Ok(Py::new(py, rejected)?.into_any()) +} + +fn into_py<'py>(py: Python<'py>, pass: Pass<'_, Bound<'py, PyAny>>) -> PyResult> { + match pass { + Pass::Swapped { .. } => Ok(Py::new(py, NativeTableSwapped)?.into_any()), + Pass::Finished { unmatched, .. } => { + let table = PyDict::new(py); + for (key, registration) in unmatched { + table.set_item(key, registration)?; + } + let registered = NativeRegistered { + unmatched: table.unbind(), + }; + Ok(Py::new(py, registered)?.into_any()) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::mpsc; + use std::time::Duration; + + use pyo3::ffi::c_str; + use rstest::rstest; + + use super::*; + + const HOST: &std::ffi::CStr = c_str!( + r#" +from types import SimpleNamespace + +class Row(dict): + names = () + +found = [] + +def resolve(rows, key, registration, custom_llm_provider): + found.append(cost_map.find_key(rows, key)) + return SimpleNamespace(key=key, base={}, matched_builtin=True) + +def merged(current, base, registration): + return Row({**(current or {}), **base, **registration}) + +def write(rows, written): + rows.update(written) +"# + ); + + #[rstest] + fn a_host_lookup_that_reaches_find_key_on_the_same_map_finishes() { + Python::initialize(); + let (sender, receiver) = mpsc::channel(); + std::thread::spawn(move || { + let outcome = Python::attach(|py| -> PyResult<(Vec>, f64)> { + let cost_map = Py::new(py, NativeModelCostMap::bundled(0, 0.0)?)?; + let module_name = std::ffi::CString::new(host::MODULE).unwrap(); + let host = + PyModule::from_code(py, HOST, c_str!("host_model_cost.py"), &module_name)?; + host.setattr("cost_map", &cost_map)?; + let modules = py.import("sys")?.getattr("modules")?; + let packages = ["litellm", "litellm.rust_bridge", "litellm.rust_bridge.host"]; + for package in packages { + modules.set_item(package, PyModule::new(py, package)?)?; + } + modules.set_item(host::MODULE, &host)?; + let rows = PyDict::new(py); + rows.set_item("Lit-Reentry", PyDict::new(py))?; + let entry = PyDict::new(py); + entry.set_item("input_cost_per_token", 1.5)?; + let entries = PyDict::new(py); + entries.set_item("lit-reentry", entry)?; + let registered = cost_map + .bind(py) + .call_method1("register", (&rows, &entries, false)); + for name in [host::MODULE].into_iter().chain(packages) { + modules.del_item(name)?; + } + registered?; + let price = rows + .get_item("lit-reentry")? + .expect("the registration landed") + .get_item("input_cost_per_token")? + .extract()?; + Ok((host.getattr("found")?.extract()?, price)) + }); + sender + .send(outcome.map_err(|error| error.to_string())) + .unwrap(); + }); + + let (found, price) = receiver + .recv_timeout(Duration::from_secs(30)) + .expect("register finished instead of deadlocking on its own lock") + .unwrap(); + assert_eq!(found, [Some("Lit-Reentry".to_owned())]); + assert!((price - 1.5).abs() < f64::EPSILON); + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_cost/table.rs b/litellm-rust/crates/python-bridge/src/model_cost/table.rs new file mode 100644 index 00000000000..67b5f581b49 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_cost/table.rs @@ -0,0 +1,166 @@ +//! Builds `litellm.model_cost` from an installed document + +use std::borrow::Cow; +use std::collections::HashMap; + +use indexmap::IndexSet; +use litellm_host_python::to_py; +use litellm_model_catalog::wire::Fields; +use litellm_model_catalog::{CostMap, FALLBACK_GENERALIZATIONS, Written}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::Value; + +use super::host::{self, Landing}; + +const ALIASES: &str = "aliases"; + +enum Served<'a> { + Model, + AliasOf(&'a str), + /// A top-level value that is not a model row, such as `sample_spec` + Section(&'a Value), +} + +/// Every served key in dict order. An `alias_of` row repeats among the aliases and keeps its document position +fn served(cost_map: &CostMap) -> impl Iterator)> { + let catalog = cost_map.catalog(); + let documented = catalog + .written() + .filter(|(key, _)| *key != FALLBACK_GENERALIZATIONS) + .filter_map(|(key, written)| { + let served = match written { + Written::Section(value) => Served::Section(value), + Written::Entry(_) => match catalog.aliases().get(key) { + Some(owner) => Served::AliasOf(owner), + None => catalog.get(key).map(|_| Served::Model)?, + }, + }; + Some((key, served)) + }); + let listed = cost_map.alias_groups().flat_map(|(owner, names)| { + names + .iter() + .skip(1) + .map(move |alias| (alias.as_str(), Served::AliasOf(owner))) + }); + documented.chain(listed) +} + +/// The installed document as a new dict +pub(crate) fn build<'py>(py: Python<'py>, cost_map: &CostMap) -> PyResult> { + let table = PyDict::new(py); + let mut rows = HashMap::new(); + for (key, served) in served(cost_map) { + let value = match served { + Served::Model => row(py, cost_map, key, &mut rows)?, + Served::AliasOf(owner) => row(py, cost_map, owner, &mut rows)?, + Served::Section(value) => to_py(py, value)?.into_bound(py), + }; + table.set_item(key, value)?; + } + Ok(table) +} + +fn row<'a, 'py>( + py: Python<'py>, + cost_map: &'a CostMap, + owner: &'a str, + rows: &mut HashMap<&'a str, Bound<'py, PyAny>>, +) -> PyResult> { + if let Some(row) = rows.get(owner) { + return Ok(row.clone()); + } + let fields = cost_map + .catalog() + .get(owner) + .expect("a served model or alias owner is a catalog row"); + let fields = to_py(py, without_alias_list(fields).as_ref())?.into_bound(py); + let row = match cost_map.names(owner) { + [] => fields, + names => host::entry(fields, names)?, + }; + rows.insert(owner, row.clone()); + Ok(row) +} + +fn without_alias_list(fields: &Fields) -> Cow<'_, Fields> { + if fields.get(ALIASES).is_none_or(Value::is_null) { + return Cow::Borrowed(fields); + } + Cow::Owned( + fields + .iter() + .filter(|(name, _)| *name != ALIASES) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ) +} + +/// The installed row under `key` with `applied` written over it in order, `None` when neither has one +pub(crate) fn registered_row<'py>( + py: Python<'py>, + cost_map: &CostMap, + key: &str, + applied: &[Bound<'py, PyDict>], +) -> PyResult>> { + let catalog = cost_map.catalog(); + let installed = catalog + .get(key) + .or_else(|| match catalog.section(key) { + Some(Value::Object(fields)) if key != FALLBACK_GENERALIZATIONS => Some(fields), + _ => None, + }) + .map(without_alias_list); + let row = match (installed, applied) { + (None, []) => return Ok(None), + (Some(fields), _) => to_py(py, fields.as_ref())? + .into_bound(py) + .cast_into::()?, + (None, _) => PyDict::new(py), + }; + for registration in applied { + row.update(registration.as_mapping())?; + } + Ok(Some(row)) +} + +/// The keys `registered_row` answers for, installed keys first +pub(crate) fn registered_keys<'a>( + cost_map: &'a CostMap, + registered: impl Iterator, +) -> Vec { + let served: IndexSet<&str> = served(cost_map).map(|(key, _)| key).collect(); + let only_registered: Vec<&str> = registered.filter(|key| !served.contains(key)).collect(); + served + .into_iter() + .chain(only_registered) + .map(str::to_owned) + .collect() +} + +/// Writes `registration` over the row under `landing.key`, and under every name that shares that row object +pub(crate) fn write<'py>( + rows: &Bound<'py, PyDict>, + landing: &Landing<'py>, + registration: &Bound<'py, PyAny>, +) -> PyResult<()> { + let current = rows.get_item(&landing.key)?; + let row = host::merged(current.as_ref(), landing, registration)?; + let names: Vec = row.getattr("names")?.extract()?; + let sharing: Vec<&String> = names + .iter() + .filter(|name| { + let served = rows.get_item(name.as_str()).ok().flatten(); + served + .zip(current.as_ref()) + .is_some_and(|(served, current)| served.is(current)) + }) + .collect(); + let written = PyDict::new(rows.py()); + written.set_item(&landing.key, &row)?; + sharing + .into_iter() + .try_for_each(|name| written.set_item(name, &row))?; + host::write(rows, &written) +} diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 0152fe8a5b9..93377abb7e9 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -56,7 +56,7 @@ impl PythonSettings { } } - #[cfg(test)] + /// A snapshot of a value Python handed in, read with the same field rules as the accessor's pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> { Snapshot { group: self, value } } diff --git a/litellm/litellm_core_utils/AGENTS.md b/litellm/litellm_core_utils/AGENTS.md new file mode 100644 index 00000000000..e817eca0ac4 --- /dev/null +++ b/litellm/litellm_core_utils/AGENTS.md @@ -0,0 +1,15 @@ +## get_model_cost_map.py + +- This module decides when to fetch and install. Fetching, parsing, integrity checks and alias expansion run in Rust +- The install order is publish, `add_known_models`, router replay, then `model_cost_inputs.reapply` + - Registrations land by lookup in the published table, so they are written after publishing +- Startup and `register_model(url)` fetch with `http_settings=None`, which reads env vars only like main's bare httpx. Reload passes `http_settings()` +- Startup falls back to the bundled map on any failure. Reload never falls back and keeps the current map +- `register_model(url)` makes one attempt and registers nothing on failure +- `utils._register_on_current_map` reruns the whole pass on the new table after `NativeTableSwapped`, and its last attempt passes `rows=None`, so no entry is lost +- Compat names exist only so code written against main keeps working. New code never calls them, and they never gain behavior of their own + - `_finalize_model_cost_map`, use `model_cost_inputs.install` + - `GetModelCostMap.validate_model_cost_map` and `_check_*`, since Rust checks fetched documents + - `ModelCostMapReloaded.model_cost_map`, use `GetModelCostMap.loaded_model_cost_map()` + - the `sleep`, `rng` and `client` parameters, which are ignored with a `DeprecationWarning` + - `AliasedCostEntry` re-exported here, which now lives in `rust_bridge/host/model_cost.py` diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index d2045c947fe..80393a72631 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -9,27 +9,23 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True ``` """ -import asyncio -import hashlib import json import os -import random import sys import threading -import time -import urllib.parse -from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass, replace +import warnings +from collections.abc import Mapping +from dataclasses import dataclass from datetime import datetime, timezone +from functools import lru_cache from importlib.resources import files from itertools import chain from pathlib import Path from types import MappingProxyType -from typing import Final, Protocol, cast # noqa: TID251 # adapts boto3 and botocore's dynamic client responses +from typing import Final, Literal -import httpx -from pydantic import ConfigDict, TypeAdapter -from typing_extensions import NotRequired, ReadOnly, TypedDict +from pydantic import JsonValue, TypeAdapter +from typing_extensions import ReadOnly, TypedDict from litellm import verbose_logger from litellm.constants import ( @@ -40,42 +36,58 @@ from litellm.litellm_core_utils.fallback_generalizations import ( set_fallback_generalizations, ) from litellm.litellm_core_utils.offline_mode import is_offline_mode +from litellm.rust_bridge._native import ( + NativeCostMapSource, + NativeFetchFailed, + NativeFetchPolicy, + NativeInstalled, + NativeModelCostMap, + NativeRejected, +) +from litellm.rust_bridge.host import model_cost as model_cost_host +from litellm.rust_bridge.host.model_cost import ( + AliasedCostEntry as AliasedCostEntry, # noqa: PLC0414 # main exported the row type from here +) +from litellm.rust_bridge.host.settings import HttpSettings -FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" -_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) -_DECODED_JSON: Final = TypeAdapter(object) -_BUNDLED_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) _CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"}) +_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) +MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 +MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0 +MODEL_COST_MAP_FETCH_MIN_DELAY_SECONDS: Final = 2.0 +_litellm_import_complete = threading.Event() + +model_cost_inputs: Final = NativeModelCostMap.bundled( + min_model_count=MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio=MODEL_COST_MAP_MAX_SHRINK_RATIO +) def _is_cli_process() -> bool: return Path(sys.argv[0]).stem in _CLI_ENTRYPOINT_NAMES -# Reserved top-level keys that are not model entries. They must be excluded -# from the model-count integrity check so a real upstream shrink can't be masked. -RESERVED_TOP_LEVEL_KEYS: Final = frozenset({"sample_spec", FALLBACK_GENERALIZATIONS_KEY}) +def _uses_bundled_map_only() -> bool: + # Can't use get_secret_bool here: this runs during litellm.__init__ before litellm._key_management_settings is set + return is_offline_mode() or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" -def _count_model_entries(model_cost: dict) -> int: - """Count actual model entries, excluding reserved meta keys.""" - return sum(1 for key in model_cost if key not in RESERVED_TOP_LEVEL_KEYS) +@dataclass(frozen=True, slots=True) +class _LoadedSnapshot: + rows: Mapping[str, Mapping[str, object]] + fields: frozenset[str] -def git_blob_id(body: bytes) -> str: - return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest() +@lru_cache(maxsize=1) +def _loaded_snapshot(revision: str) -> _LoadedSnapshot: + raw: Final = _CATALOG_ADAPTER.validate_python(model_cost_inputs.loaded()) + return _LoadedSnapshot( + rows=MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()}), + fields=frozenset(chain.from_iterable(entry.keys() for entry in raw.values())), + ) class GetModelCostMap: - """ - Handles fetching, validating, and loading the model cost map. - - Only the backup model *count* is cached (a single int). The full - backup dict is never held in memory — it is only parsed when it - needs to be *returned* as a fallback. - """ - - _backup_model_count: int = -1 # -1 = not yet loaded + """The bundled document, and the installed one as it was loaded before any registration.""" @staticmethod def read_local_model_cost_map_bytes() -> bytes: @@ -85,137 +97,72 @@ class GetModelCostMap: def read_local_model_cost_map_text() -> str: return GetModelCostMap.read_local_model_cost_map_bytes().decode("utf-8") - @staticmethod - def load_local_model_cost_map_with_revision() -> "ModelCostMapReloaded": - body: Final = GetModelCostMap.read_local_model_cost_map_bytes() - content: Final = _BUNDLED_CATALOG_ADAPTER.validate_python(json.loads(body)) - return ModelCostMapReloaded(model_cost_map=content, revision=git_blob_id(body)) - @staticmethod def load_local_model_cost_map() -> dict: """Load the local backup model cost map bundled with the package.""" - return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map + return json.loads(GetModelCostMap.read_local_model_cost_map_bytes()) - _loaded_catalog: Mapping[str, Mapping[str, object]] = MappingProxyType({}) - _loaded_catalog_fields: frozenset[str] = frozenset() - - @classmethod - def loaded_model_cost_map(cls) -> Mapping[str, Mapping[str, object]]: + @staticmethod + def loaded_model_cost_map() -> Mapping[str, Mapping[str, object]]: """The catalog as last loaded (bundled or remote), untouched by ``register_model`` or router registrations.""" - return cls._loaded_catalog + return _loaded_snapshot(model_cost_inputs.provenance.revision).rows - @classmethod - def loaded_model_cost_map_fields(cls) -> frozenset[str]: + @staticmethod + def loaded_model_cost_map_fields() -> frozenset[str]: """Every field name that at least one entry of the loaded catalog carries.""" - return cls._loaded_catalog_fields + return _loaded_snapshot(model_cost_inputs.provenance.revision).fields - @classmethod - def _snapshot_loaded_catalog(cls, model_cost: Mapping[str, object]) -> None: - raw: Final = _CATALOG_ADAPTER.validate_python(model_cost) - cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()}) - cls._loaded_catalog_fields = frozenset(chain.from_iterable(entry.keys() for entry in raw.values())) - - @classmethod - def get_backup_model_count(cls) -> int: - """Return the number of models in the local backup (cached int).""" - if cls._backup_model_count < 0: - backup: Final = cls.load_local_model_cost_map() - cls._backup_model_count = _count_model_entries(backup) - return cls._backup_model_count + @staticmethod + def get_backup_model_count() -> int: + """The number of models in the bundled document.""" + return _bundled_model_count() _get_backup_model_count = get_backup_model_count @staticmethod - def _check_is_valid_dict(fetched_map: dict) -> bool: - """Check 1: fetched map is a non-empty dict.""" - if not isinstance(fetched_map, dict): - verbose_logger.warning( - "LiteLLM: Fetched model cost map is not a dict (type=%s). Falling back to local backup.", - type(fetched_map).__name__, - ) + def _check_is_valid_dict(fetched_map: object) -> bool: + if not isinstance(fetched_map, dict) or not fetched_map: + verbose_logger.warning("LiteLLM: Fetched model cost map is not a non-empty dict.") return False - - if len(fetched_map) == 0: - verbose_logger.warning( - "LiteLLM: Fetched model cost map is empty. Falling back to local backup.", - ) - return False - return True - @classmethod + @staticmethod def _check_model_count_not_reduced( - cls, - fetched_map: dict, + fetched_map: Mapping[str, object], backup_model_count: int, min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, ) -> bool: - """Check 2: model count has not reduced significantly vs backup.""" fetched_count: Final = _count_model_entries(fetched_map) - if fetched_count < min_model_count: verbose_logger.warning( - "LiteLLM: Fetched model cost map has only %d models (minimum=%d). " - "This may indicate a corrupted upstream file. " - "Falling back to local backup.", - fetched_count, - min_model_count, + "LiteLLM: Fetched model cost map has only %d models (minimum=%d).", fetched_count, min_model_count ) return False - if backup_model_count > 0 and fetched_count < backup_model_count * max_shrink_ratio: verbose_logger.warning( - "LiteLLM: Fetched model cost map shrank significantly " - "(fetched=%d, backup=%d, threshold=%.0f%%). " - "This may indicate a corrupted upstream file. " - "Falling back to local backup.", + "LiteLLM: Fetched model cost map shrank significantly (fetched=%d, backup=%d, threshold=%.0f%%).", fetched_count, backup_model_count, max_shrink_ratio * 100, ) return False - return True - @classmethod + @staticmethod def validate_model_cost_map( - cls, - fetched_map: dict, + fetched_map: object, backup_model_count: int, min_model_count: int = MODEL_COST_MAP_MIN_MODEL_COUNT, max_shrink_ratio: float = MODEL_COST_MAP_MAX_SHRINK_RATIO, ) -> bool: - """ - Validate the integrity of a fetched model cost map. - - Runs each check in order and returns False on the first failure. - - Checks: - 1. ``_check_is_valid_dict`` -- fetched map is a non-empty dict. - 2. ``_check_model_count_not_reduced`` -- model count meets minimum - and has not shrunk >``max_shrink_ratio`` vs backup. - - Returns True if all checks pass, False otherwise. - """ - if not cls._check_is_valid_dict(fetched_map): - return False - - if not cls._check_model_count_not_reduced( - fetched_map=fetched_map, - backup_model_count=backup_model_count, - min_model_count=min_model_count, - max_shrink_ratio=max_shrink_ratio, - ): - return False - - return True - - -RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) -MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3 -MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0 -_litellm_import_complete = threading.Event() + """Compat for main's callers. Rust checks fetched documents in ``model_cost_inputs.install``.""" + return isinstance(fetched_map, dict) and ( + GetModelCostMap._check_is_valid_dict(fetched_map) + and GetModelCostMap._check_model_count_not_reduced( + fetched_map, backup_model_count, min_model_count, max_shrink_ratio + ) + ) def mark_litellm_import_complete() -> None: @@ -224,9 +171,22 @@ def mark_litellm_import_complete() -> None: @dataclass(frozen=True, slots=True) class ModelCostMapReloaded: - model_cost_map: dict # mutable-ok: adopted as litellm.model_cost, whose consumer contract is a plain mutable dict - revision: str | None = None - etag: str | None = None + """A reload that installed a new document as ``litellm.model_cost``.""" + + model_count: int + + @property + def model_cost_map(self) -> dict[str, dict[str, object]]: + """Compat for main's callers. Use ``GetModelCostMap.loaded_model_cost_map()``.""" + return model_cost_inputs.loaded() + + @property + def revision(self) -> str: + return model_cost_inputs.provenance.revision + + @property + def etag(self) -> str | None: + return model_cost_inputs.provenance.etag @dataclass(frozen=True, slots=True) @@ -238,274 +198,18 @@ ModelCostMapReloadResult = ModelCostMapReloaded | ModelCostMapReloadUnavailable @dataclass(frozen=True, slots=True) -class _FetchAttemptRetryable: - reason: str - retry_after_seconds: float | None - - -def _parse_retry_after_seconds(response: httpx.Response) -> float | None: - try: - header: Final = response.headers["Retry-After"] - except KeyError: - return None - try: - seconds: Final = float(header) - except ValueError: - return None - return seconds if seconds >= 0 else None - - -def _retry_wait_seconds(outcome: _FetchAttemptRetryable, attempt: int, rng: random.Random) -> float: - if outcome.retry_after_seconds is not None: - return min(outcome.retry_after_seconds, MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS) - return min(2.0**attempt + rng.uniform(0.0, 1.0), MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS) - - -class _AsyncGetClient(Protocol): - def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ... - - -class _SyncGetClient(Protocol): - def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ... - - -class _S3Body(Protocol): - def read(self) -> bytes: ... - - -class _S3GetObjectResponse(TypedDict): - Body: ReadOnly[_S3Body] - ETag: NotRequired[ReadOnly[str]] - - -class _S3Client(Protocol): - def get_object(self, *, Bucket: str, Key: str) -> _S3GetObjectResponse: ... - - -_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable - - -def _default_reload_client() -> _AsyncGetClient: - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client - from litellm.types.llms.custom_http import httpxSpecialProvider - - return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap) - - -def _default_s3_client(timeout: int) -> _S3Client: - import boto3 - from botocore.config import Config - - return cast( # cast-ok: boto3's generic client factory does not expose its generated S3 response type - _S3Client, - boto3.client( # pyright: ignore[reportUnknownMemberType] # generic factory returns Unknown service clients - "s3", - region_name=os.getenv("AWS_REGION_NAME") or None, - config=Config(connect_timeout=timeout, read_timeout=timeout, retries={"total_max_attempts": 1}), - ), - ) - - -def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome: - reason: Final = f"{type(error).__name__} fetching {url}: {error}" - if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)): - return ModelCostMapReloadUnavailable(reason=reason) - return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None) - - -def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _FetchAttemptOutcome: - try: - parsed: Final = _DECODED_JSON.validate_python(json.loads(content)) - except ValueError as e: - return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}") - if not isinstance(parsed, dict): - return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}") - model_cost_map: Final = _BUNDLED_CATALOG_ADAPTER.validate_python(parsed) - return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag) - - -def _split_url(url: str) -> urllib.parse.SplitResult | ModelCostMapReloadUnavailable: - try: - return urllib.parse.urlsplit(url) - except ValueError as error: - return ModelCostMapReloadUnavailable(reason=f"invalid URL {url}: {error}") - - -def _attempt_fetch_s3(parts: urllib.parse.SplitResult, url: str, timeout: int) -> _FetchAttemptOutcome: - from botocore.exceptions import BotoCoreError, ClientError - - bucket: Final = parts.netloc - key: Final = parts.path.lstrip("/") - if not bucket or not key: - return ModelCostMapReloadUnavailable(reason=f"invalid S3 URL {url}: expected s3://bucket/key") - - try: - client: Final = _default_s3_client(timeout) - response: Final = client.get_object(Bucket=bucket, Key=key) - content: Final = response["Body"].read() - except (BotoCoreError, ClientError) as error: - return ModelCostMapReloadUnavailable(reason=f"S3 {type(error).__name__} fetching {url}: {error}") - return _classify_fetch_content(content, url, response.get("ETag")) - - -async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = _split_url(url) - if isinstance(parts, ModelCostMapReloadUnavailable): - return parts - if parts.scheme == "s3": - return await asyncio.to_thread(_attempt_fetch_s3, parts, url, timeout) - try: - response: Final = await client.get(url, timeout=timeout) - except (httpx.HTTPError, httpx.InvalidURL) as e: - return _classify_fetch_error(e, url) - return _classify_fetch_response(response, url) - - -def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = _split_url(url) - if isinstance(parts, ModelCostMapReloadUnavailable): - return parts - if parts.scheme == "s3": - return _attempt_fetch_s3(parts, url, timeout) - try: - response: Final = client.get(url, timeout=timeout) - except (httpx.HTTPError, httpx.InvalidURL) as e: - return _classify_fetch_error(e, url) - return _classify_fetch_response(response, url) - - -def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome: - if response.status_code in RETRYABLE_FETCH_STATUS_CODES: - return _FetchAttemptRetryable( - reason=f"HTTP {response.status_code} from {url}", - retry_after_seconds=_parse_retry_after_seconds(response), - ) - if response.is_error: - return ModelCostMapReloadUnavailable(reason=f"HTTP {response.status_code} from {url}") - etag_values: Final = response.headers.get_list("etag") - return _classify_fetch_content(response.content, url, ", ".join(etag_values) if etag_values else None) - - -def _next_retry_wait( - outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random -) -> float | ModelCostMapReloadUnavailable: - if attempt == max_attempts: - return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)") - wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng) - verbose_logger.warning( - "LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs", - attempt, - max_attempts, - outcome.reason, - wait_seconds, - ) - return wait_seconds - - -async def _fetch_remote_model_cost_map_with_retry( - url: str, - timeout: int, - max_attempts: int, - sleep: Callable[[float], Awaitable[None]], - rng: random.Random, - client: _AsyncGetClient, -) -> ModelCostMapReloadResult: - for attempt in range(1, max_attempts + 1): - outcome = await _attempt_fetch(client=client, url=url, timeout=timeout) - if not isinstance(outcome, _FetchAttemptRetryable): - return outcome - wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) - if isinstance(wait_seconds, ModelCostMapReloadUnavailable): - return wait_seconds - await sleep(wait_seconds) - return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") - - -def _fetch_remote_model_cost_map_with_retry_sync( - url: str, - timeout: int, - attempts: range, - sleep: Callable[[float], None], - rng: random.Random, - client: _SyncGetClient, -) -> ModelCostMapReloadResult: - max_attempts: Final = attempts.stop - 1 - for attempt in attempts: - outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout) - if not isinstance(outcome, _FetchAttemptRetryable): - return outcome - wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng) - if isinstance(wait_seconds, ModelCostMapReloadUnavailable): - return wait_seconds - sleep(wait_seconds) - return ModelCostMapReloadUnavailable(reason="model cost map fetch failed") - - -async def refetch_model_cost_map( - url: str, - timeout: int = 5, - max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, - sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, - rng: random.Random | None = None, - client: "_AsyncGetClient | None" = None, -) -> ModelCostMapReloadResult: - """ - Re-fetch the model cost map for a runtime reload, retrying transient HTTP - errors (429/5xx/transport) with Retry-After-aware backoff. - - Unlike ``get_model_cost_map`` this never falls back to the packaged backup: - on failure it returns ``ModelCostMapReloadUnavailable`` so callers keep the - map they already have. - """ - if is_offline_mode() or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true": - _cost_map_source_info.loaded_at = datetime.now(timezone.utc) - _cost_map_source_info.source = "local" - _cost_map_source_info.url = None - _cost_map_source_info.is_env_forced = True - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()) - - result: Final = await _fetch_remote_model_cost_map_with_retry( - url=url, - timeout=timeout, - max_attempts=max_attempts, - sleep=sleep, - rng=rng if rng is not None else random.Random(), - client=client if client is not None else _default_reload_client(), - ) - if isinstance(result, ModelCostMapReloadUnavailable): - verbose_logger.warning( - "LiteLLM: model cost map reload failed: %s. Keeping the currently loaded map.", - result.reason, - ) - return result - if not GetModelCostMap.validate_model_cost_map( - fetched_map=result.model_cost_map, - backup_model_count=GetModelCostMap.get_backup_model_count(), - ): - return ModelCostMapReloadUnavailable(reason=f"model cost map from {url} failed integrity validation") - _cost_map_source_info.loaded_at = datetime.now(timezone.utc) - _cost_map_source_info.source = "remote" - _cost_map_source_info.url = url - _cost_map_source_info.is_env_forced = False - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(result) - - class ModelCostMapSourceInfo: - """Tracks the source of the currently loaded model cost map.""" + """Where the current cost map came from, replaced whole by every load so readers never see two loads mixed.""" - source: str = "local" # "local" or "remote" + source: Literal["local", "remote"] = "local" url: str | None = None is_env_forced: bool = False fallback_reason: str | None = None - loaded_at: "datetime | None" = None - source_revision: str | None = None - etag: str | None = None + loaded_at: datetime | None = None -# Module-level singleton tracking the source of the current cost map -_cost_map_source_info: Final = ModelCostMapSourceInfo() +class _LastLoad: + current: ModelCostMapSourceInfo = ModelCostMapSourceInfo() class CostMapProvenance(TypedDict): @@ -522,10 +226,8 @@ class CostMapSourceInfo(CostMapProvenance): def get_model_cost_map_provenance() -> CostMapProvenance: - return { - "source_revision": _cost_map_source_info.source_revision, - "etag": _cost_map_source_info.etag, - } + provenance: Final = model_cost_inputs.provenance + return {"source_revision": provenance.revision, "etag": provenance.etag} def get_model_cost_map_source_info() -> CostMapSourceInfo: @@ -541,177 +243,147 @@ def get_model_cost_map_source_info() -> CostMapSourceInfo: - source_revision: git blob id of the loaded file's bytes - etag: the ETag of the remote fetch (None for the bundled backup) """ - loaded_at: Final = _cost_map_source_info.loaded_at + last: Final = _LastLoad.current return { - "source": _cost_map_source_info.source, - "url": _cost_map_source_info.url, - "is_env_forced": _cost_map_source_info.is_env_forced, - "fallback_reason": _cost_map_source_info.fallback_reason, - "loaded_at": loaded_at.isoformat() if loaded_at is not None else None, - "source_revision": _cost_map_source_info.source_revision, - "etag": _cost_map_source_info.etag, + "source": last.source, + "url": last.url, + "is_env_forced": last.is_env_forced, + "fallback_reason": last.fallback_reason, + "loaded_at": last.loaded_at.isoformat() if last.loaded_at is not None else None, + **get_model_cost_map_provenance(), } def get_model_cost_map_loaded_at() -> "datetime | None": """When this process last loaded its cost map, stamped at the start of every load""" - return _cost_map_source_info.loaded_at + return _LastLoad.current.loaded_at -class AliasedCostEntry(dict[str, object]): - __slots__ = ("names", "registered_fields") - - def __init__( - self, entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset() - ) -> None: - super().__init__(entry) - self.names: tuple[str, ...] = names - self.registered_fields: frozenset[str] = registered_fields +def cost_map_source(url: str, http_settings: HttpSettings | None) -> NativeCostMapSource: + return NativeCostMapSource(url, http_settings=http_settings) -def _alias_groups(alias_owners: Mapping[str, str]) -> Mapping[str, tuple[str, ...]]: - return MappingProxyType( - { - owner: (owner, *(alias for alias, alias_owner in alias_owners.items() if alias_owner == owner)) - for owner in dict.fromkeys(alias_owners.values()) - } +def fetch_policy(timeout: float, max_attempts: int) -> NativeFetchPolicy: + return NativeFetchPolicy( + max_attempts=max_attempts, + max_wait_seconds=MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS, + timeout_seconds=timeout, + min_delay_seconds=MODEL_COST_MAP_FETCH_MIN_DELAY_SECONDS, ) -def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]: - """ - Expand alias rows into top-level entries that share their owner's ``AliasedCostEntry``. - - A row whose ``alias_of`` names another row is served as that row, and an - ``aliases`` list on a row adds further names for it. Every name in a group - references the **same** ``AliasedCostEntry``, which carries every name it is - served under, and neither key survives into the served entries. - - An alias that collides with a canonical entry or another alias, or a row whose - ``alias_of`` names nothing, is skipped and a warning is logged. - """ - pointer_rows: Final = { - name: target for name, info in model_cost.items() if isinstance(target := info.get("alias_of"), str) - } - alias_owners: Final[dict[str, str]] = {} - keys_with_aliases: Final[list[str]] = [] - - for model_name, model_info in model_cost.items(): - aliases: object = model_info.get("aliases") - if aliases is None or model_name in pointer_rows: - continue - keys_with_aliases.append(model_name) - if not isinstance(aliases, list): - verbose_logger.warning( - "LiteLLM model alias field for '%s' is not a list (got %s) — skipping.", - model_name, - type(aliases).__name__, - ) - continue - if not aliases: - continue - for alias in aliases: - if alias in model_cost and alias not in pointer_rows: - verbose_logger.warning( - "LiteLLM model alias conflict: alias '%s' (from '%s') " - "already exists as a canonical entry — skipping.", - alias, - model_name, - ) - continue - if alias in alias_owners: - verbose_logger.warning( - "LiteLLM model alias conflict: alias '%s' (from '%s') " - "was already claimed by another entry — skipping.", - alias, - model_name, - ) - continue - alias_owners[alias] = model_name - - for alias, target in pointer_rows.items(): - if target not in model_cost or target in pointer_rows: - verbose_logger.warning( - "LiteLLM model alias '%s' points at '%s', which is not a model entry, dropping it.", alias, target - ) - elif alias in alias_owners: - verbose_logger.warning( - "LiteLLM model alias conflict: alias '%s' (pointing at '%s') " - "was already claimed by another entry — skipping.", - alias, - target, - ) - else: - alias_owners[alias] = target - - # Remove the ``aliases`` key from entries so it doesn't pollute model info - for key in keys_with_aliases: - model_cost[key].pop("aliases", None) - for alias in pointer_rows.keys() - alias_owners.keys(): - model_cost.pop(alias) - - for owner, names in _alias_groups(alias_owners).items(): - model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names))) - return model_cost +def _set_fallback_rules(block: JsonValue) -> None: + rules: Final = block.get("rules") if isinstance(block, dict) else None + set_fallback_generalizations(rules if isinstance(rules, list) else None) -def _finalize_model_cost_map(model_cost: dict) -> dict: - """Extract fallback generalizations out of the raw map, then expand aliases. - - The ``fallback_generalizations`` block is installed into the generalizations - module and removed from the map so it is never treated as a model entry. - """ - raw: Final = model_cost.pop(FALLBACK_GENERALIZATIONS_KEY, None) - rules: Final = raw.get("rules") if isinstance(raw, dict) else None - set_fallback_generalizations(rules) - return _expand_model_aliases(model_cost) +def _rows_with_rules_set(installed: NativeInstalled) -> dict: + """The installed document's rows, after making its ``fallback_generalizations`` rules the active ones.""" + _set_fallback_rules(installed.fallback_generalizations) + return installed.rows -def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded: - _cost_map_source_info.source_revision = loaded.revision - _cost_map_source_info.etag = loaded.etag - finalized: Final = _finalize_model_cost_map(loaded.model_cost_map) - GetModelCostMap._snapshot_loaded_catalog(finalized) # pyright: ignore[reportPrivateUsage] # same module - return replace(loaded, model_cost_map=finalized) - - -def adopt_model_cost_map( - new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract -) -> int: +def adopt_model_cost_map(new_model_cost_map: dict) -> int: # mutable-ok: main's signature, the dict goes live + """Publishes the table, then writes router rows and kept registrations onto it.""" import litellm from litellm import utils - litellm.model_cost = new_model_cost_map - utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + model_cost_host.publish(new_model_cost_map) litellm.add_known_models(model_cost_map=new_model_cost_map) - fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 utils.reapply_runtime_model_cost_registrations() - return fetched_model_count + return len(new_model_cost_map) + + +def _install(installed: NativeInstalled) -> int: + adopt_model_cost_map(_rows_with_rules_set(installed)) + return installed.model_count + + +def install_model_cost_map(body: bytes, etag: str | None = None) -> ModelCostMapReloadResult: + """Installs a cost map document as ``litellm.model_cost``, every kept ``register_model`` call written on top.""" + installed: Final = model_cost_inputs.install(body, etag) + if isinstance(installed, NativeRejected): + return ModelCostMapReloadUnavailable(reason=installed.reason) + return ModelCostMapReloaded(model_count=_install(installed)) + + +def _mark_rows_with_rules_set( + source: Literal["local", "remote"], + url: str | None, + is_env_forced: bool, + fallback_reason: str | None, + loaded_at: datetime, +) -> None: + _LastLoad.current = ModelCostMapSourceInfo(source, url, is_env_forced, fallback_reason, loaded_at) + + +def _rejected_reason(rejected: NativeRejected, url: str) -> str: + if rejected.integrity: + return f"model cost map from {url} failed integrity validation" + return f"invalid JSON from {url}: {rejected.reason}" + + +def _warn_rejected(rejected: NativeRejected, url: str) -> None: + verbose_logger.warning( + "LiteLLM: Fetched model cost map from %s was rejected (%s). Using local backup instead.", url, rejected.reason + ) + + +def _warn_ignored_fetch_hooks(sleep: object, rng: object, client: object) -> None: + if (sleep, rng, client) != (None, None, None): + warnings.warn( + "The sleep, rng and client arguments of the model cost map fetch are ignored, the fetch runs natively", + DeprecationWarning, + stacklevel=3, + ) + + +async def refetch_model_cost_map( + url: str, + timeout: int = 5, + max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, + sleep: object = None, + rng: object = None, + client: object = None, +) -> ModelCostMapReloadResult: + """ + Re-fetch the model cost map for a runtime reload, retrying transient HTTP + errors (429/5xx/transport) with Retry-After-aware backoff, and install it. + + Unlike ``get_model_cost_map`` this never falls back to the packaged backup: + on failure it returns ``ModelCostMapReloadUnavailable`` and the map already + installed stays. + """ + from litellm.rust_bridge.host.settings import http_settings + + _warn_ignored_fetch_hooks(sleep, rng, client) + + if _uses_bundled_map_only(): + _mark_rows_with_rules_set("local", None, True, None, datetime.now(timezone.utc)) + return ModelCostMapReloaded(model_count=_install(model_cost_inputs.install_bundled())) + + fetched: Final = await cost_map_source(url, http_settings()).fetch(fetch_policy(timeout, max_attempts)) + if isinstance(fetched, NativeFetchFailed): + verbose_logger.warning( + "LiteLLM: model cost map reload failed: %s. Keeping the currently loaded map.", fetched.reason + ) + return ModelCostMapReloadUnavailable(reason=fetched.reason) + installed: Final = model_cost_inputs.install(fetched.body, fetched.etag) + if isinstance(installed, NativeRejected): + verbose_logger.warning( + "LiteLLM: model cost map reload failed: %s. Keeping the currently loaded map.", installed.reason + ) + return ModelCostMapReloadUnavailable(reason=_rejected_reason(installed, url)) + _mark_rows_with_rules_set("remote", url, False, None, datetime.now(timezone.utc)) + return ModelCostMapReloaded(model_count=_install(installed)) def _retry_remote_fetch_in_background( - url: str, - timeout: int, - max_attempts: int, - sleep: Callable[[float], None], - rng: random.Random, - client: _SyncGetClient, - first_outcome: _FetchAttemptRetryable, + url: str, timeout: int, max_attempts: int, first_failure: NativeFetchFailed ) -> None: try: - first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng) - if isinstance(first_wait, ModelCostMapReloadUnavailable): - return - sleep(first_wait) - result: Final = _fetch_remote_model_cost_map_with_retry_sync( - url=url, - timeout=timeout, - attempts=range(2, max_attempts + 1), - sleep=sleep, - rng=rng, - client=client, - ) - if isinstance(result, ModelCostMapReloadUnavailable): + fetched: Final = cost_map_source(url, None).fetch_blocking(fetch_policy(timeout, max_attempts), first_failure) + if isinstance(fetched, NativeFetchFailed): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", url, @@ -719,20 +391,12 @@ def _retry_remote_fetch_in_background( ) return _litellm_import_complete.wait() - if not GetModelCostMap.validate_model_cost_map( - fetched_map=result.model_cost_map, - backup_model_count=GetModelCostMap.get_backup_model_count(), - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", - url, - ) + installed: Final = model_cost_inputs.install(fetched.body, fetched.etag) + if isinstance(installed, NativeRejected): + _warn_rejected(installed, url) return - finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map - _cost_map_source_info.source = "remote" - _cost_map_source_info.fallback_reason = None - _cost_map_source_info.loaded_at = datetime.now(timezone.utc) - adopt_model_cost_map(finalized) + _mark_rows_with_rules_set("remote", url, False, None, datetime.now(timezone.utc)) + _install(installed) except Exception as e: # noqa: BLE001 # a failed background retry must not kill the task; the backup stays verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e) @@ -741,79 +405,91 @@ def get_model_cost_map( url: str, timeout: int = 5, max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, - sleep: Callable[[float], None] = time.sleep, - rng: random.Random | None = None, - client: "_SyncGetClient | None" = None, + sleep: object = None, + rng: object = None, + client: object = None, ) -> dict: """ Public entry point — returns the model cost map dict. 1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set or this is a ``lite`` / ``litellm-proxy`` CLI process, uses the local backup only. - 2. Otherwise fetches from ``url``, retrying transient HTTP errors + 2. Otherwise fetches from ``url`` once, retrying transient HTTP errors (429/5xx/transport) with Retry-After-aware backoff in a background thread, validates integrity, and falls back to the local backup on any failure. - - Only the backup model count is cached (a single int) for validation. - The full backup dict is only parsed when it must be *returned* as a - fallback — it is never held in memory long-term. """ - _cost_map_source_info.loaded_at = datetime.now(timezone.utc) - # Note: can't use get_secret_bool here — this runs during litellm.__init__ - # before litellm._key_management_settings is set. - if is_offline_mode() or os.getenv("LITELLM_LOCAL_MODEL_COST_MAP", "").lower() == "true" or _is_cli_process(): - _cost_map_source_info.source = "local" - _cost_map_source_info.url = None - _cost_map_source_info.is_env_forced = True - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + _warn_ignored_fetch_hooks(sleep, rng, client) + started: Final = datetime.now(timezone.utc) + if _uses_bundled_map_only() or _is_cli_process(): + _mark_rows_with_rules_set("local", None, True, None, started) + return _rows_with_rules_set(model_cost_inputs.install_bundled()) - _cost_map_source_info.url = url - _cost_map_source_info.is_env_forced = False - - fetch_client: Final = client if client is not None else httpx - fetch_rng: Final = rng if rng is not None else random.Random() - outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout) - if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1: - threading.Thread( - target=_retry_remote_fetch_in_background, - kwargs={ - "url": url, - "timeout": timeout, - "max_attempts": max_attempts, - "sleep": sleep, - "rng": fetch_rng, - "client": fetch_client, - "first_outcome": outcome, - }, - name="litellm-model-cost-map-retry", - daemon=True, - ).start() - if not isinstance(outcome, ModelCostMapReloaded): + fetched: Final = cost_map_source(url, None).fetch_blocking(fetch_policy(timeout, 1)) + if isinstance(fetched, NativeFetchFailed): + if fetched.retryable and max_attempts > 1: + threading.Thread( + target=_retry_remote_fetch_in_background, + args=(url, timeout, max_attempts, fetched), + name="litellm-model-cost-map-retry", + daemon=True, + ).start() verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, - outcome.reason, + fetched.reason, ) - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}" - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map - content: Final = outcome.model_cost_map + _mark_rows_with_rules_set("local", url, False, f"Remote fetch failed: {fetched.reason}", started) + return _rows_with_rules_set(model_cost_inputs.install_bundled()) - # Validate using cached count (cheap int comparison, no file I/O) - if not GetModelCostMap.validate_model_cost_map( - fetched_map=content, - backup_model_count=GetModelCostMap.get_backup_model_count(), - ): - verbose_logger.warning( - "LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s", - url, + installed: Final = model_cost_inputs.install(fetched.body, fetched.etag) + if isinstance(installed, NativeRejected): + _warn_rejected(installed, url) + fallback_reason: Final = ( + "Remote data failed integrity validation" + if installed.integrity + else f"Remote fetch failed: {_rejected_reason(installed, url)}" ) - _cost_map_source_info.source = "local" - _cost_map_source_info.fallback_reason = "Remote data failed integrity validation" - return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map + _mark_rows_with_rules_set("local", url, False, fallback_reason, started) + return _rows_with_rules_set(model_cost_inputs.install_bundled()) - _cost_map_source_info.source = "remote" - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(outcome).model_cost_map + _mark_rows_with_rules_set("remote", url, False, None, started) + return _rows_with_rules_set(installed) + + +def fetch_registration_document(url: str) -> dict[str, dict[str, object]] | None: + """The document behind ``register_model(url)``, fetched in one attempt, or ``None`` when the fetch failed.""" + fetched: Final = cost_map_source(url, None).fetch_blocking(fetch_policy(5, 1)) + if isinstance(fetched, NativeFetchFailed): + verbose_logger.warning("LiteLLM: register_model could not fetch %s: %s", url, fetched.reason) + return None + try: + document: Final = _CATALOG_ADAPTER.validate_json(fetched.body) + except ValueError as e: + verbose_logger.warning("LiteLLM: register_model got an invalid cost map from %s: %s", url, e) + return None + return {key: entry for key, entry in document.items() if key != "fallback_generalizations"} + + +def _finalize_model_cost_map( # pyright: ignore[reportUnusedFunction] # main's private name, downstream code calls it + model_cost: dict, # mutable-ok: main's signature, pops the rules block +) -> dict: + """Compat for main's callers. Use ``model_cost_inputs.install``.""" + match NativeModelCostMap.finalize(json.dumps(model_cost).encode()): + case NativeRejected(reason=reason): + raise ValueError(reason) + case installed: + model_cost.pop("fallback_generalizations", None) + return _rows_with_rules_set(installed) + + +_RESERVED_TOP_LEVEL_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"}) + + +def _count_model_entries(model_cost: Mapping[str, object]) -> int: + return sum(1 for key in model_cost if key not in _RESERVED_TOP_LEVEL_KEYS) + + +@lru_cache(maxsize=1) +def _bundled_model_count() -> int: + return _count_model_entries(GetModelCostMap.load_local_model_cost_map()) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 5a9e5745c40..9db1d50d730 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.reasoning_effort_utils import OPAQUE_VALUE, warn_dropped_reasoning from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.rust_bridge.host.model_cost import ModelCostVersion, model_cost_version from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -36,7 +37,6 @@ from litellm.types.utils import ( ProviderSpecificModelInfo, ) from litellm.utils import ( - get_model_cost_mutation_generation, supports_function_calling, supports_reasoning, supports_tool_choice, @@ -542,21 +542,15 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): for message in messages ] - # Cached index of fireworks_ai/* entries from litellm.model_cost. Building - # this index requires a full scan of model_cost (tens of thousands of - # entries), so we memoize it. The cache key is (id(model_cost), - # mutation_generation): the generation counter is bumped on every - # register_model / reload path, so add+remove or in-place value - # replacement (which can leave id and len unchanged) still invalidates. - _fireworks_index_cache: tuple[int, int, list[tuple[str, dict]]] | None = None + _fireworks_index_cache: tuple[ModelCostVersion, list[tuple[str, dict]]] | None = None @classmethod def _get_fireworks_index(cls) -> list[tuple[str, dict]]: model_cost: Final = litellm.model_cost - signature: Final = (id(model_cost), get_model_cost_mutation_generation()) + version: Final = model_cost_version() cached: Final = cls._fireworks_index_cache - if cached is not None and cached[0] == signature[0] and cached[1] == signature[1]: - return cached[2] + if cached is not None and cached[0] == version: + return cached[1] index: Final[list[tuple[str, dict]]] = [] for key, model_info in (*model_cost.items(),): @@ -570,7 +564,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): continue index.append((key_short, model_info)) - cls._fireworks_index_cache = (signature[0], signature[1], index) + cls._fireworks_index_cache = (version, index) return index @staticmethod diff --git a/litellm/proxy/common_utils/model_deprecation.py b/litellm/proxy/common_utils/model_deprecation.py index 8a91c7b0a0c..b50890089e4 100644 --- a/litellm/proxy/common_utils/model_deprecation.py +++ b/litellm/proxy/common_utils/model_deprecation.py @@ -10,7 +10,8 @@ from typing import TYPE_CHECKING, Final, TypeAlias import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.rust_bridge.host.model_cost import AliasedCostEntry from litellm.types.proxy.model_deprecation import ( DEFAULT_DEPRECATION_WARN_DAYS, DeprecationStatus, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 37dc135d4d7..55540820651 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5682,12 +5682,6 @@ def resolve_classifier_plugin( return resolved -def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: - from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map - - return adopt_model_cost_map(new_model_cost_map) - - def _websearch_handler_params(stored: Mapping[str, object]) -> dict[str, object]: """ Translate stored web search interception settings into handler kwargs. @@ -9515,7 +9509,7 @@ class ProxyConfig: ) return - models_count = _swap_in_model_cost_map(reload_result.model_cost_map) + models_count = reload_result.model_count self.model_cost_map_loaded_at = current_time await record_reload_run(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME, current_time) # Adopted last, so neither a failed fetch nor a failed status write is recorded @@ -21979,7 +21973,7 @@ async def reload_model_cost_map( detail=f"Failed to reload model cost map: {reload_result.reason}. Current pricing data was kept.", ) - models_count = _swap_in_model_cost_map(reload_result.model_cost_map) + models_count = reload_result.model_count current_time = utc_now() proxy_config.model_cost_map_loaded_at = current_time provenance: Final = get_model_cost_map_provenance() diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 114d2df026e..221968e07ed 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -8,6 +8,7 @@ from typing_extensions import disjoint_base from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings from litellm.rust_bridge.execution import Execution +from litellm.rust_bridge.host.settings import HttpSettings from litellm.rust_bridge.public_call import NativeCall _T_co = TypeVar("_T_co", covariant=True) @@ -22,6 +23,102 @@ def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... +@final +class NativeCostMapSource: + def __new__(cls, url: str, *, http_settings: HttpSettings | None) -> NativeCostMapSource: ... + def fetch(self, policy: NativeFetchPolicy) -> Future[NativeFetched | NativeFetchFailed]: ... + def fetch_blocking( + self, policy: NativeFetchPolicy, resume_from: NativeFetchFailed | None = None + ) -> NativeFetched | NativeFetchFailed: ... + +@final +class NativeFetchPolicy: + def __new__( + cls, + *, + max_attempts: int, + max_wait_seconds: float, + timeout_seconds: float, + min_delay_seconds: float, + ) -> NativeFetchPolicy: ... + +@final +class NativeFetched: + @property + def body(self) -> bytes: ... + @property + def etag(self) -> str | None: ... + +@final +class NativeFetchFailed: + @property + def reason(self) -> str: ... + @property + def retryable(self) -> bool: ... + @property + def retry_after(self) -> float | None: ... + +@final +class NativeInstalled: + @property + def rows(self) -> dict[str, dict[str, object]]: ... + @property + def model_count(self) -> int: ... + @property + def fallback_generalizations(self) -> JsonValue: ... + +@final +class NativeRegistered: + @property + def unmatched(self) -> dict[str, Mapping[str, object]]: ... + +@final +class NativeTableSwapped: ... + +@final +class NativeRejected: + @property + def reason(self) -> str: ... + @property + def integrity(self) -> bool: ... + +@final +class NativeProvenance: + @property + def revision(self) -> str: ... + @property + def etag(self) -> str | None: ... + +@final +class NativeModelCostMap: + @staticmethod + def bundled(min_model_count: int, max_shrink_ratio: float) -> NativeModelCostMap: ... + def install(self, body: bytes, etag: str | None = None) -> NativeInstalled | NativeRejected: ... + def install_bundled(self) -> NativeInstalled: ... + @staticmethod + def finalize(body: bytes) -> NativeInstalled | NativeRejected: ... + def register( + self, + rows: dict[str, dict[str, object]] | None, + entries: Mapping[str, Mapping[str, object]], + persist: bool, + custom_llm_provider: str | None = None, + ) -> NativeRegistered | NativeTableSwapped: ... + def reapply(self, rows: dict[str, dict[str, object]]) -> NativeRegistered | NativeTableSwapped: ... + def registered_row(self, key: str) -> dict[str, object] | None: ... + def registered_keys(self) -> list[str]: ... + def loaded(self) -> dict[str, dict[str, object]]: ... + def find_key(self, rows: dict[str, dict[str, object]], name: str) -> str | None: ... + def keys_in(self, rows: dict[str, dict[str, object]], provider: str) -> list[str]: ... + def invalidate(self) -> None: ... + def snapshot(self) -> NativeModelCostMapState: ... + def restore(self, state: NativeModelCostMapState) -> None: ... + @property + def provenance(self) -> NativeProvenance: ... + +@final +class NativeModelCostMapState: ... + @final class NativeClickHouseSpendConfig: def __new__(cls, database: str, url: str, retention_days: int) -> NativeClickHouseSpendConfig: ... @@ -589,19 +686,30 @@ __all__ = [ "NativeAuthCache", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", + "NativeCostMapSource", "NativeDatabase", "NativeDiagnosticProcessor", + "NativeFetchFailed", + "NativeFetchPolicy", + "NativeFetched", + "NativeInstalled", + "NativeModelCostMap", + "NativeModelCostMapState", "NativePending", + "NativeProvenance", "NativeRateLimits", "NativeRedisBreaker", "NativeRedisClient", "NativeRedisStoreError", + "NativeRegistered", + "NativeRejected", "NativeRequestRedis", "NativeRouterCooldowns", "NativeRouterUsage", "NativeSend", "NativeSensitiveRoutePins", "NativeSpendCounters", + "NativeTableSwapped", "PostCallSettlement", "ProcessReservedForForking", "ResponsesWebSocketConnection", diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 81463065acd..4cb87a85c28 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -2,10 +2,23 @@ Python state Rust needs but does not own, behind one module per question. Rust imports a module here by path and calls a function by name -- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with) +- One module answers one question: `model_capabilities` (what a model can do), `callbacks_legacy_python` (the legacy callback registries), `settings` (configured snapshots), `cache` (the live configured cache), `diagnostics` (legacy failure logging), `secret_manager` (the configured secret manager and native config capture), `http_transport` (which async transport is in use), `timeouts` (timeout to seconds), `redis_events` (reporting native Redis calls and breaker transitions to `ServiceLogging`), `redis_credentials` (what a caller's redis-py `CredentialProvider` authenticates with), `model_cost` (the live `litellm.model_cost` table, every write to it and its version) - A module is the only place Rust reads that Python state. Rust never does `py.import("litellm").getattr(...)` on a global; it asks here - Each module reuses the Python implementation the Python route runs (`AnthropicModelInfo`, `Logging`). Do not re-derive those rules here or in Rust; one resolver, two callers - The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green - Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals - Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays - `routes/.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead + +## model_cost + +- `publish`, `write` and `changed_in_place` are the only writers. Each moves `ModelCostVersion` +- The version is `(id(table), len(table), writes)`, so replacing the dict or changing its size moves it without the host + - An in-place value edit does not, so code that edits rows in place calls `changed_in_place()` +- Anything derived from the table is cached under the version (`@VersionCached`, `model_cost_version()`), never cleared by hand +- `resolve`, `merged` and `kept` are shims until `get_model_info` and a merge generic over Python values move to Rust. Then their bodies become native calls and the signatures stay +- `resolve` returns `None` when it was handed a table that is no longer published. That is how a pass learns a reload swapped the table +- Compat names exist only so code written against main keeps working. New code never calls them, and they never gain behavior of their own + - `utils._invalidate_model_cost_lowercase_map` forwards to `changed_in_place` + - `utils.get_model_cost_mutation_generation` forwards to `model_cost_version().writes` + - `utils._model_cost_lowercase_map` is always `None` diff --git a/litellm/rust_bridge/host/model_capabilities.py b/litellm/rust_bridge/host/model_capabilities.py index efb8e10022a..6fd8b42520e 100644 --- a/litellm/rust_bridge/host/model_capabilities.py +++ b/litellm/rust_bridge/host/model_capabilities.py @@ -6,6 +6,7 @@ from typing import Final, TypeAlias import litellm from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.rust_bridge.host.model_cost import ModelCostVersion, model_cost_version class _Held: @@ -21,7 +22,7 @@ class _Held: return isinstance(other, _Held) and other.value is self.value -_Fingerprint: TypeAlias = tuple[str, str | None, _Held, int, int, tuple[_Held, ...]] +_Fingerprint: TypeAlias = tuple[str, str | None, _Held, ModelCostVersion, int, tuple[_Held, ...]] def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]: @@ -33,12 +34,11 @@ def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: str, provider: str) -> _Fingerprint: - """Changes whenever an answer could: ``register_model`` bumps the generation, and replacing + """Changes whenever an answer could: every write moves the cost map version, and replacing any entry the probes look up (``litellm.model_cost[key] = ...``) changes that entry's identity. The key holds the entries it names, so a cached answer's entries are never freed and their addresses never reused.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - from litellm.utils import get_model_cost_mutation_generation model_cost: Final = as_str_mapping(litellm.model_cost) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # the model map is an untyped dict of rows keys: Final = dict.fromkeys( @@ -52,7 +52,7 @@ def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: st model, custom_llm_provider, _Held(model_cost), - get_model_cost_mutation_generation(), + model_cost_version(), len(model_cost), tuple(_Held(model_cost.get(key)) for key in keys), ) diff --git a/litellm/rust_bridge/host/model_cost.py b/litellm/rust_bridge/host/model_cost.py new file mode 100644 index 00000000000..328e0c6252e --- /dev/null +++ b/litellm/rust_bridge/host/model_cost.py @@ -0,0 +1,148 @@ +"""The live ``litellm.model_cost`` table, every write to it, and its version.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from functools import lru_cache +from typing import Final, Generic, ParamSpec, TypeVar + +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE + +P = ParamSpec("P") +R = TypeVar("R") + + +class AliasedCostEntry(dict[str, object]): + """A ``litellm.model_cost`` row, the one object every alias name of a model is served under.""" + + __slots__ = ("names", "registered_fields") + + def __init__( + self, entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset() + ) -> None: + super().__init__(entry) + self.names: tuple[str, ...] = names + self.registered_fields: frozenset[str] = registered_fields + + +entry: Final = AliasedCostEntry + + +@dataclass(frozen=True, slots=True) +class ModelCostVersion: + table: int + size: int + writes: int + + +class _Writes: + __slots__ = ("count",) + + def __init__(self) -> None: + self.count: int = 0 + + def record(self) -> None: + self.count += 1 + + +_WRITES: Final = _Writes() + + +def model_cost_version() -> ModelCostVersion: + import litellm + + table: Final = litellm.model_cost + return ModelCostVersion(table=id(table), size=len(table), writes=_WRITES.count) + + +def publish(table: dict[str, object]) -> None: + """Makes ``table`` the live ``litellm.model_cost``.""" + import litellm + + litellm.model_cost = table + _WRITES.record() + + +def write(rows: dict[str, object], written: Mapping[str, object]) -> None: + """Writes one registration's rows into ``rows``, the published table.""" + rows.update(written) + _WRITES.record() + + +def changed_in_place() -> None: + """Moves the version and drops the name index after an in-place edit of ``litellm.model_cost``.""" + from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs + + model_cost_inputs.invalidate() + _WRITES.record() + + +class VersionCached(Generic[P, R]): + """Decorates ``function`` with an ``lru_cache`` keyed on the ``litellm.model_cost`` version each call read, and + exposes that cache's ``cache_clear`` and ``cache_info``.""" + + __slots__ = ("_at", "cache_clear", "cache_info") + + def __init__(self, function: Callable[P, R]) -> None: + def at( + _version: ModelCostVersion, + *args: P.args, + **kwargs: P.kwargs, # kwargs-ok: forwards a ParamSpec unchanged + ) -> R: + return function(*args, **kwargs) + + cached: Final = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(at) + self._at: Final = cached + self.cache_clear: Final = cached.cache_clear + self.cache_info: Final = cached.cache_info + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: # kwargs-ok: forwards a ParamSpec unchanged + return self._at(model_cost_version(), *args, **kwargs) + + +@dataclass(frozen=True, slots=True) +class Landing: + key: str + base: Mapping[str, object] + matched_builtin: bool + + +def published() -> dict[str, object]: + """The table published as ``litellm.model_cost`` now.""" + import litellm + + return litellm.model_cost + + +def resolve( + rows: Mapping[str, Mapping[str, object]] | None, + key: str, + registration: Mapping[str, object], + custom_llm_provider: str | None, +) -> Landing | None: + """Where a ``register_model`` entry lands in ``rows`` and the row it starts from, ``None`` when ``rows`` is no + longer the published table. ``rows`` of ``None`` is the published table.""" + import litellm + from litellm.utils import registration_base + + if rows is not None and rows is not litellm.model_cost: + return None + return registration_base(key, registration, custom_llm_provider) + + +def merged( + current: Mapping[str, object] | None, base: Mapping[str, object], registration: Mapping[str, object] +) -> AliasedCostEntry: + """The row a registration writes over ``current``, the row now under its landing key.""" + from litellm.utils import model_cost_registration_write + + return model_cost_registration_write(current, base, registration) + + +def kept(previous: dict[str, object] | None, registration: Mapping[str, object]) -> dict[str, object]: + """The registration ``register_model(persist_across_reloads=True)`` keeps for a key: ``registration`` merged + into the one kept before it.""" + from litellm.utils import update_dictionary + + return update_dictionary(dict(previous or {}), registration) diff --git a/litellm/utils.py b/litellm/utils.py index 66b5b274f97..5a1765692ae 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -102,7 +102,11 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) -from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap +from litellm.litellm_core_utils.get_model_cost_map import ( + GetModelCostMap, + fetch_registration_document, + model_cost_inputs, +) from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired from litellm.litellm_core_utils.offline_mode import is_offline_mode @@ -114,8 +118,16 @@ from litellm.litellm_core_utils.tool_fields import ( tools_without_hosted_mcp, ) from litellm.rust_bridge import tokenizer as tokenizer_dispatch +from litellm.rust_bridge._native import NativeModelCostMap, NativeRegistered from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision +from litellm.rust_bridge.host.model_cost import ( + AliasedCostEntry, + Landing, + VersionCached, + changed_in_place, + model_cost_version, +) _CachingHandlerResponse = None _LLMCachingHandler = None @@ -3887,22 +3899,6 @@ def _builtin_cost_map() -> dict[str, Mapping[str, object]]: return cast(dict[str, Mapping[str, object]], litellm.model_cost) -def _replace_cost_map_entry( - cost_map: dict[str, Mapping[str, object]], key: str, entry: Mapping[str, object], registration: Mapping[str, object] -) -> None: - current: Final = cost_map.get(key) - merged: Final = {**(current or {}), **entry} - names: Final = current.names if isinstance(current, AliasedCostEntry) else () - previous_fields: Final[frozenset[str]] = ( - current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset() - ) - registered_fields: Final = previous_fields | frozenset( - name for name, value in registration.items() if value is not None - ) - sharing_names: Final = tuple(name for name in names if cost_map.get(name) is current) - cost_map.update(dict.fromkeys((key, *sharing_names), AliasedCostEntry(merged, names, registered_fields))) - - def _without_synthesized_costs( model_info: Mapping[str, object], key: str, value: Mapping[str, object] ) -> Mapping[str, object]: @@ -4143,48 +4139,47 @@ def _skips_builtin_registration_lookup(key: str, provider: object) -> bool: ) -@dataclass(frozen=True, slots=True) -class _RegistrationWrite: - key: str - entry: Mapping[str, object] - - def model_cost_registration_row(key: str, provider: object) -> Mapping[str, object]: """The row a ``register_model`` write of a ``key`` entry with this ``litellm_provider`` merges onto: the ``litellm.model_cost`` row it resolves to (``openai/o3-pro`` lands on ``o3-pro``), else the row it seeds ``key`` with from the built-in entry it falls back to (an OpenAI fine-tune takes its base model's row).""" - base: Final = _registration_base( - key, {"litellm_provider": provider}, custom_llm_provider=None, warning_display_name=None - ) - return _builtin_cost_map().get(base.key) or base.entry + base: Final = registration_base(key, {"litellm_provider": provider}, custom_llm_provider=None) + return _builtin_cost_map().get(base.key) or base.base -def _registration_base( - key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None -) -> _RegistrationWrite: +def registration_base(key: str, value: Mapping[str, object], custom_llm_provider: str | None) -> Landing: provider: Final = str(value.get("litellm_provider") or "") if _skips_builtin_registration_lookup(key, provider): - return _RegistrationWrite(key=key, entry=dict(_builtin_cost_map().get(key, {}))) + return Landing(key=key, base=dict(_builtin_cost_map().get(key, {})), matched_builtin=True) builtin_model_info: Final = _get_builtin_model_info_for_registration( model=key, custom_llm_provider=custom_llm_provider ) if not isinstance(builtin_model_info, ModelInfoLookupMiss): - return _RegistrationWrite( + return Landing( key=builtin_model_info["key"], - entry=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)), + base=dict(_without_synthesized_costs(builtin_model_info, key=key, value=value)), + matched_builtin=True, ) fallback: Final = _fallback_registration(key=key, provider=provider, value=value) - if not fallback.matched_builtin and ( + return Landing(key=key, base=dict(fallback.entry), matched_builtin=fallback.matched_builtin) + + +def _prices_tokens_without_cache_pricing(value: Mapping[str, object]) -> bool: + return ( value.get("cache_creation_input_token_cost") is None and value.get("cache_read_input_token_cost") is None and value.get("tiered_pricing") is None and (value.get("input_cost_per_token") is not None or value.get("output_cost_per_token") is not None) - ): - verbose_logger.warning( - "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info", - warning_display_name or key, - ) - return _RegistrationWrite(key=key, entry=dict(fallback.entry)) + ) + + +def _warn_unpriced_cache(unmatched: Mapping[str, Mapping[str, object]], warning_display_name: str | None) -> None: + for key, registration in unmatched.items(): + if _prices_tokens_without_cache_pricing(registration): + verbose_logger.warning( + "register_model: model=%s has custom pricing but not in built-in cost map and no prefix/region variant matched; cache_creation_input_token_cost and cache_read_input_token_cost will default to 0 for this model (input/output cost tracking is unaffected). To track cache cost, add them to model_info", + warning_display_name or key, + ) def native_structured_output_support(model_info: Mapping[str, object]) -> bool | None: @@ -4203,104 +4198,65 @@ def _base_fields_replaced_by(value: Mapping[str, object]) -> frozenset[str]: ) -def _registration_write( - key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None -) -> _RegistrationWrite: - base: Final = _registration_base(key, value, custom_llm_provider, warning_display_name) +def model_cost_registration_write( + current: Mapping[str, object] | None, base: Mapping[str, object], registration: Mapping[str, object] +) -> AliasedCostEntry: + """The row a ``register_model`` entry writes over ``current``.""" # ``get_model_info`` returns ``litellm_provider: None`` when the # provider is unknown (e.g. custom deployments registered via # ``Router.add_deployment``). Persisting that None into # ``litellm.model_cost`` causes ``_check_provider_match`` to drop # custom pricing on subsequent cost lookups. existing_model: Final = { - name: field_value - for name, field_value in base.entry.items() - if name != "litellm_provider" or field_value is not None + name: field_value for name, field_value in base.items() if name != "litellm_provider" or field_value is not None } - return _RegistrationWrite( - key=base.key, - entry={**update_dictionary(existing_model, value), **dict.fromkeys(_base_fields_replaced_by(value))}, + entry: Final = { + **update_dictionary(existing_model, registration), + **dict.fromkeys(_base_fields_replaced_by(registration)), + } + names: Final = current.names if isinstance(current, AliasedCostEntry) else () + previous_fields: Final[frozenset[str]] = ( + current.registered_fields if isinstance(current, AliasedCostEntry) else frozenset() ) - - -def _registration_attempt( - key: str, - value: Mapping[str, object], - custom_llm_provider: str | None, - warning_display_name: str | None, - final_attempt: bool, -) -> _RegistrationWrite | None: - source: Final = _builtin_cost_map() - write: Final = _registration_write(key, value, custom_llm_provider, warning_display_name) - _replace_cost_map_entry(_builtin_cost_map() if final_attempt else source, write.key, write.entry, value) - return write if final_attempt or litellm.model_cost is source else None + registered_fields: Final = previous_fields | frozenset( + name for name, value in registration.items() if value is not None + ) + return AliasedCostEntry({**(current or {}), **entry}, names, registered_fields) def _register_on_current_map( - key: str, value: Mapping[str, object], custom_llm_provider: str | None, warning_display_name: str | None -) -> _RegistrationWrite: - """A price data reload that swaps the map while the entry is being built or written would otherwise get the old - map's prices written over its own, or drop the entry with the old map, so each attempt writes into the map it was - built from and the entry is rebuilt when a reload replaced that map before the write finished. The last attempt - writes into the current map, so the entry is never lost.""" - attempts: Final = ( - _registration_attempt( - key, value, custom_llm_provider, warning_display_name, final_attempt=attempt == _REGISTRATION_ATTEMPTS - 1 - ) - for attempt in range(_REGISTRATION_ATTEMPTS) - ) - return next(write for write in attempts if write is not None) - - -_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload -_runtime_registration_order: Final[dict[str, int]] = {} # mutable-ok: tells which of a row's registered names came last -_runtime_registered_names: Final[dict[str, tuple[str, ...]]] = {} # mutable-ok: names registered onto each live row -_REGISTRATION_TICKS: Final = itertools.count() + entries: Mapping[str, Mapping[str, object]], persist: bool, custom_llm_provider: str | None +) -> Mapping[str, Mapping[str, object]]: + """Writes ``entries`` into the live table and returns those no built-in row matched.""" + for _ in range(_REGISTRATION_ATTEMPTS - 1): + live = litellm.model_cost + written = model_cost_inputs.register(live, entries, persist, custom_llm_provider) + if isinstance(written, NativeRegistered) and litellm.model_cost is live: + return written.unmatched + last: Final = model_cost_inputs.register(None, entries, persist, custom_llm_provider) + return last.unmatched if isinstance(last, NativeRegistered) else {} @dataclass(frozen=True, slots=True) -class _RegistrationsOverCatalog(Mapping[str, Mapping[str, object]]): - catalog: Mapping[str, Mapping[str, object]] - registrations: Mapping[str, Mapping[str, object]] - registration_order: Mapping[str, int] - registered_names: Mapping[str, tuple[str, ...]] +class _RegisteredCatalog(Mapping[str, Mapping[str, object]]): + inputs: NativeModelCostMap def __getitem__(self, key: str) -> Mapping[str, object]: - catalog_entry: Final = self.catalog.get(key) - names: Final = dict.fromkeys((*self.registered_names.get(key, ()), key)) - present: Final = tuple(name for name in names if name in self.registrations) - if not present: - return self.catalog[key] - oldest_first: Final = sorted(present, key=lambda name: self.registration_order.get(name, -1)) - registered: Final = (self.registrations[name] for name in oldest_first) - return MappingProxyType( - dict(itertools.chain.from_iterable(entry.items() for entry in (catalog_entry or {}, *registered))) - ) + row: Final = self.inputs.registered_row(key) + if row is None: + raise KeyError(key) + return MappingProxyType(row) def __iter__(self) -> Iterator[str]: - return iter({**dict.fromkeys(self.catalog), **dict.fromkeys(self.registrations)}) + return iter(self.inputs.registered_keys()) def __len__(self) -> int: - return len(self.catalog.keys() | self.registrations.keys()) - - -def _record_registered_name(written_key: str, registered_key: str) -> None: - names: Final = _runtime_registered_names.get(written_key, ()) - if registered_key not in names: - _runtime_registered_names[written_key] = (*names, registered_key) + return len(self.inputs.registered_keys()) def catalog_with_runtime_registrations() -> Mapping[str, Mapping[str, object]]: """The loaded catalog plus ``register_model`` calls, without what routers register for their deployments.""" - catalog: Final = GetModelCostMap.loaded_model_cost_map() - if not _runtime_registered_model_cost: - return catalog - return _RegistrationsOverCatalog( - catalog, - MappingProxyType(_runtime_registered_model_cost), - MappingProxyType(_runtime_registration_order), - MappingProxyType(_runtime_registered_names), - ) + return _RegisteredCatalog(model_cost_inputs) def runtime_model_cost_registrations( @@ -4374,15 +4330,9 @@ def reapply_runtime_model_cost_registrations() -> None: _LiveDeploymentReplay.callback() for extra_replay in _LiveDeploymentReplay.extra: extra_replay() - if _runtime_registered_model_cost: - register_model( - model_cost=dict( - sorted( - _runtime_registered_model_cost.items(), - key=lambda registration: _runtime_registration_order.get(registration[0], -1), - ) - ) - ) + reapplied: Final = model_cost_inputs.reapply(litellm.model_cost) + if isinstance(reapplied, NativeRegistered): + _warn_unpriced_cache(reapplied.unmatched, None) def cost_map_omits_token_price(*keys: object) -> bool: @@ -4442,34 +4392,19 @@ def register_model( keys are opaque ids rather than model names (e.g. router deployment ids). """ - loaded_model_cost = {} - if isinstance(model_cost, dict): - # Convert stringified numbers to appropriate numeric types - loaded_model_cost = model_cost - elif isinstance(model_cost, str): - loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1) - - if persist_across_reloads: - _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost - for _registered_key, _registered_value in _registrations.items(): - _runtime_registered_model_cost[_registered_key] = update_dictionary( - dict(_runtime_registered_model_cost.get(_registered_key, {})), - _RATE_TABLE.validate_python(_registered_value), - ) - _runtime_registration_order[_registered_key] = next(_REGISTRATION_TICKS) - - for key, value in loaded_model_cost.items(): - write = _register_on_current_map(str(key), value, custom_llm_provider, warning_display_name) - if persist_across_reloads: - _record_registered_name(write.key, str(key)) - - # Invalidate case-insensitive lookup map since model_cost was modified - _invalidate_model_cost_lowercase_map() - - verbose_logger.debug("added/updated model=%s in litellm.model_cost: %s", write.key, write.key) - if add_to_provider_model_sets and value.get("litellm_provider") == "openrouter": - litellm.openrouter_models.add(str(key).split("/", 1)[-1]) + loaded_model_cost: Final = ( + model_cost if isinstance(model_cost, dict) else (fetch_registration_document(model_cost) or {}) + ) + entries: Final = { + str(key): dict(_RATE_TABLE.validate_python(value) if persist_across_reloads else value) + for key, value in loaded_model_cost.items() + } + unmatched: Final = _register_on_current_map(entries, persist_across_reloads, custom_llm_provider) + _warn_unpriced_cache(unmatched, warning_display_name) if add_to_provider_model_sets: + litellm.openrouter_models.update( + key.split("/", 1)[-1] for key, value in entries.items() if value.get("litellm_provider") == "openrouter" + ) litellm.add_known_models(loaded_model_cost) return model_cost @@ -6789,123 +6724,24 @@ def _strip_model_name(model: str, custom_llm_provider: str | None) -> str: return _strip_dated_snapshot_suffix(model_name=model) -# Global case-insensitive lookup map for model_cost (built eagerly at module import) -_model_cost_lowercase_map: dict[str, str] | None = None +_model_cost_lowercase_map: Final[None] = None -# Monotonic counter bumped on every model_cost mutation. Consumers that -# memoize derived state (e.g. provider-specific indices) can include this -# value in their cache key so they invalidate even when key add+remove or -# in-place value replacement leaves len/id unchanged. -_model_cost_mutation_generation: int = 0 + +def _invalidate_model_cost_lowercase_map() -> None: # pyright: ignore[reportUnusedFunction] # downstream code calls main's name + """Compat for main's callers. Use ``changed_in_place()``.""" + changed_in_place() def get_model_cost_mutation_generation() -> int: - return _model_cost_mutation_generation - - -def _invalidate_model_cost_lowercase_map() -> None: - """Invalidate the case-insensitive lookup map for model_cost. - - Call this whenever litellm.model_cost is modified to ensure the map is rebuilt. - Also clears related LRU caches that depend on model_cost data. - """ - global _model_cost_lowercase_map, _model_cost_mutation_generation - _model_cost_lowercase_map = None - _model_cost_mutation_generation += 1 - - # Clear LRU caches that depend on model_cost data - _cached_get_model_info.cache_clear() - cached_get_model_info_helper.cache_clear() - cached_cost_map_key.cache_clear() - - -def _rebuild_model_cost_lowercase_map() -> dict[str, str]: - """Rebuild the case-insensitive lookup map from the current model_cost. - - Returns: - The rebuilt map (guaranteed to be not None). - """ - global _model_cost_lowercase_map - _model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)} - return _model_cost_lowercase_map - - -def _handle_stale_map_entry_rebuild( - potential_key_lower: str, -) -> str | None: - """ - Handle stale _model_cost_lowercase_map entry (key was popped). - - Rebuilds the map and retries the lookup. - - Returns: - The matched key if found after rebuild, None otherwise. - """ - matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower) - if matched_key is not None and matched_key in litellm.model_cost: - return matched_key - return None - - -def _handle_new_key_with_scan( - potential_key_lower: str, -) -> str | None: - """ - Handle new key added to model_cost without invalidating _model_cost_lowercase_map. - - Scans model_cost for case-insensitive match and rebuilds the map if found. - - Returns: - The matched key if found, None otherwise. - """ - global _model_cost_lowercase_map - for key in tuple(_builtin_cost_map()): - if key.lower() == potential_key_lower: - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() - return key - return None + """Compat for main's callers. Use ``model_cost_version()``.""" + return model_cost_version().writes def get_model_cost_key(potential_key: str) -> str | None: - """ - Get the actual key from model_cost, with case-insensitive fallback. - - WARNING: Only O(1) lookup operations are acceptable. O(n) lookups will cause severe - CPU overhead. This function is called frequently during router operations. - - ALLOWED HELPER FUNCTIONS (conditionally called, O(n) operations are acceptable): - - _rebuild_model_cost_lowercase_map: Rebuilds the lookup map (only when map is None) - - _handle_stale_map_entry_rebuild: Rebuilds map when stale entry detected (rare case) - - If you need to add a new helper function with O(n) operations that is conditionally - called and confirmed not to cause performance issues, add it to the allowed_helpers - list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py - """ - # Exact match (O(1)) + """The key itself when ``litellm.model_cost`` has it, otherwise the key that differs from it only in case.""" if potential_key in litellm.model_cost: return potential_key - - # Case-insensitive lookup via map (O(1)); bound locally because another thread's - # registration can invalidate the module-level map at any point - cached_lowercase_map: Final = _model_cost_lowercase_map - lowercase_map: Final = ( - cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map() - ) - - potential_key_lower: Final = potential_key.lower() - matched_key = lowercase_map.get(potential_key_lower) - - # Verify key exists (O(1) - handles model_cost.pop() case) - if matched_key is not None and matched_key in litellm.model_cost: - return matched_key - - # Rebuild map if stale entry detected (O(n) rebuild, but only when stale entry found) - if matched_key is not None: - matched_key = _handle_stale_map_entry_rebuild(potential_key_lower) - if matched_key is not None: - return matched_key - - return None + return model_cost_inputs.find_key(litellm.model_cost, potential_key) _get_model_cost_key = get_model_cost_key @@ -7169,7 +7005,7 @@ def _get_max_position_embeddings(model_name: str) -> int | None: return None -@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +@VersionCached def cached_get_model_info_helper( model: str, custom_llm_provider: str | None, @@ -7177,11 +7013,6 @@ def cached_get_model_info_helper( *, preserve_unset_costs: bool = False, ) -> ModelInfoBase: - """ - _get_model_info_helper wrapped with lru_cache - - Speed Optimization to hit high RPS - """ return get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, @@ -7193,7 +7024,7 @@ def cached_get_model_info_helper( _cached_get_model_info_helper = cached_get_model_info_helper -@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +@VersionCached def cached_cost_map_key(model: str, custom_llm_provider: str | None) -> str | None: try: return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)["key"] @@ -7753,7 +7584,7 @@ def _build_model_info( return ModelInfo(**_model_info, supported_openai_params=supported_openai_params) -@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +@VersionCached def _cached_get_model_info( model: str, custom_llm_provider: str | None = None, diff --git a/tests/code_coverage_tests/check_get_model_cost_key_performance.py b/tests/code_coverage_tests/check_get_model_cost_key_performance.py index 3ccb9b2469d..90daebc752e 100644 --- a/tests/code_coverage_tests/check_get_model_cost_key_performance.py +++ b/tests/code_coverage_tests/check_get_model_cost_key_performance.py @@ -151,12 +151,7 @@ def check_get_model_cost_key_performance(): # Check for for loops if re.search(r"\bfor\s+\w+\s+in\s+", line): - # Allow helper function calls (they're conditional) - if not re.search( - r"(_rebuild_model_cost_lowercase_map|_handle_stale_map_entry_rebuild|_handle_new_key_with_scan)", - line, - ): - problematic_lines.append((i, "for loop", line_stripped)) + problematic_lines.append((i, "for loop", line_stripped)) # Check for while loops if re.search(r"\bwhile\s+", line): @@ -187,11 +182,7 @@ def check_get_model_cost_key_performance(): # Check for calls to functions that might have O(n) operations # Allow known helper functions that are conditional - allowed_helpers = [ - "_rebuild_model_cost_lowercase_map", - "_handle_stale_map_entry_rebuild", - "_handle_new_key_with_scan", - ] + allowed_helpers: list[str] = [] # Check for function calls (pattern: function_name(...), but not function definitions) # Skip function definitions (def function_name(...)) diff --git a/tests/integration/pricing/test_price_data_reload_prices.py b/tests/integration/pricing/test_price_data_reload_prices.py index aa0addb7673..79a1ab52777 100644 --- a/tests/integration/pricing/test_price_data_reload_prices.py +++ b/tests/integration/pricing/test_price_data_reload_prices.py @@ -5,16 +5,22 @@ import threading import uuid from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor +from contextlib import AbstractContextManager from dataclasses import dataclass from pathlib import Path from types import MappingProxyType -from typing import Final +from typing import Final, TypeAlias import httpx import pytest -from integration._support.client import Gateway, eventually, object_value, string_value -from integration._support.database import read_rows -from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.client import Gateway, eventually, held, object_value, string_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import ( + OwnedProxy, + graceful_stop_seconds, + owned_proxy, + owned_proxy_process, +) from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue @@ -84,6 +90,25 @@ _CATALOG: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( ) +@dataclass(frozen=True, slots=True) +class _Published: + scale: float + added: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Failing: + status: int + + +_Served: TypeAlias = _Published | _Failing + + +class _CostMapServer: + def __init__(self, serving: _Served = _Published(1.0)) -> None: + self.serving: _Served = serving + + @dataclass(frozen=True, slots=True) class _Usage: prompt: int @@ -123,8 +148,9 @@ class _Charge: expected: float -def _catalog(scale: float) -> dict[str, JsonValue]: +def _catalog(scale: float, added: tuple[str, ...] = ()) -> dict[str, JsonValue]: return { + **{model: {"litellm_provider": "openai", "mode": "chat"} for model in added}, **{ model: {name: value * scale if isinstance(value, float) else value for name, value in row.items()} for model, row in _CATALOG.items() @@ -189,8 +215,13 @@ def _connection(model: str, wire_url: str) -> dict[str, JsonValue]: return {"api_base": f"{wire_url}/v1", "api_key": "synthetic-openai-key"} -def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str) -> Path: - config: Final = directory / "price_data_reload_config.yaml" +def _config( + directory: Path, + deployments: tuple[_Deployment, ...], + wire_url: str, + general: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + config: Final = directory / f"price_data_reload_config_{uuid.uuid4().hex}.yaml" config.write_text( json.dumps( { @@ -208,6 +239,7 @@ def _config(directory: Path, deployments: tuple[_Deployment, ...], wire_url: str "general_settings": { "master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL", + **general, }, "router_settings": {"disable_cooldowns": True}, } @@ -276,10 +308,18 @@ def _reply(request: Request) -> Reply: ) -def _responder(published: threading.Event) -> Callable[[Request], Reply]: +def _cost_map_reply(served: _Served) -> Reply: + match served: + case _Published(scale=scale, added=added): + return Reply(body=json.dumps(_catalog(scale, added)).encode()) + case _Failing(status=status): + return Reply(status=status, body=b"{}") + + +def _responder(server: _CostMapServer) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: if request.target == "/model_prices.json": - return Reply(body=json.dumps(_catalog(_PUBLISHED_SCALE if published.is_set() else 1.0)).encode()) + return _cost_map_reply(server.serving) if request.method == "GET": return Reply(body=json.dumps({"object": "list", "data": []}).encode()) return _reply(request) @@ -340,9 +380,9 @@ def _proxy_environment(wire_url: str) -> dict[str, str]: def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments: tuple[_Deployment, ...]) -> None: - published: Final = threading.Event() + server: Final = _CostMapServer() with ( - wire_server(_responder(published)) as wire, + wire_server(_responder(server)) as wire, owned_proxy( gateway, tmp_path, @@ -352,7 +392,7 @@ def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments ) as proxy, ): before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in deployments) - published.set() + server.serving = _Published(_PUBLISHED_SCALE) reload: Final = proxy.request("POST", "/reload/model_cost_map") assert reload.status_code == 200, reload.text after: Final = tuple(_charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in deployments) @@ -361,72 +401,87 @@ def _assert_priced_across_a_reload(gateway: Gateway, tmp_path: Path, deployments @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) -def test_price_data_reload_reprices_catalog_deployments_and_keeps_custom_prices( +def test_custom_prices_on_one_deployment_never_reach_a_sibling_on_the_same_backend_across_a_reload( gateway: Gateway, tmp_path: Path ) -> None: _assert_priced_across_a_reload( gateway, tmp_path, ( - _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)), _Deployment("mini-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates), + _Deployment("mini-catalog", f"openai/{_MINI}", {}, _Usage(1000, 0, 500), _catalog_rates(_MINI)), ), ) +def _catalog_tier_rates(model: str) -> Callable[[float], _Rates]: + return lambda scale: _Rates( + input=_rate(model, "input_cost_per_token_above_200k_tokens", scale), + output=_rate(model, "output_cost_per_token_above_200k_tokens", scale), + cache_read=_rate(model, "cache_read_input_token_cost", scale), + ) + + @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) -@pytest.mark.parametrize( - "deployment", - ( - _Deployment( - "region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW) - ), - _Deployment( - "finetune-unpriced", - f"openai/{_FINETUNE}:org:suffix:y", - {}, - _Usage(2000, 1000, 1000), - _catalog_rates(_FINETUNE), - ), - _Deployment( - "finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates - ), - _Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced), - _Deployment("catalog-free", f"openai/{_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)), - _Deployment("rule-free", f"openai/{_RULE_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)), - _Deployment( - "fill-rule-free", - f"openai/{_FILL_FREE}", - {}, - _Usage(2000, 1000, 1000), - _catalog_rates(_FREE), - listed_rates=_unpriced, - ), - _Deployment("rule-unpriced", f"openai/{_RULE_UNPRICED}", {}, _Usage(2000, 1000, 1000), _unpriced), - _Deployment( - "provider-reported-free", f"openrouter/{_REPORTED_FREE}", {}, _Usage(2000, 1000, 1000), _unpriced, 0.0 - ), - _Deployment( - "cross-provider-free", - f"azure/{_CLAUDE}", - {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}, - _Usage(2000, 1000, 1000), - _catalog_rates(_FREE), - ), - _Deployment( - "region-custom-above-threshold", - f"bedrock/ap-northeast-1.{_ROW}", - _CUSTOM, - _Usage(250_000, 0, 1000), - _custom_rates, - ), - ), - ids=lambda deployment: deployment.name, -) -def test_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload( - gateway: Gateway, tmp_path: Path, deployment: _Deployment +def test_every_deployment_bills_the_catalog_row_its_registration_resolves_to_across_a_reload( + gateway: Gateway, tmp_path: Path ) -> None: - _assert_priced_across_a_reload(gateway, tmp_path, (deployment,)) + _assert_priced_across_a_reload( + gateway, + tmp_path, + ( + _Deployment( + "region-unpriced", f"bedrock/ap-northeast-1.{_ROW}", {}, _Usage(1000, 0, 500), _catalog_rates(_ROW) + ), + _Deployment( + "finetune-unpriced", + f"openai/{_FINETUNE}:org:suffix:y", + {}, + _Usage(2000, 1000, 1000), + _catalog_rates(_FINETUNE), + ), + _Deployment( + "finetune-custom", f"openai/{_FINETUNE}:org:suffix:z", _CUSTOM, _Usage(2000, 0, 1000), _custom_rates + ), + _Deployment("cross-provider-unpriced", f"azure/{_CLAUDE}", {}, _Usage(2000, 1000, 1000), _unpriced), + _Deployment("catalog-free", f"openai/{_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)), + _Deployment("rule-free", f"openai/{_RULE_FREE}", {}, _Usage(2000, 1000, 1000), _catalog_rates(_FREE)), + _Deployment( + "fill-rule-free", + f"openai/{_FILL_FREE}", + {}, + _Usage(2000, 1000, 1000), + _catalog_rates(_FREE), + listed_rates=_unpriced, + ), + _Deployment("rule-unpriced", f"openai/{_RULE_UNPRICED}", {}, _Usage(2000, 1000, 1000), _unpriced), + _Deployment( + "provider-reported-free", f"openrouter/{_REPORTED_FREE}", {}, _Usage(2000, 1000, 1000), _unpriced, 0.0 + ), + _Deployment( + "cross-provider-free", + f"azure/{_CLAUDE}", + {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}, + _Usage(2000, 1000, 1000), + _catalog_rates(_FREE), + ), + _Deployment( + "region-custom-above-threshold", + f"bedrock/ap-northeast-1.{_ROW}", + _CUSTOM, + _Usage(250_000, 0, 1000), + _custom_rates, + ), + _Deployment( + "region-catalog-above-threshold", + f"bedrock/ap-northeast-1.{_ROW}", + {}, + _Usage(250_000, 0, 1000), + _catalog_tier_rates(_ROW), + listed_rates=_catalog_rates(_ROW), + ), + ), + ) def _storm_deployment(index: int) -> _Deployment: @@ -467,7 +522,7 @@ def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_pr geo: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)) finished: Final = threading.Event() with ( - wire_server(_responder(threading.Event())) as wire, + wire_server(_responder(_CostMapServer())) as wire, owned_proxy( gateway, tmp_path, @@ -492,3 +547,153 @@ def test_price_data_reloads_under_custom_priced_traffic_keep_every_deployment_pr settled: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (geo, *storm)) for charge in settled: assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6) + + +_GEO_CATALOG: Final = _Deployment("geo-unpriced", f"bedrock/{_GEO}", {}, _Usage(1000, 0, 500), _catalog_rates(_GEO)) +_POLL_SECONDS: Final = 1 +_LIFECYCLE_SETTINGS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "proxy_config_reload_interval_seconds": _POLL_SECONDS, + } +) + + +def _lifecycle_proxy( + gateway: Gateway, + tmp_path: Path, + wire_url: str, + deployments: tuple[_Deployment, ...], + database: Mapping[str, str] = MappingProxyType({}), +) -> AbstractContextManager[OwnedProxy]: + return owned_proxy_process( + gateway, + tmp_path, + {**_proxy_environment(wire_url), **database}, + config=_config(tmp_path, deployments, wire_url, _LIFECYCLE_SETTINGS), + remove_environment=("LITELLM_LOCAL_MODEL_COST_MAP",), + ) + + +def _bill(proxy: Gateway, deployment: _Deployment) -> float | None: + return _billed(proxy.client, proxy.key, deployment)[1] + + +def _reload(proxy: Gateway) -> httpx.Response: + return proxy.request("POST", "/reload/model_cost_map") + + +def _billed_at(deployment: _Deployment, scale: float) -> Callable[[float | None], bool]: + expected: Final = _expected_cost(deployment.rates(scale), deployment.usage) + return lambda charged: charged == pytest.approx(expected, rel=1e-6) + + +@pytest.mark.timeout(4 * graceful_stop_seconds() + 180) +def test_price_data_changes_only_through_a_successful_reload_or_a_restart(gateway: Gateway, tmp_path: Path) -> None: + server: Final = _CostMapServer() + with wire_server(_responder(server)) as wire, scratch_database("reload_lifecycle") as database_url: + database: Final = {"DATABASE_URL": database_url} + with _lifecycle_proxy(gateway, tmp_path, wire.url, (_GEO_CATALOG,), database=database) as first: + _charged_as_listed(first.gateway, _GEO_CATALOG, 1.0) + server.serving = _Published(_PUBLISHED_SCALE) + held( + lambda: _bill(first.gateway, _GEO_CATALOG), + _billed_at(_GEO_CATALOG, 1.0), + holding=3 * _POLL_SECONDS, + ) + server.serving = _Failing(404) + failed: Final = _reload(first.gateway) + assert failed.status_code == 502, failed.text + _charged_as_listed(first.gateway, _GEO_CATALOG, 1.0) + server.serving = _Published(_PUBLISHED_SCALE) + with _lifecycle_proxy(gateway, tmp_path, wire.url, (_GEO_CATALOG,), database=database) as restarted: + _charged_as_listed(restarted.gateway, _GEO_CATALOG, _PUBLISHED_SCALE) + + +def _add_model(proxy: Gateway, deployment: _Deployment, wire_url: str) -> None: + created: Final = proxy.post( + "/model/new", + { + "model_name": deployment.name, + "litellm_params": { + "model": deployment.model, + **_connection(deployment.model, wire_url), + **deployment.prices, + }, + "model_info": {}, + }, + ) + assert object_value(created["model_info"])["id"], created + + +def _custom_with_catalog_cache(model: str) -> Callable[[float], _Rates]: + return lambda scale: _Rates( + input=_CUSTOM["input_cost_per_token"], + output=_CUSTOM["output_cost_per_token"], + cache_read=_rate(model, "cache_read_input_token_cost", scale), + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) +def test_a_model_added_through_the_api_keeps_its_prices_and_inherits_the_reloaded_ones( + gateway: Gateway, tmp_path: Path +) -> None: + added: Final = _Deployment( + "added-finetune-custom", + f"openai/{_FINETUNE}", + _CUSTOM, + _Usage(2000, 1000, 1000), + _custom_with_catalog_cache(_FINETUNE), + ) + server: Final = _CostMapServer() + with ( + wire_server(_responder(server)) as wire, + _lifecycle_proxy(gateway, tmp_path, wire.url, (_GEO_CATALOG,)) as owned, + ): + proxy: Final = owned.gateway + _add_model(proxy, added, wire.url) + before: Final = tuple(_charged_as_listed(proxy, deployment, 1.0) for deployment in (_GEO_CATALOG, added)) + server.serving = _Published(_PUBLISHED_SCALE) + reload: Final = _reload(proxy) + assert reload.status_code == 200, reload.text + after: Final = tuple( + _charged_as_listed(proxy, deployment, _PUBLISHED_SCALE) for deployment in (_GEO_CATALOG, added) + ) + for charge in (*before, *after): + assert _spend(charge.request_id) == pytest.approx(charge.expected, rel=1e-6) + + +_WILDCARD: Final = _Deployment("openai/*", "openai/*", {}, _Usage(1000, 0, 500), _unpriced) +_LISTED_AFTER_RELOAD: Final = "gpt-audit-reload-listed" + + +def _listed_models(proxy: Gateway) -> list[str]: + data: Final = proxy.get("/v1/models")["data"] + assert isinstance(data, list), data + return [string_value(object_value(entry)["id"]) for entry in data] + + +def _deployment_ids(proxy: Gateway) -> list[str]: + entries: Final = proxy.get("/model/info")["data"] + assert isinstance(entries, list), entries + return [string_value(object_value(object_value(entry)["model_info"])["id"]) for entry in entries] + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 180) +def test_a_wildcard_lists_the_models_a_reload_adds_and_never_a_deployment_id(gateway: Gateway, tmp_path: Path) -> None: + custom: Final = _Deployment("listed-custom", f"openai/{_MINI}", _CUSTOM, _Usage(1000, 0, 500), _custom_rates) + server: Final = _CostMapServer() + with ( + wire_server(_responder(server)) as wire, + _lifecycle_proxy(gateway, tmp_path, wire.url, (_WILDCARD, custom)) as owned, + ): + proxy: Final = owned.gateway + _charged_as_listed(proxy, custom, 1.0) + assert f"openai/{_LISTED_AFTER_RELOAD}" not in _listed_models(proxy) + server.serving = _Published(1.0, added=(_LISTED_AFTER_RELOAD,)) + reload: Final = _reload(proxy) + assert reload.status_code == 200, reload.text + _charged_as_listed(proxy, custom, 1.0) + listed: Final = _listed_models(proxy) + deployment_ids: Final = _deployment_ids(proxy) + assert f"openai/{_LISTED_AFTER_RELOAD}" in listed, listed + assert [model for model in listed if any(id_ in model for id_ in deployment_ids)] == [], listed diff --git a/tests/integration/sdk/_model_cost_sdk_child.py b/tests/integration/sdk/_model_cost_sdk_child.py new file mode 100644 index 00000000000..fa86fa80f2e --- /dev/null +++ b/tests/integration/sdk/_model_cost_sdk_child.py @@ -0,0 +1,203 @@ +"""One SDK session per case, in a fresh interpreter booted from the cost map at LITELLM_MODEL_COST_MAP_URL. + +Each case does what SDK code does to ``litellm.model_cost`` and bills a real ``completion`` after every step, then +prints the charges as one JSON line. The steps use only names main exposes, so a case runs on main and the branch. +""" + +import asyncio +import json +import os +import sys +from collections.abc import Callable +from typing import Final + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import ( + ModelCostMapReloadUnavailable, + _finalize_model_cost_map, # pyright: ignore[reportPrivateUsage] # downstream code calls main's finalizer + adopt_model_cost_map, + refetch_model_cost_map, +) +from litellm.utils import ( + _cached_get_model_info_helper, # pyright: ignore[reportPrivateUsage] # downstream code clears it + _invalidate_model_cost_lowercase_map, # pyright: ignore[reportPrivateUsage] # downstream code calls it +) + +MODEL: Final = "lit-sdk-chat" +ALIAS: Final = "lit-sdk-chat-alias" +ADDED: Final = "lit-sdk-added" +FINETUNE: Final = f"ft:{MODEL}:org:custom:id" +CUSTOM_INPUT: Final = 9e-06 +CUSTOM_OUTPUT: Final = 7e-06 +BASE: Final = os.environ["MODEL_COST_SDK_BASE_URL"] +DOCUMENT_P2: Final = os.environ["MODEL_COST_SDK_DOCUMENT"] + + +def bill(model: str) -> float | None: + response: Final = litellm.completion( + model=f"openai/{model}", + api_base=f"{BASE}/v1", + api_key="sk-sdk-probe", + messages=[{"role": "user", "content": "bill"}], + ) + try: + return float(litellm.completion_cost(completion_response=response)) + except Exception: # noqa: BLE001 # an unpriced model is an observation here, not a failure + return None + + +def reload(document: str) -> bool: + result: Final = asyncio.run(refetch_model_cost_map(url=f"{BASE}/{document}")) + if isinstance(result, ModelCostMapReloadUnavailable): + return False + adopt_model_cost_map(result.model_cost_map) + return True + + +def row(scale: float) -> dict[str, object]: + return { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06 * scale, + "output_cost_per_token": 2e-06 * scale, + "cache_read_input_token_cost": 5e-07 * scale, + } + + +def reload_reprices() -> list[object]: + return [bill(MODEL), reload("p2.json"), bill(MODEL)] + + +def failed_reload_keeps_prices() -> list[object]: + return [bill(MODEL), reload("missing.json"), bill(MODEL)] + + +def billed_and_listed(model: str) -> list[object]: + info: Final = litellm.get_model_info(f"openai/{model}") + return [bill(model), info["input_cost_per_token"], info["output_cost_per_token"]] + + +def every_name_bills_and_lists_the_row_across_a_reload() -> list[object]: + names: Final = (MODEL, ALIAS, MODEL.upper()) + before: Final = [value for name in names for value in billed_and_listed(name)] + return [*before, reload("p2.json"), *(value for name in names for value in billed_and_listed(name))] + + +def registration_survives_reload() -> list[object]: + litellm.register_model({ADDED: {**row(3.0)}}) + litellm.register_model({"lit-sdk-request-scoped": row(3.0)}, persist_across_reloads=False) + before: Final = [bill(ADDED), bill("lit-sdk-request-scoped")] + return [*before, reload("p2.json"), bill(ADDED), bill("lit-sdk-request-scoped")] + + +def partial_registration_follows_reload() -> list[object]: + litellm.register_model({MODEL: {"input_cost_per_token": CUSTOM_INPUT}}) + return [bill(MODEL), reload("p2.json"), bill(MODEL)] + + +def partial_registrations_merge_and_follow_reload() -> list[object]: + litellm.register_model({MODEL: {"input_cost_per_token": CUSTOM_INPUT}}) + litellm.register_model({MODEL: {"output_cost_per_token": CUSTOM_OUTPUT}}) + return [bill(MODEL), reload("p2.json"), bill(MODEL)] + + +def finetune_registration_bills_the_same_after_a_reload() -> list[object]: + litellm.register_model({FINETUNE: {"input_cost_per_token": CUSTOM_INPUT, "litellm_provider": "openai"}}) + before: Final = bill(FINETUNE) + return [before is not None, reload("p1.json"), bill(FINETUNE) == before] + + +def finetune_registration_inherits_the_base_prices() -> list[object]: + litellm.register_model({FINETUNE: {"input_cost_per_token": CUSTOM_INPUT, "litellm_provider": "openai"}}) + return [bill(FINETUNE)] + + +def new_row_after_a_miss_without_invalidating() -> list[object]: + before: Final = bill(ADDED) + litellm.model_cost[ADDED] = row(3.0) + return [before, bill(ADDED)] + + +def value_edited_in_place_then_invalidate() -> list[object]: + before: Final = bill(MODEL) + litellm.model_cost[MODEL]["input_cost_per_token"] = CUSTOM_INPUT + _invalidate_model_cost_lowercase_map() + return [before, bill(MODEL)] + + +def table_swapped_then_invalidate() -> list[object]: + before: Final = bill(MODEL) + litellm.model_cost = {**litellm.model_cost, MODEL: row(2.0)} + _invalidate_model_cost_lowercase_map() + return [before, bill(MODEL)] + + +def table_swapped_then_helper_cache_cleared() -> list[object]: + before: Final = bill(MODEL) + litellm.model_cost = {**litellm.model_cost, MODEL: row(2.0)} + _cached_get_model_info_helper.cache_clear() + litellm.get_model_info.cache_clear() + return [before, bill(MODEL)] + + +def contents_replaced_then_invalidate() -> list[object]: + before: Final = bill(MODEL) + saved: Final = {key: dict(value) for key, value in litellm.model_cost.items()} + litellm.model_cost.clear() + litellm.model_cost.update({**saved, MODEL: row(2.0)}) + _invalidate_model_cost_lowercase_map() + return [before, bill(MODEL)] + + +def finalized_snapshot_replaces_contents() -> list[object]: + with open(DOCUMENT_P2) as handle: + finalized: Final = _finalize_model_cost_map(json.load(handle)) + litellm.model_cost.clear() + litellm.model_cost.update(finalized) + litellm.add_known_models(finalized) + return [bill(MODEL)] + + +def alias_registration_after_an_adopted_reload() -> list[object]: + reloaded: Final = reload("p2.json") + litellm.register_model({ALIAS: {"input_cost_per_token": CUSTOM_INPUT}}) + return [reloaded, bill(MODEL), bill(ALIAS)] + + +def small_map_finalizes_and_leaves_the_live_map() -> list[object]: + finalized: Final = _finalize_model_cost_map({ADDED: row(3.0)}) + return [list(finalized) == [ADDED], bill(MODEL)] + + +def popped_row_then_registered() -> list[object]: + before: Final = bill(MODEL) + litellm.model_cost.pop(MODEL) + litellm.register_model({MODEL: row(3.0)}) + return [before, bill(MODEL)] + + +CASES: Final[dict[str, Callable[[], list[object]]]] = { + case.__name__: case + for case in ( + reload_reprices, + failed_reload_keeps_prices, + every_name_bills_and_lists_the_row_across_a_reload, + registration_survives_reload, + partial_registration_follows_reload, + partial_registrations_merge_and_follow_reload, + finetune_registration_bills_the_same_after_a_reload, + finetune_registration_inherits_the_base_prices, + new_row_after_a_miss_without_invalidating, + value_edited_in_place_then_invalidate, + table_swapped_then_invalidate, + table_swapped_then_helper_cache_cleared, + contents_replaced_then_invalidate, + finalized_snapshot_replaces_contents, + alias_registration_after_an_adopted_reload, + small_map_finalizes_and_leaves_the_live_map, + popped_row_then_registered, + ) +} + +if __name__ == "__main__": + sys.stdout.write(json.dumps(CASES[sys.argv[1]]()) + "\n") diff --git a/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py index 1c419e5b11c..921552455b7 100644 --- a/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py +++ b/tests/integration/sdk/test_bedrock_native_output_register_model_sdk.py @@ -10,7 +10,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm import ModelResponse, Router -from litellm.utils import _invalidate_model_cost_lowercase_map # pyright: ignore[reportPrivateUsage] # required cache invalidation +from litellm.rust_bridge.host.model_cost import publish _OPUS_4_8: Final = "us.anthropic.claude-opus-4-8" _JSON_TOOL: Final = "json_tool_call" @@ -87,7 +87,7 @@ def opus_4_8_opted_in_for_the_process(monkeypatch: pytest.MonkeyPatch) -> Iterat ) yield monkeypatch.undo() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _model(route: str) -> str: diff --git a/tests/integration/sdk/test_model_cost_fetch_sdk.py b/tests/integration/sdk/test_model_cost_fetch_sdk.py new file mode 100644 index 00000000000..206ffa99909 --- /dev/null +++ b/tests/integration/sdk/test_model_cost_fetch_sdk.py @@ -0,0 +1,594 @@ +import asyncio +import hashlib +import json +import logging +import os +import ssl +import sys +import threading +from collections.abc import Generator, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field +from pathlib import Path +from typing import Final, Literal +from urllib.parse import urlsplit + +import pytest +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, wire_server + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import ( + GetModelCostMap, + ModelCostMapReloaded, + ModelCostMapReloadUnavailable, + get_model_cost_map_loaded_at, + get_model_cost_map_provenance, + get_model_cost_map_source_info, + model_cost_inputs, + refetch_model_cost_map, +) +from litellm.rust_bridge.host.model_cost import publish +from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + +FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_BUNDLED: Final = GetModelCostMap.read_local_model_cost_map_bytes() +_IMPORTED_MODEL_COST_INPUTS: Final = model_cost_inputs.snapshot() + + +def _blob_id(body: bytes) -> str: + return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest() + + +def _document_with(**rows: Mapping[str, object]) -> bytes: + """The bundled document with some rows replaced, so it passes the integrity limits.""" + return json.dumps({**json.loads(_BUNDLED), **rows}).encode() + + +@dataclass(frozen=True, slots=True) +class _Reply: + status: int = 200 + body: bytes = _BUNDLED + headers: Mapping[str, str] = field(default_factory=dict) + + +@dataclass(frozen=True, slots=True) +class _Server: + url: str + received: Sequence[Request] + + @property + def targets(self) -> Sequence[str]: + return [request.target for request in self.received] + + @property + def authorizations(self) -> Sequence[str]: + return [request.headers.get("authorization", "") for request in self.received] + + @property + def requests(self) -> int: + return len(self.received) + + +@contextmanager +def _serving(*replies: _Reply, tls: ssl.SSLContext | None = None) -> Generator[_Server]: + """Answers each GET with the next scripted reply, repeating the last one.""" + received: Final[list[Request]] = [] # mutable-ok: respond records each request the peer answers + + def respond(request: Request) -> Reply: + received.append(request) + reply: Final = replies[min(len(received), len(replies)) - 1] + return Reply(status=reply.status, body=reply.body, headers=reply.headers) + + with wire_server(respond, tls=tls) as wire: + yield _Server(url=f"{wire.url}/cost.json", received=received) + + +@pytest.fixture(autouse=True) +def isolated_model_cost() -> Iterator[None]: + inputs: Final = model_cost_inputs.snapshot() + rows: Final = litellm.model_cost + entries: Final = dict(rows) + model_cost_inputs.restore(_IMPORTED_MODEL_COST_INPUTS) + publish(rows) + yield + model_cost_inputs.restore(inputs) + rows.clear() + rows.update(entries) + publish(rows) + + +@pytest.fixture +def remote_source(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr(sys, "argv", ["python"]) + + +@pytest.fixture +def own_model_cost(monkeypatch: pytest.MonkeyPatch) -> None: + """Restores the binding of ``litellm.model_cost`` after a test that installs a new table.""" + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + + +@pytest.fixture +def isolated_openai_model_sets(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "open_ai_chat_completion_models", set(litellm.open_ai_chat_completion_models)) + monkeypatch.setattr(litellm, "open_ai_text_completion_models", set(litellm.open_ai_text_completion_models)) + + +@pytest.mark.usefixtures("remote_source") +def test_boot_load_installs_the_served_document_and_records_its_provenance(): + body: Final = _document_with(**{"lit-boot-model": {"litellm_provider": "openai", "mode": "chat"}}) + with _serving(_Reply(body=body, headers={"ETag": '"v1"'})) as server: + rows: Final = litellm.get_model_cost_map(url=server.url) + + assert "lit-boot-model" in rows + assert FALLBACK_GENERALIZATIONS_KEY not in rows + assert get_model_cost_map_provenance() == {"source_revision": _blob_id(body), "etag": '"v1"'} + source: Final = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["fallback_reason"]) == ("remote", server.url, None) + + +@pytest.mark.usefixtures("remote_source") +@pytest.mark.parametrize( + ("reply", "fallback_reason"), + [ + (_Reply(status=404, body=b"missing"), "Remote fetch failed: HTTP 404 from "), + (_Reply(body=b"[1, 2]"), "Remote fetch failed: invalid JSON from "), + ( + _Reply(body=json.dumps({"only": {"litellm_provider": "openai"}}).encode()), + "Remote data failed integrity validation", + ), + ], +) +def test_boot_load_falls_back_to_the_bundled_document(reply: _Reply, fallback_reason: str): + with _serving(reply) as server: + rows: Final = litellm.get_model_cost_map(url=server.url, max_attempts=1) + + assert "gpt-4o" in rows + assert get_model_cost_map_provenance() == {"source_revision": _blob_id(_BUNDLED), "etag": None} + source: Final = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["is_env_forced"]) == ("local", server.url, False) + assert source["fallback_reason"] is not None and source["fallback_reason"].startswith(fallback_reason) + + +_MOVED: Final = _Reply(status=302, body=b"", headers={"Location": "/moved.json"}) +_REDIRECTED_DOCUMENT: Final = _document_with(**{"lit-redirected-model": {"litellm_provider": "openai", "mode": "chat"}}) + + +def _rows_loaded_through(load: Literal["boot", "reload", "register_model"], url: str) -> Mapping[str, object]: + """The rows ``load`` produced from ``url``: the boot load's return value, or the live table after the others.""" + match load: + case "boot": + return litellm.get_model_cost_map(url=url, max_attempts=1) + case "reload": + assert isinstance(asyncio.run(refetch_model_cost_map(url=url, max_attempts=1)), ModelCostMapReloaded) + return litellm.model_cost + case "register_model": + litellm.register_model(url, persist_across_reloads=False) + return litellm.model_cost + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +@pytest.mark.parametrize("load", ["boot", "reload", "register_model"]) +def test_every_load_follows_a_redirect_to_the_document(load: Literal["boot", "reload", "register_model"]): + with _serving(_MOVED, _Reply(body=_REDIRECTED_DOCUMENT)) as server: + rows: Final = _rows_loaded_through(load, server.url) + + assert server.targets == ["/cost.json", "/moved.json"] + assert "lit-redirected-model" in rows + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_boot_load_transient_failure_serves_the_bundled_map_then_installs_the_remote_one_in_the_background(): + body: Final = _document_with(**{"lit-retried-model": {"litellm_provider": "openai", "mode": "chat"}}) + with _serving(_Reply(status=503, body=b"busy", headers={"Retry-After": "0"}), _Reply(body=body)) as server: + rows: Final = litellm.get_model_cost_map(url=server.url) + assert "lit-retried-model" not in rows + booted_at: Final = get_model_cost_map_loaded_at() + retry: Final = next(t for t in threading.enumerate() if t.name == "litellm-model-cost-map-retry") + retry.join(timeout=30) + + assert server.requests == 2 + assert "lit-retried-model" in litellm.model_cost + source: Final = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["is_env_forced"], source["fallback_reason"]) == ( + "remote", + server.url, + False, + None, + ) + retried_at: Final = get_model_cost_map_loaded_at() + assert booted_at is not None and retried_at is not None and retried_at > booted_at + + +@pytest.mark.usefixtures("remote_source") +def test_boot_load_does_not_retry_a_final_failure(): + with _serving(_Reply(status=404, body=b"missing")) as server: + litellm.get_model_cost_map(url=server.url) + + assert not any(t.name == "litellm-model-cost-map-retry" for t in threading.enumerate()) + assert server.requests == 1 + + +@pytest.mark.parametrize( + ("argv0", "environment"), + [ + ("/usr/bin/litellm-proxy", {}), + ("/usr/bin/lite", {}), + ("python", {"LITELLM_LOCAL_MODEL_COST_MAP": "True"}), + ("python", {"LITELLM_OFFLINE_MODE": "true"}), + ], + ids=["proxy-cli", "lite-cli", "local-map", "offline"], +) +def test_boot_load_serves_the_bundled_map_without_fetching( + monkeypatch: pytest.MonkeyPatch, argv0: str, environment: Mapping[str, str] +): + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr(sys, "argv", [argv0]) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + with _serving(_Reply()) as server: + rows: Final = litellm.get_model_cost_map(url=server.url) + + assert server.requests == 0 + assert "gpt-4o" in rows + assert get_model_cost_map_provenance() == {"source_revision": _blob_id(_BUNDLED), "etag": None} + assert get_model_cost_map_source_info()["is_env_forced"] is True + + +def test_importing_litellm_loads_the_map_litellm_model_cost_map_url_names(): + script: Final = """ +import os +import litellm +assert litellm.model_cost_map_url == os.environ["LITELLM_MODEL_COST_MAP_URL"], litellm.model_cost_map_url +assert "lit-env-url-model" in litellm.model_cost +""" + body: Final = _document_with(**{"lit-env-url-model": {"litellm_provider": "openai", "mode": "chat"}}) + environment: Final = { + key: value + for key, value in os.environ.items() + if key not in {"LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"} + } + with _serving(_Reply(body=body)) as server: + result: Final = run_child_interpreter( + script, env={**environment, "LITELLM_MODEL_COST_MAP_URL": server.url}, timeout=120 + ) + + assert result.returncode == 0, result.stderr + assert server.targets == ["/cost.json"] + + +@pytest.mark.skipif(sys.platform != "linux", reason="fork is the default start method only on Linux") +def test_importing_litellm_with_a_remote_source_leaves_the_process_free_to_fork(): + script: Final = """ +import os +import litellm +from litellm.rust_bridge._native import process_state_started, reserve_process_for_forking +assert "lit-fork-model" in litellm.model_cost +assert not process_state_started() +reserve_process_for_forking() +pid = os.fork() +if pid == 0: + os._exit(0 if "lit-fork-model" in litellm.model_cost else 1) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +""" + body: Final = _document_with(**{"lit-fork-model": {"litellm_provider": "openai", "mode": "chat"}}) + environment: Final = {key: value for key, value in os.environ.items() if key != "LITELLM_LOCAL_MODEL_COST_MAP"} + with _serving(_Reply(body=body)) as server: + result: Final = run_child_interpreter( + script, env={**environment, "LITELLM_MODEL_COST_MAP_URL": server.url}, timeout=120 + ) + + assert result.returncode == 0, result.stderr + + +def test_register_model_leaves_the_native_runtime_unstarted() -> None: + script: Final = """ +import litellm +from litellm.rust_bridge._native import process_state_started +litellm.register_model({"lit-fork-safe-model": {"input_cost_per_token": 1e-06, "litellm_provider": "openai"}}) +assert litellm.model_cost["lit-fork-safe-model"]["input_cost_per_token"] == 1e-06 +assert not process_state_started() +""" + result: Final = run_child_interpreter( + script, env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, timeout=120 + ) + + assert result.returncode == 0, result.stderr + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_reload_retries_a_429_after_its_retry_after(): + with _serving(_Reply(status=429, body=b"slow", headers={"Retry-After": "0"}), _Reply()) as server: + result: Final = asyncio.run(refetch_model_cost_map(url=server.url)) + + assert isinstance(result, ModelCostMapReloaded) + assert server.requests == 2 + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +@pytest.mark.parametrize( + ("reply", "reason", "requests"), + [ + (_Reply(status=503, body=b"busy", headers={"Retry-After": "0"}), "HTTP 503 from {url} (after 3 attempts)", 3), + (_Reply(body=b"not json"), "invalid JSON from {url}", 1), + ( + _Reply(body=json.dumps({"only": {"litellm_provider": "openai"}}).encode()), + "model cost map from {url} failed integrity validation", + 1, + ), + ], + ids=["gives-up-after-its-attempts", "not-json", "below-integrity-limits"], +) +def test_a_failed_reload_keeps_the_installed_map(reply: _Reply, reason: str, requests: int): + installed: Final = litellm.model_cost + provenance: Final = get_model_cost_map_provenance() + with _serving(reply) as server: + result: Final = asyncio.run(refetch_model_cost_map(url=server.url, max_attempts=3)) + + assert isinstance(result, ModelCostMapReloadUnavailable) + assert result.reason.startswith(reason.format(url=server.url)), result.reason + assert server.requests == requests + assert litellm.model_cost is installed + assert get_model_cost_map_provenance() == provenance + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_a_reload_rejected_by_the_integrity_limits_logs_the_counts(caplog: pytest.LogCaptureFixture): + with ( + _serving(_Reply(body=json.dumps({"only": {"litellm_provider": "openai"}}).encode())) as server, + caplog.at_level(logging.WARNING, logger="LiteLLM"), + ): + asyncio.run(refetch_model_cost_map(url=server.url)) + + assert "has 1 models, below minimum" in caplog.text + + +@pytest.mark.usefixtures("own_model_cost") +@pytest.mark.parametrize("argv0", ["/usr/bin/litellm-proxy", "/usr/bin/lite"]) +def test_a_cli_process_that_booted_on_the_bundled_map_still_fetches_on_reload( + monkeypatch: pytest.MonkeyPatch, argv0: str +): + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr(sys, "argv", [argv0]) + with _serving(_Reply()) as server: + litellm.get_model_cost_map(url=server.url) + booted_requests: Final = server.requests + result: Final = asyncio.run(refetch_model_cost_map(url=server.url)) + + assert booted_requests == 0 + assert isinstance(result, ModelCostMapReloaded) + assert server.requests == 1 + assert get_model_cost_map_source_info()["source"] == "remote" + + +@pytest.mark.usefixtures("own_model_cost") +def test_reload_with_the_local_override_installs_the_bundled_map_without_fetching(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + with _serving(_Reply()) as server: + result: Final = asyncio.run(refetch_model_cost_map(url=server.url)) + + assert isinstance(result, ModelCostMapReloaded) + assert server.requests == 0 + assert get_model_cost_map_provenance() == {"source_revision": _blob_id(_BUNDLED), "etag": None} + source: Final = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["is_env_forced"]) == ("local", None, True) + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_only_a_reload_that_installs_moves_the_source_info_and_load_time(): + with _serving(_Reply()) as first: + asyncio.run(refetch_model_cost_map(url=first.url)) + installed: Final = get_model_cost_map_source_info() + loaded_at: Final = get_model_cost_map_loaded_at() + + with _serving(_Reply(status=404, body=b"missing")) as failing: + asyncio.run(refetch_model_cost_map(url=failing.url)) + assert get_model_cost_map_source_info() == installed + assert get_model_cost_map_loaded_at() is loaded_at + + with _serving(_Reply()) as second: + asyncio.run(refetch_model_cost_map(url=second.url)) + source: Final = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["is_env_forced"], source["fallback_reason"]) == ( + "remote", + second.url, + False, + None, + ) + assert get_model_cost_map_loaded_at() is not loaded_at + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_reload_records_the_served_bytes_and_etag(): + body: Final = _document_with(**{"lit-revision-model": {"litellm_provider": "openai", "mode": "chat"}}) + with _serving(_Reply(body=body, headers={"ETag": 'W/"etag"'})) as server: + asyncio.run(refetch_model_cost_map(url=server.url)) + + assert get_model_cost_map_provenance() == {"source_revision": _blob_id(body), "etag": 'W/"etag"'} + + +def _self_signed(directory: Path) -> tuple[ssl.SSLContext, Path]: + """A server context for 127.0.0.1 and the path of its certificate, which nothing trusts by default.""" + certificate, key = write_self_signed_cert(directory) + return server_context(certificate, key), certificate + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_a_reload_rejects_an_untrusted_certificate_by_default(tmp_path: Path): + context, _ = _self_signed(tmp_path) + with _serving(_Reply(), tls=context) as server: + result: Final = asyncio.run(refetch_model_cost_map(url=server.url, max_attempts=1)) + + assert isinstance(result, ModelCostMapReloadUnavailable) + assert result.reason.startswith(f"error fetching {server.url}: ") + assert server.requests == 0 + + +@pytest.fixture +def no_tls_environment(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("SSL_VERIFY", "SSL_CERT_FILE", "SSL_CERTIFICATE"): + monkeypatch.delenv(name, raising=False) + + +@pytest.mark.usefixtures("remote_source", "no_tls_environment", "own_model_cost") +@pytest.mark.parametrize( + ("load", "setting", "trust"), + [ + ("reload", "litellm.ssl_verify", "off"), + ("reload", "litellm.ssl_verify", "certificate"), + ("boot", "SSL_VERIFY", "off"), + ("boot", "SSL_CERT_FILE", "certificate"), + ], +) +def test_a_load_trusts_the_certificate_its_settings_allow( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + load: Literal["boot", "reload"], + setting: str, + trust: Literal["off", "certificate"], +): + context, certificate = _self_signed(tmp_path) + if setting == "litellm.ssl_verify": + monkeypatch.setattr(litellm, "ssl_verify", {"off": False, "certificate": str(certificate)}[trust]) + else: + monkeypatch.setenv(setting, {"off": "False", "certificate": str(certificate)}[trust]) + with _serving(_Reply(), tls=context) as server: + _rows_loaded_through(load, server.url) + + assert server.requests == 1 + assert get_model_cost_map_source_info()["source"] == "remote" + + +@pytest.mark.usefixtures("remote_source", "no_tls_environment") +def test_a_boot_load_ignores_litellm_ssl_verify_because_it_runs_before_settings_exist( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + context, _ = _self_signed(tmp_path) + monkeypatch.setattr(litellm, "ssl_verify", False) + with _serving(_Reply(), tls=context) as server: + litellm.get_model_cost_map(url=server.url, max_attempts=1) + + assert server.requests == 0 + source: Final = get_model_cost_map_source_info() + assert source["source"] == "local" + assert source["fallback_reason"] is not None + assert source["fallback_reason"].startswith(f"Remote fetch failed: error fetching {server.url}: ") + + +_UNREACHABLE_ENDPOINT: Final = "http://127.0.0.1:9" + + +@pytest.mark.usefixtures("remote_source") +@pytest.mark.parametrize( + ("environment", "endpoint_variable", "region"), + [ + ( + {"AWS_REGION_NAME": "eu-west-3", "AWS_REGION": "ap-south-1", "AWS_DEFAULT_REGION": "sa-east-1"}, + "AWS_ENDPOINT_URL_S3", + "eu-west-3", + ), + ({"AWS_REGION": "ap-south-1", "AWS_DEFAULT_REGION": "sa-east-1"}, "AWS_ENDPOINT_URL_S3", "ap-south-1"), + ({"AWS_DEFAULT_REGION": "sa-east-1"}, "AWS_ENDPOINT_URL_S3", "sa-east-1"), + ({"AWS_ENDPOINT_URL": _UNREACHABLE_ENDPOINT}, "AWS_ENDPOINT_URL_S3", "us-east-1"), + ({}, "AWS_ENDPOINT_URL", "us-east-1"), + ({"AWS_CONFIG_FILE": "profile_region"}, "AWS_ENDPOINT_URL_S3", "eu-north-1"), + ], +) +def test_an_s3_load_signs_for_the_region_and_reaches_the_endpoint_the_environment_names( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + environment: Mapping[str, str], + endpoint_variable: str, + region: str, +): + for name in ( + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", + "AWS_PROFILE", + "AWS_SESSION_TOKEN", + "AWS_ENDPOINT_URL_S3", + "AWS_ENDPOINT_URL", + ): + monkeypatch.delenv(name, raising=False) + profile: Final = tmp_path / "profile_region" + profile.write_text("[default]\nregion = eu-north-1\n") + monkeypatch.setenv("AWS_CONFIG_FILE", str(tmp_path / "no_config")) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "no_credentials")) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIDMODELCOST") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "model-cost-secret") + for name, value in environment.items(): + monkeypatch.setenv(name, str(tmp_path / value) if name == "AWS_CONFIG_FILE" else value) + with _serving(_Reply()) as server: + monkeypatch.setenv(endpoint_variable, server.url.removesuffix("/cost.json")) + litellm.get_model_cost_map(url="s3://lit-prices/maps/cost.json", max_attempts=1) + + assert [urlsplit(target).path for target in server.targets] == ["/lit-prices/maps/cost.json"] + assert f"/{region}/s3/aws4_request" in server.authorizations[0] + assert get_model_cost_map_source_info()["source"] == "remote" + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +@pytest.mark.parametrize("load", ["boot", "reload"]) +def test_both_loads_fetch_through_the_proxy_the_environment_names( + monkeypatch: pytest.MonkeyPatch, load: Literal["boot", "reload"] +): + for name in ("http_proxy", "HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy", "NO_PROXY", "no_proxy"): + monkeypatch.delenv(name, raising=False) + url: Final = "http://lit-cost-map.invalid/cost.json" + with _serving(_Reply()) as proxy: + monkeypatch.setenv("HTTP_PROXY", proxy.url.removesuffix("/cost.json")) + rows: Final = _rows_loaded_through(load, url) + + assert "gpt-4o" in rows + assert get_model_cost_map_source_info()["source"] == "remote" + assert proxy.targets == [url] + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +def test_register_model_from_a_url_registers_the_served_rows_without_the_fallback_rules(): + document: Final = { + "lit-url-model": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 4e-06}, + FALLBACK_GENERALIZATIONS_KEY: {"rules": []}, + } + with _serving(_Reply(body=json.dumps(document).encode())) as server: + litellm.register_model(server.url) + + assert server.requests == 1 + assert litellm.model_cost["lit-url-model"]["input_cost_per_token"] == 4e-06 + assert FALLBACK_GENERALIZATIONS_KEY not in litellm.model_cost + assert model_cost_inputs.registered_keys() == [*GetModelCostMap.loaded_model_cost_map(), "lit-url-model"] + + +@pytest.mark.usefixtures("remote_source", "own_model_cost") +@pytest.mark.parametrize( + "reply", [_Reply(status=500, body=b"down"), _Reply(body=b"not json")], ids=["server_error", "invalid_json"] +) +def test_register_model_from_a_url_that_fails_registers_nothing(reply: _Reply): + with _serving(reply) as server: + litellm.register_model(server.url) + + assert model_cost_inputs.registered_keys() == list(GetModelCostMap.loaded_model_cost_map()) + + +@pytest.mark.usefixtures("isolated_openai_model_sets") +def test_register_model_url_fetch_uses_single_attempt(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + before: Final = dict(litellm.model_cost) + with wire_server(lambda _: Reply(status=503, body=b"")) as peer: + threads_before: Final = {thread.name for thread in threading.enumerate()} + litellm.register_model(model_cost=f"{peer.url}/custom_pricing.json") + + threads_after: Final = {thread.name for thread in threading.enumerate()} + requests: Final = [request.target for request in peer.drain()] + assert requests == ["/custom_pricing.json"] + assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"} + assert not any( + thread.name == "litellm-model-cost-map-retry" and thread.is_alive() for thread in threading.enumerate() + ) + assert litellm.model_cost.keys() >= before.keys() diff --git a/tests/integration/sdk/test_model_cost_sdk.py b/tests/integration/sdk/test_model_cost_sdk.py new file mode 100644 index 00000000000..6515e84de22 --- /dev/null +++ b/tests/integration/sdk/test_model_cost_sdk.py @@ -0,0 +1,188 @@ +import json +import os +import subprocess +import sys +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from importlib.resources import files +from pathlib import Path +from typing import Final + +import pytest +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_MODEL: Final = "lit-sdk-chat" +_ALIAS: Final = "lit-sdk-chat-alias" +_PROMPT: Final = 1000 +_CACHED: Final = 400 +_COMPLETION: Final = 500 +_CUSTOM_INPUT: Final = 9e-06 +_CUSTOM_OUTPUT: Final = 7e-06 +_CHARGES: Final = TypeAdapter(list[float | bool | None]) +_BUNDLED: Final = TypeAdapter(dict[str, JsonValue]).validate_json( + files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes() +) + + +class _Requested(BaseModel): + model_config = ConfigDict(frozen=True) + + model: str + + +_REQUEST: Final = TypeAdapter(_Requested) + + +@dataclass(frozen=True, slots=True) +class _Rates: + input: float + output: float + cache_read: float + + def cost(self) -> float: + return (_PROMPT - _CACHED) * self.input + _CACHED * self.cache_read + _COMPLETION * self.output + + +def _catalog(scale: float) -> _Rates: + return _Rates(input=1e-06 * scale, output=2e-06 * scale, cache_read=5e-07 * scale) + + +def _document(scale: float) -> bytes: + rates: Final = _catalog(scale) + return json.dumps( + { + **_BUNDLED, + _MODEL: { + "aliases": [_ALIAS], + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": rates.input, + "output_cost_per_token": rates.output, + "cache_read_input_token_cost": rates.cache_read, + }, + } + ).encode() + + +def _respond(request: Request) -> Reply: + match request.target: + case "/p1.json": + return Reply(body=_document(1.0)) + case "/p2.json": + return Reply(body=_document(2.0)) + case "/missing.json": + return Reply(status=404, body=b"{}") + case _: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": _REQUEST.validate_json(request.body).model, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "billed"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": _PROMPT, + "completion_tokens": _COMPLETION, + "total_tokens": _PROMPT + _COMPLETION, + "prompt_tokens_details": {"cached_tokens": _CACHED}, + }, + } + ).encode() + ) + + +@pytest.fixture(scope="module") +def cost_map_peer() -> Iterator[Wire]: + with wire_server(_respond) as peer: + yield peer + + +def _run(case: str, peer: Wire, directory: Path) -> list[float | bool | None]: + document: Final = directory / "p2.json" + document.write_bytes(_document(2.0)) + environment: Final = { + **{name: value for name, value in os.environ.items() if name != "LITELLM_LOCAL_MODEL_COST_MAP"}, + "LITELLM_MODEL_COST_MAP_URL": f"{peer.url}/p1.json", + "MODEL_COST_SDK_BASE_URL": peer.url, + "MODEL_COST_SDK_DOCUMENT": str(document), + } + result: Final = subprocess.run( # test-quality-ok: each case is one SDK session in its own interpreter + (sys.executable, "-m", "integration.sdk._model_cost_sdk_child", case), + env=environment, + capture_output=True, + text=True, + timeout=180, + check=False, + ) + assert result.returncode == 0, result.stderr[-4000:] + return _CHARGES.validate_json(result.stdout.strip().splitlines()[-1]) + + +def _custom_input(rates: _Rates) -> _Rates: + return _Rates(input=_CUSTOM_INPUT, output=rates.output, cache_read=rates.cache_read) + + +def _custom_input_and_output(rates: _Rates) -> _Rates: + return _Rates(input=_CUSTOM_INPUT, output=_CUSTOM_OUTPUT, cache_read=rates.cache_read) + + +def _charged(rates: _Rates) -> object: + return pytest.approx(rates.cost(), rel=1e-9) + + +def _billed_and_listed(rates: _Rates) -> list[object]: + return [_charged(rates), pytest.approx(rates.input, rel=1e-9), pytest.approx(rates.output, rel=1e-9)] * 3 + + +_P1: Final = _catalog(1.0) +_P2: Final = _catalog(2.0) +_P3: Final = _catalog(3.0) +_FINETUNE_BUG: Final = pytest.mark.xfail( + strict=True, + reason="BUG: a fine-tune registered with only an input price bills output at 0 and cache reads at input", +) + + +@pytest.mark.parametrize( + ("case", "expected"), + [ + ("reload_reprices", [_charged(_P1), True, _charged(_P2)]), + ("failed_reload_keeps_prices", [_charged(_P1), False, _charged(_P1)]), + ( + "every_name_bills_and_lists_the_row_across_a_reload", + [*_billed_and_listed(_P1), True, *_billed_and_listed(_P2)], + ), + ("registration_survives_reload", [_charged(_P3), _charged(_P3), True, _charged(_P3), None]), + ("partial_registration_follows_reload", [_charged(_custom_input(_P1)), True, _charged(_custom_input(_P2))]), + ( + "partial_registrations_merge_and_follow_reload", + [_charged(_custom_input_and_output(_P1)), True, _charged(_custom_input_and_output(_P2))], + ), + ("finetune_registration_bills_the_same_after_a_reload", [True, True, True]), + pytest.param( + "finetune_registration_inherits_the_base_prices", [_charged(_custom_input(_P1))], marks=_FINETUNE_BUG + ), + ("new_row_after_a_miss_without_invalidating", [None, _charged(_P3)]), + ("value_edited_in_place_then_invalidate", [_charged(_P1), _charged(_custom_input(_P1))]), + ("table_swapped_then_invalidate", [_charged(_P1), _charged(_P2)]), + ("table_swapped_then_helper_cache_cleared", [_charged(_P1), _charged(_P2)]), + ("contents_replaced_then_invalidate", [_charged(_P1), _charged(_P2)]), + ("finalized_snapshot_replaces_contents", [_charged(_P2)]), + ( + "alias_registration_after_an_adopted_reload", + [True, _charged(_custom_input(_P2)), _charged(_custom_input(_P2))], + ), + ("small_map_finalizes_and_leaves_the_live_map", [True, _charged(_P1)]), + ("popped_row_then_registered", [_charged(_P1), _charged(_P3)]), + ], + ids=lambda value: value if isinstance(value, str) else "", +) +def test_sdk_code_that_reads_or_writes_the_cost_map_gets_what_it_set_up( + case: str, expected: list[object], cost_map_peer: Wire, tmp_path: Path +) -> None: + assert _run(case, cost_map_peer, tmp_path) == expected diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index 465ba2f9150..be0c5e311da 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -28,7 +28,7 @@ from litellm.llms.custom_httpx.http_handler import( HTTPHandler, ) from litellm.types.utils import LlmProviders -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats @@ -408,7 +408,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 21772593947..fe21cb23caf 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -134,11 +134,14 @@ os.environ["PYTHON_DOTENV_DISABLED"] = "1" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import -import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency +from litellm.rust_bridge.host.model_cost import publish # noqa: E402 # same import-time dependency from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency from litellm.anthropic_beta_headers_manager import ( reload_beta_headers_config, # noqa: E402 # same import-time dependency ) +from litellm.litellm_core_utils.get_model_cost_map import ( # noqa: E402 # same import-time dependency + model_cost_inputs, +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER # noqa: E402 # same import-time dependency from litellm.litellm_core_utils.prompt_templates import ( factory as prompt_factory_module, # noqa: E402 # same import-time dependency @@ -430,22 +433,21 @@ def isolate_litellm_globals() -> Iterator[None]: tool_registry_writer_module._tool_policy_registry = original_tool_policy_registry +_IMPORTED_MODEL_COST_INPUTS: Final = model_cost_inputs.snapshot() + + @pytest.fixture(autouse=True) def isolate_router_model_cost_state() -> Iterator[None]: original_live_registries: Final = frozenset(live_router_registries) - original_runtime_registered_model_cost: Final = { - model_key: dict(model_value) - for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() - } - litellm_utils_module._invalidate_model_cost_lowercase_map() + original_model_cost_inputs: Final = model_cost_inputs.snapshot() + model_cost_inputs.restore(_IMPORTED_MODEL_COST_INPUTS) + publish(litellm.model_cost) live_router_registries.clear() - litellm_utils_module._runtime_registered_model_cost.clear() yield live_router_registries.clear() live_router_registries.update(original_live_registries) - litellm_utils_module._runtime_registered_model_cost.clear() - litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) - litellm_utils_module._invalidate_model_cost_lowercase_map() + model_cost_inputs.restore(original_model_cost_inputs) + publish(litellm.model_cost) litellm.get_model_info.cache_clear() @@ -469,7 +471,7 @@ def isolate_model_catalog() -> Iterator[None]: for models, entries in provider_model_sets: models.clear() models.update(entries) - litellm_utils_module._invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) litellm.get_model_info.cache_clear() diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 0e2813e0a0d..297039fb771 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -26,7 +26,7 @@ import litellm from litellm.constants import AWS_SIGNING_MAX_THREADS from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded, refetch_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded, install_model_cost_map from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR @@ -3963,13 +3963,8 @@ CatalogLoader: TypeAlias = Callable[[Mapping[str, Mapping[str, object]]], Awaita def load_catalog_with(monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None) -> Iterator[CatalogLoader]: async def load(rows: Mapping[str, Mapping[str, object]]) -> None: catalog: Final = {**litellm.get_model_cost_map(url=""), **rows} - transport: Final = httpx.MockTransport(lambda _: httpx.Response(200, json=catalog)) - monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP") - async with httpx.AsyncClient(transport=transport) as client: - loaded: Final = await refetch_model_cost_map(url="https://catalog.example/prices.json", client=client) - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - assert isinstance(loaded, ModelCostMapReloaded) - monkeypatch.setattr(litellm, "model_cost", loaded.model_cost_map) + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + assert isinstance(install_model_cost_map(json.dumps(catalog).encode()), ModelCostMapReloaded) litellm.get_model_info.cache_clear() yield load diff --git a/tests/unit/integrations/arize/test_arize.py b/tests/unit/integrations/arize/test_arize.py index 6accfdc9f93..42dd5e01ce9 100644 --- a/tests/unit/integrations/arize/test_arize.py +++ b/tests/unit/integrations/arize/test_arize.py @@ -16,7 +16,7 @@ from litellm.integrations.arize.arize import ArizeConfig, ArizeLogger from litellm.integrations.opentelemetry import OpenTelemetryConfig from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -346,7 +346,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/integrations/arize/test_arize_phoenix.py b/tests/unit/integrations/arize/test_arize_phoenix.py index 5edf6e6d379..73e6b62ec4c 100644 --- a/tests/unit/integrations/arize/test_arize_phoenix.py +++ b/tests/unit/integrations/arize/test_arize_phoenix.py @@ -17,7 +17,7 @@ from litellm.integrations.arize.arize_phoenix import ( ) from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -1044,7 +1044,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/integrations/opik/test_opik.py b/tests/unit/integrations/opik/test_opik.py index f7f5e252ed4..0bf35ffdabf 100644 --- a/tests/unit/integrations/opik/test_opik.py +++ b/tests/unit/integrations/opik/test_opik.py @@ -12,7 +12,7 @@ import pytest import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -341,7 +341,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 766038d5c4f..4b6f59a78ab 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -19,6 +19,7 @@ from openai.types.chat import ChatCompletionMessage from pydantic import BaseModel, ConfigDict import litellm +from litellm.rust_bridge.host.model_cost import publish from litellm import Router from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( @@ -4353,9 +4354,9 @@ class TestPromptCacheBreakpointCapability: bundled = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(bundled) as handle: monkeypatch.setattr(litellm, "model_cost", json.load(handle)) - litellm.utils.cached_get_model_info_helper.cache_clear() + publish(litellm.model_cost) yield - litellm.utils.cached_get_model_info_helper.cache_clear() + publish(litellm.model_cost) def test_listed_model_uses_the_model_map_flag(self, monkeypatch): flagged = {**litellm.model_cost["gpt-4.1"], "supports_prompt_cache_breakpoint": True} @@ -4587,9 +4588,9 @@ class TestBedrockMantleGptShipsTheOpenAIDialect: bundled = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json") with open(bundled) as handle: monkeypatch.setattr(litellm, "model_cost", json.load(handle)) - litellm.utils.cached_get_model_info_helper.cache_clear() + publish(litellm.model_cost) yield - litellm.utils.cached_get_model_info_helper.cache_clear() + publish(litellm.model_cost) def test_shipped_entry_makes_the_deployment_eligible(self): assert supports_openai_prompt_cache_breakpoint(self.MANTLE_MODEL) is True diff --git a/tests/unit/integrations/test_helicone.py b/tests/unit/integrations/test_helicone.py index a4c23a35b40..90fc2f97ee8 100644 --- a/tests/unit/integrations/test_helicone.py +++ b/tests/unit/integrations/test_helicone.py @@ -7,7 +7,7 @@ from collections.abc import Iterator from pathlib import Path from typing import Final from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from unittest.mock import MagicMock @@ -157,7 +157,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 3270cad80c9..fafc098ed5d 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -83,10 +83,10 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse, TextCompletionResponse from litellm.utils import ( Rules, - _invalidate_model_cost_lowercase_map, function_setup, validate_and_fix_openai_messages, ) +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}] @@ -6190,7 +6190,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index b4e3ef2c566..ea270fbf7c8 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -21,7 +21,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases from litellm.types.router import LiteLLM_Params -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from unittest.mock import patch @@ -209,7 +209,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 6267b841a95..b295a432a77 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -1,167 +1,54 @@ """ -Tests for model-cost-map loading: the model-count integrity check (which must -count actual model entries, not reserved meta keys) and the extraction of the -``fallback_generalizations`` block out of the raw map. +The Python host of the native model cost map: installing a document as ``litellm.model_cost``, the +``register_model`` overlay across reloads, and the lookups over the dict. """ -import asyncio, importlib, importlib.resources as importlib_get_model, json, litellm -import os -import sys +import asyncio +import json +import logging +import re import threading -from pathlib import Path +from collections.abc import Mapping from typing import Final -import boto3 import pytest -from pydantic import ValidationError -from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.litellm_core_utils import get_model_cost_map as get_model_cost_map_module +import litellm from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, - match_capability_generalizations, match_routing_generalization, set_fallback_generalizations, ) from litellm.litellm_core_utils.get_model_cost_map import ( - FALLBACK_GENERALIZATIONS_KEY, + AliasedCostEntry, GetModelCostMap, - _count_model_entries, - _finalize_model_cost_map, + ModelCostMapReloaded, + ModelCostMapReloadUnavailable, get_model_cost_map_provenance, - git_blob_id, + install_model_cost_map, + refetch_model_cost_map, ) +from litellm.rust_bridge._native import NativeInstalled, NativeModelCostMap, NativeRegistered +from litellm.utils import catalog_with_runtime_registrations, get_model_cost_key, runtime_model_cost_registrations + +FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" +_BUNDLED: Final = GetModelCostMap.read_local_model_cost_map_bytes() -def _load_root_cost_map() -> dict: - path = os.path.join(os.path.dirname(__file__), "../../../model_prices_and_context_window.json") - with open(path) as f: - return json.load(f) +def _document_with(**rows: Mapping[str, object]) -> bytes: + """The bundled document with some rows replaced, so it passes the integrity limits.""" + return json.dumps({**json.loads(_BUNDLED), **rows}).encode() -def _bundled_blob_id() -> str: - path = os.path.join(os.path.dirname(__file__), "../../../litellm/model_prices_and_context_window_backup.json") - with open(path, "rb") as f: - return git_blob_id(f.read()) +@pytest.fixture +def own_model_cost(monkeypatch: pytest.MonkeyPatch) -> None: + """Restores the binding of ``litellm.model_cost`` after a test that installs a new table.""" + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) -def test_git_blob_id_is_what_git_hash_object_prints(): - assert git_blob_id(b'{"gpt-5.4-mini": {"mode": "chat"}}\n') == "18b9a8381e13a3b38a2128f184f631f95829e987" - - -def _make_models(n: int) -> dict: - return {f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)} - - -def test_count_model_entries_excludes_reserved_keys(): - m = _make_models(3) - m["sample_spec"] = {"foo": "bar"} - m[FALLBACK_GENERALIZATIONS_KEY] = {"rules": []} - assert _count_model_entries(m) == 3 - - -def test_validation_rejects_truly_shrunk_file_even_with_meta_keys(): - """A file with only a handful of real models must be rejected as corrupt, - and the extra meta keys must not inflate the count past the minimum.""" - shrunk = _make_models(5) - shrunk["sample_spec"] = {"foo": "bar"} - shrunk[FALLBACK_GENERALIZATIONS_KEY] = {"rules": [{"name": "x"}]} - - assert ( - GetModelCostMap.validate_model_cost_map( - fetched_map=shrunk, - backup_model_count=2000, - min_model_count=50, - ) - is False - ) - - -def test_validation_accepts_healthy_file_with_meta_keys(): - healthy = _make_models(2000) - healthy["sample_spec"] = {"foo": "bar"} - healthy[FALLBACK_GENERALIZATIONS_KEY] = {"rules": []} - - assert ( - GetModelCostMap.validate_model_cost_map( - fetched_map=healthy, - backup_model_count=2000, - min_model_count=50, - ) - is True - ) - - -def test_validation_rejects_significant_shrink_vs_backup(): - # 600 real models vs a 2000-model backup is below the 50% shrink threshold. - shrunk = _make_models(600) - shrunk[FALLBACK_GENERALIZATIONS_KEY] = {"rules": []} - assert ( - GetModelCostMap.validate_model_cost_map( - fetched_map=shrunk, - backup_model_count=2000, - min_model_count=50, - max_shrink_ratio=0.5, - ) - is False - ) - - -def test_finalize_pops_key_and_installs_rules(): - previous = list(get_fallback_generalization_rules()) - try: - raw = _make_models(2) - raw[FALLBACK_GENERALIZATIONS_KEY] = { - "rules": [ - { - "name": "rule", - "pattern": r"^widget-", - "model_info": {"litellm_provider": "openai"}, - } - ] - } - finalized = _finalize_model_cost_map(raw) - - # The reserved key is removed from the returned model map ... - assert FALLBACK_GENERALIZATIONS_KEY not in finalized - # ... and its rules are installed into the generalizations module. - assert match_routing_generalization("widget-9") == "openai" - finally: - set_fallback_generalizations(previous) - - -def test_alias_of_rows_are_served_as_their_target_and_dangling_ones_are_dropped( - caplog: pytest.LogCaptureFixture, -) -> None: - raw: Final = { - "real": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 1e-6}, - "other": {"litellm_provider": "openai", "mode": "chat", "aliases": ["twin-latest"]}, - "real-latest": {"alias_of": "real"}, - "twin-latest": {"alias_of": "real"}, - "ghost-latest": {"alias_of": "ghost"}, - "chain-latest": {"alias_of": "real-latest"}, - "listed-latest": {"alias_of": "missing", "aliases": ["extra"]}, - } - with caplog.at_level("WARNING"): - finalized: Final = _finalize_model_cost_map(raw) - assert finalized["real-latest"] == finalized["real"] - assert "alias_of" not in finalized["real"] - assert set(finalized) == {"real", "other", "real-latest", "twin-latest"} - assert finalized["real-latest"]["litellm_provider"] == "openai" - assert finalized["twin-latest"] == finalized["other"] - assert {"ghost-latest", "chain-latest", "listed-latest", "twin-latest"} == { - record.getMessage().split("'")[1] for record in caplog.records - } - - -def test_finalize_with_no_block_clears_rules(): - previous = list(get_fallback_generalization_rules()) - try: - set_fallback_generalizations([{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}]) - _finalize_model_cost_map(_make_models(2)) - assert match_capability_generalizations("x-1") is None - finally: - set_fallback_generalizations(previous) +# --------------------------------------------------------------------------- +# The shipped document +# --------------------------------------------------------------------------- def test_shipped_backup_carries_the_claude_routing_rules(): @@ -253,1053 +140,420 @@ def test_shipped_backup_marks_claude_4_6_plus_adaptive_not_4_0(): assert "supports_adaptive_thinking" not in backup[non_adaptive], non_adaptive -def test_get_model_cost_map_stamps_loaded_at(): - """The load time feeds each pod's reload-due decision; a load that does not stamp it - would make manual reload requests race the proxy's startup""" - from datetime import datetime, timezone - - from litellm.litellm_core_utils import get_model_cost_map as module - - client, _calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) - - before = datetime.now(timezone.utc) - module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) - loaded_at = module.get_model_cost_map_loaded_at() - - assert loaded_at is not None - assert before <= loaded_at <= datetime.now(timezone.utc) - - # --------------------------------------------------------------------------- -# refetch_model_cost_map: retry/backoff behavior for runtime reloads +# Runtime reload # --------------------------------------------------------------------------- -import functools -import random -from datetime import datetime, timezone - -import httpx - -from litellm.litellm_core_utils.get_model_cost_map import ( - ModelCostMapReloaded, - ModelCostMapReloadUnavailable, - refetch_model_cost_map, -) - -_URL = "https://example.invalid/model_prices.json" - - -@functools.lru_cache(maxsize=1) -def _real_map_bytes() -> bytes: - return json.dumps(_load_root_cost_map()).encode() - - -class _SleepRecorder: - """Injected in place of asyncio.sleep so tests assert waits without real delay.""" - - def __init__(self): - self.waits = [] - - async def __call__(self, seconds: float) -> None: - self.waits.append(seconds) - - -@pytest.fixture(autouse=True) -def _unset_local_cost_map_env(monkeypatch): - """CI exports LITELLM_LOCAL_MODEL_COST_MAP=True; clear it so fetch behavior is deterministic.""" - monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) - - -def _mock_client(outcomes, client_cls=httpx.AsyncClient): - """httpx client over a MockTransport serving one outcome per request; an exception instance is raised.""" - calls = {"count": 0} - - def handler(request): - idx = min(calls["count"], len(outcomes) - 1) - calls["count"] += 1 - outcome = outcomes[idx] - if isinstance(outcome, Exception): - raise outcome - return outcome - - return client_cls(transport=httpx.MockTransport(handler)), calls - - -@pytest.mark.asyncio -async def test_refetch_retries_429_honoring_retry_after(): - """Two 429s with Retry-After then success: waits follow the header, not backoff.""" - client, calls = _mock_client( - [ - httpx.Response(429, headers={"Retry-After": "7"}), - httpx.Response(429, headers={"Retry-After": "7"}), - httpx.Response(200, content=_real_map_bytes()), - ] - ) - sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloaded) - assert len(result.model_cost_map) > 100 - assert calls["count"] == 3 - assert sleeper.waits == [7.0, 7.0] - - -@pytest.mark.asyncio -async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff(): - """All 429 without Retry-After: exponential backoff waits, then a failure value.""" - client, calls = _mock_client([httpx.Response(429)]) - sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloadUnavailable) - assert "429" in result.reason - assert "after 3 attempts" in result.reason - assert calls["count"] == 3 - assert len(sleeper.waits) == 2 - assert 2.0 <= sleeper.waits[0] < 3.0 - assert 4.0 <= sleeper.waits[1] < 5.0 - - -@pytest.mark.asyncio -async def test_refetch_caps_retry_after_wait(): - """A hostile/huge Retry-After is capped so reloads never sleep unbounded.""" - client, _calls = _mock_client( - [ - httpx.Response(429, headers={"Retry-After": "9999"}), - httpx.Response(200, content=_real_map_bytes()), - ] - ) - sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloaded) - assert sleeper.waits == [30.0] - - -@pytest.mark.asyncio -async def test_refetch_retries_transport_errors(): - """Connection failures are transient: retried like 5xx, succeeding when the network heals.""" - client, calls = _mock_client( - [ - httpx.ConnectError("connection refused"), - httpx.Response(200, content=_real_map_bytes()), - ] - ) - sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloaded) - assert calls["count"] == 2 - assert len(sleeper.waits) == 1 - - -@pytest.mark.asyncio -async def test_refetch_non_retryable_status_fails_immediately(): - """A 404 is permanent: one attempt, no sleeps, failure value.""" - client, calls = _mock_client([httpx.Response(404)]) - sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloadUnavailable) - assert "404" in result.reason - assert calls["count"] == 1 - assert sleeper.waits == [] - - -@pytest.mark.asyncio -async def test_refetch_invalid_json_fails_immediately(): - client, calls = _mock_client([httpx.Response(200, content=b"not json")]) - sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloadUnavailable) - assert "invalid JSON" in result.reason - assert calls["count"] == 1 - assert sleeper.waits == [] - - -@pytest.mark.asyncio -async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(): - """A drastically shrunk upstream file is rejected instead of being adopted.""" - tiny = json.dumps(_make_models(60)).encode() - client, _calls = _mock_client([httpx.Response(200, content=tiny)]) - result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) - assert isinstance(result, ModelCostMapReloadUnavailable) - assert "integrity validation" in result.reason - - -@pytest.mark.asyncio -@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"]) -async def test_refetch_respects_local_env_override(monkeypatch, flag): - """LITELLM_LOCAL_MODEL_COST_MAP=True or LITELLM_OFFLINE_MODE=True short-circuits to the bundled backup, zero HTTP.""" - monkeypatch.setenv(flag, "True") - - def _fail(request): - raise AssertionError("no HTTP request should be made when local map is forced") - - result = await refetch_model_cost_map( - url=_URL, - sleep=_SleepRecorder(), - rng=random.Random(0), - client=httpx.AsyncClient(transport=httpx.MockTransport(_fail)), - ) - assert isinstance(result, ModelCostMapReloaded) - assert len(result.model_cost_map) > 100 - - -@pytest.mark.asyncio -async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(): - body = _real_map_bytes() - client, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"abc123"'}, content=body)]) - - result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) - - assert isinstance(result, ModelCostMapReloaded) - assert result.revision == git_blob_id(body) - assert result.etag == 'W/"abc123"' - assert get_model_cost_map_provenance() == {"source_revision": git_blob_id(body), "etag": 'W/"abc123"'} - - -@pytest.mark.asyncio -async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later_registrations(monkeypatch): - import litellm - - edited = json.loads(_real_map_bytes()) - edited["gpt-5.4-mini"]["max_input_tokens"] = 777 - client, _ = _mock_client([httpx.Response(200, content=json.dumps(edited).encode())]) - - result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) - - assert isinstance(result, ModelCostMapReloaded) - monkeypatch.setattr(litellm, "model_cost", result.model_cost_map) - litellm.register_model({"gpt-5.4-mini": {"max_input_tokens": 2048}}, persist_across_reloads=False) - assert litellm.model_cost["gpt-5.4-mini"]["max_input_tokens"] == 2048 - assert GetModelCostMap.loaded_model_cost_map()["gpt-5.4-mini"]["max_input_tokens"] == 777 - - -@pytest.mark.asyncio -async def test_loaded_catalog_fields_follow_each_fetched_map(): - added = json.loads(_real_map_bytes()) - added["gpt-5.4-mini"]["supports_litellm_qa_capability"] = True - with_field, _ = _mock_client([httpx.Response(200, content=json.dumps(added).encode())]) - without_field, _ = _mock_client([httpx.Response(200, content=_real_map_bytes())]) - - await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=with_field) - fields_after_adding: Final = GetModelCostMap.loaded_model_cost_map_fields() - await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=without_field) - - assert "supports_litellm_qa_capability" in fields_after_adding - assert "supports_litellm_qa_capability" not in GetModelCostMap.loaded_model_cost_map_fields() - assert "max_input_tokens" in GetModelCostMap.loaded_model_cost_map_fields() - - -@pytest.mark.asyncio -async def test_refetch_revision_follows_the_bytes_not_the_url(): - edited = json.loads(_real_map_bytes()) - edited["gpt-5.4-mini"]["input_cost_per_token"] = 0.5 - client, _ = _mock_client( - [httpx.Response(200, content=_real_map_bytes()), httpx.Response(200, content=json.dumps(edited).encode())] - ) - - first = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) - second = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) - - assert isinstance(first, ModelCostMapReloaded) and isinstance(second, ModelCostMapReloaded) - assert first.revision != second.revision - assert get_model_cost_map_provenance()["source_revision"] == second.revision - - -@pytest.mark.asyncio -async def test_refetch_local_override_reports_the_bundled_blob_id_without_an_etag(monkeypatch): - remote, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"remote"'}, content=_real_map_bytes())]) - await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=remote) - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - - result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0)) - - assert isinstance(result, ModelCostMapReloaded) - assert result.revision == _bundled_blob_id() - assert get_model_cost_map_provenance() == {"source_revision": _bundled_blob_id(), "etag": None} - - -@pytest.mark.asyncio -async def test_refetch_stamps_loaded_at_on_remote_and_local_reloads(monkeypatch): - from litellm.litellm_core_utils import get_model_cost_map as module - - client, _ = _mock_client([httpx.Response(200, content=_real_map_bytes())]) - before_remote = datetime.now(timezone.utc) - await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) - remote_loaded_at = module.get_model_cost_map_loaded_at() - assert remote_loaded_at is not None - assert before_remote <= remote_loaded_at <= datetime.now(timezone.utc) - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - before_local = datetime.now(timezone.utc) - await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0)) - local_loaded_at = module.get_model_cost_map_loaded_at() - assert local_loaded_at is not None - assert before_local <= local_loaded_at <= datetime.now(timezone.utc) - - -# --------------------------------------------------------------------------- -# get_model_cost_map: the boot-time load retries transient failures like a reload does -# --------------------------------------------------------------------------- - -from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map, - get_model_cost_map_source_info, -) -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map -from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from unittest.mock import MagicMock, patch - - -class _FakeS3Body: - def __init__(self, content: bytes) -> None: - self.content: Final = content - - def read(self) -> bytes: - return self.content - - -class _FakeS3GetObjectResponse(TypedDict): - Body: ReadOnly[_FakeS3Body] - ETag: NotRequired[ReadOnly[str]] - - -class _FakeS3Client: - def __init__(self, content: bytes) -> None: - self.content: Final = content - self.calls: list[tuple[str, str]] = [] - - def get_object(self, *, Bucket: str, Key: str) -> _FakeS3GetObjectResponse: - self.calls.append((Bucket, Key)) - return {"Body": _FakeS3Body(self.content), "ETag": '"fake-etag"'} - - -def test_get_model_cost_map_fetches_s3_catalog(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.litellm_core_utils import get_model_cost_map as module - - source_info = module._cost_map_source_info - for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): - monkeypatch.setattr(source_info, name, getattr(source_info, name)) - - s3_content: Final = _make_models(GetModelCostMap.get_backup_model_count()) - s3_client: Final = _FakeS3Client(json.dumps(s3_content).encode()) - monkeypatch.setattr(boto3, "client", lambda *args, **kwargs: s3_client) - - cost_map: Final = get_model_cost_map("s3://my-bucket/prices/model_prices.json") - - assert cost_map == s3_content - assert s3_client.calls == [("my-bucket", "prices/model_prices.json")] - assert get_model_cost_map_source_info()["source"] == "remote" - - -class _SyncSleepRecorder: - """Injected in place of time.sleep so the boot path's waits are asserted without delay.""" - - def __init__(self, block=False): - self.waits = [] - self.block = block - self.started = threading.Event() - self.release = threading.Event() - - def __call__(self, seconds: float) -> None: - if self.block: - self.started.set() - self.release.wait(timeout=10) - self.waits.append(seconds) - - -def _retry_threads(): - return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"] - - -def test_boot_load_success_does_not_start_background_retry(): - client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - - cost_map = get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - ) - assert calls["count"] == 1 - assert sleeper.waits == [] - assert _retry_threads() == [] - assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} - assert get_model_cost_map_source_info()["source"] == "remote" - - -def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch): - import litellm - from litellm import utils as litellm_utils - from litellm.litellm_core_utils import get_model_cost_map as module - - original_model_cost = litellm.model_cost - monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost)) - for name, provider_models in tuple(vars(litellm).items()): - if name.endswith("_models") and isinstance(provider_models, set): - monkeypatch.setattr(litellm, name, set(provider_models)) - monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider)) - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) - source_info = module._cost_map_source_info - for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): - monkeypatch.setattr(source_info, name, getattr(source_info, name)) - - remote_map = _load_root_cost_map() - remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} - client, calls = _mock_client( - [ - httpx.ConnectError("connection refused"), - httpx.Response(200, content=json.dumps(remote_map).encode()), - ], - client_cls=httpx.Client, - ) - sleeper = _SyncSleepRecorder(block=True) - litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}) - - cost_map = get_model_cost_map( - url=_URL, - max_attempts=3, - sleep=sleeper, - rng=random.Random(0), - client=client, - ) - - assert calls["count"] == 1 - assert sleeper.waits == [] - assert sleeper.started.wait(timeout=10) - threads = _retry_threads() - try: - assert len(threads) == 1 - assert "claude-remote-only-test" not in cost_map - assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"].startswith("Remote fetch failed:") - sleeper.release.set() - for thread in threads: - thread.join(timeout=10) - assert all(not thread.is_alive() for thread in threads) - assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0 - assert calls["count"] == 2 - assert "claude-remote-only-test" in litellm.model_cost - assert "claude-remote-only-test" in litellm.anthropic_models - assert "my-runtime-model" in litellm.model_cost - source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["fallback_reason"] is None - finally: - sleeper.release.set() - for thread in _retry_threads(): - thread.join(timeout=10) - - -def test_boot_load_does_not_retry_non_retryable_failure(): - client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - ) - assert calls["count"] == 1 - assert sleeper.waits == [] - assert _retry_threads() == [] - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"] is not None - - -def test_boot_load_falls_back_when_urlsplit_rejects_the_url(): - client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) - sleeper = _SyncSleepRecorder() - - cost_map = get_model_cost_map( - url="http://[bad/model_prices.json", - sleep=sleeper, - rng=random.Random(0), - client=client, - ) - assert calls["count"] == 0 - assert sleeper.waits == [] - assert _retry_threads() == [] - assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert "invalid URL" in (source["fallback_reason"] or "") - - -@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"]) -def test_boot_load_respects_local_env_override(monkeypatch, flag): - """LITELLM_LOCAL_MODEL_COST_MAP=True or LITELLM_OFFLINE_MODE=True still short-circuits to the backup with zero HTTP.""" - monkeypatch.setenv(flag, "True") - - def _fail(request): - raise AssertionError("no HTTP request should be made when local map is forced") - - cost_map = get_model_cost_map( - url=_URL, - sleep=_SyncSleepRecorder(), - client=httpx.Client(transport=httpx.MockTransport(_fail)), - ) - assert len(cost_map) > 100 - assert get_model_cost_map_source_info()["is_env_forced"] is True - - -def test_boot_load_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(): - body = _real_map_bytes() - client, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=body)], client_cls=httpx.Client) - - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=client) - - source = get_model_cost_map_source_info() - assert source["source"] == "remote" - assert source["etag"] == 'W/"boot"' - assert source["source_revision"] == git_blob_id(body) - assert source["loaded_at"] is not None - - -def test_boot_load_fallback_to_the_backup_reports_its_blob_id_and_drops_the_remote_etag(): - remote, _ = _mock_client( - [httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=_real_map_bytes())], client_cls=httpx.Client - ) - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote) - failing, _ = _mock_client([httpx.Response(404)], client_cls=httpx.Client) - - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=failing) - - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["etag"] is None - assert source["source_revision"] == _bundled_blob_id() - - -def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rejected_fetch(): - remote, _ = _mock_client( - [httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=_real_map_bytes())], client_cls=httpx.Client - ) - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote) - shrunk_body = b'{"gpt-5.4-mini": {"mode": "chat", "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}' - shrunk, _ = _mock_client( - [httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client - ) - - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=shrunk) - - source = get_model_cost_map_source_info() - assert source["source"] == "local" - assert source["fallback_reason"] == "Remote data failed integrity validation" - assert source["etag"] is None - assert source["source_revision"] == _bundled_blob_id() - assert source["source_revision"] != git_blob_id(shrunk_body) - @pytest.mark.parametrize( - ("argv0", "request_count"), + ("limits", "message"), [ - ("/some/venv/bin/lite", 0), - ("/some/venv/bin/lite.exe", 0), - ("/some/venv/bin/python", 1), + ({"min_model_count": -1, "max_shrink_ratio": 0.5}, "min_model_count must be 0 or more, got -1"), + ({"min_model_count": 0, "max_shrink_ratio": 1.5}, "max_shrink_ratio must be between 0 and 1, got 1.5"), ], ) -def test_boot_load_skips_remote_fetch_for_cli_processes( - monkeypatch: pytest.MonkeyPatch, argv0: str, request_count: int -) -> None: - monkeypatch.setattr(sys, "argv", [argv0, "--version"]) - monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) - client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) +def test_integrity_limits_out_of_range_name_the_limit_and_value(limits: dict[str, float], message: str): + with pytest.raises(ValueError, match=re.escape(message)): + NativeModelCostMap.bundled(**limits) - cost_map = get_model_cost_map(url=_URL, client=client) - assert calls["count"] == request_count - assert cost_map - source = get_model_cost_map_source_info() - if request_count == 0: - assert source["source"] == "local" - else: - assert source["source"] == "remote" +# --------------------------------------------------------------------------- +# What litellm.model_cost promises +# --------------------------------------------------------------------------- -@pytest.fixture() -def _vcr_outcome_gate(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) +@pytest.mark.usefixtures("own_model_cost") +def test_model_cost_is_a_dict_that_reflects_each_install(): + first_body: Final = _document_with(**{"lit-first": {"litellm_provider": "openai"}}) + second_body: Final = _document_with(**{"lit-second": {"litellm_provider": "openai"}}) + install_model_cost_map(first_body) + first: Final = litellm.model_cost + install_model_cost_map(second_body) -@pytest.fixture(scope="session") -def event_loop(): - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - yield loop - loop.close() + assert type(litellm.model_cost) is dict + assert litellm.model_cost is not first + assert "lit-second" in litellm.model_cost + assert "lit-first" not in litellm.model_cost -@pytest.fixture(scope="function") -def setup_and_teardown(event_loop): - import litellm +@pytest.mark.usefixtures("own_model_cost") +def test_registrations_survive_a_reload_in_call_order(): + litellm.register_model( + {"lit-registered": {"litellm_provider": "openai", "input_cost_per_token": 1, "mode": "chat"}} + ) + litellm.register_model({"lit-registered": {"input_cost_per_token": 2, "output_cost_per_token": 3}}) - original_state = {} - for attr in ( - "callbacks", - "success_callback", - "failure_callback", - "_async_success_callback", - "_async_failure_callback", - ): - if hasattr(litellm, attr): - val = getattr(litellm, attr) - original_state[attr] = val.copy() if val else [] - for attr in _SCALAR_DEFAULTS: - if hasattr(litellm, attr): - original_state[attr] = getattr(litellm, attr) - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + install_model_cost_map(_BUNDLED) - asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) - importlib.reload(litellm) - asyncio.set_event_loop(event_loop) - yield - for attr, original_value in original_state.items(): - if hasattr(litellm, attr): - setattr(litellm, attr, original_value) - pending = asyncio.all_tasks(event_loop) - for task in pending: - task.cancel() - if pending: - event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + assert litellm.model_cost["lit-registered"]["input_cost_per_token"] == 2 + assert litellm.model_cost["lit-registered"]["output_cost_per_token"] == 3 + assert litellm.model_cost["lit-registered"]["mode"] == "chat" + assert litellm.get_model_info("lit-registered")["input_cost_per_token"] == 2 -_SCALAR_DEFAULTS = { - "num_retries": getattr(litellm, "num_retries", None), - "set_verbose": getattr(litellm, "set_verbose", False), - "cache": getattr(litellm, "cache", None), - "allowed_fails": getattr(litellm, "allowed_fails", 3), - "disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False), - "force_ipv4": getattr(litellm, "force_ipv4", False), - "drop_params": getattr(litellm, "drop_params", None), - "modify_params": getattr(litellm, "modify_params", False), - "api_base": getattr(litellm, "api_base", None), - "api_key": getattr(litellm, "api_key", None), - "cohere_key": getattr(litellm, "cohere_key", None), -} +@pytest.mark.usefixtures("own_model_cost") +def test_the_latest_registration_of_a_row_wins_in_the_registered_catalog_and_after_a_reload(): + litellm.register_model({"gpt-4o": {"input_cost_per_token": 1.0}}) + litellm.register_model({"openai/gpt-4o": {"input_cost_per_token": 2.0}}) + assert catalog_with_runtime_registrations()["gpt-4o"]["input_cost_per_token"] == 2.0 + litellm.register_model({"gpt-4o": {"input_cost_per_token": 3.0}}) + install_model_cost_map(_BUNDLED) -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestCheckIsValidDict: - """Unit tests for _check_is_valid_dict.""" + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 3.0 + assert catalog_with_runtime_registrations()["gpt-4o"]["input_cost_per_token"] == 3.0 - def test_should_reject_non_dict(self): - """Non-dict should fail.""" - assert GetModelCostMap._check_is_valid_dict("not a dict") is False - def test_should_reject_empty_dict(self): - """Empty dict should fail.""" - assert GetModelCostMap._check_is_valid_dict({}) is False +@pytest.mark.usefixtures("own_model_cost") +def test_a_registration_moves_to_the_row_a_reload_adds_for_it(): + litellm.register_model({"openai/lit-arrives": {"litellm_provider": "openai", "input_cost_per_token": 5.0}}) - def test_should_reject_list(self): - """List should fail.""" - assert GetModelCostMap._check_is_valid_dict([1, 2, 3]) is False + install_model_cost_map(_document_with(**{"lit-arrives": {"litellm_provider": "openai", "mode": "chat"}})) - def test_should_reject_none(self): - """None should fail.""" - assert GetModelCostMap._check_is_valid_dict(None) is False + assert litellm.model_cost["lit-arrives"]["input_cost_per_token"] == 5.0 + assert catalog_with_runtime_registrations()["lit-arrives"]["input_cost_per_token"] == 5.0 - def test_should_accept_non_empty_dict(self): - """Non-empty dict should pass.""" - assert GetModelCostMap._check_is_valid_dict({"model": {}}) is True +@pytest.mark.usefixtures("own_model_cost") +def test_a_reload_counts_the_rows_it_serves_and_not_the_registrations(): + litellm.register_model({"lit-only-registered": {"litellm_provider": "openai", "mode": "chat"}}) -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestCheckModelCountNotReduced: - """Unit tests for _check_model_count_not_reduced.""" + result: Final = install_model_cost_map(_BUNDLED) - def test_should_reject_too_few_models(self): - """Fetched map with fewer models than min_model_count should fail.""" - small_map = {f"model-{i}": {} for i in range(5)} - assert ( - GetModelCostMap._check_model_count_not_reduced( - fetched_map=small_map, backup_model_count=0, min_model_count=10 - ) - is False - ) + assert isinstance(result, ModelCostMapReloaded) + assert result.model_count == len(GetModelCostMap.loaded_model_cost_map()) == len(litellm.model_cost) - 1 - def test_should_reject_significant_shrinkage(self): - """Fetched map that shrunk >50% vs backup should fail.""" - fetched = {f"model-{i}": {} for i in range(40)} # 40% of 100 - assert ( - GetModelCostMap._check_model_count_not_reduced( - fetched_map=fetched, backup_model_count=100, min_model_count=10 - ) - is False - ) - def test_should_accept_when_above_threshold(self): - """Fetched map at 60% of backup (above 50% threshold) should pass.""" - fetched = {f"model-{i}": {} for i in range(60)} - assert ( - GetModelCostMap._check_model_count_not_reduced( - fetched_map=fetched, backup_model_count=100, min_model_count=10 - ) - is True - ) +@pytest.mark.usefixtures("own_model_cost") +def test_a_reload_adds_the_models_it_serves_to_the_provider_sets(): + install_model_cost_map(_document_with(**{"lit-reloaded-chat": {"litellm_provider": "openai", "mode": "chat"}})) - def test_should_accept_growth(self): - """Fetched map larger than backup should pass.""" - fetched = {f"model-{i}": {} for i in range(120)} - assert ( - GetModelCostMap._check_model_count_not_reduced( - fetched_map=fetched, backup_model_count=100, min_model_count=10 - ) - is True - ) + assert "lit-reloaded-chat" in litellm.models_by_provider["openai"] - def test_should_accept_with_empty_backup(self): - """When backup is empty, only min_model_count matters.""" - fetched = {f"model-{i}": {} for i in range(15)} - assert ( - GetModelCostMap._check_model_count_not_reduced( - fetched_map=fetched, backup_model_count=0, min_model_count=10 - ) - is True - ) - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestValidateModelCostMap: - """Unit tests for validate_model_cost_map (combines both checks).""" - - def test_should_reject_non_dict(self): - """Non-dict should fail at check 1.""" - assert GetModelCostMap.validate_model_cost_map(fetched_map="not a dict", backup_model_count=0) is False - - def test_should_reject_empty_map(self): - """Empty dict should fail at check 1.""" - assert GetModelCostMap.validate_model_cost_map(fetched_map={}, backup_model_count=0) is False - - def test_should_reject_significant_shrinkage(self): - """Should fail at check 2 (shrinkage).""" - fetched = {f"model-{i}": {} for i in range(40)} - assert ( - GetModelCostMap.validate_model_cost_map(fetched_map=fetched, backup_model_count=100, min_model_count=10) - is False - ) - - def test_should_accept_valid_map(self): - """Should pass both checks.""" - fetched = {f"model-{i}": {} for i in range(120)} - assert ( - GetModelCostMap.validate_model_cost_map(fetched_map=fetched, backup_model_count=100, min_model_count=10) - is True - ) - - def test_should_accept_equal_size_map(self): - """Equal size should pass both checks.""" - fetched = {f"model-{i}": {} for i in range(100)} - assert ( - GetModelCostMap.validate_model_cost_map(fetched_map=fetched, backup_model_count=100, min_model_count=10) - is True - ) - - -def _fetch_with_single_outcome(monkeypatch, outcome): - from litellm.litellm_core_utils import get_model_cost_map as module - - monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) - source_info = module._cost_map_source_info - for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"): - monkeypatch.setattr(source_info, name, getattr(source_info, name)) - client, calls = _mock_client([outcome], client_cls=httpx.Client) - result = get_model_cost_map(url=_URL, max_attempts=1, client=client) - return result, calls["count"], get_model_cost_map_source_info() - - -def _backup_keys(): - return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys() - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestGetModelCostMapFallback: - """Tests for get_model_cost_map fallback behavior with bad upstream.""" - - def test_should_fallback_to_backup_on_invalid_json(self, monkeypatch): - """When upstream returns invalid JSON, should fall back to local backup.""" - result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.Response(200, content=b"not json")) - - assert fetches == 1 - assert result.keys() == _backup_keys() - assert source["source"] == "local" - assert source["fallback_reason"].startswith("Remote fetch failed") - - def test_should_fallback_to_backup_on_network_error(self, monkeypatch): - """When upstream is unreachable, should fall back to local backup.""" - result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.ConnectError("Connection refused")) - - assert fetches == 1 - assert result.keys() == _backup_keys() - assert source["source"] == "local" - assert source["fallback_reason"].startswith("Remote fetch failed") - - def test_should_fallback_when_fetched_map_is_empty(self, monkeypatch): - """When upstream returns valid JSON but empty dict, should fall back.""" - result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.Response(200, content=b"{}")) - - assert fetches == 1 - assert result.keys() == _backup_keys() - assert source["source"] == "local" - - def test_should_fallback_when_fetched_json_is_not_object(self, monkeypatch): - result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.Response(200, content=b"[]")) - - assert fetches == 1 - assert result.keys() == _backup_keys() - assert source["source"] == "local" - assert source["fallback_reason"] == f"Remote fetch failed: expected a JSON object from {_URL}, got list" - - def test_should_fallback_when_fetched_map_shrinks_dramatically(self, monkeypatch): - """When upstream returns far fewer models than backup, should fall back.""" - tiny_map = {f"model-{i}": {"litellm_provider": "test"} for i in range(11)} - result, fetches, source = _fetch_with_single_outcome( - monkeypatch, httpx.Response(200, content=json.dumps(tiny_map).encode()) - ) - - assert fetches == 1 - assert result.keys() == _backup_keys() - assert source == {**source, "source": "local", "fallback_reason": "Remote data failed integrity validation"} - - def test_should_use_local_map_when_env_var_set(self): - """LITELLM_LOCAL_MODEL_COST_MAP=True should skip remote fetch entirely.""" - with patch.dict(os.environ, {"LITELLM_LOCAL_MODEL_COST_MAP": "True"}): - with patch("httpx.get") as mock_get: - result = get_model_cost_map("https://fake-url.com/model_prices.json") - mock_get.assert_not_called() - - assert isinstance(result, dict) - assert len(result) > 0 - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestBackupModelCostMapExists: - """Validates the local backup file is always present and valid.""" - - def test_should_have_backup_file(self): - """The backup model cost map must exist and be loadable.""" - backup = GetModelCostMap.load_local_model_cost_map() - assert isinstance(backup, dict) - assert len(backup) > 0, "Backup model cost map is empty" - - def test_should_have_minimum_models_in_backup(self): - """The backup must contain a reasonable number of models.""" - backup = GetModelCostMap.load_local_model_cost_map() - assert len(backup) > 100, f"Backup has only {len(backup)} models, expected > 100" - - -@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") -class TestBadHostedModelCostMap: - """ - Simulates the hosted model cost map being bad (invalid JSON / corrupted). - - When the hosted map is bad, get_model_cost_map() falls back to the local - backup. These tests verify that after fallback: - - get_model_info() still works for models in the backup - - litellm.completion() still works - """ - - def test_should_model_info_pass_after_bad_hosted_map(self): - """ - If the hosted map is bad, get_model_cost_map falls back to the local - backup. get_model_info should still work for models in the backup. - """ - mock_response = MagicMock() - mock_response.raise_for_status = MagicMock() - mock_response.json.side_effect = json.JSONDecodeError("bad json", "", 0) - - with patch("httpx.get", return_value=mock_response): - fallback_map = get_model_cost_map("https://fake-url.com/bad.json") - - original = litellm.model_cost - litellm.model_cost = fallback_map - try: - # gpt-4o is in every backup — should work fine - info = litellm.get_model_info("gpt-4o") - assert info is not None - assert info["input_cost_per_token"] > 0 - finally: - litellm.model_cost = original - - def test_should_completion_pass_after_bad_hosted_map(self): - """ - If the hosted map is bad, litellm.completion() should still work. - - Uses litellm's built-in mock_response param so the real completion - path is exercised (routing, cost calculator, logging) without - needing API credentials. - """ - # Simulate bad hosted map → fallback to backup - mock_http = MagicMock() - mock_http.raise_for_status = MagicMock() - mock_http.json.side_effect = json.JSONDecodeError("bad json", "", 0) - - with patch("httpx.get", return_value=mock_http): - fallback_map = get_model_cost_map("https://fake-url.com/bad.json") - - original = litellm.model_cost - litellm.model_cost = fallback_map - try: - # mock_response goes through the real completion path — - # routing, cost calculator, logging — but skips the HTTP call - response = litellm.completion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "say hi"}], - mock_response="hello from mock", - ) - assert response is not None - assert response.choices[0].message.content == "hello from mock" - finally: - litellm.model_cost = original - - -@pytest.fixture() -def _vcr_outcome_gate_local_testing(request, vcr): - install_live_call_probe(request, vcr) - yield - record_vcr_outcome(request, vcr) - - -@pytest.fixture(scope="function") -def isolate_litellm_state(): - """ - Per-function isolation fixture. - - Resets litellm globals to their true defaults before each test and - restores them afterward, so tests don't leak side effects. - Works safely under pytest-xdist parallel execution. - """ - original_state = {} - for attr in ( - "callbacks", - "success_callback", - "failure_callback", - "_async_success_callback", - "_async_failure_callback", - ): - if hasattr(litellm, attr): - val = getattr(litellm, attr) - original_state[attr] = val.copy() if val else [] - for attr in ("pre_call_rules", "post_call_rules"): - if hasattr(litellm, attr): - val = getattr(litellm, attr) - original_state[attr] = val.copy() if val else [] - for attr in _SCALAR_DEFAULTS_local_testing: - if hasattr(litellm, attr): - original_state[attr] = getattr(litellm, attr) - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - for attr in ( - "callbacks", - "success_callback", - "failure_callback", - "_async_success_callback", - "_async_failure_callback", - "pre_call_rules", - "post_call_rules", - ): - if hasattr(litellm, attr): - setattr(litellm, attr, []) - for attr, default_val in _SCALAR_DEFAULTS_local_testing.items(): - if hasattr(litellm, attr): - setattr(litellm, attr, default_val) - yield - asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - for attr, original_value in original_state.items(): - if hasattr(litellm, attr): - setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() - - -_SCALAR_DEFAULTS_local_testing = { - "num_retries": getattr(litellm, "num_retries", None), - "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), - "request_timeout": getattr(litellm, "request_timeout", None), - "set_verbose": getattr(litellm, "set_verbose", False), - "cache": getattr(litellm, "cache", None), - "allowed_fails": getattr(litellm, "allowed_fails", 3), - "default_fallbacks": getattr(litellm, "default_fallbacks", None), - "enable_azure_ad_token_refresh": getattr(litellm, "enable_azure_ad_token_refresh", None), - "tag_budget_config": getattr(litellm, "tag_budget_config", None), - "model_cost": getattr(litellm, "model_cost", None), - "token_counter": getattr(litellm, "token_counter", None), - "disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False), - "force_ipv4": getattr(litellm, "force_ipv4", False), - "drop_params": getattr(litellm, "drop_params", None), - "modify_params": getattr(litellm, "modify_params", False), - "api_base": getattr(litellm, "api_base", None), - "api_key": getattr(litellm, "api_key", None), -} - - -@pytest.fixture(scope="module") -def setup_and_teardown_local_testing(): - """ - Module-scoped setup. Reloads litellm only in single-process mode - (skipped under xdist to avoid cross-worker interference). - """ - import litellm - - worker_id = os.environ.get("PYTEST_XDIST_WORKER", None) - if worker_id is None: - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - import litellm.proxy.proxy_server - - importlib.reload(litellm.proxy.proxy_server) - except Exception: - pass - if hasattr(litellm, "in_memory_llm_clients_cache"): - litellm.in_memory_llm_clients_cache.flush_cache() - yield - - -@pytest.mark.usefixtures( - "_vcr_outcome_gate_local_testing", - "isolate_litellm_state", - "setup_and_teardown_local_testing", -) -def test_get_model_cost_map(): - try: - print(litellm.get_model_cost_map(url="fake-url")) - except Exception as e: - pytest.fail(f"An exception occurred: {e}") - - -@pytest.mark.usefixtures( - "_vcr_outcome_gate_local_testing", - "isolate_litellm_state", - "setup_and_teardown_local_testing", -) -def test_get_backup_model_cost_map(): - with importlib_get_model.open_text("litellm", "model_prices_and_context_window_backup.json") as f: - print("inside backup") - content = json.load(f) - print("content", content) - - -@pytest.mark.parametrize("backup", [b'["only-in-the-backup"]', b'"only-in-the-backup"', b"null"]) -def test_a_bundled_backup_that_is_not_a_json_object_is_rejected_without_quoting_it( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path, backup: bytes +@pytest.mark.usefixtures("own_model_cost") +def test_a_reload_repeats_the_cache_pricing_warning_for_a_registration_no_catalog_row_matches( + caplog: pytest.LogCaptureFixture, ): - (tmp_path / "model_prices_and_context_window_backup.json").write_bytes(backup) - monkeypatch.setattr(get_model_cost_map_module, "files", lambda _package: tmp_path) + litellm.register_model({"lit-unmatched": {"litellm_provider": "openai", "input_cost_per_token": 1e-06}}) + caplog.clear() - with pytest.raises(ValidationError) as rejected: - GetModelCostMap.load_local_model_cost_map_with_revision() + with caplog.at_level(logging.WARNING, logger=litellm.verbose_logger.name): + install_model_cost_map(_BUNDLED) - assert "only-in-the-backup" not in str(rejected.value) + assert any( + "lit-unmatched" in record.message and "cache_creation_input_token_cost" in record.message + for record in caplog.records + ), caplog.text + + +@pytest.mark.usefixtures("own_model_cost") +def test_a_registration_that_does_not_persist_is_gone_after_a_reload(): + litellm.register_model( + {"lit-request-scoped": {"litellm_provider": "openai", "input_cost_per_token": 1}}, persist_across_reloads=False + ) + assert "lit-request-scoped" in litellm.model_cost + + install_model_cost_map(_BUNDLED) + + assert "lit-request-scoped" not in litellm.model_cost + + +@pytest.mark.usefixtures("own_model_cost") +def test_a_registration_wins_the_fields_it_shares_with_a_request_scoped_row_after_a_reload(): + litellm.register_model({"gpt-4o": {"input_cost_per_token": 7.0}}) + litellm.register_model( + {"gpt-4o": {"input_cost_per_token": 9.0, "output_cost_per_token": 8.0}}, persist_across_reloads=False + ) + + install_model_cost_map(_BUNDLED) + + assert litellm.model_cost["gpt-4o"]["input_cost_per_token"] == 7.0 + assert ( + litellm.model_cost["gpt-4o"]["output_cost_per_token"] == json.loads(_BUNDLED)["gpt-4o"]["output_cost_per_token"] + ) + + +@pytest.mark.usefixtures("own_model_cost") +def test_a_rejected_document_changes_nothing(): + installed: Final = litellm.model_cost + rows: Final = dict(installed) + provenance: Final = get_model_cost_map_provenance() + + result: Final = install_model_cost_map(json.dumps({"only": {"litellm_provider": "openai"}}).encode()) + + assert isinstance(result, ModelCostMapReloadUnavailable) + assert litellm.model_cost is installed + assert litellm.model_cost == rows + assert get_model_cost_map_provenance() == provenance + + +@pytest.mark.usefixtures("own_model_cost") +def test_a_reader_iterating_the_old_table_finishes_while_another_thread_installs(): + old: Final = litellm.model_cost + installer: Final = threading.Thread(target=install_model_cost_map, args=(_BUNDLED,)) + walked: Final = iter(old.items()) + first: Final = next(walked) + installer.start() + installer.join() + + assert len((first, *walked)) == len(old) + assert litellm.model_cost is not old + + +@pytest.mark.usefixtures("own_model_cost") +def test_get_model_info_follows_a_reload_and_a_registration(): + litellm.get_model_info("gpt-4o") + install_model_cost_map(_document_with(**{"gpt-4o": {"litellm_provider": "openai", "input_cost_per_token": 0.5}})) + assert litellm.get_model_info("gpt-4o")["input_cost_per_token"] == 0.5 + + litellm.register_model({"gpt-4o": {"input_cost_per_token": 0.25}}) + assert litellm.get_model_info("gpt-4o")["input_cost_per_token"] == 0.25 + + +@pytest.mark.usefixtures("own_model_cost") +def test_the_registered_catalog_reads_the_document_under_registrations_and_never_request_scoped_rows(): + litellm.register_model({"lit-kept": {"litellm_provider": "openai", "input_cost_per_token": 1}}) + litellm.register_model({"lit-request-row": {"litellm_provider": "openai"}}, persist_across_reloads=False) + + for _ in range(2): + catalog = catalog_with_runtime_registrations() + assert catalog["lit-kept"]["input_cost_per_token"] == 1 + assert "lit-request-row" not in catalog + assert catalog["gpt-4o"] == GetModelCostMap.loaded_model_cost_map()["gpt-4o"] + install_model_cost_map(_BUNDLED) + + +@pytest.mark.usefixtures("own_model_cost") +def test_the_loaded_catalog_follows_installs_and_ignores_registrations(): + install_model_cost_map(_document_with(**{"lit-loaded": {"litellm_provider": "openai", "lit_field": 1}})) + litellm.register_model({"lit-loaded": {"input_cost_per_token": 5}}) + + assert GetModelCostMap.loaded_model_cost_map()["lit-loaded"] == {"litellm_provider": "openai", "lit_field": 1} + assert "lit_field" in GetModelCostMap.loaded_model_cost_map_fields() + + install_model_cost_map(_document_with(**{"lit-loaded": {"litellm_provider": "openai", "lit_other_field": 2}})) + + assert GetModelCostMap.loaded_model_cost_map()["lit-loaded"] == {"litellm_provider": "openai", "lit_other_field": 2} + assert "lit_field" not in GetModelCostMap.loaded_model_cost_map_fields() + assert "lit_other_field" in GetModelCostMap.loaded_model_cost_map_fields() + + +@pytest.mark.usefixtures("own_model_cost") +@pytest.mark.parametrize("add_to_provider_model_sets", [True, False]) +def test_register_model_adds_names_to_the_provider_sets_only_when_asked_and_a_reload_keeps_it_so( + add_to_provider_model_sets: bool, +): + chat_model: Final = f"lit-registered-chat-{add_to_provider_model_sets}".lower() + litellm.register_model( + { + "openrouter/lit-vendor/lit-model": {"litellm_provider": "openrouter", "mode": "chat"}, + chat_model: {"litellm_provider": "openai", "mode": "chat"}, + }, + add_to_provider_model_sets=add_to_provider_model_sets, + ) + install_model_cost_map(_BUNDLED) + + assert ("lit-vendor/lit-model" in litellm.openrouter_models) is add_to_provider_model_sets + assert (chat_model in litellm.models_by_provider["openai"]) is add_to_provider_model_sets + + +@pytest.mark.usefixtures("own_model_cost") +def test_install_hands_the_fallback_rules_to_the_generalization_lookup(): + rules: Final = [{"name": "lit-rule", "pattern": "^lit-family-", "model_info": {"litellm_provider": "openai"}}] + previous: Final = list(get_fallback_generalization_rules()) + try: + install_model_cost_map(_document_with(**{FALLBACK_GENERALIZATIONS_KEY: {"rules": rules}})) + assert match_routing_generalization("lit-family-1") == "openai" + assert FALLBACK_GENERALIZATIONS_KEY not in litellm.model_cost + finally: + set_fallback_generalizations(previous) + + +@pytest.mark.usefixtures("own_model_cost") +def test_alias_names_share_one_entry_after_an_install_and_a_registration(): + install_model_cost_map( + _document_with(**{"lit-owner": {"litellm_provider": "openai", "mode": "chat", "aliases": ["lit-alias"]}}) + ) + assert litellm.model_cost["lit-alias"] is litellm.model_cost["lit-owner"] + + litellm.register_model({"lit-alias": {"input_cost_per_token": 3e-06}}) + + entry: Final = litellm.model_cost["lit-owner"] + assert litellm.model_cost["lit-alias"] is entry + assert isinstance(entry, AliasedCostEntry) + assert entry.registered_fields == frozenset({"input_cost_per_token"}) + assert runtime_model_cost_registrations("lit-owner", "openai") == (("lit-owner", {"input_cost_per_token": 3e-06}),) + + litellm.register_model({"lit-owner": {"output_cost_per_token": 4e-06}}) + + assert runtime_model_cost_registrations("lit-owner", "openai") == ( + ("lit-owner", {"input_cost_per_token": 3e-06, "output_cost_per_token": 4e-06}), + ) + + +def test_without_registrations_the_registered_catalog_is_the_loaded_catalog(): + catalog: Final = catalog_with_runtime_registrations() + + assert {key: dict(row) for key, row in catalog.items()} == { + key: dict(row) for key, row in GetModelCostMap.loaded_model_cost_map().items() + } + + +# --------------------------------------------------------------------------- +# A second instance over a dict of its own +# --------------------------------------------------------------------------- + + +def _own_instance(document: Mapping[str, Mapping[str, object]]) -> tuple[NativeModelCostMap, dict]: + instance: Final = NativeModelCostMap.bundled(0, 0.0) + installed: Final = instance.install(json.dumps(document).encode()) + assert isinstance(installed, NativeInstalled) + return instance, installed.rows + + +def test_alias_of_rows_are_served_as_their_target_and_dangling_ones_are_dropped(): + _, rows = _own_instance( + { + "real": {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 1e-6}, + "other": {"litellm_provider": "openai", "mode": "chat", "aliases": ["twin-latest"]}, + "real-latest": {"alias_of": "real"}, + "twin-latest": {"alias_of": "real"}, + "ghost-latest": {"alias_of": "ghost"}, + "chain-latest": {"alias_of": "real-latest"}, + "listed-latest": {"alias_of": "missing", "aliases": ["extra"]}, + } + ) + + assert rows["real-latest"] is rows["real"] + assert "alias_of" not in rows["real"] + assert set(rows) == {"real", "other", "real-latest", "twin-latest"} + assert rows["twin-latest"] is rows["other"] + + +def test_an_instance_writes_only_the_published_dict_it_is_handed(monkeypatch: pytest.MonkeyPatch): + instance, rows = _own_instance({"lit-own": {"litellm_provider": "openai", "mode": "chat"}}) + previous: Final = litellm.model_cost + monkeypatch.setattr(litellm, "model_cost", rows) + + written: Final = instance.register(rows, {"lit-own-new": {"litellm_provider": "openai"}}, True) + + assert isinstance(written, NativeRegistered) + assert "lit-own-new" in rows + assert "lit-own-new" not in previous + assert instance.find_key(rows, "LIT-OWN-NEW") == "lit-own-new" + + +def test_the_name_index_is_a_cache(): + instance, rows = _own_instance( + { + "lit-a": {"litellm_provider": "p1"}, + "LIT-B": {"litellm_provider": "p2"}, + "lit-c": {"litellm_provider": "p1"}, + } + ) + + def queries() -> tuple[list[str | None], list[list[str]]]: + return ( + [instance.find_key(rows, name) for name in ("LIT-A", "lit-b", "missing")], + [instance.keys_in(rows, provider) for provider in ("p1", "p2", "p3")], + ) + + before: Final = queries() + instance.invalidate() + + assert queries() == before + assert before == (["lit-a", "LIT-B", None], [["lit-a", "lit-c"], ["LIT-B"], []]) + + +def test_a_key_written_straight_to_the_dict_is_found_until_the_next_install(): + instance, rows = _own_instance({"lit-a": {"litellm_provider": "p"}}) + rows.setdefault("Lit-Direct", {"litellm_provider": "p"}) + assert instance.find_key(rows, "lit-direct") == "Lit-Direct" + + reinstalled: Final = instance.install(json.dumps({"lit-a": {"litellm_provider": "p"}}).encode()) + assert isinstance(reinstalled, NativeInstalled) + assert instance.find_key(reinstalled.rows, "lit-direct") is None + + +def test_a_key_swapped_in_for_a_popped_one_replaces_it_in_lookups(): + instance, rows = _own_instance({"lit-a": {"litellm_provider": "p"}, "Lit-Gone": {"litellm_provider": "p"}}) + assert instance.find_key(rows, "lit-gone") == "Lit-Gone" + rows.pop("Lit-Gone") + rows["Lit-New"] = {"litellm_provider": "p"} + + assert instance.find_key(rows, "lit-new") == "Lit-New" + assert instance.find_key(rows, "lit-gone") is None + + +def test_a_rebound_dict_is_what_lookups_read_next(): + instance, rows = _own_instance({"lit-a": {"litellm_provider": "p"}}) + rebound: Final = {"Lit-Rebound": {"litellm_provider": "p"}} + + assert instance.find_key(rows, "lit-rebound") is None + assert instance.find_key(rebound, "lit-rebound") == "Lit-Rebound" + + +def test_case_insensitive_lookups_go_through_the_live_dict(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, "Lit-Mixed-Case": {"litellm_provider": "p"}}) + + assert get_model_cost_key("GPT-4O") == "gpt-4o" + assert get_model_cost_key("lit-mixed-case") == "Lit-Mixed-Case" + assert get_model_cost_key("lit-unknown") is None + + +def test_main_cost_map_names_downstream_code_calls_stay_callable_with_their_main_signatures(): + # Not an API we would design today. Downstream open source SDK code calls these names from main directly, so + # removing one breaks it. Their behavior is pinned in tests/integration/sdk/test_model_cost_sdk.py + import inspect + + from litellm.litellm_core_utils import get_model_cost_map as cost_map_module + from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetModelCostMap, ModelCostMapReloaded + + assert litellm.utils._model_cost_lowercase_map is None + generation: Final = litellm.utils.get_model_cost_mutation_generation() + litellm.utils._invalidate_model_cost_lowercase_map() + assert litellm.utils.get_model_cost_mutation_generation() > generation + for cached in ( + litellm.utils._cached_get_model_info, + litellm.utils._cached_get_model_info_helper, + litellm.utils.cached_cost_map_key, + litellm.get_model_info, + ): + cached.cache_clear() + assert cached.cache_info().currsize == 0 + assert litellm.utils._cached_get_model_info("gpt-4o")["key"] == "gpt-4o" + assert GetModelCostMap.get_backup_model_count() == GetModelCostMap._get_backup_model_count() > 0 + assert GetModelCostMap.validate_model_cost_map(fetched_map={}, backup_model_count=1) is False + assert GetModelCostMap.validate_model_cost_map( + fetched_map=GetModelCostMap.load_local_model_cost_map(), + backup_model_count=GetModelCostMap.get_backup_model_count(), + ) + for function, arguments in ( + (cost_map_module.refetch_model_cost_map, {"url": "u", "timeout": 5, "max_attempts": 1, "client": None}), + (cost_map_module.get_model_cost_map, {"url": "u", "timeout": 5, "max_attempts": 1, "sleep": None, "rng": None}), + (cost_map_module.adopt_model_cost_map, {"new_model_cost_map": {}}), + (cost_map_module._finalize_model_cost_map, {"model_cost": {}}), + ): + inspect.signature(function).bind(**arguments) + assert {"model_cost_map", "revision", "etag"} <= set(dir(ModelCostMapReloaded)) + assert issubclass(AliasedCostEntry, dict) + + +@pytest.mark.usefixtures("own_model_cost") +def test_main_cost_map_helpers_keep_their_side_effects(monkeypatch: pytest.MonkeyPatch): + from litellm.litellm_core_utils.get_model_cost_map import ( + _finalize_model_cost_map, # pyright: ignore[reportPrivateUsage] # downstream code calls main's finalizer + ) + + document: Final = json.loads(_BUNDLED) + assert FALLBACK_GENERALIZATIONS_KEY in document + + rows: Final = _finalize_model_cost_map(document) + + assert FALLBACK_GENERALIZATIONS_KEY not in document + assert FALLBACK_GENERALIZATIONS_KEY not in rows + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + with pytest.warns(DeprecationWarning, match="ignored"): + asyncio.run(refetch_model_cost_map(url="http://127.0.0.1:9/unused.json", sleep=asyncio.sleep)) diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index a83e98afb89..b4bf1120935 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -32,7 +32,7 @@ from litellm.types.utils import ( StandardLoggingPayloadErrorInformation, StreamingChoices, ) -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish import unittest.mock @@ -1413,7 +1413,7 @@ async def test_messages_honors_a_chat_pin_on_a_prefixed_responses_deployment( "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], } litellm.model_cost["perplexity/pin-test-model"] = {"litellm_provider": "perplexity", "mode": "chat"} - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() @@ -2838,7 +2838,7 @@ async def test_azure_messages_replay_encrypted_reasoning_only_to_responses( }, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) litellm.in_memory_llm_clients_cache.flush_cache() chat_reply: Final = { "id": "chatcmpl-azure", diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index 8571a0c3cbc..6358d6976c4 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -26,7 +26,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig from litellm.router import Router -from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params +from litellm.utils import get_optional_params +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome _MAPPED_PARAMS: Final = TypeAdapter(dict[str, object]) @@ -598,7 +599,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index fdd21c87732..5ead4ad3986 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py @@ -15,17 +15,18 @@ from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) from litellm.types.utils import ImageObject, ImageResponse -from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params_image_gen +from litellm.utils import get_optional_params_image_gen +from litellm.rust_bridge.host.model_cost import publish @pytest.fixture(autouse=True) def use_local_model_cost_map(monkeypatch): monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) yield litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize( diff --git a/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py index c8365e7b7c0..927833b530d 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_fw_models_metadata.py @@ -19,7 +19,7 @@ def use_local_model_cost_map(): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") import litellm - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish original_model_cost = litellm.model_cost litellm.model_cost = json.loads( @@ -28,13 +28,13 @@ def use_local_model_cost_map(): .read_text(encoding="utf-8") ) litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) try: yield litellm finally: litellm.model_cost = original_model_cost litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) monkeypatch.undo() diff --git a/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py index 4756773aa3d..e23be01df55 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py @@ -14,7 +14,7 @@ def use_local_model_cost_map(): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") import litellm - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish original_model_cost = litellm.model_cost litellm.model_cost = json.loads( @@ -23,13 +23,13 @@ def use_local_model_cost_map(): .read_text(encoding="utf-8") ) litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) try: yield litellm finally: litellm.model_cost = original_model_cost litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) monkeypatch.undo() diff --git a/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py index 122dd5b555a..d17c2569b92 100644 --- a/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py +++ b/tests/unit/llms/bedrock/image_edit/test_amazon_nova_canvas_image_edit.py @@ -27,7 +27,7 @@ def ensure_nova_canvas_image_edit_model_cost_flags(monkeypatch): Full ``model_prices_and_context_window.json`` includes these flags, but CI or alternate cost maps may omit them—merge minimal entries so tests match production. """ - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish for key in ( "amazon.nova-canvas-v1:0", @@ -46,11 +46,11 @@ def ensure_nova_canvas_image_edit_model_cost_flags(monkeypatch): "supports_nova_canvas_image_edit": True, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) yield - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_get_config_class_nova_canvas(): diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 0de8efa2662..b0071287746 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -42,7 +42,7 @@ from litellm.types.llms.custom_http import VerifyTypes from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from litellm.exceptions import Timeout as LitellmTimeout from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -2183,7 +2183,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 53122246de4..bdeff226af2 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -42,7 +42,7 @@ from litellm.types.utils import ( PromptTokensDetails, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from openai.types.beta.assistant import Assistant from openai.types.beta.assistant_deleted import AssistantDeleted from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -1830,7 +1830,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/llms/parallel_ai/test_parallel_ai_search.py b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py index 2bb07ecca75..1400fcc7d4e 100644 --- a/tests/unit/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/unit/llms/parallel_ai/test_parallel_ai_search.py @@ -57,13 +57,13 @@ def bundled_cost_map(monkeypatch): litellm caches model-info lookups, so swapping ``model_cost`` only takes effect once those caches are invalidated -- on the way in and back out. """ - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) yield monkeypatch.undo() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) class TestParallelAISearch: diff --git a/tests/unit/llms/test_custom_llm.py b/tests/unit/llms/test_custom_llm.py index 409cb42bf3b..e731d79ed99 100644 --- a/tests/unit/llms/test_custom_llm.py +++ b/tests/unit/llms/test_custom_llm.py @@ -35,7 +35,8 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, ) -from litellm.utils import ModelResponseIterator, _invalidate_model_cost_lowercase_map +from litellm.utils import ModelResponseIterator +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -729,7 +730,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 03bc308a89d..8c2cb643943 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -36,7 +36,8 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage -from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper +from litellm.utils import CustomStreamWrapper +from litellm.rust_bridge.host.model_cost import publish from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger from litellm.llms.vertex_ai.gemini.transformation import ( @@ -8161,7 +8162,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { @@ -9676,7 +9677,7 @@ async def test_gemini_pro_json_schema_args_sent_httpx( monkeypatch.setenv("GEMINI_API_KEY", "test-gemini-api-key") monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) monkeypatch.setattr(litellm, "set_verbose", True) messages = [{"role": "user", "content": "List 5 cookie recipes"}] from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -9782,7 +9783,7 @@ async def test_gemini_pro_json_schema_args_sent_httpx_openai_schema( monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) monkeypatch.setattr(litellm, "set_verbose", True) messages = [{"role": "user", "content": "List 5 cookie recipes"}] from litellm.llms.custom_httpx.http_handler import HTTPHandler diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index cbb37c964c6..feda96c46b7 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -1086,7 +1086,7 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None: from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache from litellm.proxy.utils import litellm_config_cache - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish litellm_config_cache.flush_cache() prisma_client = _reload_config_prisma_client() @@ -1100,13 +1100,12 @@ async def test_model_cost_map_reload_does_not_publish_config_change() -> None: with patch( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new=AsyncMock( - return_value=ModelCostMapReloaded(model_cost_map={"gpt-5.2": {"input_cost_per_token": 0.001}}) + return_value=ModelCostMapReloaded(model_count=len({"gpt-5.2": {"input_cost_per_token": 0.001}})) ), ): await ProxyConfig()._check_and_reload_model_cost_map(prisma_client=prisma_client) finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) _set_redis_usage_cache(previous_cache) prisma_client.db.litellm_config.update_many.assert_awaited_once() diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index e4e32eeeac9..588411cd57a 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -34,7 +34,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from fastapi import Request as Request_http_parsing from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from starlette.types import Message from starlette.websockets import WebSocket from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -1561,7 +1561,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/proxy/common_utils/test_model_deprecation.py b/tests/unit/proxy/common_utils/test_model_deprecation.py index a0fe560e6fb..57839d65cfb 100644 --- a/tests/unit/proxy/common_utils/test_model_deprecation.py +++ b/tests/unit/proxy/common_utils/test_model_deprecation.py @@ -15,7 +15,7 @@ import pytest import litellm from litellm import Router from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.rust_bridge.host.model_cost import AliasedCostEntry from litellm.proxy.common_utils.model_deprecation import ( _classify, _parse_deprecation_date, diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py index cbb3e5f981c..c24d04b3a8b 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter.py @@ -14,7 +14,7 @@ from litellm import DualCache as DualCache_dynamic_rate, Router from litellm._uuid import uuid from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import UserAPIKeyAuth -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from typing import Optional @@ -139,7 +139,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/proxy/hooks/test_prompt_injection_detection.py b/tests/unit/proxy/hooks/test_prompt_injection_detection.py index 3c871ca7333..1b50de938b6 100644 --- a/tests/unit/proxy/hooks/test_prompt_injection_detection.py +++ b/tests/unit/proxy/hooks/test_prompt_injection_detection.py @@ -17,7 +17,7 @@ from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm import Router as Router_prompt_injection from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -279,7 +279,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), diff --git a/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py b/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py index 7678e1c4d2d..0b488618636 100644 --- a/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py +++ b/tests/unit/proxy/management_endpoints/test_sso_helper_utils.py @@ -9,7 +9,7 @@ import os import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from litellm.proxy._types import LitellmUserRoles @@ -87,7 +87,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py index 5e31906a4b9..e3a499c7afc 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py +++ b/tests/unit/proxy/pass_through_endpoints/test_deepgram_ws_passthrough_routes.py @@ -2,6 +2,7 @@ target URL.""" import asyncio +import json from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType, SimpleNamespace @@ -17,7 +18,7 @@ from starlette.websockets import WebSocketDisconnect import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import WEBSOCKET_CLOSE_REASON_MAX_BYTES -from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded, install_model_cost_map from litellm.llms.deepgram.common_utils import DeepgramSpeakMeter from litellm.proxy._lazy_features import LAZY_FEATURES from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth @@ -546,7 +547,9 @@ async def test_deepgram_speak_bills_at_the_connect_time_rate_after_a_cost_map_re relay: Final = await _serve_speak(_FakeWebSocket("/deepgram/v1/speak", "model=aura-2-thalia-en")) meters: Final = [call.meter for call in relay.calls if isinstance(call.meter, DeepgramSpeakMeter)] - adopt_model_cost_map({key: row for key, row in litellm.model_cost.items() if key != "deepgram/aura-2"}) + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + reloaded: Final = {key: row for key, row in litellm.model_cost.items() if key != "deepgram/aura-2"} + assert isinstance(install_model_cost_map(json.dumps(reloaded).encode()), ModelCostMapReloaded) assert [meter.price(47).response_cost for meter in meters] == [47 * rate] diff --git a/tests/unit/proxy/proxy_server/test_routes_model_cost_map.py b/tests/unit/proxy/proxy_server/test_routes_model_cost_map.py index 31f69d686d4..0c4c5d82c18 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_cost_map.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_cost_map.py @@ -10,7 +10,12 @@ Routes covered: from __future__ import annotations +import hashlib import json +import threading +from collections.abc import Iterator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from unittest.mock import AsyncMock, MagicMock @@ -32,6 +37,35 @@ _ROOT_COST_MAP = Path(__file__).resolve().parents[4] / "model_prices_and_context # --------------------------------------------------------------------------- +def _git_blob_id(body: bytes) -> str: + return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest() + + +@contextmanager +def _serving(body: bytes, etag: str) -> Iterator[str]: + """A local HTTP server that answers every GET with ``body`` and ``etag``, so the native fetch runs for real.""" + + class _Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(200) + self.send_header("ETag", etag) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return None + + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}/model_prices_and_context_window.json" + finally: + server.shutdown() + server.server_close() + + def _attach_litellm_config(mock_prisma): """Attach a litellm_config table mock (not in conftest's _PRISMA_TABLES).""" table = MagicMock() @@ -62,17 +96,9 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): table = _attach_litellm_config(mock_prisma) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - fake_cost_map = {"gpt-4": {"input_cost": 0.03}, "gpt-3.5": {"input_cost": 0.002}} monkeypatch.setattr( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - AsyncMock(return_value=ModelCostMapReloaded(model_cost_map=fake_cost_map)), - ) - monkeypatch.setattr("litellm.add_known_models", lambda model_cost_map=None: None) - monkeypatch.setattr("litellm.model_cost", {}, raising=False) - monkeypatch.setattr( - "litellm.proxy.proxy_server._invalidate_model_cost_lowercase_map", - lambda: None, - raising=False, + AsyncMock(return_value=ModelCostMapReloaded(model_count=2)), ) async def _fake_invalidate(name): @@ -100,10 +126,7 @@ def test_reload_model_cost_map_happy(client, auth_as, monkeypatch, mock_prisma): def test_reload_model_cost_map_surfaces_the_blob_id_of_the_bytes_served_on_every_status_surface( client, auth_as, monkeypatch, mock_prisma ): - import httpx - import litellm - from litellm.litellm_core_utils.get_model_cost_map import _finalize_model_cost_map, git_blob_id from litellm.proxy import proxy_server as ps from litellm.proxy._types import LitellmUserRoles @@ -111,12 +134,7 @@ def test_reload_model_cost_map_surfaces_the_blob_id_of_the_bytes_served_on_every monkeypatch.setattr(ps, "prisma_client", mock_prisma) monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) body = _ROOT_COST_MAP.read_bytes() - expected = {"source_revision": git_blob_id(body), "etag": _SERVED_ETAG} - served = httpx.Response(200, headers={"ETag": _SERVED_ETAG}, content=body) - monkeypatch.setattr( - "litellm.litellm_core_utils.get_model_cost_map._default_reload_client", - lambda: httpx.AsyncClient(transport=httpx.MockTransport(lambda request: served)), - ) + expected = {"source_revision": _git_blob_id(body), "etag": _SERVED_ETAG} monkeypatch.setattr("litellm.add_known_models", lambda model_cost_map=None: None) monkeypatch.setattr("litellm.model_cost", {}, raising=False) @@ -125,7 +143,8 @@ def test_reload_model_cost_map_surfaces_the_blob_id_of_the_bytes_served_on_every monkeypatch.setattr(ps, "invalidate_config_param", _fake_invalidate) - with auth_as(LitellmUserRoles.PROXY_ADMIN): + with _serving(body, _SERVED_ETAG) as url, auth_as(LitellmUserRoles.PROXY_ADMIN): + monkeypatch.setattr(litellm, "model_cost_map_url", url) reload_response = client.post("/reload/model_cost_map") source_response = client.get("/model/cost_map/source") status_response = client.get("/schedule/model_cost_map_reload/status") @@ -142,7 +161,7 @@ def test_reload_model_cost_map_surfaces_the_blob_id_of_the_bytes_served_on_every assert {key: status_response.json()[key] for key in expected} == expected assert public_response.status_code == 200 assert "gpt-4o" in public_response.json() - assert reload_body["models_count"] == len(_finalize_model_cost_map(json.loads(body))) + assert reload_body["models_count"] == len(litellm.model_cost) def test_reload_model_cost_map_fetch_failure_502_keeps_map( diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index f3fe1e2520a..a988250159f 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -27,7 +27,7 @@ from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server from litellm.proxy._types import LiteLLM_ProxyModelTable, UserAPIKeyAuth -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from .conftest import normalize, route_entry_points_to_stubs # type: ignore[import-not-found] @@ -76,7 +76,7 @@ async def test_discovery_preserves_model_info_fallbacks( "fallback-deployment": {"litellm_provider": "openai", "mode": "chat"}, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) monkeypatch.setattr(proxy_server, "llm_router", router) handler: Final = AsyncHTTPHandler() await handler.client.aclose() @@ -117,7 +117,7 @@ async def test_discovery_preserves_model_info_fallbacks( assert info["input_cost_per_token"] == builtin["input_cost_per_token"] assert info["output_cost_per_token"] == builtin["output_cost_per_token"] assert info["mode"] is None - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) async def test_upstream_limits_reach_model_info_routes( @@ -163,7 +163,7 @@ async def test_upstream_limits_reach_model_info_routes( group_response: Final = client.get("/model_group/info") assert group_response.status_code == 200, group_response.text assert group_response.json()["data"][0]["max_input_tokens"] == 4096 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) # --------------------------------------------------------------------------- @@ -582,7 +582,7 @@ def test_model_info_id_lookup_reports_the_same_cost_as_the_list( "vllm-priced-id": (declared["input_cost_per_token"], declared["output_cost_per_token"]), } assert by_id == listed - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _discoverable_models(monkeypatch: pytest.MonkeyPatch, models: Mapping[str, list[str]]) -> None: @@ -707,7 +707,7 @@ def test_model_info_prices_expanded_wildcard_rows_as_the_model_their_calls_route }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ { @@ -763,7 +763,7 @@ def test_model_info_prices_expanded_wildcard_rows_as_the_model_their_calls_route for model_name, call in calls.items() } ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_prices_a_wildcard_row_as_the_other_deployment_its_calls_reach( @@ -815,7 +815,7 @@ def test_model_info_prices_a_wildcard_row_as_the_other_deployment_its_calls_reac }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ { @@ -890,7 +890,7 @@ def test_model_info_prices_a_wildcard_row_as_the_other_deployment_its_calls_reac for key, (input_cost, output_cost) in listed.items() } ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _catch_all_and_shadowed_wildcard_router(monkeypatch: pytest.MonkeyPatch) -> litellm.Router: @@ -928,7 +928,7 @@ def _catch_all_and_shadowed_wildcard_router(monkeypatch: pytest.MonkeyPatch) -> }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ { @@ -997,7 +997,7 @@ def test_model_info_prices_a_wildcard_row_aliased_to_a_sibling_name_as_the_alias shown["input_cost_per_token"] * call.usage.prompt_tokens + shown["output_cost_per_token"] * call.usage.completion_tokens ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_prices_a_wildcard_row_the_proxy_renames_as_the_model_its_calls_reach( @@ -1033,7 +1033,7 @@ def test_model_info_prices_a_wildcard_row_the_proxy_renames_as_the_model_its_cal "my-fw/acme-pro": (5e-08, 6e-08, "fireworks_ai"), "gpt-5.6-acme": (7e-07, 8e-07, "deepseek"), } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_prices_a_wildcard_row_whose_team_alias_targets_a_deleted_team_deployment( @@ -1071,7 +1071,7 @@ def test_model_info_prices_a_wildcard_row_whose_team_alias_targets_a_deleted_tea 5e-08, "fireworks_ai", ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _fireworks_wildcard_rows(client: TestClient) -> dict[str, dict[str, object]]: @@ -1122,7 +1122,7 @@ def test_model_info_prices_a_wildcard_row_as_the_deployment_the_callers_team_pub assert team_rows["my-fw/acme-mini"].get("input_cost_per_token") == 7e-07 assert other_rows["my-fw/acme-mini"].get("input_cost_per_token") == 5e-08 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _call_cost_at_row_rates(call: litellm.ModelResponse, row: Mapping[str, float]) -> float: @@ -1165,7 +1165,7 @@ def test_model_info_keeps_the_price_and_limits_of_a_row_another_wildcard_serves( "mode": "chat", } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, "gpt-5.6-acme": entry}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ { @@ -1194,7 +1194,7 @@ def test_model_info_keeps_the_price_and_limits_of_a_row_another_wildcard_serves( model=rows[0]["model_name"], messages=[{"role": "user", "content": "hi"}], mock_response="hello there" ) assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, rows[0]["model_info"])) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_prices_an_aliased_wildcard_row_as_the_model_the_other_wildcard_bills( @@ -1220,7 +1220,7 @@ def test_model_info_prices_an_aliased_wildcard_row_as_the_model_the_other_wildca }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ { @@ -1247,7 +1247,7 @@ def test_model_info_prices_an_aliased_wildcard_row_as_the_model_the_other_wildca ) assert aliased["model_info"]["input_cost_per_token"] == 3e-08 assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, aliased["model_info"])) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_keeps_the_price_of_a_wildcard_row_whose_listed_name_no_route_serves( @@ -1265,7 +1265,7 @@ def test_model_info_keeps_the_price_of_a_wildcard_row_whose_listed_name_no_route "mode": "chat", } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, "gpt-5.6-acme": entry}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ { @@ -1287,7 +1287,7 @@ def test_model_info_keeps_the_price_of_a_wildcard_row_whose_listed_name_no_route assert { row["model_info"].get("input_cost_per_token") for row in response.json()["data"] if "*" not in row["model_name"] } == {3e-07} - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize( @@ -1335,7 +1335,7 @@ def test_model_info_shows_no_wildcard_price_when_the_cost_map_cannot_say_what_ca }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[_FIREWORKS_WILDCARD, *({"model_name": "acme-target", **target} for target in targets)], model_group_alias={"my-fw/acme-mini": "acme-target"}, @@ -1349,7 +1349,7 @@ def test_model_info_shows_no_wildcard_price_when_the_cost_map_cannot_say_what_ca rows: Final = _fireworks_wildcard_rows(client) assert rows["my-fw/acme-mini"].get("input_cost_per_token") is None - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_prices_a_wildcard_row_aliased_to_an_aliased_group_as_the_group_its_calls_reach( @@ -1375,7 +1375,7 @@ def test_model_info_prices_a_wildcard_row_aliased_to_an_aliased_group_as_the_gro }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ _FIREWORKS_WILDCARD, @@ -1405,7 +1405,7 @@ def test_model_info_prices_a_wildcard_row_aliased_to_an_aliased_group_as_the_gro ) assert rows["my-fw/acme-mini"].get("input_cost_per_token") is not None assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, rows["my-fw/acme-mini"])) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_model_info_prices_a_wildcard_row_aliased_to_another_deployment_as_that_deployment_not_the_wildcards_base_model( @@ -1431,7 +1431,7 @@ def test_model_info_prices_a_wildcard_row_aliased_to_another_deployment_as_that_ }, } monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = litellm.Router( model_list=[ {**_FIREWORKS_WILDCARD, "model_info": {"id": "fireworks-wildcard", "base_model": "deepseek/acme-base"}}, @@ -1456,7 +1456,7 @@ def test_model_info_prices_a_wildcard_row_aliased_to_another_deployment_as_that_ ) assert rows["my-fw/acme-mini"].get("input_cost_per_token") is not None assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, rows["my-fw/acme-mini"])) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_saving_a_replayed_model_info_wildcard_row_stores_none_of_its_shown_metadata( @@ -1506,7 +1506,7 @@ def test_saving_a_replayed_model_info_wildcard_row_stores_none_of_its_shown_meta ] assert {name: info["access_groups"] for name, info in stored.items()} == dict.fromkeys(rows, ["prod"]) assert {name: shown_metadata[name] & set(info) for name, info in stored.items()} == dict.fromkeys(rows, set()) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index f93c23c8ce4..551c4b76766 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -17,6 +17,7 @@ import pytest import respx import litellm +from litellm.rust_bridge.host.model_cost import publish from litellm.litellm_core_utils import get_llm_provider_logic from litellm.proxy import proxy_server @@ -256,10 +257,10 @@ def lookup_fixture_model(monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setitem(litellm.model_cost, "lookup-fixture-model", entry) litellm.get_model_info.cache_clear() - litellm.utils.cached_get_model_info_helper.cache_clear() + publish(litellm.model_cost) yield entry litellm.get_model_info.cache_clear() - litellm.utils.cached_get_model_info_helper.cache_clear() + publish(litellm.model_cost) def test_model_info_lookup_returns_full_cost_map_entry_for_unregistered_model(client, auth_as, lookup_fixture_model): diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index c62e73232d0..1688dab6b5f 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -39,7 +39,12 @@ from litellm._logging import turn_on_json, verbose_proxy_logger from litellm.caching.caching import RedisCache from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded +from litellm.litellm_core_utils.get_model_cost_map import ( + GetModelCostMap, + ModelCostMapReloaded, + ModelCostMapReloadResult, + install_model_cost_map, +) from litellm.proxy._types import ( AwsErrorTypeProxyException, LitellmUserRoles, @@ -54,7 +59,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from litellm.proxy.common_utils.banner import LITELLM_BANNER from tests._master_key import MASTER_KEY from litellm.redis.stores.connection import connect @@ -4699,7 +4704,9 @@ class TestPriceDataReloadAPI: with patch( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new=AsyncMock( - return_value=ModelCostMapReloaded(model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}) + return_value=ModelCostMapReloaded( + model_count=1 + ) ), ): # Mock the database connection @@ -4721,8 +4728,7 @@ class TestPriceDataReloadAPI: assert data["models_count"] == 1 finally: # Restore the full model cost map so subsequent tests are not affected - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_reload_model_cost_map_non_admin_access(self, client_with_auth, monkeypatch): """Test that non-admin users cannot access the reload endpoint""" @@ -4955,22 +4961,13 @@ class TestPriceDataReloadIntegration: def test_complete_reload_flow(self, client_with_auth): """Test the complete reload flow from API to model cost update""" - # Mock the model cost map - mock_cost_map = { - "gpt-3.5-turbo": { - "input_cost_per_token": 0.001, - "output_cost_per_token": 0.002, - }, - "gpt-4": {"input_cost_per_token": 0.03, "output_cost_per_token": 0.06}, - } - from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded original_model_cost = litellm.model_cost.copy() try: with patch( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock(return_value=ModelCostMapReloaded(model_cost_map=mock_cost_map)), + new=AsyncMock(return_value=ModelCostMapReloaded(model_count=2)), ): # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: @@ -4986,8 +4983,7 @@ class TestPriceDataReloadIntegration: response = client_with_auth.get("/public/litellm_model_cost_map") assert response.status_code == 200 finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_pod_data_clock_seeded_from_actual_cost_map_load(self): """Regression: seeding from ProxyConfig construction time instead of the real @@ -5044,12 +5040,11 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} + model_count=1 ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - assert litellm.model_cost["gpt-3.5-turbo"] == {"input_cost_per_token": 0.001} assert proxy_config.model_cost_map_loaded_at == frozen_now assert mock_prisma.db.litellm_config.update_many.call_args[1] == { "data": {"last_run_at": frozen_now}, @@ -5058,8 +5053,7 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.upsert.assert_not_called() assert proxy_config.model_cost_map_applied_revision == 4 finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_distributed_reload_ignores_already_applied_request(self): """ @@ -5097,8 +5091,7 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.update_many.assert_not_called() assert proxy_config.model_cost_map_loaded_at == pod_data_loaded_at finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_periodic_reload_uses_pod_local_data_age(self): """ @@ -5129,12 +5122,11 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-4-test": {"input_cost_per_token": 0.5}} + model_count=1 ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - assert litellm.model_cost["gpt-4-test"] == {"input_cost_per_token": 0.5} assert proxy_config.model_cost_map_loaded_at == frozen_now assert mock_prisma.db.litellm_config.update_many.call_args[1]["data"] == {"last_run_at": frozen_now} @@ -5148,8 +5140,7 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.update_many.assert_not_called() assert proxy_config.model_cost_map_loaded_at == frozen_now - timedelta(hours=1) finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_every_pod_applies_a_manual_revision_exactly_once(self): """The fleet property: no pod clears the revision, so each one reloads on the tick @@ -5174,7 +5165,7 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + model_count=1 ) for _ in range(3): @@ -5184,8 +5175,7 @@ class TestPriceDataReloadIntegration: assert mock_get_map.call_count == len(pods) assert all(p.model_cost_map_applied_revision == 1 for p in pods) finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) @pytest.mark.parametrize( "published_revision, expect_reload", @@ -5218,7 +5208,7 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + model_count=1 ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -5227,8 +5217,7 @@ class TestPriceDataReloadIntegration: assert mock_get_map.call_count == (1 if expect_reload else 0) assert proxy_config.model_cost_map_applied_revision == published_revision finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_distributed_reload_stamps_last_run_without_creating_row(self): """ @@ -5254,7 +5243,7 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + model_count=1 ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -5266,8 +5255,7 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.upsert.assert_not_called() mock_prisma.db.litellm_config.create.assert_not_called() finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_distributed_reload_leaves_request_unserved_when_status_write_fails(self): """ @@ -5296,15 +5284,14 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.1}} + model_count=1 ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) assert proxy_config.model_cost_map_applied_revision == 0 finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_distributed_reload_keeps_current_map_when_fetch_fails(self): """Fetch failure during a periodic reload must not downgrade the pod or count the @@ -5360,8 +5347,6 @@ class TestPriceDataReloadIntegration: The swap happens early in the handler, so a failure in the bookkeeping after it is swallowed by the surrounding except and would otherwise leave the metadata correct while the path is quietly broken""" - from litellm import utils as litellm_utils - from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() @@ -5374,20 +5359,18 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) original_model_cost = litellm.model_cost - original_registry = dict(litellm_utils._runtime_registered_model_cost) try: litellm.register_model( model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} ) + async def reload_the_bundled_document(url: str) -> ModelCostMapReloadResult: + return install_model_cost_map(GetModelCostMap.read_local_model_cost_map_bytes()) + with ( patch( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock( - return_value=ModelCostMapReloaded( - model_cost_map={"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} - ) - ), + new=reload_the_bundled_document, ), patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), patch("litellm.proxy.proxy_server.verbose_proxy_logger") as mock_logger, @@ -5395,39 +5378,29 @@ class TestPriceDataReloadIntegration: asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) mock_logger.exception.assert_not_called() + assert litellm.model_cost is not original_model_cost assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 assert "gpt-4o" in litellm.model_cost assert proxy_config.model_cost_map_applied_revision == 7 finally: - litellm.model_cost = original_model_cost - litellm_utils._runtime_registered_model_cost.clear() - litellm_utils._runtime_registered_model_cost.update(original_registry) - _invalidate_model_cost_lowercase_map() - - def test_swap_in_model_cost_map_counts_the_fetched_catalog_only(self): - """The count the reload endpoints report describes the price data, so it - is taken before the runtime registrations are written back into the same - dict. Counting after would inflate it by however many deployments and - overrides this pod happens to be carrying""" - from litellm import utils as litellm_utils - from litellm.proxy.proxy_server import _swap_in_model_cost_map + publish(original_model_cost) + def test_install_counts_the_fetched_catalog_only(self): + """The count the reload endpoints report describes the price data, so the + runtime registrations written back into the same dict are not counted. + Counting them would inflate it by however many deployments and overrides + this pod happens to be carrying""" original_model_cost = litellm.model_cost - original_registry = dict(litellm_utils._runtime_registered_model_cost) try: litellm.register_model( model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} ) + reloaded = install_model_cost_map(GetModelCostMap.read_local_model_cost_map_bytes()) - models_count = _swap_in_model_cost_map({"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}) - - assert models_count == 1 + assert reloaded == ModelCostMapReloaded(model_count=len(litellm.model_cost) - 1) assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 finally: - litellm.model_cost = original_model_cost - litellm_utils._runtime_registered_model_cost.clear() - litellm_utils._runtime_registered_model_cost.update(original_registry) - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_manual_reload_preserves_interval_hours(self, monkeypatch): """ @@ -5460,7 +5433,7 @@ class TestPriceDataReloadIntegration: patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = ModelCostMapReloaded( - model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + model_count=1 ) mock_prisma.db.litellm_config.upsert = AsyncMock( return_value=_reload_schedule_row({}, reload_revision=9) @@ -5485,8 +5458,7 @@ class TestPriceDataReloadIntegration: "the serving pod must adopt the revision it published, not reload again" ) finally: - litellm.model_cost = original_model_cost - _invalidate_model_cost_lowercase_map() + publish(original_model_cost) def test_anthropic_beta_headers_reload_preserves_interval_hours(self): """Test that _check_and_reload_anthropic_beta_headers preserves interval_hours after reload. diff --git a/tests/unit/responses/test_responses_utils.py b/tests/unit/responses/test_responses_utils.py index 49ff7918849..e3eb9961228 100644 --- a/tests/unit/responses/test_responses_utils.py +++ b/tests/unit/responses/test_responses_utils.py @@ -13,7 +13,7 @@ from litellm.responses.litellm_completion_transformation.transformation import L from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIOptionalRequestParams, ResponsesAPIResponse from litellm.types.utils import ModelInfo, Usage -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish class TestResponsesAPIRequestUtils: @@ -984,7 +984,7 @@ def test_responses_object_bills_an_explicit_cache_hit_at_the_explicit_rate(): "cache_read_input_token_cost": 5e-07, "explicit_cache_read_input_token_cost": 2.5e-07, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) try: chat_usage = Usage( prompt_tokens=1000, @@ -1012,4 +1012,4 @@ def test_responses_object_bills_an_explicit_cache_hit_at_the_explicit_rate(): litellm.model_cost.pop(model, None) else: litellm.model_cost[model] = previous_row - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) diff --git a/tests/unit/router_strategy/test_lowest_cost.py b/tests/unit/router_strategy/test_lowest_cost.py index e3fdcd69499..f71d1ad9144 100644 --- a/tests/unit/router_strategy/test_lowest_cost.py +++ b/tests/unit/router_strategy/test_lowest_cost.py @@ -7,7 +7,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome DEPLOYMENT_ID = "9876" @@ -163,7 +163,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index f04ade1ce2d..1139de6c0fc 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -11,7 +11,7 @@ import traceback from typing import Dict, Final, Iterator from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from unittest.mock import patch @@ -724,7 +724,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/rust_bridge/host/test_model_cost.py b/tests/unit/rust_bridge/host/test_model_cost.py new file mode 100644 index 00000000000..82a52a941a1 --- /dev/null +++ b/tests/unit/rust_bridge/host/test_model_cost.py @@ -0,0 +1,80 @@ +from typing import Final + +import pytest + +import litellm +from litellm.cost_calculator import cost_per_token +from litellm.litellm_core_utils.get_model_cost_map import model_cost_inputs +from litellm.rust_bridge.host.model_cost import ( + changed_in_place, + entry, + merged, + publish, + resolve, + write, +) + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +def test_resolve_lands_on_the_row_a_lookup_in_the_published_table_resolves() -> None: + landing: Final = resolve(litellm.model_cost, "openai/gpt-4o", {"input_cost_per_token": 1.0}, None) + + assert landing is not None + assert (landing.key, landing.matched_builtin) == ("gpt-4o", True) + + +def test_merged_keeps_the_alias_names_and_adds_only_the_fields_a_registration_set() -> None: + current: Final = entry({"mode": "chat", "x": 1}, ("owner", "alias"), frozenset({"mode"})) + + row: Final = merged(current, {"x": 2}, {"litellm_provider": "openai", "output_cost_per_token": None}) + + assert row == {"mode": "chat", "x": 2, "litellm_provider": "openai"} + assert (row.names, row.registered_fields) == (("owner", "alias"), frozenset({"mode", "litellm_provider"})) + + +_MODEL: Final = "lit-table-model" + + +def _row(price: float) -> dict[str, object]: + return {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": price, "output_cost_per_token": price} + + +def _billed() -> float: + prompt, completion = cost_per_token( + model=_MODEL, custom_llm_provider="openai", prompt_tokens=1, completion_tokens=0 + ) + return prompt + completion + + +def test_a_published_table_is_read_in_place_of_what_was_cached_for_the_one_before() -> None: + publish({**litellm.model_cost, _MODEL: _row(1.0)}) + assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 1.0 + assert _billed() == 1.0 + + publish({**litellm.model_cost, _MODEL: _row(2.0)}) + + assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 2.0 + assert _billed() == 2.0 + + +def test_rows_written_into_the_live_table_replace_a_cached_value_of_the_same_size() -> None: + publish({**litellm.model_cost, _MODEL: _row(1.0)}) + assert _billed() == 1.0 + + write(litellm.model_cost, {_MODEL: _row(3.0)}) + + assert litellm.get_model_info(_MODEL)["input_cost_per_token"] == 3.0 + assert _billed() == 3.0 + + +def test_a_provider_changed_in_place_regroups_the_key_once_the_change_is_reported( + monkeypatch: pytest.MonkeyPatch, +) -> None: + assert "gpt-4o" not in model_cost_inputs.keys_in(litellm.model_cost, "lit-moved") + + monkeypatch.setitem(litellm.model_cost["gpt-4o"], "litellm_provider", "lit-moved") + changed_in_place() + + assert model_cost_inputs.keys_in(litellm.model_cost, "lit-moved") == ["gpt-4o"] + assert "gpt-4o" not in model_cost_inputs.keys_in(litellm.model_cost, "openai") diff --git a/tests/unit/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py index e9a153a3f8d..35cac056b38 100644 --- a/tests/unit/test_conftest_isolation.py +++ b/tests/unit/test_conftest_isolation.py @@ -1,4 +1,5 @@ import asyncio +import json import sys from collections.abc import Iterator from dataclasses import dataclass @@ -9,8 +10,13 @@ import pytest import litellm from litellm import Router -from litellm import utils as litellm_utils_module from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.get_model_cost_map import ( + GetModelCostMap, + ModelCostMapReloaded, + install_model_cost_map, + model_cost_inputs, +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from tests.unit import conftest as unit_harness @@ -80,11 +86,11 @@ class _CanaryRouterHolder: def test_register_model_ledger_entry_is_scoped_to_this_test(): litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) - assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost + assert CANARY_MODEL in model_cost_inputs.registered_keys() def test_register_model_ledger_entry_was_rolled_back(): - assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost + assert CANARY_MODEL not in model_cost_inputs.registered_keys() def test_register_model_catalog_writes_are_scoped_to_this_test(): @@ -158,7 +164,7 @@ def test_aws_cache_reset_sees_a_replaced_module(monkeypatch: pytest.MonkeyPatch) @pytest.fixture(scope="module", autouse=True) def registered_collection_models() -> Iterator[None]: - original_registrations: Final = dict(litellm_utils_module._runtime_registered_model_cost) + original_inputs: Final = model_cost_inputs.snapshot() original_routers: Final = frozenset(unit_harness.live_router_registries) litellm.register_model( {"collection-isolation-ledger": {"litellm_provider": "openai", "input_cost_per_token": 0.001}} @@ -172,17 +178,18 @@ def registered_collection_models() -> Iterator[None]: ] ) yield - litellm_utils_module._runtime_registered_model_cost.clear() - litellm_utils_module._runtime_registered_model_cost.update(original_registrations) + model_cost_inputs.restore(original_inputs) unit_harness.live_router_registries.discard(collection_router._core._hosted) unit_harness.live_router_registries.update(original_routers) -def test_collection_model_metadata_is_not_replayed_into_a_test_catalog() -> None: - from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map - - catalog: Final = {"unit-test-catalog-model": {"litellm_provider": "openai", "input_cost_per_token": 0.001}} - adopt_model_cost_map(catalog) +def test_collection_model_metadata_is_not_replayed_into_a_test_catalog(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + catalog: Final = { + **GetModelCostMap.load_local_model_cost_map(), + "unit-test-catalog-model": {"litellm_provider": "openai", "input_cost_per_token": 0.001}, + } + assert isinstance(install_model_cost_map(json.dumps(catalog).encode()), ModelCostMapReloaded) assert "collection-isolation-ledger" not in litellm.model_cost assert "collection-isolation-router" not in litellm.model_cost assert "unit-test-catalog-model" in litellm.model_cost diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 1d9f79aafc6..9f3bfc2de86 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -37,7 +37,8 @@ from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject -from litellm.utils import _invalidate_model_cost_lowercase_map, cached_get_model_info_helper, supports_prompt_caching +from litellm.utils import cached_get_model_info_helper, supports_prompt_caching +from litellm.rust_bridge.host.model_cost import publish import os import threading from unittest.mock import MagicMock, patch @@ -724,7 +725,7 @@ def test_completion_cost_bills_a_requested_model_priced_after_an_unpriced_lookup requested_model, {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) cost: Final = completion_cost(completion_response=display_name_response()) @@ -11060,7 +11061,6 @@ def test_bedrock_converse_profile_preserves_registered_price_override( import litellm.utils as model_utils from litellm.cost_calculator import pricing_entry_for_cost_calc - monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) catalog: Final = dict(litellm.model_cost) model: Final = "openai.gpt-5.6-sol" litellm.register_model( @@ -11125,7 +11125,7 @@ def test_bedrock_converse_profile_preserves_registered_price_override( ) == pytest.approx(10 * updated_rate + 5 * rate) monkeypatch.setattr(litellm, "model_cost", catalog) - model_utils._invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) model_utils.reapply_runtime_model_cost_registrations() profile_prices: Final = litellm.get_model_info(f"us.{model}", custom_llm_provider="bedrock") expected_after_reload: Final = ( @@ -11144,9 +11144,6 @@ def test_bedrock_converse_profile_preserves_registered_price_override( def test_bedrock_converse_profile_ignores_registered_capabilities_without_prices( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: - import litellm.utils as model_utils - - monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) model: Final = "openai.gpt-6-sol" litellm.register_model({model: {"supports_function_calling": True}}) response: Final = ModelResponse(model=model, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)) @@ -11179,7 +11176,6 @@ def test_bedrock_converse_profile_preserves_cache_only_registered_prices( import litellm.utils as model_utils from litellm.cost_calculator import pricing_entry_for_cost_calc - monkeypatch.setattr(model_utils, "_runtime_registered_model_cost", {}) model: Final = "openai.cache-overlay-test" profile: Final = f"{profile_prefix}.{model}" catalog: Final = { @@ -11200,7 +11196,7 @@ def test_bedrock_converse_profile_preserves_cache_only_registered_prices( } profile_before: Final = dict(catalog[profile]) monkeypatch.setattr(litellm, "model_cost", dict(catalog)) - model_utils._invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) litellm.register_model({model: {cache_field: rate}}, persist_across_reloads=persist_across_reloads) response: Final = ModelResponse( model=model, @@ -11251,7 +11247,7 @@ def test_bedrock_converse_profile_preserves_cache_only_registered_prices( assert litellm.model_cost[profile] == profile_before monkeypatch.setattr(litellm, "model_cost", dict(catalog)) - model_utils._invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) model_utils.reapply_runtime_model_cost_registrations() expected_after_reload: Final = ( 10 * 0.02 + 3 * 0.04 + 5 * (rate if persist_across_reloads else catalog[profile][cache_field]) @@ -11269,8 +11265,6 @@ def test_bedrock_converse_profile_preserves_cache_only_registered_prices( def test_bedrock_cache_registration_with_unpriced_profile_keeps_registered_rates( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, profile: str ) -> None: - import litellm.utils as model_utils - model: Final = "openai.registered-cache-test" monkeypatch.setattr( litellm, @@ -11286,7 +11280,7 @@ def test_bedrock_cache_registration_with_unpriced_profile_keeps_registered_rates "us.openai.unpriced-profile-test": {"litellm_provider": "bedrock", "mode": "chat"}, }, ) - model_utils._invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) litellm.register_model({model: {"cache_read_input_token_cost": 0.001}}, persist_across_reloads=False) response: Final = ModelResponse( model=model, @@ -11378,7 +11372,6 @@ def test_bedrock_registered_cache_prices_preserve_dynamic_threshold_rates( cache_field: str, tier_override: float | None, ) -> None: - import litellm.utils as model_utils from litellm.cost_calculator import pricing_entry_for_cost_calc model: Final = "openai.threshold-overlay-test" @@ -11400,7 +11393,7 @@ def test_bedrock_registered_cache_prices_preserve_dynamic_threshold_rates( }, }, ) - model_utils._invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) litellm.register_model( { model: { @@ -11683,7 +11676,7 @@ def test_per_call_custom_prices_keep_the_regional_and_speed_multipliers( multiplier: float, ) -> None: monkeypatch.setitem(litellm.model_cost, call_kwargs["model"], {**row, **_PER_CALL_RATES}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) response: Final = litellm.completion( messages=[{"role": "user", "content": "hello"}], diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 0509b8dcf30..dbd9ba8a32e 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -5572,21 +5572,14 @@ def _priced_at(prompt_tokens, completion_tokens): @pytest.fixture def local_cost_map(monkeypatch): """The prices these tests assert are the checked-in ones. Setting the environment - variable alone does not reload the map, so pin the map itself. - - Prices are read through two separate lru_caches, so pinning ``model_cost`` is not - enough on its own: an entry warmed against the network-fetched map keeps its old - prices and billing reads those while the assertions read the pinned map. - ``_invalidate_model_cost_lowercase_map`` clears both caches, where - ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out - so entries never leak across tests in either direction.""" - from litellm.utils import _invalidate_model_cost_lowercase_map + variable alone does not reload the map, so pin the map itself.""" + from litellm.rust_bridge.host.model_cost import publish monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) yield - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): diff --git a/tests/unit/test_model_cost_aliases.py b/tests/unit/test_model_cost_aliases.py index f9f92a85cb2..c01e4d07e3e 100644 --- a/tests/unit/test_model_cost_aliases.py +++ b/tests/unit/test_model_cost_aliases.py @@ -1,14 +1,32 @@ """ Tests for the ``aliases`` feature in the model cost map. -The ``_expand_model_aliases`` function processes ``aliases`` lists from model -entries, creating shared dict references for alias entries at load time. +An entry's ``aliases`` list serves the entry under every alias name, all of them +one shared ``AliasedCostEntry`` object, when a document is installed. """ -from unittest.mock import patch +import json +import logging +from typing import Final -from litellm import verbose_logger -from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases +import pytest + +from litellm.rust_bridge._native import NativeInstalled, NativeModelCostMap, NativeRejected +from litellm.rust_bridge.host.model_cost import AliasedCostEntry + + +def _install(document: dict[str, dict[str, object]]) -> NativeInstalled: + installed: Final = NativeModelCostMap.bundled(0, 0.0).install(json.dumps(document).encode()) + assert isinstance(installed, NativeInstalled) + return installed + + +def _installed(document: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]: + return _install(document).rows + + +def _row(**fields: object) -> dict[str, object]: + return {"litellm_provider": "test", "mode": "chat", **fields} # --------------------------------------------------------------------------- @@ -17,97 +35,48 @@ from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases class TestExpandModelAliases: - """Unit tests for _expand_model_aliases.""" - def test_basic_expansion(self): """Aliases are added as top-level entries in model_cost.""" - model_cost = { - "my-model-latest": { - "aliases": ["my-model-20250101"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + result = _installed({"my-model-latest": _row(aliases=["my-model-20250101"], input_cost_per_token=1e-06)}) - assert "my-model-20250101" in result assert result["my-model-20250101"]["input_cost_per_token"] == 1e-06 assert result["my-model-20250101"]["litellm_provider"] == "test" def test_multiple_aliases(self): """A single entry can declare multiple aliases.""" - model_cost = { - "provider/model-latest": { - "aliases": ["provider/model-v1", "provider/model-v2"], - "input_cost_per_token": 5e-06, - "litellm_provider": "provider", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + result = _installed({"provider/model-latest": _row(aliases=["provider/model-v1", "provider/model-v2"])}) - assert "provider/model-v1" in result - assert "provider/model-v2" in result + assert list(result) == ["provider/model-latest", "provider/model-v1", "provider/model-v2"] - def test_shared_dict_reference(self): - """Alias entries share the same dict object as the canonical entry (no copy).""" - model_cost = { - "canonical-model": { - "aliases": ["alias-model"], - "input_cost_per_token": 2e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + def test_shared_entry_object(self): + """Alias entries are the same object as the canonical entry, which knows every name it is served under.""" + result = _installed({"canonical-model": _row(aliases=["alias-model"], input_cost_per_token=2e-06)}) - assert result["alias-model"] is result["canonical-model"] + entry: Final = result["canonical-model"] + assert result["alias-model"] is entry + assert isinstance(entry, AliasedCostEntry) + assert entry.names == ("canonical-model", "alias-model") + assert entry.registered_fields == frozenset() def test_aliases_key_removed(self): """The ``aliases`` key is removed from the entry after expansion.""" - model_cost = { - "my-model": { - "aliases": ["my-model-alias"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + result = _installed({"my-model": _row(aliases=["my-model-alias"])}) assert "aliases" not in result["my-model"] assert "aliases" not in result["my-model-alias"] - def test_entries_without_aliases_unchanged(self): - """Entries with no ``aliases`` key are left untouched.""" - model_cost = { - "plain-model": { - "input_cost_per_token": 3e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + def test_entries_without_aliases_stay_plain_dicts(self): + """Entries with no ``aliases`` key are served as written, as plain dicts.""" + result = _installed({"plain-model": _row(input_cost_per_token=3e-06)}) - assert "plain-model" in result - assert result["plain-model"]["input_cost_per_token"] == 3e-06 - assert len(result) == 1 + assert result == {"plain-model": _row(input_cost_per_token=3e-06)} + assert type(result["plain-model"]) is dict def test_empty_aliases_list(self): """An empty ``aliases`` list is treated the same as no aliases.""" - model_cost = { - "model-a": { - "aliases": [], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + result = _installed({"model-a": _row(aliases=[])}) - assert len(result) == 1 - assert "model-a" in result + assert list(result) == ["model-a"] assert "aliases" not in result["model-a"] @@ -116,81 +85,119 @@ class TestExpandModelAliases: # --------------------------------------------------------------------------- -class TestAliasConflicts: - """Tests for alias conflict detection and handling.""" - - def test_alias_conflicts_with_canonical_entry(self): - """Alias that matches an existing canonical entry is skipped with a warning.""" - model_cost = { - "model-latest": { - "aliases": ["model-dated"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", +@pytest.mark.parametrize( + ("document", "name", "row", "other", "shared"), + [ + pytest.param( + { + "model-latest": _row(aliases=["model-dated"], input_cost_per_token=1e-06), + "model-dated": _row(input_cost_per_token=5e-06), }, - "model-dated": { - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", + "model-dated", + _row(input_cost_per_token=5e-06), + "model-latest", + False, + id="alias-names-a-later-entry", + ), + pytest.param( + { + "existing-model": _row(input_cost_per_token=9.99e-06), + "other-model": _row(aliases=["existing-model"], input_cost_per_token=1e-06), }, - } - with patch.object(verbose_logger, "warning") as mock_warn: - result = _expand_model_aliases(model_cost) - - # The canonical "model-dated" entry is preserved, not overwritten - assert "model-dated" in result - # Verify a warning about the alias conflict was logged - mock_warn.assert_called() - warning_messages = " ".join(str(c) for c in mock_warn.call_args_list) - assert "alias conflict" in warning_messages.lower() - - def test_duplicate_alias_across_entries(self): - """Same alias claimed by two different entries: second one is skipped.""" - model_cost = { - "model-a": { - "aliases": ["shared-alias"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", + "existing-model", + _row(input_cost_per_token=9.99e-06), + "other-model", + False, + id="alias-names-an-earlier-entry", + ), + pytest.param( + { + "model-a": _row(aliases=["shared-alias"], input_cost_per_token=1e-06), + "model-b": _row(aliases=["shared-alias"], input_cost_per_token=2e-06), }, - "model-b": { - "aliases": ["shared-alias"], - "input_cost_per_token": 2e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - with patch.object(verbose_logger, "warning") as mock_warn: - result = _expand_model_aliases(model_cost) + "shared-alias", + _row(input_cost_per_token=1e-06), + "model-a", + True, + id="two-entries-claim-one-alias", + ), + ], +) +def test_an_alias_never_replaces_a_model_name_and_the_first_claim_wins( + document: dict[str, dict[str, object]], name: str, row: dict[str, object], other: str, shared: bool +): + result: Final = _installed(document) - # "shared-alias" should point to model-a (first one wins) - assert "shared-alias" in result - assert result["shared-alias"]["input_cost_per_token"] == 1e-06 - # Verify a warning about the alias conflict was logged - mock_warn.assert_called() - warning_messages = " ".join(str(c) for c in mock_warn.call_args_list) - assert "alias conflict" in warning_messages.lower() + assert result[name] == row + assert (result[name] is result[other]) is shared - def test_canonical_entry_not_overwritten_by_alias(self): - """An alias must never overwrite an existing canonical entry's data.""" - original_cost = 9.99e-06 - model_cost = { - "existing-model": { - "input_cost_per_token": original_cost, - "litellm_provider": "test", - "mode": "chat", - }, - "other-model": { - "aliases": ["existing-model"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) - # Original entry must be preserved - assert result["existing-model"]["input_cost_per_token"] == original_cost +class TestAliasOfRows: + def test_an_alias_of_row_is_served_as_its_target_where_the_document_put_it(self): + """An ``alias_of`` row is its target's object, even when written before the target.""" + result = _installed( + { + "model-latest": {"alias_of": "model"}, + "model": _row(input_cost_per_token=1e-06), + } + ) + + assert list(result) == ["model-latest", "model"] + assert result["model-latest"] is result["model"] + assert result["model"] == _row(input_cost_per_token=1e-06) + + @pytest.mark.parametrize( + "pointer", [{"alias_of": "ghost"}, {"alias_of": "model", "mode": "chat"}], ids=["no-target", "other-fields"] + ) + def test_an_alias_of_row_the_catalog_cannot_serve_is_dropped(self, pointer: dict[str, object]): + result = _installed({"model": _row(), "model-latest": pointer}) + + assert list(result) == ["model"] + + def test_the_model_count_is_every_key_the_document_serves(self): + installed: Final = _install( + { + "model": _row(aliases=["model-dated"]), + "model-latest": {"alias_of": "model"}, + "ghost-latest": {"alias_of": "ghost"}, + "sample_spec": {"mode": "chat"}, + } + ) + + assert list(installed.rows) == ["model", "model-latest", "sample_spec", "model-dated"] + assert installed.model_count == 4 + + +class TestAliasWarnings: + @pytest.mark.parametrize( + ("document", "warned"), + [ + ({"model": _row(), "other": _row(aliases=["model"])}, 'alias "model" is already a model name'), + ( + {"a": _row(aliases=["shared"]), "b": _row(aliases=["shared"])}, + 'alias "shared" already belongs to a', + ), + ({"model": _row(), "ghost-latest": {"alias_of": "ghost"}}, 'alias_of "ghost" is not a model name'), + ( + {"model": _row(), "model-latest": {"alias_of": "model", "mode": "chat"}}, + "model-latest.mode: an alias_of row carries no other field", + ), + ({"model": _row(aliases="model-dated")}, "model.aliases: invalid value"), + ], + ) + def test_an_alias_the_document_cannot_serve_is_logged( + self, caplog: pytest.LogCaptureFixture, document: dict[str, dict[str, object]], warned: str + ): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + _installed(document) + + assert any(warned in record.getMessage() for record in caplog.records), caplog.text + + def test_a_document_with_servable_aliases_logs_nothing(self, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + _installed({"model": _row(aliases=["model-dated"]), "model-latest": {"alias_of": "model"}}) + + assert caplog.records == [] # --------------------------------------------------------------------------- @@ -199,47 +206,24 @@ class TestAliasConflicts: class TestAliasIntegration: - """Higher-level tests verifying aliases work with the model_cost dict.""" - def test_mutation_through_alias_visible_on_canonical(self): """Since alias is a shared reference, mutations are visible on both.""" - model_cost = { - "canonical": { - "aliases": ["alias"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + result = _installed({"canonical": _row(aliases=["alias"], input_cost_per_token=1e-06)}) - # Mutate via alias result["alias"]["input_cost_per_token"] = 999 assert result["canonical"]["input_cost_per_token"] == 999 def test_mixed_entries_with_and_without_aliases(self): - """A model_cost dict with a mix of aliased and plain entries.""" - model_cost = { - "model-with-alias": { - "aliases": ["alias-1", "alias-2"], - "input_cost_per_token": 1e-06, - "litellm_provider": "test", - "mode": "chat", - }, - "plain-model": { - "input_cost_per_token": 2e-06, - "litellm_provider": "test", - "mode": "chat", - }, - } - result = _expand_model_aliases(model_cost) + """Aliases follow every document key, grouped by owner.""" + result = _installed( + { + "model-with-alias": _row(aliases=["alias-1", "alias-2"]), + "plain-model": _row(), + } + ) - assert len(result) == 4 # 2 canonical + 2 aliases - assert "alias-1" in result - assert "alias-2" in result - assert "plain-model" in result - assert "model-with-alias" in result + assert list(result) == ["model-with-alias", "plain-model", "alias-1", "alias-2"] - def test_expand_on_empty_dict(self): - """Expanding an empty dict returns an empty dict.""" - assert _expand_model_aliases({}) == {} + def test_an_empty_document_is_rejected(self): + """An empty document is never installed.""" + assert isinstance(NativeModelCostMap.bundled(0, 0.0).install(b"{}"), NativeRejected) diff --git a/tests/unit/test_no_top_level_test_invocations.py b/tests/unit/test_no_top_level_test_invocations.py index 0a7c3600752..386c71d779e 100644 --- a/tests/unit/test_no_top_level_test_invocations.py +++ b/tests/unit/test_no_top_level_test_invocations.py @@ -8,7 +8,7 @@ import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome LOCAL_TESTING_DIR = Path(__file__).parents[1] / "local_testing" @@ -101,7 +101,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index fb42f6229e0..6ceea94d843 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1383,13 +1383,6 @@ ALIAS_CASES: Final = ( "is_non_openai_azure_model", False, ), - ( - "litellm.litellm_core_utils.get_model_cost_map", - "GetModelCostMap", - "_get_backup_model_count", - "get_backup_model_count", - False, - ), ( "litellm.litellm_core_utils.health_check_helpers", "HealthCheckHelpers", diff --git a/tests/unit/test_pydantic_namespaces.py b/tests/unit/test_pydantic_namespaces.py index 81d075077db..fc21ea501db 100644 --- a/tests/unit/test_pydantic_namespaces.py +++ b/tests/unit/test_pydantic_namespaces.py @@ -7,7 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -79,7 +79,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/test_register_model_custom_pricing.py b/tests/unit/test_register_model_custom_pricing.py index afff995aa69..d4f591e7291 100644 --- a/tests/unit/test_register_model_custom_pricing.py +++ b/tests/unit/test_register_model_custom_pricing.py @@ -10,15 +10,25 @@ calculations for DB-sourced models with prompt caching pricing. """ import copy +import datetime import itertools +import json import os -from collections.abc import Callable, Mapping +from collections.abc import Callable +from decimal import Decimal from typing import Final, NoReturn import pytest import litellm -from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases, adopt_model_cost_map +from litellm.litellm_core_utils.get_model_cost_map import ( + GetModelCostMap, + ModelCostMapReloaded, + install_model_cost_map, + model_cost_inputs, +) +from litellm.rust_bridge._native import NativeTableSwapped +from litellm.rust_bridge.host.model_cost import publish from litellm.types.utils import ( Choices, CompletionTokensDetailsWrapper, @@ -27,7 +37,6 @@ from litellm.types.utils import ( PromptTokensDetailsWrapper, Usage, ) -from litellm.utils import _invalidate_model_cost_lowercase_map def _snapshot_model_cost_entries(keys): @@ -40,7 +49,7 @@ def _restore_model_cost_entries(original_entries): litellm.model_cost.pop(key, None) else: litellm.model_cost[key] = value - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_receives_cache_pricing_fields(): @@ -137,10 +146,6 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp "max_tokens": 4096, } - # ``register_model`` calls ``get_model_info.cache_clear`` via - # ``_invalidate_model_cost_lowercase_map``, so the replacement must - # expose a no-op ``cache_clear`` attribute. - _fake_get_model_info.cache_clear = lambda: None monkeypatch.setattr(litellm_utils, "get_model_info", _fake_get_model_info) try: @@ -253,9 +258,9 @@ def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(monkeypa litellm.model_cost.pop(registered_key, None) litellm.model_cost = original_model_cost os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog): @@ -620,7 +625,7 @@ def test_embedding_direct_sdk_custom_pricing_prices_the_call_without_writing_the assert plain._hidden_params["response_cost"] != response._hidden_params["response_cost"] finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(monkeypatch: pytest.MonkeyPatch) -> None: @@ -649,7 +654,7 @@ def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(mo "anthropic.lit-5853-geo-model": sibling_row, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: raise RuntimeError("dictionary changed size during iteration") @@ -664,7 +669,7 @@ def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(mo assert litellm.get_model_info(geo_key)[field] == geo_row[field] assert litellm.model_cost[geo_key]["litellm_provider"] == geo_row["litellm_provider"] finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize("route", ["", "converse/", "invoke/"]) @@ -688,7 +693,7 @@ def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once( }, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) real_lookup: Final = litellm.utils.get_model_info lookups: Final = itertools.count() @@ -706,7 +711,7 @@ def test_register_model_leaves_no_shadow_key_when_the_lookup_trips_once( assert litellm.get_model_info(alias)["key"] == catalog_key assert litellm.get_model_info(alias)["input_cost_per_token"] == repriced_input_cost finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize("key_prefix", ["", "bedrock/", "bedrock/converse/"]) @@ -733,7 +738,7 @@ def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup def long_prompt_cost_after_registering_with(lookup: Callable[..., object]) -> float: monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) with pytest.MonkeyPatch.context() as lookup_patch: lookup_patch.setattr(litellm.utils, "get_model_info", lookup) litellm.register_model( @@ -760,7 +765,7 @@ def test_register_model_bills_a_catalog_rows_tier_rates_the_same_when_the_lookup long_prompt_cost_after_registering_with(litellm.utils.get_model_info) ) finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize("lookup_trips", [False, True]) @@ -781,7 +786,7 @@ def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_th monkeypatch.setattr( litellm, "open_ai_chat_completion_models", {*litellm.open_ai_chat_completion_models, openai_key} ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: raise RuntimeError("dictionary changed size during iteration") @@ -800,7 +805,11 @@ def test_register_model_prices_a_provider_less_azure_ai_key_without_repricing_th (1e-03, 2e-03) ) finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) + + +def _publish(cost_map: dict[str, dict[str, object]]) -> None: + publish(cost_map) def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monkeypatch: pytest.MonkeyPatch) -> None: @@ -816,7 +825,7 @@ def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monk builtin_cost_map: Final = litellm.get_model_cost_map(url="") monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", {**builtin_cost_map, catalog_key: dict(catalog_row)}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) reloaded_cost_map: Final = { **builtin_cost_map, catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost}, @@ -826,7 +835,7 @@ def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monk def lookup_racing_a_price_data_reload(model: str, custom_llm_provider: str | None = None) -> object: info: Final = real_lookup(model=model, custom_llm_provider=custom_llm_provider) if litellm.model_cost is not reloaded_cost_map: - adopt_model_cost_map(reloaded_cost_map) + _publish(reloaded_cost_map) return info with pytest.MonkeyPatch.context() as lookup_patch: @@ -837,7 +846,7 @@ def test_register_model_keeps_the_prices_a_reload_adopts_while_it_registers(monk assert litellm.model_cost[catalog_key]["input_cost_per_token"] == reloaded_input_cost assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes( @@ -859,19 +868,19 @@ def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes( catalog_key: {**catalog_row, "input_cost_per_token": reloaded_input_cost}, } - class CostMapAReloadReplacesAsItIsWritten(dict[str, dict[str, object]]): - def update(self, other: Mapping[str, dict[str, object]], /) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # registration writes pass one mapping - if litellm.model_cost is self: - adopt_model_cost_map(reloaded_cost_map) - super().update(other) + written_map: Final = {**builtin_cost_map, catalog_key: dict(catalog_row)} + registration_base: Final = litellm.utils.registration_base + + def reload_while_resolving(*args: object, **kwargs: object) -> object: + resolved = registration_base(*args, **kwargs) + if litellm.model_cost is written_map: + _publish(reloaded_cost_map) + return resolved monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr( - litellm, - "model_cost", - CostMapAReloadReplacesAsItIsWritten({**builtin_cost_map, catalog_key: dict(catalog_row)}), - ) - _invalidate_model_cost_lowercase_map() + monkeypatch.setattr(litellm, "model_cost", written_map) + monkeypatch.setattr(litellm.utils, "registration_base", reload_while_resolving) + publish(litellm.model_cost) litellm.register_model( {deployment_key: {"litellm_provider": "bedrock", "output_cost_per_token": registered_output_cost}}, @@ -883,7 +892,35 @@ def test_register_model_keeps_its_price_when_a_reload_adopts_a_map_as_it_writes( assert litellm.get_model_info(deployment_key)["input_cost_per_token"] == reloaded_input_cost assert litellm.get_model_info(deployment_key)["output_cost_per_token"] == registered_output_cost finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) + + +def test_register_model_lands_in_the_live_map_when_a_reload_swaps_it_on_every_attempt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + key: Final = "lit-always-swapped-model" + builtin_cost_map: Final = litellm.get_model_cost_map(url="") + registration_base: Final = litellm.utils.registration_base + + def reload_while_resolving(*args: object, **kwargs: object) -> object: + resolved = registration_base(*args, **kwargs) + _publish(dict(builtin_cost_map)) + return resolved + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", dict(builtin_cost_map)) + monkeypatch.setattr(litellm.utils, "registration_base", reload_while_resolving) + publish(litellm.model_cost) + + litellm.register_model( + {key: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 1e-06}}, + persist_across_reloads=False, + ) + + try: + assert litellm.model_cost.get(key, {}).get("input_cost_per_token") == 1e-06, sorted(litellm.model_cost)[:3] + finally: + publish(litellm.model_cost) def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: pytest.MonkeyPatch) -> None: @@ -897,7 +934,7 @@ def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: catalog_key: {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 2e-06}, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) walk: Final = iter(litellm.model_cost[catalog_key].items()) first_field: Final = next(walk) @@ -914,7 +951,7 @@ def test_register_model_leaves_the_entry_a_lookup_is_walking_intact(monkeypatch: assert litellm.get_model_info(catalog_key)["input_cost_per_token"] == 1e-06 assert litellm.get_model_info(catalog_key)["supports_vision"] is True finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _CANONICAL_KEY: Final = "lit-5853-canonical-model" @@ -923,23 +960,18 @@ _ALIAS_KEY: Final = "lit-5853-alias-model" def _install_aliased_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr( - litellm, - "model_cost", - _expand_model_aliases( - { - **litellm.get_model_cost_map(url=""), - _CANONICAL_KEY: { - "aliases": [_ALIAS_KEY], - "litellm_provider": "openai", - "mode": "chat", - "input_cost_per_token": 2e-06, - "output_cost_per_token": 8e-06, - }, - } - ), - ) - _invalidate_model_cost_lowercase_map() + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + document: Final = { + **GetModelCostMap.load_local_model_cost_map(), + _CANONICAL_KEY: { + "aliases": [_ALIAS_KEY], + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 2e-06, + "output_cost_per_token": 8e-06, + }, + } + assert isinstance(install_model_cost_map(json.dumps(document).encode()), ModelCostMapReloaded) def _token_prices(name: str) -> tuple[float | None, float | None]: @@ -967,7 +999,7 @@ def test_register_model_prices_every_alias_of_the_registered_entry( _ALIAS_KEY: (1e-06, 4e-06), } finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_leaves_an_alias_repointed_at_its_own_entry(monkeypatch: pytest.MonkeyPatch) -> None: @@ -989,13 +1021,13 @@ def test_register_model_leaves_an_alias_repointed_at_its_own_entry(monkeypatch: _ALIAS_KEY: (5e-06, 6e-06), } finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) builtin_key: Final = "us.anthropic.claude-sonnet-4-6" registered_key: Final = f"bedrock/bedrock/bedrock/{builtin_key}" builtin: Final = litellm.model_cost[builtin_key] @@ -1009,7 +1041,7 @@ def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypa assert info["output_cost_per_token"] == builtin["output_cost_per_token"] > 0 assert info["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"] finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_resolves_to( @@ -1030,7 +1062,7 @@ def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_reso } monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", {**litellm.get_model_cost_map(url=""), sibling_key: sibling_row}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) previous_rules: Final = list(get_fallback_generalization_rules()) set_fallback_generalizations( [ @@ -1057,7 +1089,7 @@ def test_register_model_keeps_the_capabilities_an_unregistered_region_alias_reso assert registered["input_cost_per_token"] == sibling_row["input_cost_per_token"] finally: set_fallback_generalizations(previous_rules) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize( @@ -1131,7 +1163,7 @@ def test_register_model_keeps_custom_flat_prices_over_inherited_rate_variants( builtin_key: {"litellm_provider": "bedrock", "mode": "chat", **builtin_prices}, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) try: litellm.register_model( @@ -1151,7 +1183,7 @@ def test_register_model_keeps_custom_flat_prices_over_inherited_rate_variants( ) assert cost == pytest.approx(expected_cost) finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pytest.MonkeyPatch) -> None: @@ -1171,7 +1203,7 @@ def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pyt }, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) registered_key: Final = f"bedrock/{openai_key}" try: @@ -1183,7 +1215,7 @@ def test_register_model_never_inherits_pricing_across_providers(monkeypatch: pyt assert "input_cost_per_token" not in registered assert "output_cost_per_token" not in registered finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_update_dictionary_merges_nested_dicts_without_aliasing(): @@ -1519,7 +1551,7 @@ def test_direct_call_with_flat_rates_is_priced_from_them_without_writing_the_cos assert model_key not in litellm.model_cost finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_direct_call_with_audio_rates_is_priced_per_modality_without_writing_the_cost_map(): @@ -1539,7 +1571,7 @@ def test_direct_call_with_audio_rates_is_priced_per_modality_without_writing_the assert model_key not in litellm.model_cost finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_direct_call_with_directional_per_second_rates_is_priced_from_them_without_writing_the_cost_map(): @@ -1559,7 +1591,7 @@ def test_direct_call_with_directional_per_second_rates_is_priced_from_them_witho assert model_key not in litellm.model_cost finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_router_call_with_request_prices_bills_them_over_the_deployment_rates(): @@ -1636,7 +1668,7 @@ def test_direct_call_flat_rates_override_the_catalogs_tiers_per_direction(): assert "input_cost_per_token" not in litellm.model_cost[model_key] finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_direct_call_with_only_a_per_second_rate_on_a_mapped_model_is_billed_per_second(): @@ -1666,7 +1698,7 @@ def test_direct_call_with_only_a_per_second_rate_on_a_mapped_model_is_billed_per assert litellm.model_cost[model_key]["input_cost_per_token"] == 1e-06 finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_direct_call_flat_rate_applies_to_the_callers_own_tiers_not_the_catalogs(): @@ -1698,4 +1730,168 @@ def test_direct_call_flat_rate_applies_to_the_callers_own_tiers_not_the_catalogs assert litellm.model_cost[model_key]["tiered_pricing"][0]["output_cost_per_token"] == 6e-07 finally: litellm.model_cost.pop(model_key, None) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) + + +_RELOAD_BASE_KEY: Final = "lit-reload-base-model" +_RELOAD_NESTED_KEY: Final = "lit-reload-nested-model" +_RELOAD_DOCUMENT: Final = { + **GetModelCostMap.load_local_model_cost_map(), + _RELOAD_BASE_KEY: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + "cache_read_input_token_cost": 1e-07, + "cache_creation_input_token_cost": 3e-06, + }, + _RELOAD_NESTED_KEY: { + "litellm_provider": "openai", + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.02, + "search_context_size_high": 0.03, + }, + }, +} + + +def _reload_cost_map() -> None: + assert isinstance(install_model_cost_map(json.dumps(_RELOAD_DOCUMENT).encode()), ModelCostMapReloaded) + + +def test_a_reload_keeps_what_a_registration_inherited_for_a_key_the_document_lacks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + _reload_cost_map() + registered_key: Final = f"bedrock/bedrock/bedrock/{_RELOAD_BASE_KEY}" + litellm.register_model( + {registered_key: {"input_cost_per_token": 5e-06, "output_cost_per_token": 6e-06, "litellm_provider": "bedrock"}} + ) + registered: Final = dict(litellm.model_cost[registered_key]) + assert ( + registered["cache_read_input_token_cost"] == _RELOAD_DOCUMENT[_RELOAD_BASE_KEY]["cache_read_input_token_cost"] + ) + + _reload_cost_map() + + assert dict(litellm.model_cost[registered_key]) == registered + + +def test_a_persisted_registration_keeps_values_json_cannot_hold_across_a_reload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + registered_key: Final = "lit-registered-with-python-values" + created_at: Final = datetime.datetime(2026, 1, 2, tzinfo=datetime.timezone.utc) + litellm.register_model( + { + registered_key: { + "litellm_provider": "openai", + "input_cost_per_token": Decimal("0.000001"), + "created_at": created_at, + } + } + ) + + _reload_cost_map() + + assert litellm.model_cost[registered_key]["created_at"] is created_at + assert litellm.model_cost[registered_key]["input_cost_per_token"] == Decimal("0.000001") + + +def test_a_reload_keeps_the_document_sub_keys_beside_a_registered_nested_field( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_cost", litellm.model_cost) + _reload_cost_map() + litellm.register_model({_RELOAD_NESTED_KEY: {"search_context_cost_per_query": {"search_context_size_low": 0.5}}}) + expected: Final = { + **_RELOAD_DOCUMENT[_RELOAD_NESTED_KEY]["search_context_cost_per_query"], + "search_context_size_low": 0.5, + } + assert litellm.model_cost[_RELOAD_NESTED_KEY]["search_context_cost_per_query"] == expected + + _reload_cost_map() + + assert litellm.model_cost[_RELOAD_NESTED_KEY]["search_context_cost_per_query"] == expected + + +def test_a_lookup_that_raises_reaches_the_caller_and_keeps_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + error: Final = LookupError("lit-lookup-failed") + + def failing_lookup(*_args: object, **_kwargs: object) -> NoReturn: + raise error + + monkeypatch.setattr(litellm.utils, "registration_base", failing_lookup) + registered_before: Final = model_cost_inputs.registered_keys() + + with pytest.raises(LookupError) as raised: + litellm.register_model({"lit-raise-model": {"input_cost_per_token": 1e-06, "litellm_provider": "openai"}}) + + assert raised.value is error + assert model_cost_inputs.registered_keys() == registered_before + assert "lit-raise-model" not in litellm.model_cost + + +def test_registering_into_a_table_that_is_not_published_reports_it_and_writes_nothing() -> None: + other: Final = {"lit-other-model": {"litellm_provider": "openai"}} + + result: Final = model_cost_inputs.register(other, {"lit-swap-model": {"input_cost_per_token": 1e-06}}, False) + + assert isinstance(result, NativeTableSwapped) + assert other == {"lit-other-model": {"litellm_provider": "openai"}} + + +def test_an_entry_landing_on_an_earlier_entry_of_the_same_call_merges_onto_that_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + key: Final = "lit-chained-model" + litellm.register_model( + { + key: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + } + }, + persist_across_reloads=False, + ) + assert litellm.get_model_info(f"openai/{key}")["max_input_tokens"] is None + + litellm.register_model( + {key: {"max_input_tokens": 4096}, f"openai/{key}": {"input_cost_per_token": 5e-06}}, + persist_across_reloads=False, + ) + + assert litellm.model_cost[key]["max_input_tokens"] == 4096 + assert litellm.model_cost[key]["input_cost_per_token"] == 5e-06 + + +def test_a_router_deployment_registers_a_datetime_in_its_model_info(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm import Router + + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + created_at: Final = datetime.datetime(2026, 1, 2, 3, 4, 5, tzinfo=datetime.timezone.utc) + router: Final = Router( + model_list=[ + { + "model_name": "lit-datetime-model", + "litellm_params": { + "model": "openai/lit-datetime-model", + "api_key": "fake-key-for-registration", + "input_cost_per_token": 1e-06, + "output_cost_per_token": 2e-06, + }, + "model_info": {"id": "lit-datetime-deployment", "created_at": created_at}, + } + ] + ) + + assert litellm.model_cost["lit-datetime-deployment"]["created_at"] == created_at + assert litellm.model_cost["lit-datetime-deployment"]["input_cost_per_token"] == 1e-06 + del router diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 24ade9b48fb..6557db7be32 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -101,7 +101,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) from litellm.types.utils import CallTypes, CredentialItem, Usage -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE from tests.large_text import text @@ -14262,7 +14262,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index cd71c3f9bee..f5c71668e63 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -42,9 +42,9 @@ from litellm.types.router import ( ModelInfo, ) from litellm.utils import ( - _invalidate_model_cost_lowercase_map, reapply_runtime_model_cost_registrations, ) +from litellm.rust_bridge.host.model_cost import publish _COPILOT_TOKEN_EXPIRES_AT: Final = 4_102_444_800 @@ -53,15 +53,14 @@ _COPILOT_TOKEN_EXPIRES_AT: Final = 4_102_444_800 def _simulate_price_data_reload(fetched_catalog): """Drive what a price data reload does to this process's litellm state. - Mirrors `litellm.proxy.proxy_server._swap_in_model_cost_map`, which is the + Mirrors `litellm.litellm_core_utils.get_model_cost_map._install`, which is the one place both reload paths adopt a freshly fetched catalog; that wiring is covered in the proxy's own tests, so these exercise the replay itself without dragging the proxy in. The provider model sets that helper also repopulates are left alone, since nothing here reads them and rebuilding them from a two-entry catalog would outlive the test. """ - litellm.model_cost = fetched_catalog - _invalidate_model_cost_lowercase_map() + publish(fetched_catalog) reapply_runtime_model_cost_registrations() @@ -81,7 +80,7 @@ def _restore_model_cost_entries(original_entries): litellm.model_cost.pop(key, None) else: litellm.model_cost[key] = value - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize("initial_count", (1, DEFAULT_MAX_LRU_CACHE_SIZE + 1)) @@ -127,7 +126,7 @@ async def test_discovered_limits_survive_deployment_growth_and_removal( assert all( router.get_configured_token_limits(deployment.model_name) == (4096, 4096) for deployment in deployments[:-2] ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch: pytest.MonkeyPatch) -> None: @@ -172,7 +171,7 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch router.cache.in_memory_cache.flush_cache() await router.arefresh_model_info(client=handler) assert router.get_configured_token_limits("local") == (2048, 2048) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) async def test_discovery_sends_the_secret_behind_an_oidc_api_key( @@ -207,7 +206,7 @@ async def test_discovery_sends_the_secret_behind_an_oidc_api_key( handler.client = client await router.arefresh_model_info(client=handler) assert authorizations == ["Bearer local-secret"] - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None: @@ -261,7 +260,7 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: await first.arefresh_model_info(client=handler) assert first.get_configured_token_limits("local") == (None, None) assert second.get_configured_token_limits("local") == (2048, 2048) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None: @@ -297,7 +296,7 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa assert router.get_configured_token_limits("first") == (2048, 2048) assert router.get_configured_token_limits("second") == (None, None) assert router.get_configured_token_limits("third") == (2048, 2048) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize("provider", ("hosted_vllm", "openai", "openai_like", "text-completion-openai")) @@ -354,7 +353,7 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( router.get_router_model_info(id="two", deployment=None, received_model_name="local")["max_output_tokens"] == 512 ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypatch: pytest.MonkeyPatch) -> None: @@ -408,7 +407,7 @@ async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypa assert router.get_configured_token_limits("byok") == (None, None) assert next(responses, None) is None assert responder.call_count == 2 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_should_not_pollute_shared_key_with_zero_cost_pricing(): @@ -844,7 +843,7 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): backend_entry["litellm_provider"] = "chatgpt" backend_entry["mode"] = "responses" litellm.model_cost[backend_model] = backend_entry - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) Router( model_list=[ @@ -1160,7 +1159,7 @@ def test_partial_custom_pricing_bills_explicit_cache_hits( "cache_read_input_token_cost": 5e-07, "explicit_cache_read_input_token_cost": 2.5e-07, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) try: Router( model_list=[ @@ -1234,13 +1233,13 @@ _AZURE_TIER_BACKEND_ENTRY: Final = { def _register_tier_backend() -> None: litellm.model_cost[_TIER_BACKEND_KEY] = copy.deepcopy(_TIER_BACKEND_ENTRY) litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _register_azure_tier_backend() -> None: litellm.model_cost[_AZURE_TIER_BACKEND_KEY] = copy.deepcopy(_AZURE_TIER_BACKEND_ENTRY) litellm.get_model_info.cache_clear() - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None: @@ -2014,13 +2013,6 @@ def test_price_data_reload_preserves_router_registered_model_info(monkeypatch): max_input_tokens / max_output_tokens from every custom model group and /model_group/info starts reporting nulls. """ - from litellm import utils as litellm_utils - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) router = Router( model_list=[ @@ -2052,8 +2044,7 @@ def test_price_data_reload_preserves_router_registered_model_info(monkeypatch): assert after.max_input_tokens == 128000 assert after.max_output_tokens == 16384 finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_price_data_reload_preserves_custom_override_of_a_catalog_model(monkeypatch): @@ -2062,13 +2053,6 @@ def test_price_data_reload_preserves_custom_override_of_a_catalog_model(monkeypa the same bug: the reload does not blank the metadata, it reverts the operator's model_info override to the upstream catalog values. """ - from litellm import utils as litellm_utils - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) router = Router( model_list=[ @@ -2102,8 +2086,7 @@ def test_price_data_reload_preserves_custom_override_of_a_catalog_model(monkeypa assert after.max_input_tokens == 12345 assert after.max_output_tokens == 678 finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_deleted_deployments_are_not_replayed_onto_later_reloads(monkeypatch): @@ -2114,13 +2097,6 @@ def test_deleted_deployments_are_not_replayed_onto_later_reloads(monkeypatch): key that another live deployment still points at must survive the same deletion. """ - from litellm import utils as litellm_utils - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) router = Router( model_list=[ @@ -2159,8 +2135,7 @@ def test_deleted_deployments_are_not_replayed_onto_later_reloads(monkeypatch): assert surviving["max_input_tokens"] == 222 assert "hosted_vllm/shared-backend" in litellm.model_cost finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_deleting_a_deployment_leaves_catalog_pricing_for_its_backend_model(monkeypatch): @@ -2168,13 +2143,6 @@ def test_deleting_a_deployment_leaves_catalog_pricing_for_its_backend_model(monk A backend key is shared with the fetched catalog, so withdrawing the entries a deleted deployment owns must not take real upstream pricing down with it. """ - from litellm import utils as litellm_utils - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) backend_model = "gemini/gemini-2.5-pro" catalog_entry = litellm.get_model_info(model=backend_model) @@ -2203,8 +2171,7 @@ def test_deleting_a_deployment_leaves_catalog_pricing_for_its_backend_model(monk assert "doomed-gemini-id" not in litellm.model_cost assert litellm.model_cost[backend_model]["input_cost_per_token"] == catalog_input_cost finally: - litellm.model_cost = saved_catalog - _invalidate_model_cost_lowercase_map() + publish(saved_catalog) def test_repointing_a_deployment_drops_its_previous_backend_key(monkeypatch): @@ -2213,13 +2180,6 @@ def test_repointing_a_deployment_drops_its_previous_backend_key(monkeypatch): old backend key behind, and a replayed registry would re-assert it onto every later catalog for the life of the process. """ - from litellm import utils as litellm_utils - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) router = Router( model_list=[ @@ -2249,8 +2209,7 @@ def test_repointing_a_deployment_drops_its_previous_backend_key(monkeypatch): assert "hosted_vllm/new-backend" in litellm.model_cost assert "moving-target-id" in litellm.model_cost finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_mantle_chat_deployment_does_not_rewrite_the_shared_responses_row(): @@ -2274,7 +2233,7 @@ def test_mantle_chat_deployment_does_not_rewrite_the_shared_responses_row(): "litellm_provider": "bedrock_mantle", "mode": "responses", } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) Router( model_list=[ @@ -2314,7 +2273,7 @@ def _register_routing_test_row(monkeypatch: pytest.MonkeyPatch, mode: str) -> No "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize("registered_mode, pinned_mode", [("responses", "chat"), ("chat", "responses")]) @@ -2588,7 +2547,7 @@ def test_a_github_copilot_chat_pin_opts_a_dual_endpoint_model_out_of_the_respons "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], "max_input_tokens": 1000, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) assert github_copilot_supports_responses_api("github_copilot/routing-test-model") is True Router( @@ -2630,7 +2589,7 @@ def test_a_github_copilot_responses_pin_still_opts_the_model_into_the_responses_ "supported_endpoints": ["/v1/responses"] if registered_mode == "responses" else ["/v1/chat/completions"], "max_input_tokens": 1000, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) Router( model_list=[ @@ -2685,8 +2644,7 @@ def test_a_discarded_router_stops_contributing_to_later_reloads(monkeypatch): assert "hosted_vllm/kept-backend" in litellm.model_cost assert kept.model_list # keep the live router referenced for the duration finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: @@ -2727,8 +2685,7 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: assert "member_auto_router" not in litellm.model_cost["gpt-4o"] assert router.model_list finally: - litellm.model_cost = saved_catalog - _invalidate_model_cost_lowercase_map() + publish(saved_catalog) def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): @@ -2883,8 +2840,7 @@ def test_reload_keeps_custom_pricing_configured_on_litellm_params_for_a_db_model assert litellm.model_cost["db-priced-id"]["input_cost_per_token"] == 0.000123 assert litellm.model_cost["db-priced-id"]["output_cost_per_token"] == 0.000456 finally: - litellm.model_cost = saved_catalog - _invalidate_model_cost_lowercase_map() + publish(saved_catalog) def test_reload_keeps_custom_pricing_of_a_db_model_on_a_router_built_without_a_model_list( @@ -2949,8 +2905,7 @@ def test_replay_live_router_model_cost_rebuilds_every_live_router(): assert litellm.model_cost["second-id"]["max_input_tokens"] == 222 assert first.model_list and second.model_list finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_strategy_router_alias_pricing_never_enters_model_cost(monkeypatch): @@ -2962,13 +2917,6 @@ def test_strategy_router_alias_pricing_never_enters_model_cost(monkeypatch): strip must also survive a price-data reload, which rebuilds entries by walking the live routers. """ - from litellm import utils as litellm_utils - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) router = Router( model_list=[ @@ -3008,8 +2956,7 @@ def test_strategy_router_alias_pricing_never_enters_model_cost(monkeypatch): _assert_alias_unpriced() assert router.model_list finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) # --- a config.yaml PTU deployment must not also bill per token ------------------ @@ -3605,9 +3552,9 @@ def test_every_declaring_deployment_is_named(caplog): def _simulate_price_data_reload_with_provider_sets(monkeypatch, fetched_catalog): """Like `_simulate_price_data_reload`, plus the provider model-set refresh the proxy's - `_swap_in_model_cost_map` does before replaying, so bare names in the new catalog resolve.""" + `_install` does before replaying, so bare names in the new catalog resolve.""" monkeypatch.setattr(litellm, "model_cost", fetched_catalog) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) litellm.add_known_models(model_cost_map=fetched_catalog) reapply_runtime_model_cost_registrations() @@ -3717,7 +3664,7 @@ def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_ monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") catalog: Final = litellm.get_model_cost_map(url="") monkeypatch.setattr(litellm, "model_cost", dict(catalog)) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) backends: Final = tuple( f"bedrock/au.anthropic.{model}" for model in ("claude-sonnet-5", "claude-opus-5", "claude-haiku-4-5-20251001-v1:0", "claude-sonnet-4-6") @@ -3780,7 +3727,7 @@ def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_ assert litellm.get_model_info(backend)["input_cost_per_token"] == canonical_prices[backend] assert tuple(router.get_model_names()) == tuple(backend.rsplit(".", 1)[-1] for backend in backends) finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _COLLIDING_BUILTIN_KEY: Final = "baseten/zai-org/GLM-5.2" @@ -3951,12 +3898,12 @@ def test_router_capacity_widens_to_a_larger_base_model_window( base: _capacity_test_metadata(provider, 8192), }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = _capacity_test_router(route, base, base_location, explicit_limit) expected: Final = explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192) assert _effective_max_input_tokens(router) == expected assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) @@ -3973,11 +3920,11 @@ def test_router_capacity_never_narrows_to_a_smaller_base_model_window( base: _capacity_test_metadata(provider, 1024), }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = _capacity_test_router(route, base, base_location, None) assert _effective_max_input_tokens(router) == 8192 assert _deployments_allowed_for_2048_input_tokens(router) == 1 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) @@ -3985,10 +3932,10 @@ def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_is_ monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str ) -> None: monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, base: _capacity_test_metadata(provider, 1024)}) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = _capacity_test_router(f"{route}-not-in-cost-map", base, "model_info", None) assert _deployments_allowed_for_2048_input_tokens(router) == 1 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) @@ -4005,11 +3952,11 @@ def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_dec "model_cost", {**litellm.model_cost, route: route_without_window, base: _capacity_test_metadata(provider, 1024)}, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = _capacity_test_router(route, base, "model_info", None) assert _effective_max_input_tokens(router) is None assert _deployments_allowed_for_2048_input_tokens(router) == 1 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SAME_MODEL_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9" @@ -4057,11 +4004,11 @@ def test_router_capacity_widens_only_for_an_exact_base_model_entry_of_the_same_m "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended": larger_window, }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = _capacity_test_router(route, base_model, "model_info", None) assert _effective_max_input_tokens(router) == expected assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_one( @@ -4077,7 +4024,7 @@ def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_o "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": _capacity_test_metadata("bedrock", 8192), }, ) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) router: Final = Router( model_list=[ { @@ -4096,7 +4043,7 @@ def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_o ) assert _effective_max_input_tokens(router) == 1024 assert _deployments_allowed_for_2048_input_tokens(router) == 0 - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_deployment_id_equal_to_another_providers_catalog_key_bills_its_own_prices_under_pricing_provider() -> None: diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index fbc3c0e47c3..5856caf7740 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -21,7 +21,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -444,7 +444,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index a8421d8afbd..13a45ceeaf3 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -88,7 +88,6 @@ from litellm.secret_managers.main import oidc_cache from litellm.types.videos.main import VideoObject from litellm.utils import ( acredential_updates, - _invalidate_model_cost_lowercase_map, add_custom_logger_callback_to_specific_event, check_valid_key, CustomStreamWrapper, @@ -121,6 +120,7 @@ from litellm.utils import ( validate_environment, validate_chat_completion_tool_choice, ) +from litellm.rust_bridge.host.model_cost import publish @pytest.mark.parametrize("root_base_model", [42, ["gemini-3.5-flash"], {"model": "gemini-3.5-flash"}]) @@ -2524,11 +2524,9 @@ def test_register_model_with_scientific_notation(): test_model_name = f"test-scientific-notation-model-{uuid.uuid4().hex[:12]}" # Clear LRU caches that might have stale data - from litellm.utils import ( - _invalidate_model_cost_lowercase_map, - ) + from litellm.rust_bridge.host.model_cost import publish - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) model_cost_dict = { test_model_name: { "max_tokens": 8192, @@ -2551,27 +2549,7 @@ def test_register_model_with_scientific_notation(): # Clean up after test if test_model_name in litellm.model_cost: del litellm.model_cost[test_model_name] - _invalidate_model_cost_lowercase_map() - - -@pytest.mark.usefixtures("isolated_openai_model_sets") -@respx.mock -def test_register_model_url_fetch_uses_single_attempt(monkeypatch): - monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False) - monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) - before = dict(litellm.model_cost) - threads_before = {thread.name for thread in threading.enumerate()} - route = respx.get("https://example.invalid/custom_pricing.json").mock(return_value=httpx.Response(503)) - - litellm.register_model(model_cost="https://example.invalid/custom_pricing.json") - - threads_after = {thread.name for thread in threading.enumerate()} - assert route.call_count == 1 - assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"} - assert not any( - thread.name == "litellm-model-cost-map-retry" and thread.is_alive() for thread in threading.enumerate() - ) - assert litellm.model_cost.keys() >= before.keys() + publish(litellm.model_cost) def test_register_model_openrouter_without_slash(): @@ -5188,15 +5166,10 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( """ from litellm import utils as litellm_utils from litellm.utils import ( - _invalidate_model_cost_lowercase_map, reapply_runtime_model_cost_registrations, ) + from litellm.rust_bridge.host.model_cost import publish - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) # Only the recorded half is under test here; the live-router rebuild is covered # in test_router_model_cost_isolation.py. Routers built by earlier tests in this # process stay in the weak set until they are collected, so leaving the callback @@ -5230,7 +5203,7 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( "max_input_tokens": 4242, } } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) reapply_runtime_model_cost_registrations() expected_rate: Final = rate if updated_rate is None else updated_rate @@ -5238,8 +5211,7 @@ def test_reapply_runtime_registrations_replays_register_model_overrides( assert litellm.model_cost["openai/gpt-4o"]["supports_function_calling"] is False assert litellm.model_cost["openai/gpt-4o"]["max_input_tokens"] == 4242 finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) @pytest.mark.usefixtures("isolated_openai_model_sets", "local_model_cost_map") @@ -5250,17 +5222,10 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey the catalog generation it was applied to and silently beat fresh upstream pricing forever, while a durable override registered alongside it survives. """ - from litellm import utils as litellm_utils from litellm.utils import ( - _invalidate_model_cost_lowercase_map, reapply_runtime_model_cost_registrations, ) - - monkeypatch.setattr( - litellm_utils, - "_runtime_registered_model_cost", - dict(litellm_utils._runtime_registered_model_cost), - ) + from litellm.rust_bridge.host.model_cost import publish saved_model_cost = litellm.model_cost try: @@ -5277,23 +5242,17 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey "openai/gpt-4o": {"litellm_provider": "openai", "input_cost_per_token": 0.000999}, "openai/gpt-4o-mini": {"litellm_provider": "openai", "input_cost_per_token": 0.000888}, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) reapply_runtime_model_cost_registrations() assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000111 assert litellm.model_cost["openai/gpt-4o-mini"]["input_cost_per_token"] == 0.000888 finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() + publish(saved_model_cost) def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.utils import _invalidate_model_cost_lowercase_map, _rebuild_model_cost_lowercase_map - - class _KeyRegisteredMidWalk(str): - def lower(self) -> str: - litellm.model_cost.setdefault("lit-5853-registered-during-the-rebuild", {"mode": "chat"}) - return super().lower() + from litellm.rust_bridge.host.model_cost import publish class _EntryRegisteredMidWalk(dict): def get(self, field: str, default: object = None) -> object: @@ -5301,35 +5260,12 @@ def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch: p return super().get(field, default) monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) - litellm.model_cost[_KeyRegisteredMidWalk("LIT-5853-WALKED-LAST")] = _EntryRegisteredMidWalk(mode="chat") + litellm.model_cost["LIT-5853-WALKED-LAST"] = _EntryRegisteredMidWalk(mode="chat") try: - assert _rebuild_model_cost_lowercase_map()["lit-5853-walked-last"] == "LIT-5853-WALKED-LAST" litellm.add_known_models(model_cost_map=litellm.model_cost) - assert "lit-5853-registered-during-the-rebuild" in litellm.model_cost assert "lit-5853-registered-during-add-known-models" in litellm.model_cost finally: - _invalidate_model_cost_lowercase_map() - - -class _KeyWhoseLoweringLosesTheRace(str): - """Lowering the key is the one Python-level step between the lookup map's rebuild and - its read, so another thread's registration invalidating the map lands right there.""" - - def lower(self) -> str: - from litellm.utils import _invalidate_model_cost_lowercase_map - - _invalidate_model_cost_lowercase_map() - return super().lower() - - -def test_case_insensitive_model_cost_lookup_survives_an_invalidation_mid_lookup() -> None: - from litellm.utils import _get_model_cost_key, _invalidate_model_cost_lowercase_map - - _invalidate_model_cost_lowercase_map() - try: - assert _get_model_cost_key(_KeyWhoseLoweringLosesTheRace("GPT-4O")) == "gpt-4o" - finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def _register_a_model_mid_walk() -> None: @@ -5357,7 +5293,7 @@ class _EntryReadDuringARegistration(dict[str, object]): return super().__contains__(key) -@pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5", "case-insensitive scan"]) +@pytest.mark.parametrize("walk", ["fal_ai", "fireworks_ai", "ollama", "gpt-5"]) def test_request_path_cost_map_walks_survive_a_registration_mid_walk( monkeypatch: pytest.MonkeyPatch, walk: str ) -> None: @@ -5365,14 +5301,13 @@ def test_request_path_cost_map_walks_survive_a_registration_mid_walk( from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.ollama.common_utils import OllamaModelInfo from litellm.llms.openai.chat.gpt_reasoning_transformation import _catalogue_declares_default_effort - from litellm.utils import _handle_new_key_with_scan, _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish walks: Final[Mapping[str, Callable[[], object]]] = { "fal_ai": lambda: _keyed_rows("lit-5853-model", "standard"), "fireworks_ai": FireworksAIConfig._get_fireworks_index, "ollama": lambda: OllamaModelInfo._is_static_ollama_model("ollama/lit-5853-model"), "gpt-5": _catalogue_declares_default_effort, - "case-insensitive scan": lambda: _handle_new_key_with_scan("lit-5853-model"), } monkeypatch.setattr( litellm, @@ -5384,7 +5319,7 @@ def test_request_path_cost_map_walks_survive_a_registration_mid_walk( walks[walk]() assert "lit-5853-registered-mid-walk" in litellm.model_cost finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) class _JsonCapture(logging.Handler): @@ -8506,7 +8441,7 @@ def test_get_model_info_gemini(monkeypatch): def test_register_model_without_a_provider_never_inherits_another_providers_entry( monkeypatch: pytest.MonkeyPatch, ) -> None: - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish key: Final = "azure/claude-sonnet-4-6" monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) @@ -8521,7 +8456,7 @@ def test_register_model_without_a_provider_never_inherits_another_providers_entr assert info["input_cost_per_token"] == 1e-07 assert info["output_cost_per_token"] == 4e-07 finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize( @@ -8539,7 +8474,7 @@ def test_register_model_without_a_provider_never_inherits_another_providers_entr def test_register_model_over_a_builtin_row_bills_explicit_cache_hits( monkeypatch: pytest.MonkeyPatch, registered_prices: dict[str, float], billed_cache_read_rate: float ) -> None: - from litellm.utils import _invalidate_model_cost_lowercase_map + from litellm.rust_bridge.host.model_cost import publish key: Final = "dashscope/qwen-register-explicit-cache-test" monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) @@ -8551,7 +8486,7 @@ def test_register_model_over_a_builtin_row_bills_explicit_cache_hits( "cache_read_input_token_cost": 5e-07, "explicit_cache_read_input_token_cost": 2.5e-07, } - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) try: litellm.register_model({key: registered_prices}, persist_across_reloads=False) explicit_hit: Final = ModelResponse( @@ -8566,7 +8501,7 @@ def test_register_model_over_a_builtin_row_bills_explicit_cache_hits( 1000 * billed_cache_read_rate ) finally: - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) @pytest.mark.parametrize( @@ -8946,7 +8881,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = { @@ -11609,7 +11544,7 @@ def test_the_catalog_with_register_model_entries_lists_each_model_once_and_leave ) def test_runtime_model_cost_registrations_resolve_bedrock_routing(monkeypatch: pytest.MonkeyPatch, key: str) -> None: from litellm import utils as model_utils - from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry + from litellm.rust_bridge.host.model_cost import AliasedCostEntry rates: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} alias: Final = "openai.test-converse-alias" diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index 7548001ce4c..cde0a8978c0 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -13,7 +13,8 @@ from litellm import get_model_info from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.utils import ModelInfoBase -from litellm.utils import _invalidate_model_cost_lowercase_map, supports_native_structured_output +from litellm.utils import supports_native_structured_output +from litellm.rust_bridge.host.model_cost import publish import os from unittest.mock import MagicMock, patch @@ -23,7 +24,7 @@ def isolate_model_info_state(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) monkeypatch.setattr(litellm, "custom_provider_map", list(litellm.custom_provider_map)) yield - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) def test_get_model_info_simple_model_name(): @@ -119,7 +120,7 @@ def test_get_model_info_bedrock_region(monkeypatch): key: value for key, value in litellm.get_model_cost_map(url="").items() if key != regional_model } monkeypatch.setattr(litellm, "model_cost", model_cost_without_regional_entry) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) info = litellm.get_model_info(model=regional_model, custom_llm_provider="bedrock") print("info", info) assert info["key"] == "anthropic.claude-haiku-4-5-20251001-v1:0" diff --git a/tests/unit/types/test_files.py b/tests/unit/types/test_files.py index fffb1eecc61..5b2c82b69db 100644 --- a/tests/unit/types/test_files.py +++ b/tests/unit/types/test_files.py @@ -16,7 +16,7 @@ from litellm.types.files import ( get_file_mime_type_from_extension, get_file_type_from_extension, ) -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.rust_bridge.host.model_cost import publish from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -120,7 +120,7 @@ def isolate_litellm_state(): for attr, original_value in original_state.items(): if hasattr(litellm, attr): setattr(litellm, attr, original_value) - _invalidate_model_cost_lowercase_map() + publish(litellm.model_cost) _SCALAR_DEFAULTS = {