diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index cbeea71ad0f..c4c7417a52b 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -578,9 +578,6 @@ jobs: - name: code_qa_check_tests run: uv run --no-sync python ./tests/code_coverage_tests/code_qa_check_tests.py - - name: check_get_model_cost_key_performance - run: uv run --no-sync python ./tests/code_coverage_tests/check_get_model_cost_key_performance.py - - name: test_proxy_types_import run: uv run --no-sync python ./tests/code_coverage_tests/test_proxy_types_import.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b93a8827428..53e9329099a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3590,19 +3590,14 @@ dependencies = [ name = "litellm-cache-s3" version = "0.1.0" dependencies = [ - "aws-credential-types", - "aws-sdk-s3", - "aws-smithy-runtime-api", "aws-smithy-types", - "aws-types", "futures-util", - "http 1.4.2", "litellm-auth-aws", "litellm-cache", "litellm-cache-testing", "litellm-error", "litellm-http", - "reqwest", + "litellm-storage-s3", "rstest", "serde_json", "tokio", @@ -4250,12 +4245,14 @@ dependencies = [ "jsonschema", "litellm-cost", "litellm-error", + "litellm-python-compat", "macro_rules_attribute", "rstest", "schemars 1.2.2", "serde", "serde_json", "serde_with", + "sha1 0.10.7", "strum", "thiserror 2.0.19", "time", @@ -4265,14 +4262,19 @@ dependencies = [ name = "litellm-model-catalog-refresh" version = "0.1.0" dependencies = [ + "http 1.4.2", + "litellm-auth-aws", "litellm-error", "litellm-http", - "litellm-model-catalog", + "litellm-python-compat", + "litellm-storage-s3", "reqwest", "rstest", + "tempfile", "thiserror 2.0.19", "tokio", "url", + "wiremock", ] [[package]] @@ -4303,6 +4305,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-model-catalog-refresh", "litellm-router-types", "litellm-secret-value", "litellm-secrets", @@ -4615,6 +4618,22 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-storage-s3" +version = "0.1.0" +dependencies = [ + "aws-credential-types", + "aws-sdk-s3", + "aws-smithy-runtime-api", + "aws-smithy-types", + "http 1.4.2", + "litellm-auth-aws", + "litellm-http", + "reqwest", + "rstest", + "tokio", +] + [[package]] name = "litellm-testkit" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 943e8576ae6..8bc0b66687f 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -58,6 +58,7 @@ litellm-cache = { path = "crates/cache" } litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-s3 = { path = "crates/cache-s3" } +litellm-storage-s3 = { path = "crates/storage-s3" } litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } litellm-cache-response = { path = "crates/cache-response" } diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index d51c282d4d7..115d6526196 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -9,14 +9,9 @@ repository.workspace = true litellm-http.workspace = true litellm-cache.workspace = true litellm-auth-aws.workspace = true -aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } -aws-credential-types = "1.3.0" -aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] } -aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] } -aws-types = "1.6.0" +litellm-storage-s3.workspace = true +aws-smithy-types = "1.6.0" futures-util.workspace = true -http.workspace = true -reqwest.workspace = true tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 5369bca0f18..ecc0df035e8 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -4,21 +4,16 @@ use std::{ time::{Duration, SystemTime}, }; -use aws_sdk_s3::{ - config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation}, - error::SdkError, - primitives::ByteStream, -}; use aws_smithy_types::{DateTime, date_time::Format}; use futures_util::future::try_join_all; use litellm_auth_aws::AwsAuthConfig; use litellm_cache::{ BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache, }; +use litellm_storage_s3::aws_sdk_s3::{self, error::SdkError, primitives::ByteStream}; +use litellm_storage_s3::s3_config; use tokio::runtime::Handle; -use crate::{auth::S3Credentials, transport::ReqwestHttpClient}; - pub struct S3Endpoint { pub url: String, } @@ -49,17 +44,12 @@ impl S3Cache { runtime: Handle, ) -> Self { let endpoint_url: Option = config.endpoint.map(|endpoint| endpoint.url); - let base = aws_sdk_s3::Config::builder() - .behavior_version(BehaviorVersion::latest()) - .region(Region::new(config.region.clone())) - .http_client(ReqwestHttpClient(http)) - .credentials_provider(S3Credentials::new(config.auth)) - .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) - .response_checksum_validation(ResponseChecksumValidation::WhenRequired); - let builder = match &endpoint_url { - Some(url) => base.endpoint_url(url).force_path_style(true), - None => base, - }; + let builder = s3_config( + config.region.clone(), + endpoint_url.clone(), + http, + config.auth, + ); Self { client: aws_sdk_s3::Client::from_conf(builder.build()), codec, diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs index 9be0210aad3..43006532441 100644 --- a/litellm-rust/crates/cache-s3/src/lib.rs +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -1,6 +1,3 @@ -mod auth; mod cache; -mod transport; -pub use auth::S3Credentials; pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index 02b8d4f2dcb..21d38d9c21a 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -40,5 +40,6 @@ Terminal notifications share the selected response or exception. An ordinary not - Keep coroutine driving in the shared Python driver and the native handle - Shared driver implementation: `litellm/rust_bridge/execution.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding - Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract +- Every runtime entry in `runtime.rs` passes the fork gate except `run_detached_value`, whose current-thread runtime ends with the call and leaves no thread for a fork to lose. Use it only for work that runs before a process decides whether it forks, such as the cost map refresh `import litellm` starts - References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html) - [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html) diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index 8e37ab02669..71916074432 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -36,8 +36,8 @@ 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, - run_sync_value, runtime_started, + ready_future, reserve_process_for_forking, run_async, run_async_value, run_detached_value, + run_sync, run_sync_value, 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 e8872c9340a..154b70c10f2 100644 --- a/litellm-rust/crates/host-python/src/runtime.rs +++ b/litellm-rust/crates/host-python/src/runtime.rs @@ -103,6 +103,24 @@ where run_sync_on(py, runtime()?, future, map_error) } +/// Runs `future` with the GIL released on a current-thread runtime that is dropped before this returns, so no runtime +/// thread outlives the call and the fork gate is never claimed. For work at import time, before the process has +/// decided whether it forks workers or reserves itself for forking +pub fn run_detached_value(py: Python<'_>, future: F) -> PyResult +where + T: Send, + F: Future> + Send, +{ + if Handle::try_current().is_ok() { + return Err(Error::SyncInsideTokio.into()); + } + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|error| PyRuntimeError::new_err(Report::new(&error).to_string()))?; + release_gil(py, move || runtime.block_on(future)) +} + pub fn run_sync_value(py: Python<'_>, future: F) -> PyResult where T: Send + 'static, @@ -541,6 +559,21 @@ asyncio.run(exercise()) }); } + #[rstest] + fn a_detached_run_completes_without_claiming_the_fork_gate( + #[from(initialized_python)] python: &InitializedPython, + ) { + let started_before = runtime_started(); + let value = python.attach(|py| { + run_detached_value(py, async { + tokio::time::sleep(Duration::from_millis(1)).await; + Ok(7) + }) + }); + assert_eq!(value.unwrap(), 7); + assert_eq!(runtime_started(), started_before); + } + #[rstest] fn sync_runner_rejects_calls_from_a_tokio_context( #[from(initialized_python)] python: &InitializedPython, diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 3bdc63453ef..bfd0009a4a9 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -37,6 +37,29 @@ pub struct HttpClientPool { clients: Mutex, } +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(), + } +} + impl HttpClientPool { pub fn new(media_resolver: Arc) -> Self { Self::with_ttl(media_resolver, CLIENT_TTL) @@ -55,36 +78,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 +95,28 @@ impl HttpClientPool { Ok(Client::new(client)) } + /// A client built from the same settings as `client` that the pool never holds, for work that must not share + /// connections or take the pool's lock, such as a fetch on a runtime that ends with it or across a fork + pub fn unpooled_client( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + self.build(&effective(config, variant), variant) + .map(Client::new) + } + + fn build( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + Ok(self + .apply(variant, reqwest::ClientBuilder::try_from(config)?) + .build() + .map_err(Repr::Client)?) + } + fn lock(&self) -> MutexGuard<'_, Clients> { self.clients.lock().unwrap_or_else(PoisonError::into_inner) } @@ -224,6 +246,27 @@ mod tests { assert_eq!(connections.load(Ordering::SeqCst), 3); } + #[tokio::test] + async fn unpooled_clients_share_no_connection_with_each_other_or_the_pool() { + let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; + let url = format!("http://{address}"); + let pool = pool(); + get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await; + for expected in [2, 3] { + let response = pool + .unpooled_client(&config("a"), ClientVariant::NoRedirect) + .unwrap() + .get(&url) + .send() + .await + .unwrap(); + assert_eq!(response.status(), 204); + assert_eq!(connections.load(Ordering::SeqCst), expected); + } + get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 3); + } + #[tokio::test] async fn provider_clients_route_through_the_resolved_proxy_not_the_process_environment() { let (proxy, connections, requests) = serve("HTTP/1.1 204 No Content").await; 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..83e7ecdb6bd --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/AGENTS.md @@ -0,0 +1,5 @@ +# Model catalog refresh + +`litellm-model-catalog-refresh` fetches a cost map document from an HTTP URL, a `file://` path or an `s3://` object and retries it by main's status rule. It returns bytes and never parses them, so every rule about what a map contains stays in `litellm-model-catalog` + +`Fetch` is one source, `SourceFetch` picks it from a location, and `refresh` / `resume` drive the attempts through an injected `Pause`. They return an `Attempt` that is never `Retryable`. S3 clients come from `litellm-storage-s3` diff --git a/litellm-rust/crates/model-catalog-refresh/Cargo.toml b/litellm-rust/crates/model-catalog-refresh/Cargo.toml index 27219042c6c..9703fc36f56 100644 --- a/litellm-rust/crates/model-catalog-refresh/Cargo.toml +++ b/litellm-rust/crates/model-catalog-refresh/Cargo.toml @@ -6,14 +6,20 @@ license.workspace = true repository.workspace = true [dependencies] +http.workspace = true +litellm-auth-aws.workspace = true +litellm-storage-s3.workspace = true +litellm-error.workspace = true litellm-http.workspace = true -litellm-model-catalog.workspace = true +litellm-python-compat.workspace = true thiserror.workspace = true +tokio = { workspace = true, features = ["fs", "time"] } url.workspace = true [dev-dependencies] -litellm-error.workspace = true litellm-http = { workspace = true, features = ["test-support"] } reqwest.workspace = true rstest.workspace = true +tempfile.workspace = true tokio.workspace = true +wiremock.workspace = true diff --git a/litellm-rust/crates/model-catalog-refresh/src/error.rs b/litellm-rust/crates/model-catalog-refresh/src/error.rs index 656f32370cb..fc8dd61bd0a 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/error.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/error.rs @@ -14,14 +14,32 @@ enum Repr { Read(#[source] io::Error), #[error("invalid catalog URL")] InvalidUrl(#[source] url::ParseError), - #[error("unsupported catalog URL scheme {scheme}")] + #[error("unsupported catalog URL scheme {scheme:?}")] UnsupportedScheme { scheme: String }, + #[error("expected s3://bucket/key")] + InvalidS3Url, + #[error("S3 GetObject failed")] + S3(#[source] Box), + #[error("timed out")] + TimedOut, } impl Error { - pub fn unsupported_scheme(scheme: String) -> Self { + pub(crate) fn unsupported_scheme(scheme: String) -> Self { Self(Box::new(Repr::UnsupportedScheme { scheme })) } + + pub(crate) fn invalid_s3_url() -> Self { + Self(Box::new(Repr::InvalidS3Url)) + } + + pub(crate) fn s3(cause: impl std::error::Error + Send + Sync + 'static) -> Self { + Self(Box::new(Repr::S3(Box::new(cause)))) + } + + pub(crate) fn timed_out() -> Self { + Self(Box::new(Repr::TimedOut)) + } } impl From for Error { @@ -48,25 +66,10 @@ mod tests { use rstest::rstest; use super::*; - use crate::{DoneReason, Reply, Step}; #[rstest] - fn exhausted_transport_keeps_its_typed_cause() { - let reply = Reply::Transport { - error: io::Error::from(io::ErrorKind::ConnectionReset).into(), - }; - let Reply::Transport { error } = reply else { - panic!("transport failure") - }; - let step = Step::Done { - reason: DoneReason::Exhausted(error), - }; - let Step::Done { - reason: DoneReason::Exhausted(error), - } = step - else { - panic!("exhausted attempts") - }; + fn a_read_failure_keeps_its_typed_cause_once() { + let error = Error::from(io::Error::from(io::ErrorKind::ConnectionReset)); let cause = find::(&error).expect("typed I/O cause"); assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset); assert_eq!( diff --git a/litellm-rust/crates/model-catalog-refresh/src/fetch.rs b/litellm-rust/crates/model-catalog-refresh/src/fetch.rs new file mode 100644 index 00000000000..f487b314ba1 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/fetch.rs @@ -0,0 +1,34 @@ +use std::time::Duration; + +use crate::Error; + +/// One attempt. Where to fetch from is the `Fetch`, so the request carries no URL +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FetchRequest { + pub if_none_match: Option, + pub timeout: Duration, +} + +#[derive(Debug)] +pub enum Reply { + /// An HTTP, S3 or file answer + Response { + status: u16, + /// Every `ETag` value, joined with `", "` + etag: Option, + /// Every `Retry-After` value, joined with `", "` + retry_after: Option, + body: Vec, + }, + /// No answer reached us, and another attempt may get one + Transport { error: Error }, + /// The request itself cannot succeed, so retrying is pointless + Refused { error: Error }, +} + +/// One source of the catalog's bytes +pub trait Fetch: Sync { + /// The source as the caller wrote it, for failure reasons + fn location(&self) -> &str; + fn fetch(&self, request: &FetchRequest) -> impl Future + Send; +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/file.rs b/litellm-rust/crates/model-catalog-refresh/src/file.rs new file mode 100644 index 00000000000..2df77bec419 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/file.rs @@ -0,0 +1,73 @@ +use std::path::PathBuf; + +use crate::{Fetch, FetchRequest, Reply}; + +/// A catalog in a local file, answered as a 200 with no `ETag` +#[derive(Clone, Debug)] +pub struct FileFetch { + location: String, + path: PathBuf, +} + +impl FileFetch { + pub fn new(location: String, path: PathBuf) -> Self { + Self { location, path } + } +} + +impl Fetch for FileFetch { + fn location(&self) -> &str { + &self.location + } + + async fn fetch(&self, _: &FetchRequest) -> Reply { + match tokio::fs::read(&self.path).await { + Ok(body) => Reply::Response { + status: 200, + etag: None, + retry_after: None, + body, + }, + Err(error) => Reply::Refused { + error: error.into(), + }, + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + + use super::*; + + fn request() -> FetchRequest { + FetchRequest { + if_none_match: None, + timeout: Duration::from_secs(5), + } + } + + #[rstest] + #[tokio::test] + async fn reads_the_file_and_refuses_a_missing_one() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("map.json"); + std::fs::write(&path, b"{}").unwrap(); + let present = FileFetch::new("file://map".to_owned(), path.clone()); + assert!(matches!( + present.fetch(&request()).await, + Reply::Response { status: 200, ref body, .. } if body == b"{}" + )); + let missing = FileFetch::new( + "file://missing".to_owned(), + directory.path().join("missing.json"), + ); + assert!(matches!( + missing.fetch(&request()).await, + Reply::Refused { .. } + )); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/http.rs b/litellm-rust/crates/model-catalog-refresh/src/http.rs new file mode 100644 index 00000000000..1556e870db2 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/http.rs @@ -0,0 +1,148 @@ +use http::HeaderMap; +use http::header::{ETAG, IF_NONE_MATCH, RETRY_AFTER}; + +use crate::{Fetch, FetchRequest, Reply}; + +/// A catalog over HTTP. The client never follows a redirect, as main's `httpx.get` does not +#[derive(Clone, Debug)] +pub struct HttpFetch { + client: litellm_http::Client, + url: String, +} + +impl HttpFetch { + pub fn new(client: litellm_http::Client, url: String) -> Self { + Self { client, url } + } +} + +impl Fetch for HttpFetch { + fn location(&self) -> &str { + &self.url + } + + async fn fetch(&self, request: &FetchRequest) -> Reply { + let builder = self.client.get(&self.url).timeout(request.timeout); + let builder = match &request.if_none_match { + Some(etag) => builder.header(IF_NONE_MATCH, etag), + None => builder, + }; + let response = match builder.send().await { + Ok(response) => response, + Err(error) if error.is_builder() => { + return Reply::Refused { + error: litellm_http::Error::from(error).into(), + }; + } + Err(error) => { + return Reply::Transport { + error: litellm_http::Error::from(error).into(), + }; + } + }; + let status = response.status().as_u16(); + let etag = joined(response.headers(), &ETAG); + let retry_after = joined(response.headers(), &RETRY_AFTER); + match response.bytes().await { + Ok(body) => Reply::Response { + status, + etag, + retry_after, + body: body.to_vec(), + }, + Err(error) => Reply::Transport { + error: litellm_http::Error::from(error).into(), + }, + } + } +} + +/// Every value of `name` joined with `", "`, as httpx reads a repeated header +fn joined(headers: &HeaderMap, name: &http::HeaderName) -> Option { + let values: Vec = headers + .get_all(name) + .iter() + .map(|value| String::from_utf8_lossy(value.as_bytes()).into_owned()) + .collect(); + (!values.is_empty()).then(|| values.join(", ")) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + use wiremock::matchers::{header, method}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + + fn request(if_none_match: Option<&str>) -> FetchRequest { + FetchRequest { + if_none_match: if_none_match.map(str::to_owned), + timeout: Duration::from_secs(5), + } + } + + #[rstest] + #[tokio::test] + async fn reads_status_headers_and_body_without_following_redirects() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with( + ResponseTemplate::new(302) + .insert_header("location", "/elsewhere") + .append_header("etag", "\"a\"") + .append_header("etag", "\"b\"") + .insert_header("retry-after", "3") + .set_body_string("moved"), + ) + .mount(&server) + .await; + let fetch = HttpFetch::new(litellm_http::Client::for_test(), server.uri()); + let Reply::Response { + status, + etag, + retry_after, + body, + } = fetch.fetch(&request(None)).await + else { + panic!("a response") + }; + assert_eq!(status, 302); + assert_eq!(etag.as_deref(), Some("\"a\", \"b\"")); + assert_eq!(retry_after.as_deref(), Some("3")); + assert_eq!(body, b"moved"); + } + + #[rstest] + #[tokio::test] + async fn sends_if_none_match_when_asked() { + let server = MockServer::start().await; + Mock::given(header("if-none-match", "\"v1\"")) + .respond_with(ResponseTemplate::new(304)) + .mount(&server) + .await; + let fetch = HttpFetch::new(litellm_http::Client::for_test(), server.uri()); + let reply = fetch.fetch(&request(Some("\"v1\""))).await; + assert!(matches!(reply, Reply::Response { status: 304, .. })); + } + + #[rstest] + #[case::refused_connection("http://127.0.0.1:1/map.json", false)] + #[case::unparseable("http://[invalid/map.json", true)] + #[tokio::test] + async fn transport_failures_retry_and_builder_failures_are_final( + #[case] url: &str, + #[case] refused: bool, + ) { + let fetch = HttpFetch::new(litellm_http::Client::for_test(), url.to_owned()); + let reply = fetch.fetch(&request(None)).await; + assert_eq!(matches!(reply, Reply::Refused { .. }), refused, "{reply:?}"); + assert_eq!( + matches!(reply, Reply::Transport { .. }), + !refused, + "{reply:?}" + ); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/lib.rs b/litellm-rust/crates/model-catalog-refresh/src/lib.rs index 006c4887fad..a072d07496b 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/lib.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/lib.rs @@ -1,9 +1,21 @@ mod error; -mod policy; +mod fetch; +mod file; +mod http; +mod refresh; +mod retry; +mod s3; mod source; -mod step; +mod source_fetch; +mod status; pub use error::Error; -pub use policy::*; +pub use fetch::*; +pub use file::*; +pub use http::*; +pub use refresh::*; +pub use retry::*; +pub use s3::*; pub use source::*; -pub use step::*; +pub use source_fetch::*; +pub use status::*; diff --git a/litellm-rust/crates/model-catalog-refresh/src/policy.rs b/litellm-rust/crates/model-catalog-refresh/src/policy.rs deleted file mode 100644 index bc32e40408e..00000000000 --- a/litellm-rust/crates/model-catalog-refresh/src/policy.rs +++ /dev/null @@ -1,20 +0,0 @@ -use std::time::Duration; - -use litellm_model_catalog::IntegrityLimits; - -/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct RetryPolicy { - pub max_attempts: u32, - pub max_wait: Duration, - pub timeout: Duration, - pub backoff_base_seconds: f64, - pub jitter_min_seconds: f64, - pub jitter_max_seconds: f64, - /// Main's 50 model floor and 0.5 shrink ratio - pub limits: IntegrityLimits, -} - -/// Main's default `LITELLM_MODEL_COST_MAP_URL` -pub const DEFAULT_REGISTRY_URL: &str = - "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"; diff --git a/litellm-rust/crates/model-catalog-refresh/src/refresh.rs b/litellm-rust/crates/model-catalog-refresh/src/refresh.rs new file mode 100644 index 00000000000..52af9e3e902 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/refresh.rs @@ -0,0 +1,330 @@ +use crate::retry::{AfterRetryable, after_retryable}; +use crate::status::{Attempt, classify}; +use crate::{Fetch, FetchRequest}; + +/// A retryable failure about to be waited out +#[derive(Clone, Debug, PartialEq)] +pub struct Retrying<'a> { + pub attempt: i64, + pub max_attempts: i64, + pub reason: &'a str, + pub wait_seconds: f64, +} + +/// What a refresh does between attempts, injected so tests and hosts choose their own clock and randomness. A host +/// whose jitter or report can fail returns that failure from the `sleep` that follows it, which ends the refresh +pub trait Pause: Send { + type Stop; + + /// Seconds added to a backoff, drawn only when there is no usable `Retry-After` + fn jitter(&mut self) -> f64; + fn retrying(&mut self, retrying: &Retrying<'_>); + fn sleep(&mut self, seconds: f64) -> impl Future> + Send; +} + +/// Up to `max_attempts` attempts, as main's `_fetch_remote_model_cost_map_with_retry`. Never `Attempt::Retryable` +pub async fn refresh( + fetch: &impl Fetch, + request: &FetchRequest, + max_attempts: i64, + pause: &mut P, +) -> Result { + if max_attempts < 1 { + return Ok(Attempt::Failed { + reason: "model cost map fetch failed".to_owned(), + }); + } + let first = attempt(fetch, request).await; + resume(fetch, request, 1, max_attempts, first, pause).await +} + +/// One attempt, for a host that makes the first attempt itself and resumes only after a retryable one +pub async fn attempt(fetch: &impl Fetch, request: &FetchRequest) -> Attempt { + classify(fetch.fetch(request).await, fetch.location(), request) +} + +/// Carries on after `outcome`, the result of attempt number `attempt`. Never `Attempt::Retryable` +pub async fn resume( + fetch: &impl Fetch, + request: &FetchRequest, + attempt: i64, + max_attempts: i64, + outcome: Attempt, + pause: &mut P, +) -> Result { + let (mut attempt_number, mut outcome) = (attempt, outcome); + loop { + let (reason, retry_after) = match outcome { + Attempt::Retryable { + reason, + retry_after, + } => (reason, retry_after), + done => return Ok(done), + }; + let seconds = + match after_retryable(attempt_number, max_attempts, retry_after, || pause.jitter()) { + AfterRetryable::GiveUp => { + return Ok(Attempt::Failed { + reason: format!("{reason} (after {max_attempts} attempts)"), + }); + } + AfterRetryable::Wait { seconds } => seconds, + }; + pause.retrying(&Retrying { + attempt: attempt_number, + max_attempts, + reason: &reason, + wait_seconds: seconds, + }); + pause.sleep(seconds).await?; + attempt_number += 1; + outcome = self::attempt(fetch, request).await; + } +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + use std::io; + use std::sync::Mutex; + use std::time::Duration; + + use rstest::rstest; + + use super::*; + use crate::Reply; + + const URL: &str = "https://example.com/map.json"; + + /// Answers each attempt from a script and records the requests it saw + struct Scripted { + replies: Mutex>, + requests: Mutex>, + } + + impl Scripted { + fn new(replies: impl IntoIterator) -> Self { + Self { + replies: Mutex::new(replies.into_iter().collect()), + requests: Mutex::new(Vec::new()), + } + } + + fn attempts(&self) -> usize { + self.requests.lock().unwrap().len() + } + } + + impl Fetch for Scripted { + fn location(&self) -> &str { + URL + } + + async fn fetch(&self, request: &FetchRequest) -> Reply { + self.requests.lock().unwrap().push(request.clone()); + self.replies + .lock() + .unwrap() + .pop_front() + .expect("a scripted reply") + } + } + + #[derive(Default)] + struct Recording { + jitters: usize, + slept: Vec, + retried: Vec<(i64, i64, String, f64)>, + } + + impl Pause for Recording { + type Stop = std::convert::Infallible; + + fn jitter(&mut self) -> f64 { + self.jitters += 1; + 0.5 + } + + async fn sleep(&mut self, seconds: f64) -> Result<(), Self::Stop> { + self.slept.push(seconds); + Ok(()) + } + + fn retrying(&mut self, retrying: &Retrying<'_>) { + self.retried.push(( + retrying.attempt, + retrying.max_attempts, + retrying.reason.to_owned(), + retrying.wait_seconds, + )); + } + } + + fn request() -> FetchRequest { + FetchRequest { + if_none_match: None, + timeout: Duration::from_secs(5), + } + } + + fn status(status: u16, retry_after: Option<&str>) -> Reply { + Reply::Response { + status, + etag: None, + retry_after: retry_after.map(str::to_owned), + body: Vec::new(), + } + } + + fn ok(body: &[u8]) -> Reply { + Reply::Response { + status: 200, + etag: Some("\"v2\"".to_owned()), + retry_after: None, + body: body.to_vec(), + } + } + + fn reset() -> Reply { + Reply::Transport { + error: io::Error::from(io::ErrorKind::ConnectionReset).into(), + } + } + + #[rstest] + #[tokio::test] + async fn retries_then_returns_the_fetched_body() { + let fetch = Scripted::new([status(503, Some("1")), status(503, None), ok(b"hello\n")]); + let mut pause = Recording::default(); + let refreshed = refresh(&fetch, &request(), 3, &mut pause).await.unwrap(); + assert_eq!( + refreshed, + Attempt::Fetched { + body: b"hello\n".to_vec(), + etag: Some("\"v2\"".to_owned()), + } + ); + assert_eq!(pause.slept, vec![1.0, 4.5]); + assert_eq!(pause.jitters, 1); + assert_eq!( + pause.retried, + vec![ + (1, 3, format!("HTTP 503 from {URL}"), 1.0), + (2, 3, format!("HTTP 503 from {URL}"), 4.5), + ] + ); + } + + #[rstest] + #[tokio::test] + async fn gives_up_after_the_last_attempt_without_sleeping() { + let fetch = Scripted::new([status(503, Some("0")), reset(), status(429, Some("0"))]); + let mut pause = Recording::default(); + let refreshed = refresh(&fetch, &request(), 3, &mut pause).await.unwrap(); + assert_eq!( + refreshed, + Attempt::Failed { + reason: format!("HTTP 429 from {URL} (after 3 attempts)"), + } + ); + assert_eq!(fetch.attempts(), 3); + assert_eq!(pause.slept.len(), 2); + } + + #[rstest] + #[tokio::test] + async fn a_final_status_stops_at_once() { + let fetch = Scripted::new([status(404, None)]); + let mut pause = Recording::default(); + let refreshed = refresh(&fetch, &request(), 3, &mut pause).await.unwrap(); + assert_eq!( + refreshed, + Attempt::Failed { + reason: format!("HTTP 404 from {URL}"), + } + ); + assert!(pause.slept.is_empty()); + } + + #[rstest] + #[case::zero(0)] + #[case::negative(-1)] + #[tokio::test] + async fn no_attempts_fetch_nothing(#[case] max_attempts: i64) { + let fetch = Scripted::new([]); + let refreshed = refresh(&fetch, &request(), max_attempts, &mut Recording::default()) + .await + .unwrap(); + assert_eq!( + refreshed, + Attempt::Failed { + reason: "model cost map fetch failed".to_owned(), + } + ); + assert_eq!(fetch.attempts(), 0); + } + + #[rstest] + #[tokio::test] + async fn resume_waits_out_the_first_attempt_then_counts_from_two() { + let fetch = Scripted::new([ok(b"{}")]); + let mut pause = Recording::default(); + let first = Attempt::Retryable { + reason: "HTTP 503 from x".to_owned(), + retry_after: Some(2.0), + }; + let refreshed = resume(&fetch, &request(), 1, 3, first, &mut pause) + .await + .unwrap(); + assert!(matches!(refreshed, Attempt::Fetched { .. })); + assert_eq!( + pause.retried, + vec![(1, 3, "HTTP 503 from x".to_owned(), 2.0)] + ); + assert_eq!(fetch.attempts(), 1); + } + + /// Stops at the first sleep, as a host whose jitter or report failed does + struct Stopping; + + impl Pause for Stopping { + type Stop = &'static str; + + fn jitter(&mut self) -> f64 { + 0.0 + } + + fn retrying(&mut self, _: &Retrying<'_>) {} + + async fn sleep(&mut self, _: f64) -> Result<(), Self::Stop> { + Err("stopped") + } + } + + #[rstest] + #[tokio::test] + async fn a_stopped_pause_ends_the_refresh_before_the_next_attempt() { + let fetch = Scripted::new([status(503, None), ok(b"{}")]); + let stopped = refresh(&fetch, &request(), 3, &mut Stopping).await; + assert_eq!(stopped, Err("stopped")); + assert_eq!(fetch.attempts(), 1); + } + + #[rstest] + #[tokio::test] + async fn a_conditional_not_modified_ends_the_refresh() { + let fetch = Scripted::new([status(304, None)]); + let request = FetchRequest { + if_none_match: Some("\"v1\"".to_owned()), + ..request() + }; + let refreshed = refresh(&fetch, &request, 3, &mut Recording::default()) + .await + .unwrap(); + assert_eq!(refreshed, Attempt::NotModified); + assert_eq!( + fetch.requests.lock().unwrap()[0].if_none_match.as_deref(), + Some("\"v1\"") + ); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/retry.rs b/litellm-rust/crates/model-catalog-refresh/src/retry.rs new file mode 100644 index 00000000000..abcf4c1619f --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/retry.rs @@ -0,0 +1,79 @@ +const MAX_WAIT_SECONDS: f64 = 30.0; +const BACKOFF_BASE_SECONDS: f64 = 2.0; + +/// What follows a retryable failure of attempt `attempt` out of `max_attempts` +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum AfterRetryable { + GiveUp, + Wait { seconds: f64 }, +} + +/// A usable `Retry-After` is waited, capped at 30 s. Otherwise `2^attempt` plus `jitter()`, capped the same, so +/// `jitter` runs only when there is no usable `Retry-After` +pub fn after_retryable( + attempt: i64, + max_attempts: i64, + retry_after_seconds: Option, + jitter: impl FnOnce() -> f64, +) -> AfterRetryable { + if attempt == max_attempts { + return AfterRetryable::GiveUp; + } + let seconds = match retry_after_seconds.filter(|seconds| *seconds >= 0.0) { + Some(retry_after) => retry_after.min(MAX_WAIT_SECONDS), + None => (backoff(attempt) + jitter()).min(MAX_WAIT_SECONDS), + }; + AfterRetryable::Wait { seconds } +} + +fn backoff(attempt: i64) -> f64 { + i32::try_from(attempt).map_or(f64::INFINITY, |exponent| { + BACKOFF_BASE_SECONDS.powi(exponent) + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::last_attempt(3, 3, Some(1.0), AfterRetryable::GiveUp)] + #[case::retry_after(1, 3, Some(4.5), AfterRetryable::Wait { seconds: 4.5 })] + #[case::retry_after_capped(1, 3, Some(120.0), AfterRetryable::Wait { seconds: 30.0 })] + #[case::infinite_retry_after(1, 3, Some(f64::INFINITY), AfterRetryable::Wait { seconds: 30.0 })] + #[case::zero_retry_after(1, 3, Some(0.0), AfterRetryable::Wait { seconds: 0.0 })] + #[case::negative_retry_after_backs_off(1, 3, Some(-1.0), AfterRetryable::Wait { seconds: 2.25 })] + #[case::nan_retry_after_backs_off(1, 3, Some(f64::NAN), AfterRetryable::Wait { seconds: 2.25 })] + #[case::first_backoff(1, 3, None, AfterRetryable::Wait { seconds: 2.25 })] + #[case::second_backoff(2, 3, None, AfterRetryable::Wait { seconds: 4.25 })] + #[case::backoff_capped(5, 9, None, AfterRetryable::Wait { seconds: 30.0 })] + #[case::past_the_last_attempt_still_waits(4, 3, None, AfterRetryable::Wait { seconds: 16.25 })] + fn waits_match_main( + #[case] attempt: i64, + #[case] max_attempts: i64, + #[case] retry_after: Option, + #[case] expected: AfterRetryable, + ) { + assert_eq!( + after_retryable(attempt, max_attempts, retry_after, || 0.25), + expected + ); + } + + #[rstest] + #[case::retry_after(Some(1.0), false)] + #[case::no_retry_after(None, true)] + #[case::unusable_retry_after(Some(-1.0), true)] + #[case::give_up(None, false)] + fn jitter_is_drawn_only_for_a_backoff(#[case] retry_after: Option, #[case] drawn: bool) { + let mut draws = 0; + let max_attempts = if drawn || retry_after.is_some() { 3 } else { 1 }; + let _ = after_retryable(1, max_attempts, retry_after, || { + draws += 1; + 0.0 + }); + assert_eq!(draws, usize::from(drawn)); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/s3.rs b/litellm-rust/crates/model-catalog-refresh/src/s3.rs new file mode 100644 index 00000000000..ff6407e97ab --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/s3.rs @@ -0,0 +1,132 @@ +use litellm_storage_s3::aws_sdk_s3; + +use crate::source::S3Object; +use crate::{Error, Fetch, FetchRequest, Reply}; + +/// A catalog in an S3 object. Every failure is final, as main's boto3 fetch treats them +#[derive(Clone, Debug)] +pub struct S3Fetch { + client: aws_sdk_s3::Client, + location: String, + object: S3Object, +} + +impl S3Fetch { + /// `config` is built by `litellm_storage_s3::s3_config`, and retries are left to the refresh + pub fn new(config: aws_sdk_s3::config::Builder, location: String, object: S3Object) -> Self { + let config = config.retry_config(aws_sdk_s3::config::retry::RetryConfig::disabled()); + Self { + client: aws_sdk_s3::Client::from_conf(config.build()), + location, + object, + } + } +} + +impl Fetch for S3Fetch { + fn location(&self) -> &str { + &self.location + } + + async fn fetch(&self, request: &FetchRequest) -> Reply { + let get = async { + let output = self + .client + .get_object() + .bucket(&self.object.bucket) + .key(&self.object.key) + .send() + .await + .map_err(Error::s3)?; + let etag = output.e_tag().map(str::to_owned); + let body = output.body.collect().await.map_err(Error::s3)?; + Ok::<_, Error>((etag, body.into_bytes().to_vec())) + }; + match tokio::time::timeout(request.timeout, get).await { + Ok(Ok((etag, body))) => Reply::Response { + status: 200, + etag, + retry_after: None, + body, + }, + Ok(Err(error)) => Reply::Refused { error }, + Err(_) => Reply::Refused { + error: Error::timed_out(), + }, + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_auth_aws::AwsAuthConfig; + use rstest::rstest; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + + async fn fetch_from(status: u16) -> (Reply, Vec) { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/bucket/prices/map.json")) + .respond_with( + ResponseTemplate::new(status) + .insert_header("etag", "\"s3-etag\"") + .set_body_string("{}"), + ) + .mount(&server) + .await; + let config = litellm_storage_s3::s3_config( + "us-east-1".to_owned(), + Some(server.uri()), + litellm_http::Client::for_test(), + AwsAuthConfig { + access_key_id: Some("access".to_owned()), + secret_access_key: Some("secret".to_owned()), + ..AwsAuthConfig::default() + }, + ); + let fetch = S3Fetch::new( + config, + "s3://bucket/prices/map.json".to_owned(), + S3Object { + bucket: "bucket".to_owned(), + key: "prices/map.json".to_owned(), + }, + ); + let request = FetchRequest { + if_none_match: None, + timeout: Duration::from_secs(5), + }; + let reply = fetch.fetch(&request).await; + (reply, server.received_requests().await.unwrap_or_default()) + } + + #[rstest] + #[tokio::test] + async fn reads_the_object_and_its_etag() { + let (reply, requests) = fetch_from(200).await; + let Reply::Response { + status, etag, body, .. + } = reply + else { + panic!("an object") + }; + assert_eq!( + (status, etag.as_deref(), body.as_slice()), + (200, Some("\"s3-etag\""), b"{}".as_slice()) + ); + assert_eq!(requests.len(), 1); + } + + #[rstest] + #[tokio::test] + async fn a_server_error_is_final_and_not_retried_by_the_sdk() { + let (reply, requests) = fetch_from(503).await; + assert!(matches!(reply, Reply::Refused { .. }), "{reply:?}"); + assert_eq!(requests.len(), 1); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/source.rs b/litellm-rust/crates/model-catalog-refresh/src/source.rs index 713e0386941..b6a082dbed7 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/source.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/source.rs @@ -2,6 +2,12 @@ use std::path::PathBuf; use url::Url; +use crate::Error; + +/// Main's default `LITELLM_MODEL_COST_MAP_URL` +pub const DEFAULT_REGISTRY_URL: &str = + "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"; + /// Where a fetched catalog comes from #[derive(Clone, Debug, PartialEq, Eq)] pub enum Source { @@ -15,3 +21,86 @@ pub struct S3Object { pub bucket: String, pub key: String, } + +impl Source { + /// Main splits an `s3` URL the way `urllib.parse.urlsplit` does, the bucket is the netloc and the key is the path + /// without its leading slashes + pub fn parse(location: &str) -> Result { + let scheme = location + .split_once(':') + .map(|(scheme, _)| scheme.to_ascii_lowercase()); + match scheme.as_deref() { + Some("s3") => s3_object(location).map(Source::S3), + Some("file") => Url::parse(location)? + .to_file_path() + .map(Source::File) + .map_err(|()| Error::unsupported_scheme("file".to_owned())), + Some("http" | "https") => Ok(Source::Http(Url::parse(location)?)), + other => Err(Error::unsupported_scheme( + other.unwrap_or_default().to_owned(), + )), + } + } +} + +fn s3_object(location: &str) -> Result { + let after_scheme = location.split_once(':').map_or("", |(_, rest)| rest); + let Some(rest) = after_scheme.strip_prefix("//") else { + return Err(Error::invalid_s3_url()); + }; + let rest = rest.split(['?', '#']).next().unwrap_or_default(); + let (bucket, path) = rest.split_once('/').unwrap_or((rest, "")); + let key = path.trim_start_matches('/'); + if bucket.is_empty() || key.is_empty() { + return Err(Error::invalid_s3_url()); + } + Ok(S3Object { + bucket: bucket.to_owned(), + key: key.to_owned(), + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn s3(bucket: &str, key: &str) -> Option { + Some(Source::S3(S3Object { + bucket: bucket.to_owned(), + key: key.to_owned(), + })) + } + + #[rstest] + #[case::s3("s3://bucket/path/map.json", s3("bucket", "path/map.json"))] + #[case::s3_upper_scheme("S3://bucket/map.json", s3("bucket", "map.json"))] + #[case::s3_leading_slashes("s3://bucket//map.json", s3("bucket", "map.json"))] + #[case::s3_query_dropped("s3://bucket/map.json?versionId=1#x", s3("bucket", "map.json"))] + #[case::s3_no_key("s3://bucket/", None)] + #[case::s3_no_bucket("s3:///map.json", None)] + #[case::s3_no_slashes("s3:bucket/map.json", None)] + #[case::empty("", None)] + #[case::ftp("ftp://host/map.json", None)] + #[case::relative("map.json", None)] + fn sources_split_like_urlsplit(#[case] location: &str, #[case] expected: Option) { + assert_eq!(Source::parse(location).ok(), expected); + } + + #[rstest] + fn file_urls_are_decoded_paths() { + assert_eq!( + Source::parse("file:///tmp/a%20b.json").ok(), + Some(Source::File(PathBuf::from("/tmp/a b.json"))) + ); + } + + #[rstest] + fn http_urls_keep_their_query() { + let Ok(Source::Http(url)) = Source::parse("https://example.com/map.json?ref=main") else { + panic!("an HTTP source") + }; + assert_eq!(url.as_str(), "https://example.com/map.json?ref=main"); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs b/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs new file mode 100644 index 00000000000..ba59bc90b11 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs @@ -0,0 +1,99 @@ +use litellm_auth_aws::AwsAuthConfig; + +use crate::source::Source; +use crate::{Error, Fetch, FetchRequest, FileFetch, HttpFetch, Reply, S3Fetch}; + +const DEFAULT_S3_REGION: &str = "us-east-1"; + +/// The `Fetch` a catalog location names +#[derive(Clone, Debug)] +pub enum SourceFetch { + Http(HttpFetch), + File(FileFetch), + S3(Box), +} + +/// What a host supplies for whichever source a location names +pub struct SourceClients<'a> { + pub http: litellm_http::Client, + pub env: &'a dyn Fn(&str) -> Option, +} + +impl SourceFetch { + /// Fails for a location no source can serve, which main reports as a final failure + pub fn new(location: &str, clients: SourceClients<'_>) -> Result { + let owned = location.to_owned(); + Ok(match Source::parse(location)? { + Source::Http(_) => Self::Http(HttpFetch::new(clients.http, owned)), + Source::File(path) => Self::File(FileFetch::new(owned, path)), + Source::S3(object) => { + let region = s3_region(clients.env); + let auth = AwsAuthConfig { + region_name: Some(region.clone()), + ..AwsAuthConfig::default() + }; + let config = litellm_storage_s3::s3_config( + region, + s3_endpoint(clients.env), + clients.http, + auth, + ); + Self::S3(Box::new(S3Fetch::new(config, owned, object))) + } + }) + } +} + +/// `AWS_REGION_NAME` as main reads it, then the AWS SDK's own variables +fn s3_region(env: &dyn Fn(&str) -> Option) -> String { + ["AWS_REGION_NAME", "AWS_REGION", "AWS_DEFAULT_REGION"] + .into_iter() + .find_map(|name| env(name).filter(|region| !region.is_empty())) + .unwrap_or_else(|| DEFAULT_S3_REGION.to_owned()) +} + +/// `AWS_ENDPOINT_URL_S3` then `AWS_ENDPOINT_URL`, which boto3 also reads, for MinIO and other S3 compatible stores +fn s3_endpoint(env: &dyn Fn(&str) -> Option) -> Option { + ["AWS_ENDPOINT_URL_S3", "AWS_ENDPOINT_URL"] + .into_iter() + .find_map(|name| env(name).filter(|endpoint| !endpoint.is_empty())) +} + +impl Fetch for SourceFetch { + fn location(&self) -> &str { + match self { + Self::Http(fetch) => fetch.location(), + Self::File(fetch) => fetch.location(), + Self::S3(fetch) => fetch.location(), + } + } + + async fn fetch(&self, request: &FetchRequest) -> Reply { + match self { + Self::Http(fetch) => fetch.fetch(request).await, + Self::File(fetch) => fetch.fetch(request).await, + Self::S3(fetch) => fetch.fetch(request).await, + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::region_name(&[("AWS_REGION_NAME", "eu-west-1"), ("AWS_REGION", "us-west-2")], "eu-west-1")] + #[case::sdk_region(&[("AWS_REGION_NAME", ""), ("AWS_REGION", "us-west-2")], "us-west-2")] + #[case::sdk_default_region(&[("AWS_DEFAULT_REGION", "ap-south-1")], "ap-south-1")] + #[case::none(&[], DEFAULT_S3_REGION)] + fn s3_region_prefers_main_s_variable(#[case] env: &[(&str, &str)], #[case] expected: &str) { + let lookup = |name: &str| { + env.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + }; + assert_eq!(s3_region(&lookup), expected); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/status.rs b/litellm-rust/crates/model-catalog-refresh/src/status.rs new file mode 100644 index 00000000000..9a76ba052f1 --- /dev/null +++ b/litellm-rust/crates/model-catalog-refresh/src/status.rs @@ -0,0 +1,183 @@ +use litellm_error::Report; + +use crate::{FetchRequest, Reply}; + +/// Main's `RETRYABLE_FETCH_STATUS_CODES` +pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504]; + +/// What one attempt came to, by main's `_classify_fetch_response` and `_classify_fetch_error` +#[derive(Clone, Debug, PartialEq)] +pub enum Attempt { + Fetched { + body: Vec, + etag: Option, + }, + /// A 304 to a request that sent `If-None-Match` + NotModified, + Retryable { + reason: String, + retry_after: Option, + }, + Failed { + reason: String, + }, +} + +/// A retryable status waits out its `Retry-After`, any other 4xx or 5xx is final, and every other status is a +/// document. A transport failure is retried and a refused request is final +pub fn classify(reply: Reply, location: &str, request: &FetchRequest) -> Attempt { + match reply { + Reply::Response { status, .. } if RETRYABLE_STATUSES.contains(&status) => { + let Reply::Response { retry_after, .. } = reply else { + unreachable!("matched a response") + }; + Attempt::Retryable { + reason: format!("HTTP {status} from {location}"), + retry_after: retry_after + .as_deref() + .and_then(litellm_python_compat::number::float), + } + } + Reply::Response { status, .. } if (400..=599).contains(&status) => Attempt::Failed { + reason: format!("HTTP {status} from {location}"), + }, + Reply::Response { status: 304, .. } if request.if_none_match.is_some() => { + Attempt::NotModified + } + Reply::Response { body, etag, .. } => Attempt::Fetched { body, etag }, + Reply::Transport { error } => Attempt::Retryable { + reason: format!("could not fetch {location}: {}", Report::new(&error)), + retry_after: None, + }, + Reply::Refused { error } => refused(location, &error), + } +} + +/// The final failure for a request that cannot succeed +pub fn refused(location: &str, error: &crate::Error) -> Attempt { + Attempt::Failed { + reason: format!("could not fetch {location}: {}", Report::new(error)), + } +} + +#[cfg(test)] +mod tests { + use std::io; + use std::time::Duration; + + use rstest::rstest; + + use super::*; + + const URL: &str = "https://example.com/map.json"; + + fn request(if_none_match: Option<&str>) -> FetchRequest { + FetchRequest { + if_none_match: if_none_match.map(str::to_owned), + timeout: Duration::from_secs(5), + } + } + + fn response(status: u16, retry_after: Option<&str>) -> Reply { + Reply::Response { + status, + etag: Some("\"v1\"".to_owned()), + retry_after: retry_after.map(str::to_owned), + body: b"{}".to_vec(), + } + } + + fn retryable(status: u16, retry_after: Option) -> Attempt { + Attempt::Retryable { + reason: format!("HTTP {status} from {URL}"), + retry_after, + } + } + + fn failed(status: u16) -> Attempt { + Attempt::Failed { + reason: format!("HTTP {status} from {URL}"), + } + } + + fn fetched() -> Attempt { + Attempt::Fetched { + body: b"{}".to_vec(), + etag: Some("\"v1\"".to_owned()), + } + } + + /// The same cases as `test_fetch_status_rule_matches_main` in `test_get_model_cost_map.py`, which runs them + /// against main's `_classify_fetch_response` + #[rstest] + #[case::ok(200, None, fetched())] + #[case::created(201, None, fetched())] + #[case::redirect_is_a_document(302, None, fetched())] + #[case::past_599_is_a_document(600, None, fetched())] + #[case::not_modified_without_an_etag_is_a_document(304, None, fetched())] + #[case::too_many_requests(429, None, retryable(429, None))] + #[case::server_error(500, None, retryable(500, None))] + #[case::bad_gateway(502, None, retryable(502, None))] + #[case::unavailable(503, None, retryable(503, None))] + #[case::gateway_timeout(504, None, retryable(504, None))] + #[case::not_implemented_is_final(501, None, failed(501))] + #[case::not_found_is_final(404, None, failed(404))] + #[case::bad_request_is_final(400, None, failed(400))] + #[case::last_server_error_is_final(599, None, failed(599))] + #[case::retry_after_seconds(503, Some("4.5"), retryable(503, Some(4.5)))] + #[case::retry_after_padded(503, Some(" 2 "), retryable(503, Some(2.0)))] + #[case::retry_after_negative(503, Some("-1"), retryable(503, Some(-1.0)))] + #[case::retry_after_infinite(503, Some("inf"), retryable(503, Some(f64::INFINITY)))] + #[case::retry_after_http_date(503, Some("Wed, 21 Oct 2026 07:28:00 GMT"), retryable(503, None))] + #[case::retry_after_repeated(503, Some("1, 2"), retryable(503, None))] + #[case::retry_after_ignored_on_a_final_status(404, Some("1"), failed(404))] + fn status_rule_matches_main( + #[case] status: u16, + #[case] retry_after: Option<&str>, + #[case] expected: Attempt, + ) { + assert_eq!( + classify(response(status, retry_after), URL, &request(None)), + expected + ); + } + + #[rstest] + fn not_modified_needs_a_conditional_request() { + assert_eq!( + classify(response(304, None), URL, &request(Some("\"v1\""))), + Attempt::NotModified + ); + } + + #[rstest] + fn transport_failures_retry_and_refusals_are_final() { + let transport = classify( + Reply::Transport { + error: io::Error::from(io::ErrorKind::ConnectionReset).into(), + }, + URL, + &request(None), + ); + assert!(matches!( + transport, + Attempt::Retryable { + retry_after: None, + .. + } + )); + let refused = classify( + Reply::Refused { + error: crate::Error::unsupported_scheme("ftp".to_owned()), + }, + URL, + &request(None), + ); + assert_eq!( + refused, + Attempt::Failed { + reason: format!("could not fetch {URL}: unsupported catalog URL scheme \"ftp\""), + } + ); + } +} diff --git a/litellm-rust/crates/model-catalog-refresh/src/step.rs b/litellm-rust/crates/model-catalog-refresh/src/step.rs deleted file mode 100644 index 822d1807916..00000000000 --- a/litellm-rust/crates/model-catalog-refresh/src/step.rs +++ /dev/null @@ -1,66 +0,0 @@ -use std::time::Duration; - -use litellm_model_catalog::Catalog; - -use crate::Error; -use crate::source::Source; - -/// One attempt. Where to fetch from is the `Source`, so the request carries no URL -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct FetchRequest { - pub if_none_match: Option, - pub timeout: Duration, -} - -#[derive(Debug)] -pub enum Reply { - /// An HTTP, S3 or file answer. The body is read only on 200 - Response { - status: u16, - etag: Option, - retry_after: Option, - body: Option>, - }, - /// No answer reached us - Transport { error: Error }, -} - -/// What the fetch state machine asks its driver to do next -#[derive(Debug)] -pub enum Step { - Fetch(FetchRequest), - Sleep(Duration), - /// A parsed document to install, with its rejected row errors - Adopt { - catalog: Box, - installed: Installed, - errors: Vec, - }, - /// 304, a final status, or retries exhausted - Done { - reason: DoneReason, - }, -} - -#[derive(Debug)] -pub enum DoneReason { - NotModified, - Status(u16), - Exhausted(Error), -} - -/// What is installed now, the input to the next revalidation -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct Installed { - pub origin: Origin, - pub etag: Option, - /// `git_blob_id` of the installed bytes - pub revision: String, - pub model_count: usize, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum Origin { - Bundled, - Remote { source: Source }, -} diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index abc5d198729..cd70b2f6f61 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -1,3 +1,7 @@ +## Boundary + +This crate holds every rule about the cost map that more than one host applies, and it does no I/O and has no pyo3. Fetching and retrying a map is `litellm-model-catalog-refresh`. `python-bridge/src/model_registry/` only adapts Python objects to the traits here (`FieldValue`, `Rows`) and keeps the process statics, so a rule written there is a bug. Move it here and call it from the bridge + ## Concepts Every price in a row lowers to one `Price { meter, channel, when }` and its `UnitPrice`. These keys come from the `FlatKey::parse` tests @@ -46,7 +50,11 @@ Every price in a row lowers to one `Price { meter, channel, when }` and its `Uni `store/` is the catalog's in-memory store and the only place that indexes model names. Its rules are in [src/store/AGENTS.md](src/store/AGENTS.md) -`layers.rs`, `resolve/`, `model_info.rs` and `register.rs` hold types only. The functions over them land with their first caller +`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller + +`provenance.rs` is where the installed map came from, which hosts record and report + +`register.rs` is the `register_model` merge rule, `(base row, registration) -> row`, generic over `FieldValue` so python-bridge runs it on Python objects and a JSON host on `serde_json::Value` 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..85429f61c5f 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-python-compat.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/aliases.rs b/litellm-rust/crates/model-catalog/src/aliases.rs new file mode 100644 index 00000000000..4cae7fee51b --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/aliases.rs @@ -0,0 +1,95 @@ +use std::hash::Hash; + +use indexmap::IndexMap; + +/// Why an alias was left unclaimed +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Unclaimed { + /// The alias is already a top-level key of the map + ShadowsModel, + TakenBy(O), +} + +/// Aliases in claim order, each served as the row of the first model that claimed it +#[derive(Clone, Debug)] +pub struct AliasClaims { + owners: IndexMap, +} + +impl AliasClaims { + pub fn new() -> Self { + Self { + owners: IndexMap::new(), + } + } + + pub fn claim(&mut self, alias: &A, owner: &O, is_model: bool) -> Result<(), Unclaimed> { + if is_model { + return Err(Unclaimed::ShadowsModel); + } + if let Some(taken_by) = self.owners.get(alias) { + return Err(Unclaimed::TakenBy(taken_by.clone())); + } + self.owners.insert(alias.clone(), owner.clone()); + Ok(()) + } + + /// Each owner with its aliases, owners in the order of their first claim + pub fn groups(&self) -> IndexMap<&O, Vec<&A>> { + self.owners + .iter() + .fold(IndexMap::new(), |mut groups, (alias, owner)| { + groups.entry(owner).or_insert_with(Vec::new).push(alias); + groups + }) + } + + pub fn into_owners(self) -> IndexMap { + self.owners + } +} + +impl Default for AliasClaims { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + type Claims = AliasClaims<&'static str, &'static str>; + + #[rstest] + #[case::free("a", "m", false, Ok(()))] + #[case::shadows_a_model("a", "m", true, Err(Unclaimed::ShadowsModel))] + #[case::taken("taken", "m", false, Err(Unclaimed::TakenBy("first")))] + #[case::shadow_wins_over_taken("taken", "m", true, Err(Unclaimed::ShadowsModel))] + fn a_claim_is_refused_by_a_model_key_first_then_by_an_earlier_claim( + #[case] alias: &'static str, + #[case] owner: &'static str, + #[case] is_model: bool, + #[case] expected: Result<(), Unclaimed<&'static str>>, + ) { + let mut claims = Claims::new(); + claims.claim(&"taken", &"first", false).unwrap(); + assert_eq!(claims.claim(&alias, &owner, is_model), expected); + } + + #[rstest] + fn groups_keep_owners_in_first_claim_order_and_aliases_in_claim_order() { + let mut claims = Claims::new(); + for (alias, owner) in [("b1", "b"), ("a1", "a"), ("b2", "b"), ("a1", "c")] { + let _ = claims.claim(&alias, &owner, false); + } + let groups: Vec<(&str, Vec<&str>)> = claims + .groups() + .into_iter() + .map(|(owner, aliases)| (*owner, aliases.into_iter().copied().collect())) + .collect(); + assert_eq!(groups, vec![("b", vec!["b1", "b2"]), ("a", vec!["a1"])]); + } +} diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index 98af785b17e..0d0015ae16c 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -5,15 +5,16 @@ use indexmap::IndexMap; use serde::Deserialize; use serde_json::Value; +use crate::aliases::{AliasClaims, Unclaimed}; use crate::error::{Error, Loaded, Rejects, Repr, pointer_child, pointer_for_key}; use crate::fallback::{FallbackGeneralizations, FallbackRule}; +use crate::integrity::{ + FALLBACK_GENERALIZATIONS, SAMPLE_SPEC, SizeProblem, is_model_key, size_problem, +}; use crate::model::Model; use crate::store::{Rows, Store}; use crate::wire::{EntryWire, Fields}; -const SAMPLE_SPEC: &str = "sample_spec"; -const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; - #[derive(Clone, Copy, Debug, PartialEq)] pub struct IntegrityLimits { min_model_count: usize, @@ -80,13 +81,18 @@ impl CatalogRows { } } -/// Model names match ignoring case, and keys are grouped by their raw `litellm_provider` +/// The name two spellings of one `model_cost` key share, for every `model_cost` source. Names match ignoring case +pub fn fold_model_name(name: &str) -> String { + name.to_lowercase() +} + +/// Keys are grouped by their raw `litellm_provider` impl Rows for CatalogRows { type Row = Fields; type Stamp = (); fn fold(name: &str) -> String { - name.to_lowercase() + fold_model_name(name) } fn stamp(&self) -> Self::Stamp {} @@ -104,7 +110,7 @@ impl Rows for CatalogRows { for key in keys { let provider = self .fields(key) - .and_then(|fields| fields.get("litellm_provider")?.as_str()); + .and_then(|fields| fields.get(crate::LITELLM_PROVIDER)?.as_str()); visit(key, provider); } } @@ -172,24 +178,24 @@ impl Catalog { return Err(Error::new("", Repr::Empty)); } - let mut aliases: IndexMap = IndexMap::new(); + let mut claims = AliasClaims::new(); for (model, names) in alias_lists { for alias in names { - if entries.contains_key(&alias) { - problems.push(Error::from_key(&model, Repr::AliasShadowsModel { alias })); - } else if let Some(taken_by) = aliases.get(&alias) { - problems.push(Error::from_key( - &model, - Repr::AliasTaken { - alias, - taken_by: taken_by.clone(), - }, - )); - } else { - aliases.insert(alias, model.clone()); + match claims.claim(&alias, &model, entries.contains_key(&alias)) { + Ok(()) => {} + Err(Unclaimed::ShadowsModel) => { + problems.push(Error::from_key(&model, Repr::AliasShadowsModel { alias })); + } + Err(Unclaimed::TakenBy(taken_by)) => { + problems.push(Error::from_key( + &model, + Repr::AliasTaken { alias, taken_by }, + )); + } } } } + let aliases = claims.into_owners(); let fallback_generalizations = sections .iter() @@ -275,27 +281,30 @@ impl Catalog { } pub fn check_size(&self, reference: usize, limits: IntegrityLimits) -> Result<(), Error> { - let actual = self.model_count(); - if actual < limits.min_model_count { - return Err(Error::new( - "", - Repr::BelowMinimum { - actual, - minimum: limits.min_model_count, - }, - )); - } - if (actual as f64) < (reference as f64) * limits.min_reference_ratio { - return Err(Error::new( + let problem = size_problem( + self.model_count(), + reference, + limits.min_model_count, + limits.min_reference_ratio, + ); + match problem { + None => Ok(()), + Some(SizeProblem::BelowMinimum { actual, minimum }) => { + Err(Error::new("", Repr::BelowMinimum { actual, minimum })) + } + Some(SizeProblem::Shrunk { + actual, + reference, + ratio, + }) => Err(Error::new( "", Repr::Shrunk { actual, reference, - ratio: limits.min_reference_ratio, + ratio, }, - )); + )), } - Ok(()) } pub fn without_rejected(self, errors: &[Error]) -> (Self, Vec<&Error>) { @@ -420,7 +429,7 @@ fn read_row(fields: &Fields) -> (Option, Option>, Vec) } fn is_section(key: &str) -> bool { - key == SAMPLE_SPEC || key == FALLBACK_GENERALIZATIONS + !is_model_key(key) } fn unknown_timezone(wire: &EntryWire) -> Option { diff --git a/litellm-rust/crates/model-catalog/src/integrity.rs b/litellm-rust/crates/model-catalog/src/integrity.rs new file mode 100644 index 00000000000..93e11ac3312 --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/integrity.rs @@ -0,0 +1,95 @@ +use sha1::{Digest, Sha1}; + +pub const SAMPLE_SPEC: &str = "sample_spec"; +pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations"; + +/// Why a fetched map is too small to serve +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum SizeProblem { + BelowMinimum { + actual: usize, + minimum: usize, + }, + Shrunk { + actual: usize, + reference: usize, + ratio: f64, + }, +} + +/// Whether a top-level key is a model, as opposed to a document section +pub fn is_model_key(key: &str) -> bool { + key != SAMPLE_SPEC && key != FALLBACK_GENERALIZATIONS +} + +/// `actual` models against `minimum` and against `ratio` of the `reference` map's models, skipped for an empty reference +pub fn size_problem( + actual: usize, + reference: usize, + minimum: usize, + ratio: f64, +) -> Option { + if actual < minimum { + return Some(SizeProblem::BelowMinimum { actual, minimum }); + } + (reference > 0 && (actual as f64) < (reference as f64) * ratio).then_some(SizeProblem::Shrunk { + actual, + reference, + ratio, + }) +} + +/// The git blob id of a document's bytes, which is the revision a host reports for it +pub fn revision(body: &[u8]) -> String { + let mut hasher = Sha1::new(); + hasher.update(format!("blob {}\0", body.len()).as_bytes()); + hasher.update(body); + hasher + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::model("gpt-4o", true)] + #[case::sample_spec("sample_spec", false)] + #[case::fallback_generalizations("fallback_generalizations", false)] + #[case::case_sensitive("Sample_Spec", true)] + fn sections_are_not_models(#[case] key: &str, #[case] expected: bool) { + assert_eq!(is_model_key(key), expected); + } + + #[rstest] + #[case::fits(10, 10, 5, 0.5, None)] + #[case::below_minimum(4, 10, 5, 0.5, Some(SizeProblem::BelowMinimum { actual: 4, minimum: 5 }))] + #[case::minimum_checked_first(1, 10, 5, 0.5, Some(SizeProblem::BelowMinimum { actual: 1, minimum: 5 }))] + #[case::shrunk(5, 11, 0, 0.5, Some(SizeProblem::Shrunk { actual: 5, reference: 11, ratio: 0.5 }))] + #[case::exactly_at_ratio(5, 10, 0, 0.5, None)] + #[case::empty_reference_skips_ratio(0, 0, 0, 2.0, None)] + #[case::ratio_above_one(10, 10, 0, 1.5, Some(SizeProblem::Shrunk { actual: 10, reference: 10, ratio: 1.5 }))] + #[case::nan_ratio_never_shrinks(0, 10, 0, f64::NAN, None)] + fn size_matches_python_validate_model_cost_map( + #[case] actual: usize, + #[case] reference: usize, + #[case] minimum: usize, + #[case] ratio: f64, + #[case] expected: Option, + ) { + assert_eq!(size_problem(actual, reference, minimum, ratio), expected); + } + + #[rstest] + fn the_revision_is_the_git_blob_id_of_the_body() { + assert_eq!( + revision(b"hello\n"), + "ce013625030ba8dba906f756967f9e9ca394464a" + ); + } +} 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..6df27f94411 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -35,27 +35,31 @@ macro_rules_attribute::attribute_alias! { #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]; } +mod aliases; mod bundled; mod catalog; mod error; mod fallback; -mod layers; +mod integrity; mod lower; mod model; mod model_info; +mod provenance; mod register; mod resolve; mod store; mod vocabulary; pub mod wire; +pub use aliases::*; pub use bundled::*; pub use catalog::*; pub use error::*; pub use fallback::*; -pub use layers::*; +pub use integrity::*; pub use model::*; pub use model_info::*; +pub use provenance::*; pub use register::*; pub use resolve::*; pub use store::*; diff --git a/litellm-rust/crates/model-catalog/src/provenance.rs b/litellm-rust/crates/model-catalog/src/provenance.rs new file mode 100644 index 00000000000..fd748b6555d --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/provenance.rs @@ -0,0 +1,26 @@ +use std::time::SystemTime; + +/// Where the installed catalog came from and when, written once per install and never edited +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Provenance { + pub origin: Origin, + /// `LITELLM_LOCAL_MODEL_COST_MAP`, offline mode or a CLI process kept the bundled map + pub is_env_forced: bool, + /// Why the remote map was not installed, for a bundled map that stands in for it + pub fallback_reason: Option, + pub loaded_at: SystemTime, + /// `git_blob_id` of the installed bytes + pub revision: Option, + pub etag: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Origin { + /// The map shipped with the package, standing in for `attempted` when that is set + Bundled { + attempted: Option, + }, + Remote { + url: String, + }, +} diff --git a/litellm-rust/crates/model-catalog/src/register.rs b/litellm-rust/crates/model-catalog/src/register.rs index ed34082ae57..6b2694ffd90 100644 --- a/litellm-rust/crates/model-catalog/src/register.rs +++ b/litellm-rust/crates/model-catalog/src/register.rs @@ -1,26 +1,500 @@ -use std::collections::{BTreeSet, HashMap}; - use indexmap::IndexMap; +use litellm_python_compat::{BigInt, number}; +use serde_json::Value; -use crate::wire::Fields; +/// One row's fields in the order Python's dict holds them +pub type Fields = IndexMap; -/// 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, +const INPUT_COST_PER_TOKEN: &str = "input_cost_per_token"; +const OUTPUT_COST_PER_TOKEN: &str = "output_cost_per_token"; +const TOKEN_COST_FIELDS: [&str; 2] = [INPUT_COST_PER_TOKEN, OUTPUT_COST_PER_TOKEN]; +const CACHE_READ_INPUT_TOKEN_COST: &str = "cache_read_input_token_cost"; +const EXPLICIT_CACHE_READ_INPUT_TOKEN_COST: &str = "explicit_cache_read_input_token_cost"; +/// The field a row names its provider in, which `Store` groups keys by +pub const LITELLM_PROVIDER: &str = "litellm_provider"; + +/// Whether there are rows and none of them carries a per-token price, each row given as its field test. +/// `get_model_info` fills a missing price with `0`, so this tells an unpriced row from a declared zero +pub fn omits_token_price bool>(rows: impl IntoIterator) -> bool { + let mut rows = rows.into_iter().peekable(); + rows.peek().is_some() && rows.all(|has| !TOKEN_COST_FIELDS.into_iter().any(&has)) } -/// 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>, +/// How the merge rule sees one field value +#[derive(Clone, Debug, PartialEq)] +pub enum Shape { + None, + Str(String), + Map, + Other, +} + +/// A number parsed from a stringified registration value +#[derive(Clone, Debug, PartialEq)] +pub enum Number { + Float(f64), + Int(BigInt), +} + +/// A field value of the host that owns the rows, Python objects in python-bridge, JSON in a JSON host +pub trait FieldValue: Clone { + fn shape(&self) -> Shape; + fn none() -> Self; + /// `None` when the host cannot hold `number`, which keeps the string + fn number(number: Number) -> Option; + /// For a dict `self`, `{**base, **self}` when `base` is a dict, else `dict(self)` + fn merged_over(&self, base: Option<&Self>) -> Self; + /// The value as pydantic validates `tuple[Mapping[str, object], ...]`, `None` where it raises + fn rate_tiers(&self) -> Option>>; + /// The value as pydantic validates `Mapping[str, object]`, `None` where it raises + fn rate_table(&self) -> Option>; + fn tier_list(tiers: Vec>) -> Self; + fn table(fields: Fields) -> Self; +} + +/// Main's `update_dictionary`. A `None` value is skipped, a string becomes the number it spells, a dict is merged +/// over a dict already there, and anything else replaces the field. New fields land after the existing ones +pub fn update_dictionary(existing: Fields, new: &Fields) -> Fields { + new.iter().fold(existing, |mut merged, (name, value)| { + let updated = match value.shape() { + Shape::None => return merged, + Shape::Str(text) => stringified_number(&text) + .and_then(V::number) + .unwrap_or_else(|| value.clone()), + Shape::Map => value.merged_over(merged.get(name)), + Shape::Other => value.clone(), + }; + merged.insert(name.clone(), updated); + merged + }) +} + +/// Main's `_convert_stringified_numbers`, `float` when the text has an `e` or a `.`, else `int` +pub fn stringified_number(text: &str) -> Option { + if text.contains(['e', 'E', '.']) { + return number::float(text).map(Number::Float); + } + number::int(text).map(Number::Int) +} + +/// The row a `register_model` call writes, its `base` row with `registration` merged over it. A `litellm_provider` +/// of `None` in `base` is dropped so the registration's provider or none is kept +pub fn registration_row(base: &Fields, registration: &Fields) -> Fields { + let existing: Fields = base + .iter() + .filter(|(name, value)| *name != LITELLM_PROVIDER || value.shape() != Shape::None) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(); + let merged = update_dictionary(existing, registration); + base_fields_replaced_by(registration) + .iter() + .fold(merged, |mut row, name| { + row.insert((*name).to_owned(), V::none()); + row + }) +} + +/// Fields of the base row a registration clears, so a flat cache read price is not shadowed by the base row's +/// explicit one +pub fn base_fields_replaced_by(registration: &Fields) -> &'static [&'static str] { + if is_set(registration, CACHE_READ_INPUT_TOKEN_COST) + && !is_set(registration, EXPLICIT_CACHE_READ_INPUT_TOKEN_COST) + { + return &[EXPLICIT_CACHE_READ_INPUT_TOKEN_COST]; + } + &[] +} + +/// Main's `_without_synthesized_costs`. `get_model_info` fills a missing per-token price with `0`, so a token price +/// field survives only when the raw row has it or the registration names it +pub fn without_synthesized_costs( + model_info: &Fields, + raw_fields: impl Fn(&str) -> bool, + registration: &Fields, +) -> Fields { + model_info + .iter() + .filter(|(name, _)| { + !TOKEN_COST_FIELDS.contains(&name.as_str()) + || raw_fields(name) + || registration.contains_key(*name) + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +/// Main's `_without_overridden_price_variants`. A flat per-token price in `registration` is that direction's price +/// everywhere, so every inherited rate that would stand in for it is dropped and each `tiered_pricing` tier carries +/// the flat price +pub fn without_overridden_price_variants( + entry: &Fields, + registration: &Fields, +) -> Fields { + let flat_prices: Fields = TOKEN_COST_FIELDS + .iter() + .filter_map(|name| { + registration + .get(*name) + .filter(|value| value.shape() != Shape::None) + .map(|value| ((*name).to_owned(), value.clone())) + }) + .collect(); + let overridden: Vec<&str> = flat_prices + .keys() + .flat_map(|name| rate_fields_overridden_by(name)) + .copied() + .collect(); + if overridden.is_empty() { + return entry.clone(); + } + entry + .iter() + .filter(|(name, _)| !is_overridden_rate(name, &overridden)) + .map(|(name, value)| { + ( + name.clone(), + rate_table_with_flat_prices(name, value, &overridden, &flat_prices), + ) + }) + .collect() +} + +/// Main's `model_info_with_pricing_overrides` before its pydantic validation, `overrides` already narrowed to +/// server derived pricing keys +pub fn with_pricing_overrides( + model_info: &Fields, + overrides: &Fields, +) -> Fields { + update_dictionary( + without_overridden_price_variants(model_info, overrides), + overrides, + ) +} + +fn is_set(fields: &Fields, name: &str) -> bool { + fields + .get(name) + .is_some_and(|value| value.shape() != Shape::None) +} + +fn rate_fields_overridden_by(name: &str) -> &'static [&'static str] { + match name { + INPUT_COST_PER_TOKEN => &[INPUT_COST_PER_TOKEN], + OUTPUT_COST_PER_TOKEN => &[OUTPUT_COST_PER_TOKEN, "output_cost_per_reasoning_token"], + _ => &[], + } +} + +fn is_overridden_rate(name: &str, overridden: &[&str]) -> bool { + let is_variant = overridden.iter().any(|rate| { + name.strip_prefix(rate) + .is_some_and(|rest| rest.starts_with('_')) + }); + overridden.contains(&name) || (is_variant && !name.contains("_cache_")) +} + +fn without_overridden_rates(rates: &Fields, overridden: &[&str]) -> Fields { + rates + .iter() + .filter(|(name, _)| !is_overridden_rate(name, overridden)) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() +} + +fn rate_table_with_flat_prices( + name: &str, + table: &V, + overridden: &[&str], + flat_prices: &Fields, +) -> V { + let rewritten = match name { + "tiered_pricing" => table.rate_tiers().map(|tiers| { + V::tier_list( + tiers + .iter() + .map(|tier| { + let mut priced = without_overridden_rates(tier, overridden); + priced.extend(flat_prices.clone()); + priced + }) + .collect(), + ) + }), + "off_peak_pricing" => table + .rate_table() + .map(|rates| V::table(without_overridden_rates(&rates, overridden))), + _ => None, + }; + rewritten.unwrap_or_else(|| table.clone()) +} + +impl FieldValue for Value { + fn shape(&self) -> Shape { + match self { + Value::Null => Shape::None, + Value::String(text) => Shape::Str(text.clone()), + Value::Object(_) => Shape::Map, + _ => Shape::Other, + } + } + + fn none() -> Self { + Value::Null + } + + fn number(number: Number) -> Option { + match number { + Number::Float(value) => serde_json::Number::from_f64(value).map(Value::Number), + Number::Int(value) => i64::try_from(&value) + .map(Value::from) + .or_else(|_| u64::try_from(&value).map(Value::from)) + .ok(), + } + } + + fn merged_over(&self, base: Option<&Self>) -> Self { + match (base, self) { + (Some(Value::Object(base)), Value::Object(over)) => { + let mut merged = base.clone(); + merged.extend(over.clone()); + Value::Object(merged) + } + _ => self.clone(), + } + } + + fn rate_tiers(&self) -> Option>> { + let Value::Array(tiers) = self else { + return None; + }; + tiers.iter().map(Value::rate_table).collect() + } + + fn rate_table(&self) -> Option> { + let Value::Object(rates) = self else { + return None; + }; + Some( + rates + .iter() + .map(|(name, rate)| (name.clone(), rate.clone())) + .collect(), + ) + } + + fn tier_list(tiers: Vec>) -> Self { + Value::Array(tiers.into_iter().map(Value::table).collect()) + } + + fn table(fields: Fields) -> Self { + Value::Object(fields.into_iter().collect()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn fields(value: Value) -> Fields { + value.rate_table().expect("a JSON object") + } + + fn written(row: Fields) -> String { + Value::table(row).to_string() + } + + /// Each expected row is what main's `update_dictionary` returns for the same JSON, keys in dict order + #[rstest] + #[case::none_is_skipped(json!({"a": 1}), json!({"a": null, "b": null}), r#"{"a":1}"#)] + #[case::scientific_string(json!({}), json!({"input_cost_per_token": "3e-07"}), r#"{"input_cost_per_token":3e-7}"#)] + #[case::integer_string(json!({}), json!({"max_tokens": "8192"}), r#"{"max_tokens":8192}"#)] + #[case::padded_integer_string(json!({}), json!({"max_tokens": " 8_192 "}), r#"{"max_tokens":8192}"#)] + #[case::point_makes_a_float(json!({}), json!({"max_tokens": "8192."}), r#"{"max_tokens":8192.0}"#)] + #[case::text_stays(json!({}), json!({"mode": "chat"}), r#"{"mode":"chat"}"#)] + #[case::e_without_a_number_stays(json!({}), json!({"litellm_provider": "openai"}), r#"{"litellm_provider":"openai"}"#)] + #[case::infinity_without_e_or_point_stays(json!({}), json!({"x": "inf"}), r#"{"x":"inf"}"#)] + #[case::dict_merges_over_dict(json!({"m": {"a": 1, "b": 2}}), json!({"m": {"b": 3, "c": 4}}), r#"{"m":{"a":1,"b":3,"c":4}}"#)] + #[case::dict_replaces_a_scalar(json!({"m": 1}), json!({"m": {"b": 3}}), r#"{"m":{"b":3}}"#)] + #[case::nested_strings_are_kept(json!({}), json!({"m": {"b": "3"}}), r#"{"m":{"b":"3"}}"#)] + #[case::other_values_replace(json!({"a": 1, "b": true}), json!({"b": [1, "2"], "c": 0.5}), r#"{"a":1,"b":[1,"2"],"c":0.5}"#)] + #[case::existing_order_then_new(json!({"z": 1, "a": 2}), json!({"m": 3, "z": 4}), r#"{"z":4,"a":2,"m":3}"#)] + fn update_dictionary_matches_main( + #[case] existing: Value, + #[case] new: Value, + #[case] expected: &str, + ) { + assert_eq!( + written(update_dictionary(fields(existing), &fields(new))), + expected + ); + } + + #[rstest] + #[case::provider_none_is_dropped( + json!({"litellm_provider": null, "max_tokens": 10}), + json!({"input_cost_per_token": 1}), + r#"{"max_tokens":10,"input_cost_per_token":1}"# + )] + #[case::provider_none_gives_way_to_the_registration( + json!({"litellm_provider": null}), + json!({"litellm_provider": "openai"}), + r#"{"litellm_provider":"openai"}"# + )] + #[case::cache_read_clears_explicit( + json!({"explicit_cache_read_input_token_cost": 1, "a": 1}), + json!({"cache_read_input_token_cost": 2}), + r#"{"explicit_cache_read_input_token_cost":null,"a":1,"cache_read_input_token_cost":2}"# + )] + #[case::explicit_cache_read_is_kept_when_registered( + json!({}), + json!({"cache_read_input_token_cost": 2, "explicit_cache_read_input_token_cost": 3}), + r#"{"cache_read_input_token_cost":2,"explicit_cache_read_input_token_cost":3}"# + )] + #[case::cache_read_none_clears_nothing( + json!({"explicit_cache_read_input_token_cost": 1}), + json!({"cache_read_input_token_cost": null}), + r#"{"explicit_cache_read_input_token_cost":1}"# + )] + fn registration_row_matches_main( + #[case] base: Value, + #[case] registration: Value, + #[case] expected: &str, + ) { + assert_eq!( + written(registration_row(&fields(base), &fields(registration))), + expected + ); + } + + #[rstest] + #[case::no_flat_price( + json!({"input_cost_per_token_above_200k_tokens": 1}), + json!({"cache_read_input_token_cost": 1}), + r#"{"input_cost_per_token_above_200k_tokens":1}"# + )] + #[case::input_variants_dropped_cache_kept( + json!({"input_cost_per_token": 1, "input_cost_per_token_above_200k_tokens": 2, "input_cost_per_token_cache_hit": 3, "output_cost_per_token_priority": 4}), + json!({"input_cost_per_token": 9}), + r#"{"input_cost_per_token_cache_hit":3,"output_cost_per_token_priority":4}"# + )] + #[case::output_drops_reasoning( + json!({"output_cost_per_reasoning_token": 1, "output_cost_per_reasoning_token_batches": 2, "output_cost_per_token_batches": 3}), + json!({"output_cost_per_token": 9}), + r#"{}"# + )] + #[case::none_price_overrides_nothing( + json!({"input_cost_per_token_priority": 1}), + json!({"input_cost_per_token": null}), + r#"{"input_cost_per_token_priority":1}"# + )] + #[case::name_without_underscore_is_not_a_variant( + json!({"input_cost_per_tokens": 1}), + json!({"input_cost_per_token": 9}), + r#"{"input_cost_per_tokens":1}"# + )] + #[case::tiers_carry_the_flat_price( + json!({"tiered_pricing": [{"input_cost_per_token": 1, "output_cost_per_token": 2, "range": [0, 10]}]}), + json!({"input_cost_per_token": "9"}), + r#"{"tiered_pricing":[{"output_cost_per_token":2,"range":[0,10],"input_cost_per_token":"9"}]}"# + )] + #[case::invalid_tiers_are_kept( + json!({"tiered_pricing": [1, {"input_cost_per_token": 1}]}), + json!({"input_cost_per_token": 9}), + r#"{"tiered_pricing":[1,{"input_cost_per_token":1}]}"# + )] + #[case::off_peak_drops_overridden( + json!({"off_peak_pricing": {"input_cost_per_token": 1, "hours": [1, 2]}}), + json!({"input_cost_per_token": 9}), + r#"{"off_peak_pricing":{"hours":[1,2]}}"# + )] + #[case::invalid_off_peak_is_kept( + json!({"off_peak_pricing": [1]}), + json!({"input_cost_per_token": 9}), + r#"{"off_peak_pricing":[1]}"# + )] + fn without_overridden_price_variants_matches_main( + #[case] entry: Value, + #[case] registration: Value, + #[case] expected: &str, + ) { + assert_eq!( + written(without_overridden_price_variants( + &fields(entry), + &fields(registration) + )), + expected + ); + } + + #[rstest] + #[case::synthesized_dropped(json!({}), json!({}), r#"{"key":"m","mode":"chat"}"#)] + #[case::raw_price_kept(json!({"input_cost_per_token": 1}), json!({}), r#"{"key":"m","input_cost_per_token":0,"mode":"chat"}"#)] + #[case::registered_price_kept(json!({}), json!({"output_cost_per_token": null}), r#"{"key":"m","output_cost_per_token":0,"mode":"chat"}"#)] + fn without_synthesized_costs_matches_main( + #[case] raw: Value, + #[case] registration: Value, + #[case] expected: &str, + ) { + let model_info = fields( + json!({"key": "m", "input_cost_per_token": 0, "output_cost_per_token": 0, "mode": "chat"}), + ); + let raw = fields(raw); + assert_eq!( + written(without_synthesized_costs( + &model_info, + |name| raw.contains_key(name), + &fields(registration) + )), + expected + ); + } + + #[rstest] + fn pricing_overrides_replace_variants_then_merge() { + let model_info = fields( + json!({"input_cost_per_token": 1, "input_cost_per_token_above_200k_tokens": 2, "mode": "chat"}), + ); + let overrides = fields(json!({"input_cost_per_token": "5e-07"})); + assert_eq!( + written(with_pricing_overrides(&model_info, &overrides)), + r#"{"mode":"chat","input_cost_per_token":5e-7}"# + ); + } + + #[rstest] + fn a_number_json_cannot_hold_keeps_the_string() { + let row = update_dictionary( + Fields::new(), + &fields(json!({"a": "1e400", "b": "99999999999999999999999"})), + ); + assert_eq!( + written(row), + r#"{"a":"1e400","b":"99999999999999999999999"}"# + ); + } + + #[rstest] + #[case::no_rows(json!([]), false)] + #[case::unpriced_rows(json!([{"mode": "chat"}, {}]), true)] + #[case::declared_zero(json!([{"mode": "chat"}, {"output_cost_per_token": 0}]), false)] + #[case::input_price(json!([{"input_cost_per_token": null}]), false)] + fn a_map_omits_token_prices_only_when_no_row_carries_one( + #[case] rows: Value, + #[case] expected: bool, + ) { + let rows: Vec> = rows + .as_array() + .unwrap() + .iter() + .cloned() + .map(fields) + .collect(); + assert_eq!( + omits_token_price(rows.iter().map(|row| |name: &str| row.contains_key(name))), + expected + ); + } } diff --git a/litellm-rust/crates/model-catalog/src/store/AGENTS.md b/litellm-rust/crates/model-catalog/src/store/AGENTS.md index f1dacd2e217..f7c1258696b 100644 --- a/litellm-rust/crates/model-catalog/src/store/AGENTS.md +++ b/litellm-rust/crates/model-catalog/src/store/AGENTS.md @@ -1,6 +1,6 @@ # Store -`Store` is an in-memory store, a small database built for exactly three queries. It knows nothing about the catalog. A source picks its own `Rows::Row` type, how names fold for matching, and the group each key is listed under. The catalog's sources fold to lowercase and group by raw `litellm_provider` +`Store` is an in-memory store, a small database built for exactly three queries. It knows nothing about the catalog. A source picks its own `Rows::Row` type, how names fold for matching, and the group each key is listed under. The catalog's sources, `CatalogRows` and python-bridge's `litellm.model_cost`, fold with `fold_model_name` and group by raw `LITELLM_PROVIDER` | Query | Method | Real caller | |---|---|---| diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index fd826f549e5..0873f3a334a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -49,6 +49,7 @@ litellm-llms-types.workspace = true litellm-router-types.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 } pyo3.workspace = true prost.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index b4e20a8f750..8fa17436e5d 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -131,6 +131,14 @@ pub(crate) fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult, variant: ClientVariant) -> PyResult { + let config = call_config(py, &PyDict::new(py), true)?; + pool() + .unpooled_client(&config, variant) + .map_err(client_error) +} + pub(crate) fn client_error(error: litellm_http::Error) -> PyErr { match error.tls_source() { Some(TlsSource::ClientIdentity) => PyValueError::new_err( diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index de35419308a..9ff9fa9fb69 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -34,7 +34,15 @@ mod _native { #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] - use crate::model_registry::bundled_model_cost_map; + use crate::model_registry::{ + AliasWarning, ModelCostMapFetch, ModelCountProblem, PyProvenance, PyRefresh, + bundled_model_cost_map, cost_map_omits_token_price, count_model_entries, + expand_model_aliases, find_model_cost_key, git_blob_id, install_model_cost_map_provenance, + invalidate_model_cost_keys, live_row_or, model_cost_map_provenance, model_count_problem, + registration_row, stringified_number, take_fallback_generalization_rules, + update_dictionary, with_pricing_overrides, without_overridden_price_variants, + without_synthesized_costs, + }; #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] @@ -113,6 +121,28 @@ mod tests { "NativeDatabase", "trace_encode_error", "bundled_model_cost_map", + "count_model_entries", + "model_count_problem", + "ModelCountProblem", + "expand_model_aliases", + "AliasWarning", + "git_blob_id", + "ModelCostMapFetch", + "ModelCostMapProvenance", + "ModelCostMapRefresh", + "find_model_cost_key", + "cost_map_omits_token_price", + "live_row_or", + "invalidate_model_cost_keys", + "install_model_cost_map_provenance", + "model_cost_map_provenance", + "registration_row", + "stringified_number", + "take_fallback_generalization_rules", + "update_dictionary", + "with_pricing_overrides", + "without_overridden_price_variants", + "without_synthesized_costs", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/model_registry.rs b/litellm-rust/crates/python-bridge/src/model_registry.rs index b7b7cd1fbb3..da166b01cc5 100644 --- a/litellm-rust/crates/python-bridge/src/model_registry.rs +++ b/litellm-rust/crates/python-bridge/src/model_registry.rs @@ -1,6 +1,92 @@ +mod aliases; +mod fetch; +mod keys; +mod live; +mod merge; +mod provenance; + +use litellm_model_catalog::{FALLBACK_GENERALIZATIONS, SizeProblem, is_model_key, size_problem}; +use pyo3::intern; use pyo3::prelude::*; +use pyo3::pybacked::PyBackedBytes; use pyo3::sync::PyOnceLock; -use pyo3::types::PyBytes; +use pyo3::types::{PyBytes, PyDict, PyString}; + +pub(crate) use aliases::{AliasWarning, expand_model_aliases}; +pub(crate) use fetch::{ModelCostMapFetch, PyRefresh}; +pub(crate) use keys::{find_model_cost_key, invalidate_model_cost_keys}; +pub(crate) use live::{cost_map_omits_token_price, live_row_or}; +pub(crate) use merge::{ + registration_row, stringified_number, update_dictionary, with_pricing_overrides, + without_overridden_price_variants, without_synthesized_costs, +}; +pub(crate) use provenance::{ + PyProvenance, install_model_cost_map_provenance, model_cost_map_provenance, +}; + +/// Why a fetched map has too few models, as main's `_check_model_count_not_reduced` decides +#[pyclass(frozen, module = "litellm.rust_bridge._native")] +pub(crate) enum ModelCountProblem { + BelowMinimum { actual: usize, minimum: usize }, + Shrunk { actual: usize, reference: usize }, +} + +/// Top-level keys of `model_cost` that are models. A non-string key is counted, as Python counts it +#[pyfunction] +pub(crate) fn count_model_entries(model_cost: &Bound<'_, PyAny>) -> PyResult { + model_cost.try_iter()?.try_fold(0, |count, key| { + let key = key?; + let is_model = match key.cast::() { + Ok(key) => is_model_key(&key.to_string_lossy()), + Err(_) => true, + }; + Ok(count + usize::from(is_model)) + }) +} + +/// A negative minimum or reference passes the check that uses it, as it does in Python +#[pyfunction] +pub(crate) fn model_count_problem( + fetched_map: &Bound<'_, PyAny>, + backup_model_count: i64, + min_model_count: i64, + max_shrink_ratio: f64, +) -> PyResult> { + let problem = size_problem( + count_model_entries(fetched_map)?, + usize::try_from(backup_model_count).unwrap_or(0), + usize::try_from(min_model_count).unwrap_or(0), + max_shrink_ratio, + ); + Ok(problem.map(|problem| match problem { + SizeProblem::BelowMinimum { actual, minimum } => { + ModelCountProblem::BelowMinimum { actual, minimum } + } + SizeProblem::Shrunk { + actual, reference, .. + } => ModelCountProblem::Shrunk { actual, reference }, + })) +} + +/// Pops the `fallback_generalizations` block and returns its `rules`, `None` when the block is not a dict, as main's +/// `_finalize_model_cost_map` does before it expands aliases +#[pyfunction] +pub(crate) fn take_fallback_generalization_rules<'py>( + model_cost: &Bound<'py, PyDict>, +) -> PyResult> { + let py = model_cost.py(); + let block = + model_cost.call_method1(intern!(py, "pop"), (FALLBACK_GENERALIZATIONS, py.None()))?; + if !block.is_instance_of::() { + return Ok(py.None().into_bound(py)); + } + block.call_method1(intern!(py, "get"), (intern!(py, "rules"),)) +} + +#[pyfunction] +pub(crate) fn git_blob_id(body: PyBackedBytes) -> String { + litellm_model_catalog::revision(&body) +} #[pyfunction] pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> { diff --git a/litellm-rust/crates/python-bridge/src/model_registry/aliases.rs b/litellm-rust/crates/python-bridge/src/model_registry/aliases.rs new file mode 100644 index 00000000000..6c2df648a8e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry/aliases.rs @@ -0,0 +1,141 @@ +use std::collections::HashMap; + +use litellm_model_catalog::{AliasClaims, Unclaimed}; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList, PyTuple}; + +/// A skipped alias, carrying the objects main's warning formats +#[pyclass(frozen, module = "litellm.rust_bridge._native")] +pub(crate) enum AliasWarning { + NotAList { model: Py, type_name: String }, + ShadowsModel { alias: Py, model: Py }, + Taken { alias: Py, model: Py }, +} + +/// Expands every row's `aliases` into top-level keys sharing one `entry_type` row, as main's `_expand_model_aliases`. +/// Returns the warnings in main's order and the error main raises mid-scan, in which case `model_cost` is untouched +#[pyfunction] +pub(crate) fn expand_model_aliases<'py>( + model_cost: &Bound<'py, PyDict>, + entry_type: &Bound<'py, PyAny>, +) -> PyResult<(Vec, Option)> { + let scan = match scan(model_cost) { + Ok(scan) => scan, + Err(Failed { warnings, error }) => return Ok((warnings, Some(error))), + }; + apply(model_cost, entry_type, &scan)?; + Ok((scan.warnings, None)) +} + +struct Scan<'py> { + warnings: Vec, + with_aliases: Vec>, + owners: Vec>, + /// Claimed aliases by their class under Python equality, so `1`, `1.0` and `True` collide as they do in a dict + aliases: HashMap>, + claims: AliasClaims, +} + +struct Failed { + warnings: Vec, + error: PyErr, +} + +fn scan<'py>(model_cost: &Bound<'py, PyDict>) -> Result, Failed> { + let py = model_cost.py(); + let classes = PyDict::new(py); + let mut scan = Scan { + warnings: Vec::new(), + with_aliases: Vec::new(), + owners: Vec::new(), + aliases: HashMap::new(), + claims: AliasClaims::new(), + }; + for (name, info) in model_cost.iter() { + if let Err(error) = scan_row(model_cost, &classes, &mut scan, name, &info) { + return Err(Failed { + warnings: scan.warnings, + error, + }); + } + } + Ok(scan) +} + +fn scan_row<'py>( + model_cost: &Bound<'py, PyDict>, + classes: &Bound<'py, PyDict>, + scan: &mut Scan<'py>, + name: Bound<'py, PyAny>, + info: &Bound<'py, PyAny>, +) -> PyResult<()> { + let py = model_cost.py(); + let aliases = info.call_method1(intern!(py, "get"), (intern!(py, "aliases"),))?; + if aliases.is_none() { + return Ok(()); + } + scan.with_aliases.push(name.clone()); + let Ok(aliases) = aliases.cast::() else { + scan.warnings.push(AliasWarning::NotAList { + model: name.unbind(), + type_name: aliases.get_type().name()?.to_string(), + }); + return Ok(()); + }; + let owner = scan.owners.len(); + scan.owners.push(name.clone()); + for alias in aliases.iter() { + let is_model = model_cost.contains(&alias)?; + let class = class_of(classes, &alias)?; + match scan.claims.claim(&class, &owner, is_model) { + Ok(()) => { + scan.aliases.insert(class, alias); + } + Err(Unclaimed::ShadowsModel) => scan.warnings.push(AliasWarning::ShadowsModel { + alias: alias.unbind(), + model: name.clone().unbind(), + }), + Err(Unclaimed::TakenBy(_)) => scan.warnings.push(AliasWarning::Taken { + alias: alias.unbind(), + model: name.clone().unbind(), + }), + } + } + Ok(()) +} + +fn class_of(classes: &Bound<'_, PyDict>, alias: &Bound<'_, PyAny>) -> PyResult { + if let Some(class) = classes.get_item(alias)? { + return class.extract(); + } + let class = classes.len(); + classes.set_item(alias, class)?; + Ok(class) +} + +fn apply<'py>( + model_cost: &Bound<'py, PyDict>, + entry_type: &Bound<'py, PyAny>, + scan: &Scan<'py>, +) -> PyResult<()> { + let py = model_cost.py(); + for name in &scan.with_aliases { + model_cost + .as_any() + .get_item(name)? + .call_method1(intern!(py, "pop"), (intern!(py, "aliases"), py.None()))?; + } + for (owner, aliases) in scan.claims.groups() { + let owner = &scan.owners[*owner]; + let names: Vec<_> = std::iter::once(owner.clone()) + .chain(aliases.into_iter().map(|class| scan.aliases[class].clone())) + .collect(); + let names = PyTuple::new(py, names)?; + let entry = entry_type.call1((model_cost.as_any().get_item(owner)?, &names))?; + for name in names.iter() { + model_cost.set_item(name, &entry)?; + } + } + Ok(()) +} diff --git a/litellm-rust/crates/python-bridge/src/model_registry/fetch.rs b/litellm-rust/crates/python-bridge/src/model_registry/fetch.rs new file mode 100644 index 00000000000..07995365430 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry/fetch.rs @@ -0,0 +1,322 @@ +use std::time::Duration; + +use litellm_http::ClientVariant; +use litellm_model_catalog_refresh::{ + Attempt, FetchRequest, Pause, Retrying, SourceClients, SourceFetch, attempt, refresh, refused, + resume, +}; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +/// One fetch of the cost map, or a whole refresh +#[pyclass(frozen, module = "litellm.rust_bridge._native")] +pub(crate) enum ModelCostMapFetch { + /// The document's bytes for Python to decode, with the `git_blob_id` of those bytes + Adopted { + body: Py, + revision: String, + etag: Option, + }, + NotModified {}, + Failed { + reason: String, + }, + /// Only from `attempt`, to resume with `run` + Retryable { + reason: String, + retry_after: Option, + }, +} + +fn draw_jitter(rng: &Bound<'_, PyAny>) -> PyResult { + rng.call_method1(intern!(rng.py(), "uniform"), (0.0, 1.0))? + .extract() +} + +/// One cost map refresh, built where it is created so the thread that runs it never reads Python settings or +/// imports. `sleep` is `None` to sleep in Rust, `on_retry(attempt, max_attempts, reason, wait_seconds)` runs before +/// each wait +#[pyclass( + frozen, + module = "litellm.rust_bridge._native", + name = "ModelCostMapRefresh" +)] +pub(crate) struct PyRefresh { + run: Run, + rng: Py, + on_retry: Py, + sleep: Option>, +} + +#[pymethods] +impl PyRefresh { + #[new] + #[pyo3(signature = (url, *, timeout, max_attempts, rng, on_retry, sleep = None))] + fn new( + py: Python<'_>, + url: &str, + timeout: f64, + max_attempts: i64, + rng: Py, + on_retry: Py, + sleep: Option>, + ) -> PyResult { + let http = crate::http::host_client(py, ClientVariant::NoRedirect)?; + Ok(Self::with_client( + http, + url, + timeout, + max_attempts, + rng, + on_retry, + sleep, + )) + } + + /// The refresh `import litellm` starts, for `run_detached`. Its client is built fresh and never pooled, so its + /// connections end with its runtime and a fork mid-fetch leaves no shared lock or socket in the child + #[staticmethod] + #[pyo3(signature = (url, *, timeout, max_attempts, rng, on_retry))] + fn at_import( + py: Python<'_>, + url: &str, + timeout: f64, + max_attempts: i64, + rng: Py, + on_retry: Py, + ) -> PyResult { + let http = crate::http::unpooled_host_client(py, ClientVariant::NoRedirect)?; + Ok(Self::with_client( + http, + url, + timeout, + max_attempts, + rng, + on_retry, + None, + )) + } + + /// Attempt 1 alone, which `get_model_cost_map` makes before deciding to retry in the background + fn attempt(&self, py: Python<'_>) -> PyResult { + let run = self.run.clone(); + let first = litellm_host_python::run_sync_value(py, async move { + Ok(match &run.source { + Ok(fetch) => attempt(fetch, &run.request).await, + Err(reason) => reason.clone(), + }) + })?; + Python::attach(|py| Ok(from_attempt(py, first))) + } + + /// Every attempt with the GIL released, resuming after `first` when it is a retryable attempt 1 + #[pyo3(signature = (first = None))] + fn run( + &self, + py: Python<'_>, + first: Option<&Bound<'_, ModelCostMapFetch>>, + ) -> PyResult { + let first = first.map(|first| to_attempt(py, first.get())); + let run = self.run.clone(); + let mut pause = self.pause(py, false); + let refreshed = litellm_host_python::run_sync_value(py, async move { + run.refresh(first, &mut pause).await + })?; + Python::attach(|py| Ok(from_attempt(py, refreshed))) + } + + /// `run` on a runtime that ends with the call, for the refresh `import litellm` starts. The process starts no + /// runtime thread, so it can still fork workers or reserve itself for forking afterwards + fn run_detached(&self, py: Python<'_>) -> PyResult { + let run = self.run.clone(); + let mut pause = self.pause(py, false); + let refreshed = litellm_host_python::run_detached_value(py, run.refresh(None, &mut pause))?; + Python::attach(|py| Ok(from_attempt(py, refreshed))) + } + + /// `run` as an awaitable, where `sleep` returns an awaitable + fn arun<'py>(&self, py: Python<'py>) -> PyResult> { + let run = self.run.clone(); + let mut pause = self.pause(py, true); + litellm_host_python::run_async_value(py, async move { + let refreshed = run.refresh(None, &mut pause).await?; + Ok(Python::attach(|py| from_attempt(py, refreshed))) + }) + } +} + +impl PyRefresh { + fn with_client( + http: litellm_http::Client, + url: &str, + timeout: f64, + max_attempts: i64, + rng: Py, + on_retry: Py, + sleep: Option>, + ) -> Self { + let env = |name: &str| std::env::var(name).ok(); + let source = SourceFetch::new(url, SourceClients { http, env: &env }) + .map_err(|error| refused(url, &error)); + Self { + run: Run { + source, + request: request(timeout), + max_attempts, + }, + rng, + on_retry, + sleep, + } + } + + fn pause(&self, py: Python<'_>, asynchronous: bool) -> PythonPause { + PythonPause::new( + self.rng.clone_ref(py), + self.on_retry.clone_ref(py), + self.sleep.as_ref().map(|sleep| sleep.clone_ref(py)), + asynchronous, + ) + } +} + +#[derive(Clone)] +struct Run { + source: Result, + request: FetchRequest, + max_attempts: i64, +} + +impl Run { + async fn refresh(self, first: Option, pause: &mut PythonPause) -> PyResult { + let fetch = match self.source { + Ok(fetch) => fetch, + Err(failed) => return Ok(failed), + }; + match first { + Some(first) => resume(&fetch, &self.request, 1, self.max_attempts, first, pause).await, + None => refresh(&fetch, &self.request, self.max_attempts, pause).await, + } + } +} + +fn request(timeout: f64) -> FetchRequest { + FetchRequest { + if_none_match: None, + timeout: Duration::try_from_secs_f64(timeout).unwrap_or(Duration::MAX), + } +} + +fn from_attempt(py: Python<'_>, attempt: Attempt) -> ModelCostMapFetch { + match attempt { + Attempt::Fetched { body, etag } => ModelCostMapFetch::Adopted { + revision: litellm_model_catalog::revision(&body), + body: PyBytes::new(py, &body).unbind(), + etag, + }, + Attempt::NotModified => ModelCostMapFetch::NotModified {}, + Attempt::Retryable { + reason, + retry_after, + } => ModelCostMapFetch::Retryable { + reason, + retry_after, + }, + Attempt::Failed { reason } => ModelCostMapFetch::Failed { reason }, + } +} + +fn to_attempt(py: Python<'_>, fetch: &ModelCostMapFetch) -> Attempt { + match fetch { + ModelCostMapFetch::Adopted { body, etag, .. } => Attempt::Fetched { + body: body.bind(py).as_bytes().to_vec(), + etag: etag.clone(), + }, + ModelCostMapFetch::NotModified {} => Attempt::NotModified, + ModelCostMapFetch::Failed { reason } => Attempt::Failed { + reason: reason.clone(), + }, + ModelCostMapFetch::Retryable { + reason, + retry_after, + } => Attempt::Retryable { + reason: reason.clone(), + retry_after: *retry_after, + }, + } +} + +/// Python's `rng`, `sleep` and retry report. A failure in any of them ends the refresh at the next sleep with that +/// exception, as main's loop raised it +struct PythonPause { + rng: Py, + on_retry: Py, + sleep: Option>, + asynchronous: bool, + failed: Option, +} + +impl PythonPause { + fn new( + rng: Py, + on_retry: Py, + sleep: Option>, + asynchronous: bool, + ) -> Self { + Self { + rng, + on_retry, + sleep, + asynchronous, + failed: None, + } + } +} + +impl Pause for PythonPause { + type Stop = PyErr; + + fn jitter(&mut self) -> f64 { + Python::attach(|py| draw_jitter(self.rng.bind(py))).unwrap_or_else(|error| { + self.failed.get_or_insert(error); + 0.0 + }) + } + + fn retrying(&mut self, retrying: &Retrying<'_>) { + let reported = Python::attach(|py| { + self.on_retry + .call1( + py, + ( + retrying.attempt, + retrying.max_attempts, + retrying.reason, + retrying.wait_seconds, + ), + ) + .map(drop) + }); + if let Err(error) = reported { + self.failed.get_or_insert(error); + } + } + + async fn sleep(&mut self, seconds: f64) -> PyResult<()> { + if let Some(error) = self.failed.take() { + return Err(error); + } + let Some(sleep) = &self.sleep else { + tokio::time::sleep(Duration::try_from_secs_f64(seconds).unwrap_or_default()).await; + return Ok(()); + }; + if !self.asynchronous { + return Python::attach(|py| sleep.call1(py, (seconds,)).map(drop)); + } + let slept = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(sleep.bind(py).call1((seconds,))?) + })?; + slept.await.map(drop) + } +} diff --git a/litellm-rust/crates/python-bridge/src/model_registry/keys.rs b/litellm-rust/crates/python-bridge/src/model_registry/keys.rs new file mode 100644 index 00000000000..03cb69556cb --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry/keys.rs @@ -0,0 +1,82 @@ +use std::borrow::Cow; +use std::sync::LazyLock; + +use litellm_model_catalog::{LITELLM_PROVIDER, Rows, Store, fold_model_name}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyString}; + +use super::live::current; + +/// The bound `litellm.model_cost`, read again on every query because a reload rebinds it +struct LiveModelCost; + +impl Rows for LiveModelCost { + type Row = (); + /// The dict object and its length, which miss a delete plus an add between two reads + type Stamp = Option<(usize, usize)>; + + fn fold(name: &str) -> String { + fold_model_name(name) + } + + fn stamp(&self) -> Self::Stamp { + Python::attach(|py| { + current(py) + .ok() + .map(|dict| (dict.as_ptr() as usize, dict.len())) + }) + } + + fn contains(&self, key: &str) -> bool { + Python::attach(|py| current(py).is_ok_and(|dict| dict.contains(key).unwrap_or(false))) + } + + fn row(&self, key: &str) -> Option> { + self.contains(key).then_some(Cow::Owned(())) + } + + /// A key that is not a `str` cannot be matched by name and is left out + fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) { + Python::attach(|py| { + let Ok(dict) = current(py) else { + return; + }; + for (key, row) in dict + .copy() + .map(|copy| copy.into_iter().collect::>()) + .unwrap_or_default() + { + let Some(key) = key + .cast::() + .ok() + .and_then(|key| key.to_cow().ok()) + else { + continue; + }; + let provider = provider(&row); + visit(&key, provider.as_deref()); + } + }); + } +} + +fn provider(row: &Bound<'_, PyAny>) -> Option { + let row = row.cast::().ok()?; + let provider = row.get_item(LITELLM_PROVIDER).ok()??; + provider.extract().ok() +} + +static STORE: LazyLock> = LazyLock::new(|| Store::new(LiveModelCost)); + +/// The `litellm.model_cost` key for `name`: `name` itself when present, else the key that lowercases to the same +/// name, the later key winning a tie as main's `_model_cost_lowercase_map` did +#[pyfunction] +pub(crate) fn find_model_cost_key(name: &str) -> Option { + STORE.find_key(name) +} + +/// For a write the `(id, len)` stamp cannot see, such as a key replaced by another of the same count +#[pyfunction] +pub(crate) fn invalidate_model_cost_keys() { + STORE.invalidate(); +} diff --git a/litellm-rust/crates/python-bridge/src/model_registry/live.rs b/litellm-rust/crates/python-bridge/src/model_registry/live.rs new file mode 100644 index 00000000000..c19b4f90b08 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry/live.rs @@ -0,0 +1,46 @@ +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::{PyDict, PyString, PyTuple}; + +const MODULE: &str = "litellm.rust_bridge.host.model_cost"; + +/// The bound `litellm.model_cost` +pub(super) fn current(py: Python<'_>) -> PyResult> { + static CURRENT: PyOnceLock> = PyOnceLock::new(); + CURRENT + .import(py, MODULE, "current")? + .call0()? + .cast_into::() + .map_err(PyErr::from) +} + +/// `omits_token_price` over the dict rows of `litellm.model_cost` under the `str` keys +#[pyfunction] +pub(crate) fn cost_map_omits_token_price(keys: &Bound<'_, PyTuple>) -> PyResult { + let model_cost = current(keys.py())?; + let rows = keys + .iter() + .filter(|key| key.is_instance_of::()) + .map(|key| model_cost.get_item(&key)) + .collect::>>()?; + let rows = rows + .into_iter() + .flatten() + .filter_map(|row| row.cast_into::().ok()); + Ok(litellm_model_catalog::omits_token_price(rows.map(|row| { + move |name: &str| row.contains(name).unwrap_or(false) + }))) +} + +/// The `litellm.model_cost` row under `key` when it is truthy, else `fallback`, the row a registration of `key` +/// merges onto +#[pyfunction] +pub(crate) fn live_row_or<'py>( + key: &Bound<'py, PyAny>, + fallback: Bound<'py, PyAny>, +) -> PyResult> { + let py = key.py(); + let row = current(py)?.call_method1(intern!(py, "get"), (key,))?; + Ok(if row.is_truthy()? { row } else { fallback }) +} diff --git a/litellm-rust/crates/python-bridge/src/model_registry/merge.rs b/litellm-rust/crates/python-bridge/src/model_registry/merge.rs new file mode 100644 index 00000000000..7b98d31cf13 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry/merge.rs @@ -0,0 +1,219 @@ +use litellm_model_catalog::{FieldValue, Fields, Number, Shape}; +use pyo3::exceptions::PyTypeError; +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::{PyDict, PyFloat, PyInt, PyList, PyString, PyTuple, PyType}; + +/// A Python value under the registration merge rule. Values the rule does not rewrite keep their identity +struct PyField(Py); + +impl Clone for PyField { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl FieldValue for PyField { + fn shape(&self) -> Shape { + Python::attach(|py| { + let value = self.0.bind(py); + if value.is_none() { + return Shape::None; + } + if let Ok(text) = value.cast::() { + return Shape::Str(text.to_string_lossy().into_owned()); + } + if value.is_instance_of::() { + return Shape::Map; + } + Shape::Other + }) + } + + fn none() -> Self { + Python::attach(|py| Self(py.None())) + } + + fn number(number: Number) -> Option { + Python::attach(|py| { + let value = match number { + Number::Float(value) => PyFloat::new(py, value).into_any(), + Number::Int(value) => py.get_type::().call1((value.to_string(),)).ok()?, + }; + Some(Self(value.unbind())) + }) + } + + fn merged_over(&self, base: Option<&Self>) -> Self { + Python::attach(|py| { + let merged = PyDict::new(py); + let base = base + .map(|base| base.0.bind(py)) + .filter(|base| base.is_instance_of::()); + let copied = base.map_or(Ok(()), |base| { + merged.update(base.cast::().expect("a dict").as_mapping()) + }); + copied + .and_then(|()| { + merged.update( + self.0 + .bind(py) + .cast::() + .expect("a dict") + .as_mapping(), + ) + }) + .expect("dict.update of two dicts"); + Self(merged.into_any().unbind()) + }) + } + + fn rate_tiers(&self) -> Option>> { + Python::attach(|py| { + let tiers = self.0.bind(py); + if !(tiers.is_instance_of::() || tiers.is_instance_of::()) { + return None; + } + tiers + .try_iter() + .ok()? + .map(|tier| tier.ok().and_then(|tier| rate_table(&tier))) + .collect() + }) + } + + fn rate_table(&self) -> Option> { + Python::attach(|py| rate_table(self.0.bind(py))) + } + + fn tier_list(tiers: Vec>) -> Self { + Python::attach(|py| { + let tiers: Vec> = + tiers.into_iter().map(|tier| to_dict(py, tier)).collect(); + Self( + PyList::new(py, tiers) + .expect("a list of dicts") + .into_any() + .unbind(), + ) + }) + } + + fn table(fields: Fields) -> Self { + Python::attach(|py| Self(to_dict(py, fields).into_any().unbind())) + } +} + +/// The value as pydantic's `Mapping[str, object]` reads it, `None` where it raises +fn rate_table(value: &Bound<'_, PyAny>) -> Option> { + if !value.is_instance(mapping_type(value.py()).ok()?).ok()? { + return None; + } + let items = value.call_method0("items").ok()?; + items + .try_iter() + .ok()? + .map(|item| { + let (name, rate): (Bound<'_, PyAny>, Bound<'_, PyAny>) = item.ok()?.extract().ok()?; + Some(( + name.cast::().ok()?.to_str().ok()?.to_owned(), + PyField(rate.unbind()), + )) + }) + .collect() +} + +fn mapping_type(py: Python<'_>) -> PyResult<&Bound<'_, PyType>> { + static MAPPING: PyOnceLock> = PyOnceLock::new(); + MAPPING.import(py, "collections.abc", "Mapping") +} + +fn to_dict(py: Python<'_>, fields: Fields) -> Bound<'_, PyDict> { + let dict = PyDict::new(py); + for (name, value) in fields { + dict.set_item(name, value.0).expect("str keys"); + } + dict +} + +/// A row's fields in dict order. A field name that is not a `str` is a `TypeError` +fn fields(row: &Bound<'_, PyAny>) -> PyResult> { + row.call_method0("items")? + .try_iter()? + .map(|item| { + let (name, value): (Bound<'_, PyAny>, Bound<'_, PyAny>) = item?.extract()?; + let Ok(text) = name.cast::() else { + let type_name = name.get_type().name()?; + return Err(PyTypeError::new_err(format!( + "field names must be str, got {type_name}" + ))); + }; + let name = text.to_str()?.to_owned(); + Ok((name, PyField(value.unbind()))) + }) + .collect() +} + +/// Main's `update_dictionary`, as a new dict whose keys are `existing`'s then `new`'s +#[pyfunction] +pub(crate) fn update_dictionary<'py>( + existing: &Bound<'py, PyAny>, + new: &Bound<'py, PyAny>, +) -> PyResult> { + let merged = litellm_model_catalog::update_dictionary(fields(existing)?, &fields(new)?); + Ok(to_dict(existing.py(), merged)) +} + +/// The number a stringified registration value spells, as main's `_convert_stringified_numbers` +#[pyfunction] +pub(crate) fn stringified_number(py: Python<'_>, text: &str) -> Option> { + litellm_model_catalog::stringified_number(text) + .and_then(PyField::number) + .map(|value| value.0.clone_ref(py)) +} + +#[pyfunction] +pub(crate) fn registration_row<'py>( + base: &Bound<'py, PyAny>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + let row = litellm_model_catalog::registration_row(&fields(base)?, &fields(registration)?); + Ok(to_dict(base.py(), row)) +} + +#[pyfunction] +pub(crate) fn without_overridden_price_variants<'py>( + entry: &Bound<'py, PyAny>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + let row = litellm_model_catalog::without_overridden_price_variants( + &fields(entry)?, + &fields(registration)?, + ); + Ok(to_dict(entry.py(), row)) +} + +/// `raw_entry` is the `litellm.model_cost` row the lookup resolved to, whose fields decide which prices are real +#[pyfunction] +pub(crate) fn without_synthesized_costs<'py>( + model_info: &Bound<'py, PyAny>, + raw_entry: &Bound<'py, PyAny>, + registration: &Bound<'py, PyAny>, +) -> PyResult> { + let row = litellm_model_catalog::without_synthesized_costs( + &fields(model_info)?, + |name| raw_entry.contains(name).unwrap_or(false), + &fields(registration)?, + ); + Ok(to_dict(model_info.py(), row)) +} + +#[pyfunction] +pub(crate) fn with_pricing_overrides<'py>( + model_info: &Bound<'py, PyAny>, + overrides: &Bound<'py, PyAny>, +) -> PyResult> { + let row = + litellm_model_catalog::with_pricing_overrides(&fields(model_info)?, &fields(overrides)?); + Ok(to_dict(model_info.py(), row)) +} diff --git a/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs b/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs new file mode 100644 index 00000000000..607428017b3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs @@ -0,0 +1,110 @@ +use std::sync::{LazyLock, PoisonError, RwLock}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use litellm_model_catalog::{Origin, Provenance}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; + +/// The installed cost map's provenance, as `get_model_cost_map_source_info` reports it +#[pyclass( + frozen, + module = "litellm.rust_bridge._native", + name = "ModelCostMapProvenance" +)] +pub(crate) struct PyProvenance(Provenance); + +#[pymethods] +impl PyProvenance { + #[getter] + fn source(&self) -> &'static str { + match self.0.origin { + Origin::Bundled { .. } => "local", + Origin::Remote { .. } => "remote", + } + } + + #[getter] + fn url(&self) -> Option<&str> { + match &self.0.origin { + Origin::Bundled { attempted } => attempted.as_deref(), + Origin::Remote { url } => Some(url), + } + } + + #[getter] + fn is_env_forced(&self) -> bool { + self.0.is_env_forced + } + + #[getter] + fn fallback_reason(&self) -> Option<&str> { + self.0.fallback_reason.as_deref() + } + + /// Seconds since the Unix epoch + #[getter] + fn loaded_at(&self) -> f64 { + self.0 + .loaded_at + .duration_since(UNIX_EPOCH) + .map_or(0.0, |since| since.as_secs_f64()) + } + + #[getter] + fn revision(&self) -> Option<&str> { + self.0.revision.as_deref() + } + + #[getter] + fn etag(&self) -> Option<&str> { + self.0.etag.as_deref() + } +} + +static INSTALLED: LazyLock>>> = LazyLock::new(RwLock::default); + +/// Replaces the installed provenance, stamped now. Only the step that installs a map calls this +#[pyfunction] +#[pyo3(signature = (*, source, url, is_env_forced, fallback_reason, revision, etag))] +pub(crate) fn install_model_cost_map_provenance( + py: Python<'_>, + source: &str, + url: Option, + is_env_forced: bool, + fallback_reason: Option, + revision: Option, + etag: Option, +) -> PyResult<()> { + let origin = match (source, url) { + ("local", attempted) => Origin::Bundled { attempted }, + ("remote", Some(url)) => Origin::Remote { url }, + _ => { + return Err(PyValueError::new_err(format!( + "a model cost map comes from 'local' or a 'remote' URL, got {source:?}" + ))); + } + }; + let provenance = Py::new( + py, + PyProvenance(Provenance { + origin, + is_env_forced, + fallback_reason, + loaded_at: SystemTime::now(), + revision, + etag, + }), + )?; + *INSTALLED.write().unwrap_or_else(PoisonError::into_inner) = Some(provenance); + Ok(()) +} + +/// `None` before the first install +#[pyfunction] +pub(crate) fn model_cost_map_provenance(py: Python<'_>) -> Option> { + INSTALLED + .read() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .map(|provenance| provenance.clone_ref(py)) +} diff --git a/litellm-rust/crates/python-compat/src/lib.rs b/litellm-rust/crates/python-compat/src/lib.rs index b23aff409f2..760faaf56e7 100644 --- a/litellm-rust/crates/python-compat/src/lib.rs +++ b/litellm-rust/crates/python-compat/src/lib.rs @@ -10,6 +10,7 @@ //! | [`json`] | `json.dumps(value)`, and `json.loads(json.dumps(value))` as a JSON value | //! | [`pickle`] | `pickle.loads(data)` and `pickle.dumps(value)` for plain data | //! | [`truthy`] | `bool(value)` | +//! | [`number`] | `float(text)` and `int(text)` | //! //! [`Value`] is the closed data model these formats share. Live Python objects //! (descriptors, `__bool__`, `__str__`, callbacks) are out of scope: those belong to the @@ -24,6 +25,7 @@ mod error; pub mod json; pub mod literal; +pub mod number; pub mod pickle; pub mod repr; pub mod truthy; diff --git a/litellm-rust/crates/python-compat/src/number.rs b/litellm-rust/crates/python-compat/src/number.rs new file mode 100644 index 00000000000..e9608f1dee1 --- /dev/null +++ b/litellm-rust/crates/python-compat/src/number.rs @@ -0,0 +1,99 @@ +use num_bigint::BigInt; + +/// `float(text)` for a `str`, `None` where CPython raises `ValueError`. Non-ASCII digits, which CPython reads as their +/// decimal value, are refused +pub fn float(text: &str) -> Option { + let digits = without_underscores(text.trim_matches(char::is_whitespace))?; + digits.parse().ok() +} + +/// `int(text)` for a `str` in base 10, `None` where CPython raises `ValueError`. Non-ASCII digits, which CPython reads +/// as their decimal value, are refused +pub fn int(text: &str) -> Option { + let digits = without_underscores(text.trim_matches(char::is_whitespace))?; + let unsigned = digits.strip_prefix(['+', '-']).unwrap_or(&digits); + if unsigned.is_empty() || !unsigned.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + digits.parse().ok() +} + +/// CPython accepts a single `_` only between two digits +fn without_underscores(text: &str) -> Option { + if !text.is_ascii() { + return None; + } + let bytes = text.as_bytes(); + let misplaced = bytes.iter().enumerate().any(|(at, byte)| { + *byte == b'_' + && !(at > 0 + && bytes[at - 1].is_ascii_digit() + && bytes.get(at + 1).is_some_and(u8::is_ascii_digit)) + }); + (!misplaced).then(|| text.replace('_', "")) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + /// Each expected value is what CPython 3.14.7 returns for `float(text)` and `int(text)` + #[rstest] + #[case::scientific("3e-07", Some(3e-07), None)] + #[case::upper_exponent("3E-07", Some(3e-07), None)] + #[case::padded(" 8192 ", Some(8192.0), Some(8192))] + #[case::unicode_whitespace("\u{a0}12\u{3000}", Some(12.0), Some(12))] + #[case::next_line("\u{85}1\u{85}", Some(1.0), Some(1))] + #[case::file_separator_is_not_whitespace("\u{1c}12\u{1f}", None, None)] + #[case::zero_width_space_is_not_whitespace("\u{200b}12", None, None)] + #[case::byte_order_mark_is_not_whitespace("\u{feff}1", None, None)] + #[case::underscores("1_000", Some(1000.0), Some(1000))] + #[case::underscore_in_exponent("1.5e3_0", Some(1.5e30), None)] + #[case::underscore_before_point("1_0.5", Some(10.5), None)] + #[case::double_underscore("1__0", None, None)] + #[case::leading_underscore("_1", None, None)] + #[case::trailing_underscore("1_", None, None)] + #[case::plus("+5", Some(5.0), Some(5))] + #[case::negative_zero("-0", Some(-0.0), Some(0))] + #[case::leading_zeros("0010", Some(10.0), Some(10))] + #[case::overflow("1e400", Some(f64::INFINITY), None)] + #[case::negative_overflow("-1e400", Some(f64::NEG_INFINITY), None)] + #[case::infinity("Infinity", Some(f64::INFINITY), None)] + #[case::signed_inf("+inf", Some(f64::INFINITY), None)] + #[case::trailing_point("1.", Some(1.0), None)] + #[case::point_before_exponent("5.e3", Some(5000.0), None)] + #[case::leading_point(".5", Some(0.5), None)] + #[case::bare_point(".", None, None)] + #[case::bare_exponent("e5", None, None)] + #[case::open_exponent("1e", None, None)] + #[case::empty("", None, None)] + #[case::blank(" ", None, None)] + #[case::hex("0x10", None, None)] + #[case::complex("1.5j", None, None)] + #[case::suffix("12abc", None, None)] + #[case::non_ascii_digits_are_a_known_limit("\u{661}\u{662}", None, None)] + fn matches_cpython( + #[case] text: &str, + #[case] float_value: Option, + #[case] int_value: Option, + ) { + assert_eq!(float(text).map(f64::to_bits), float_value.map(f64::to_bits)); + assert_eq!(int(text), int_value.map(BigInt::from)); + } + + #[rstest] + #[case::nan("nan")] + #[case::signed_nan("-NaN")] + fn nan_parses(#[case] text: &str) { + assert!(float(text).is_some_and(f64::is_nan)); + assert_eq!(int(text), None); + } + + #[rstest] + fn int_has_no_size_limit() { + let digits = "9".repeat(40); + assert_eq!(int(&digits).map(|value| value.to_string()), Some(digits)); + } +} diff --git a/litellm-rust/crates/storage-s3/AGENTS.md b/litellm-rust/crates/storage-s3/AGENTS.md new file mode 100644 index 00000000000..aec1e3a507d --- /dev/null +++ b/litellm-rust/crates/storage-s3/AGENTS.md @@ -0,0 +1,5 @@ +# S3 storage + +`litellm-storage-s3` builds an `aws_sdk_s3` client config whose requests go through `litellm_http::Client` and whose credentials come from `litellm-auth-aws`. It re-exports `aws_sdk_s3` so every S3 user shares one SDK version + +It knows no buckets, keys or object formats. `litellm-cache-s3` stores cache entries and `litellm-model-catalog-refresh` reads the cost map object on top of it diff --git a/litellm-rust/crates/storage-s3/Cargo.toml b/litellm-rust/crates/storage-s3/Cargo.toml new file mode 100644 index 00000000000..332bcb1c67a --- /dev/null +++ b/litellm-rust/crates/storage-s3/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-storage-s3" +version = "0.1.0" +description = "Shared S3 client config over LiteLLM's HTTP client and AWS credentials" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-http.workspace = true +aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } +aws-credential-types = "1.3.0" +aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] } +aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] } +http.workspace = true +reqwest.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/storage-s3/src/auth.rs similarity index 97% rename from litellm-rust/crates/cache-s3/src/auth.rs rename to litellm-rust/crates/storage-s3/src/auth.rs index 659390a5659..7fc172002ad 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/storage-s3/src/auth.rs @@ -40,7 +40,7 @@ impl ProvideCredentials for S3Credentials { secret_access_key, self.config.session_token.clone(), None, - "litellm-s3-cache", + "litellm-s3", )); } self.auth diff --git a/litellm-rust/crates/storage-s3/src/config.rs b/litellm-rust/crates/storage-s3/src/config.rs new file mode 100644 index 00000000000..00e90098959 --- /dev/null +++ b/litellm-rust/crates/storage-s3/src/config.rs @@ -0,0 +1,27 @@ +use aws_sdk_s3::config::{ + BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation, +}; +use litellm_auth_aws::AwsAuthConfig; + +use crate::{auth::S3Credentials, transport::ReqwestHttpClient}; + +/// An S3 client config over litellm's HTTP client and AWS credential chain. A custom `endpoint` is addressed path +/// style, as MinIO and other S3 compatible stores expect +pub fn s3_config( + region: String, + endpoint: Option, + http: litellm_http::Client, + auth: AwsAuthConfig, +) -> aws_sdk_s3::config::Builder { + let base = aws_sdk_s3::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(region)) + .http_client(ReqwestHttpClient(http)) + .credentials_provider(S3Credentials::new(auth)) + .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) + .response_checksum_validation(ResponseChecksumValidation::WhenRequired); + match endpoint { + Some(url) => base.endpoint_url(url).force_path_style(true), + None => base, + } +} diff --git a/litellm-rust/crates/storage-s3/src/lib.rs b/litellm-rust/crates/storage-s3/src/lib.rs new file mode 100644 index 00000000000..0acb64c2429 --- /dev/null +++ b/litellm-rust/crates/storage-s3/src/lib.rs @@ -0,0 +1,7 @@ +mod auth; +mod config; +mod transport; + +pub use auth::S3Credentials; +pub use aws_sdk_s3; +pub use config::s3_config; diff --git a/litellm-rust/crates/cache-s3/src/transport.rs b/litellm-rust/crates/storage-s3/src/transport.rs similarity index 100% rename from litellm-rust/crates/cache-s3/src/transport.rs rename to litellm-rust/crates/storage-s3/src/transport.rs diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/storage-s3/tests/auth.rs similarity index 97% rename from litellm-rust/crates/cache-s3/tests/auth.rs rename to litellm-rust/crates/storage-s3/tests/auth.rs index 59c5aa2e091..b770de3bcf6 100644 --- a/litellm-rust/crates/cache-s3/tests/auth.rs +++ b/litellm-rust/crates/storage-s3/tests/auth.rs @@ -1,6 +1,6 @@ use aws_credential_types::provider::ProvideCredentials; use litellm_auth_aws::AwsAuthConfig; -use litellm_cache_s3::S3Credentials; +use litellm_storage_s3::S3Credentials; use rstest::rstest; fn explicit(session_token: Option<&str>) -> AwsAuthConfig { diff --git a/litellm/__init__.py b/litellm/__init__.py index c1f4bb6a3df..a0b363e715f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -548,10 +548,14 @@ _key_management_system: Optional["KeyManagementSystem"] = None #### PII MASKING #### output_parse_pii: bool = False ############################################# -from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete +from litellm.litellm_core_utils.get_model_cost_map import ( + get_model_cost_map as get_model_cost_map, + load_model_cost_map_at_import, + mark_litellm_import_complete, +) from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired -model_cost = get_model_cost_map(url=model_cost_map_url) +model_cost = load_model_cost_map_at_import(url=model_cost_map_url) cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount apply_search_free_tier: bool = False cost_margin_config: Dict[ diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index e452befe295..2e11931afd7 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -9,15 +9,14 @@ export LITELLM_LOCAL_MODEL_COST_MAP=True ``` """ -import asyncio -import hashlib +import atexit +import functools import json import os import random import sys import threading -import time -import urllib.parse +import warnings from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from datetime import datetime, timezone @@ -25,11 +24,10 @@ 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 -import httpx from pydantic import ConfigDict, TypeAdapter -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict, assert_never from litellm import verbose_logger from litellm.constants import ( @@ -40,6 +38,19 @@ 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 ( + AliasWarning, + ModelCostMapFetch, + ModelCostMapRefresh, + ModelCountProblem, + count_model_entries, + expand_model_aliases, + git_blob_id, + install_model_cost_map_provenance, + model_cost_map_provenance, + model_count_problem, + take_fallback_generalization_rules, +) FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations" _CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) @@ -52,20 +63,6 @@ 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 _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) - - -def git_blob_id(body: bytes) -> str: - return hashlib.sha1(b"blob %d\0" % len(body) + body, usedforsecurity=False).hexdigest() - - class GetModelCostMap: """ Handles fetching, validating, and loading the model cost map. @@ -120,7 +117,7 @@ class GetModelCostMap: """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) + cls._backup_model_count = count_model_entries(backup) return cls._backup_model_count _get_backup_model_count = get_backup_model_count @@ -152,31 +149,31 @@ class GetModelCostMap: 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, - ) - 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.", - fetched_count, - backup_model_count, - max_shrink_ratio * 100, - ) - return False - - return True + problem: Final = model_count_problem(fetched_map, backup_model_count, min_model_count, max_shrink_ratio) + match problem: + case None: + return True + case ModelCountProblem.BelowMinimum(actual=actual, minimum=minimum): + 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.", + actual, + minimum, + ) + case ModelCountProblem.Shrunk(actual=actual, reference=reference): + 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.", + actual, + reference, + max_shrink_ratio * 100, + ) + case _: + assert_never(problem) + return False @classmethod def validate_model_cost_map( @@ -212,14 +209,39 @@ class GetModelCostMap: 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() +class _PendingImportRefresh: + """Single slot for the URL ``import litellm`` refreshes from once the import has finished""" + + url: str | None = None + + +class _BackgroundAdoption: + """A background fetch adopts under ``lock`` unless the interpreter is exiting. A thread that is still importing when + finalization hangs it holds the import lock and deadlocks the exit, so exit waits for an adoption in progress and + stops the next one""" + + lock: Final = threading.Lock() + exiting: Final = threading.Event() + + +def _stop_background_adoption() -> None: + with _BackgroundAdoption.lock: + _BackgroundAdoption.exiting.set() + + +atexit.register(_stop_background_adoption) + + def mark_litellm_import_complete() -> None: _litellm_import_complete.set() + url: Final = _PendingImportRefresh.url + _PendingImportRefresh.url = None + if url is not None: + _start_import_refresh(url) @dataclass(frozen=True, slots=True) @@ -237,217 +259,108 @@ class ModelCostMapReloadUnavailable: ModelCostMapReloadResult = ModelCostMapReloaded | ModelCostMapReloadUnavailable -@dataclass(frozen=True, slots=True) -class _FetchAttemptRetryable: - reason: str - retry_after_seconds: float | None +def _warn_ignored_client(client: object) -> None: + if client is not None: + warnings.warn( + "client= is ignored: the model cost map is fetched in Rust with litellm's HTTP settings", + DeprecationWarning, + stacklevel=3, + ) -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 _warn_retrying(attempt: int, max_attempts: int, reason: str, wait_seconds: float) -> None: + verbose_logger.warning( + "LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs", + attempt, + max_attempts, + reason, + wait_seconds, ) -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: +def _decoded( + fetched: ModelCostMapFetch.Adopted + | ModelCostMapFetch.NotModified + | ModelCostMapFetch.Failed + | ModelCostMapFetch.Retryable, + url: str, +) -> ModelCostMapReloadResult: + match fetched: + case ModelCostMapFetch.Adopted(body=body, revision=revision, etag=etag): + pass + case ModelCostMapFetch.Failed(reason=reason) | ModelCostMapFetch.Retryable(reason=reason): + return ModelCostMapReloadUnavailable(reason=reason) + case ModelCostMapFetch.NotModified(): + return ModelCostMapReloadUnavailable(reason=f"{url} was not modified") + case _: + assert_never(fetched) try: - parsed: Final = _DECODED_JSON.validate_python(json.loads(content)) + parsed: Final = _DECODED_JSON.validate_python(json.loads(body)) 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) + return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=revision, 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 _uses_bundled_map_only() -> bool: + # Note: 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 _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, +def _install_bundled( + attempted_url: str | None, fallback_reason: str | None, is_env_forced: bool = False +) -> dict[str, object]: + loaded: Final = _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()) + install_model_cost_map_provenance( + source="local", + url=attempted_url, + is_env_forced=is_env_forced, + fallback_reason=fallback_reason, + revision=loaded.revision, + etag=None, ) - return wait_seconds + return loaded.model_cost_map -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 _install_remote(fetched: ModelCostMapReloaded, url: str) -> ModelCostMapReloaded: + finalized: Final = _finalize_loaded_model_cost_map(fetched) + install_model_cost_map_provenance( + source="remote", + url=url, + is_env_forced=False, + fallback_reason=None, + revision=finalized.revision, + etag=finalized.etag, + ) + return finalized -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") +def _keep_bundled_after(url: str, fallback_reason: str) -> None: + """The bundled map served at import stays, now standing in for a refresh of ``url`` that failed. A map another + load installed meanwhile keeps its own provenance""" + current: Final = model_cost_map_provenance() + if current is None or current.source != "local" or current.url != url: + return + install_model_cost_map_provenance( + source="local", + url=url, + is_env_forced=False, + fallback_reason=fallback_reason, + revision=current.revision, + etag=None, + ) 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, + sleep: Callable[[float], Awaitable[None]] | None = None, rng: random.Random | None = None, - client: "_AsyncGetClient | None" = None, + client: object = None, ) -> ModelCostMapReloadResult: """ Re-fetch the model cost map for a runtime reload, retrying transient HTTP @@ -455,24 +368,26 @@ async def refetch_model_cost_map( 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. + map they already have. ``sleep`` defaults to waiting in Rust, and ``client`` + is ignored. """ - 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()) + _warn_ignored_client(client) + if _uses_bundled_map_only(): + loaded: Final = _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()) + install_model_cost_map_provenance( + source="local", url=None, is_env_forced=True, fallback_reason=None, revision=loaded.revision, etag=None + ) + return loaded - result: Final = await _fetch_remote_model_cost_map_with_retry( - url=url, + fetched: Final = await ModelCostMapRefresh( + 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(), - ) + on_retry=_warn_retrying, + sleep=sleep, + ).arun() + result: Final = _decoded(fetched, url) if isinstance(result, ModelCostMapReloadUnavailable): verbose_logger.warning( "LiteLLM: model cost map reload failed: %s. Keeping the currently loaded map.", @@ -484,28 +399,7 @@ async def refetch_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.""" - - source: str = "local" # "local" or "remote" - 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 - - -# Module-level singleton tracking the source of the current cost map -_cost_map_source_info: Final = ModelCostMapSourceInfo() + return _install_remote(result, url) class CostMapProvenance(TypedDict): @@ -522,9 +416,10 @@ class CostMapSourceInfo(CostMapProvenance): def get_model_cost_map_provenance() -> CostMapProvenance: + provenance: Final = model_cost_map_provenance() return { - "source_revision": _cost_map_source_info.source_revision, - "etag": _cost_map_source_info.etag, + "source_revision": provenance.revision if provenance is not None else None, + "etag": provenance.etag if provenance is not None else None, } @@ -537,25 +432,36 @@ def get_model_cost_map_source_info() -> CostMapSourceInfo: - url: the remote URL attempted (or None for local-only) - is_env_forced: True if LITELLM_LOCAL_MODEL_COST_MAP=True forced local usage - fallback_reason: human-readable reason if remote failed and local was used - - loaded_at: ISO 8601 time this process last loaded the map + - loaded_at: ISO 8601 time this process last installed the map - 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 + provenance: Final = model_cost_map_provenance() + if provenance is None: + return { + "source": "local", + "url": None, + "is_env_forced": False, + "fallback_reason": None, + "loaded_at": None, + "source_revision": None, + "etag": None, + } 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": provenance.source, + "url": provenance.url, + "is_env_forced": provenance.is_env_forced, + "fallback_reason": provenance.fallback_reason, + "loaded_at": datetime.fromtimestamp(provenance.loaded_at, timezone.utc).isoformat(), + "source_revision": provenance.revision, + "etag": provenance.etag, } 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 + """When this process last installed its cost map""" + provenance: Final = model_cost_map_provenance() + return datetime.fromtimestamp(provenance.loaded_at, timezone.utc) if provenance is not None else None class AliasedCostEntry(dict[str, object]): @@ -569,13 +475,28 @@ class AliasedCostEntry(dict[str, object]): self.registered_fields: frozenset[str] = registered_fields -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 _warn_skipped_alias(warning: AliasWarning.NotAList | AliasWarning.ShadowsModel | AliasWarning.Taken) -> None: + match warning: + case AliasWarning.NotAList(model=model, type_name=type_name): + verbose_logger.warning( + "LiteLLM model alias field for '%s' is not a list (got %s) — skipping.", + model, + type_name, + ) + case AliasWarning.ShadowsModel(alias=alias, model=model): + verbose_logger.warning( + "LiteLLM model alias conflict: alias '%s' (from '%s') already exists as a canonical entry — skipping.", + alias, + model, + ) + case AliasWarning.Taken(alias=alias, model=model): + verbose_logger.warning( + "LiteLLM model alias conflict: alias '%s' (from '%s') was already claimed by another entry — skipping.", + alias, + model, + ) + case _: + assert_never(warning) def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, dict[str, object]]: @@ -589,48 +510,11 @@ def _expand_model_aliases(model_cost: dict[str, dict[str, object]]) -> dict[str, If an alias collides with an existing canonical entry the alias is skipped and a warning is logged. """ - 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: - 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: - 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 - - # 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 owner, names in _alias_groups(alias_owners).items(): - model_cost.update(dict.fromkeys(names, AliasedCostEntry(model_cost[owner], names))) + warnings, error = expand_model_aliases(model_cost, AliasedCostEntry) + for warning in warnings: + _warn_skipped_alias(warning) + if error is not None: + raise error return model_cost @@ -640,56 +524,43 @@ def _finalize_model_cost_map(model_cost: dict) -> dict: 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) + rules: Final = take_fallback_generalization_rules(model_cost) + set_fallback_generalizations(rules if isinstance(rules, list) else None) return _expand_model_aliases(model_cost) 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) +MODEL_COST_ADOPTION_LOCK: Final = threading.RLock() +"""Held while a map is adopted and while a router registers its deployments and joins the replay, so an adoption in +another thread never lands between the two and drops that router's rows""" + + def adopt_model_cost_map( new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract ) -> int: 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 - 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() + with MODEL_COST_ADOPTION_LOCK: + litellm.model_cost = new_model_cost_map + utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation + 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 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, + refresh: ModelCostMapRefresh, url: str, max_attempts: int, first_outcome: ModelCostMapFetch.Retryable ) -> 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, - ) + fetched: Final = refresh.run(first=first_outcome) + result: Final = _decoded(fetched, url) if isinstance(result, ModelCostMapReloadUnavailable): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup", @@ -698,31 +569,123 @@ 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, - ) - 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) + with _BackgroundAdoption.lock: + if _BackgroundAdoption.exiting.is_set(): + return + 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, + ) + return + adopt_model_cost_map(_install_remote(result, url).model_cost_map) 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) +_import_refresh_adopted: Final = threading.Event() +_import_refresh_idle: Final = threading.Event() +_import_refresh_idle.set() +_FORK_WAIT_SECONDS: Final = 5.0 + + +def _refresh_in_background(refresh: ModelCostMapRefresh, url: str) -> None: + try: + fetched: Final = refresh.run_detached() + with _BackgroundAdoption.lock: + if _BackgroundAdoption.exiting.is_set(): + return + _adopt_import_refresh(_decoded(fetched, url), url) + except Exception as e: # noqa: BLE001 # a failed background refresh must not kill the thread; the bundled map stays + verbose_logger.warning("LiteLLM: Background model cost map refresh failed: %s", e) + finally: + _import_refresh_idle.set() + + +def _adopt_import_refresh(result: ModelCostMapReloadResult, url: str) -> None: + if isinstance(result, ModelCostMapReloadUnavailable): + verbose_logger.warning( + "LiteLLM: Failed to fetch remote model cost map from %s: %s. Keeping the bundled map.", + url, + result.reason, + ) + _keep_bundled_after(url, f"Remote fetch failed: {result.reason}") + return + 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. Keeping the bundled map. url=%s", + url, + ) + _keep_bundled_after(url, "Remote data failed integrity validation") + return + adopt_model_cost_map(_install_remote(result, url).model_cost_map) + _import_refresh_adopted.set() + + +def _start_import_refresh(url: str) -> None: + refresh: Final = ModelCostMapRefresh.at_import( + url, + timeout=5, + max_attempts=MODEL_COST_MAP_FETCH_MAX_ATTEMPTS, + rng=random.Random(), + on_retry=_warn_retrying, + ) + _import_refresh_idle.clear() + threading.Thread( + target=_refresh_in_background, + args=(refresh, url), + name="litellm-model-cost-map-refresh", + daemon=True, + ).start() + + +def _let_the_import_refresh_finish_before_fork() -> None: + """A fork while the refresh thread is inside the system resolver leaves the child's resolver broken (macOS + segfaults on the child's next lookup), so a fork waits up to one attempt's timeout for the refresh to finish""" + _import_refresh_idle.wait(timeout=_FORK_WAIT_SECONDS) + + +def _restart_import_refresh_in_child(url: str) -> None: + """A worker forked before its parent adopted the remote map, as gunicorn and granian fork theirs, fetches it itself""" + if _litellm_import_complete.is_set() and not _import_refresh_adopted.is_set(): + _start_import_refresh(url) + + +def load_model_cost_map_at_import(url: str) -> dict: + """The bundled map, served at once while ``url`` is fetched in the background. + + The fetch starts once ``import litellm`` finishes, and a map that passes the integrity + check is adopted. It runs on a runtime that ends with it, so the process can still fork. + ``LITELLM_LOCAL_MODEL_COST_MAP``, offline mode and the ``lite`` / ``litellm-proxy`` + CLI processes never fetch. + """ + if _uses_bundled_map_only() or _is_cli_process(): + return _install_bundled(attempted_url=None, fallback_reason=None, is_env_forced=True) + bundled: Final = _install_bundled(attempted_url=url, fallback_reason=None) + os.register_at_fork( + before=_let_the_import_refresh_finish_before_fork, + after_in_child=functools.partial(_restart_import_refresh_in_child, url), + ) + if _litellm_import_complete.is_set(): + _start_import_refresh(url) + else: + _PendingImportRefresh.url = url + return bundled + + 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, + sleep: Callable[[float], None] | None = None, rng: random.Random | None = None, - client: "_SyncGetClient | None" = None, + client: object = None, ) -> dict: """ Public entry point — returns the model cost map dict. @@ -734,65 +697,46 @@ def get_model_cost_map( 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. + ``sleep`` defaults to waiting in Rust, and ``client`` is ignored. """ - _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_client(client) + if _uses_bundled_map_only() or _is_cli_process(): + return _install_bundled(attempted_url=None, fallback_reason=None, is_env_forced=True) - _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: + refresh: Final = ModelCostMapRefresh( + url, + timeout=timeout, + max_attempts=max_attempts, + rng=rng if rng is not None else random.Random(), + on_retry=_warn_retrying, + sleep=sleep, + ) + first: Final = refresh.attempt() + if isinstance(first, ModelCostMapFetch.Retryable) 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, - }, + kwargs={"refresh": refresh, "url": url, "max_attempts": max_attempts, "first_outcome": first}, name="litellm-model-cost-map-retry", daemon=True, ).start() - if not isinstance(outcome, ModelCostMapReloaded): + outcome: Final = _decoded(first, url) + if isinstance(outcome, ModelCostMapReloadUnavailable): verbose_logger.warning( "LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.", url, outcome.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 + return _install_bundled(attempted_url=url, fallback_reason=f"Remote fetch failed: {outcome.reason}") # Validate using cached count (cheap int comparison, no file I/O) if not GetModelCostMap.validate_model_cost_map( - fetched_map=content, + fetched_map=outcome.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, ) - _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 + return _install_bundled(attempted_url=url, fallback_reason="Remote data failed integrity validation") - _cost_map_source_info.source = "remote" - _cost_map_source_info.fallback_reason = None - return _finalize_loaded_model_cost_map(outcome).model_cost_map + return _install_remote(outcome, url).model_cost_map diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index e28ebf439b5..14937fcfcf9 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -144,7 +144,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( is_registered_custom_provider, resolve_model_and_provider_for_metadata, ) -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.get_model_cost_map import MODEL_COST_ADOPTION_LOCK, GetModelCostMap from litellm.litellm_core_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, ) @@ -1851,19 +1851,20 @@ class PythonCore(DeprecatedRouterMembers): self.model_group_affinity_config = model_group_affinity_config warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) - if model_list is not None: - # set_model_list will build indices automatically - self.set_model_list(model_list) - self.healthy_deployments: list = self._core_model_list - for m in model_list: - if "model" in m["litellm_params"]: - self._core_deployment_latency_map[m["litellm_params"]["model"]] = 0 - else: - self._core_model_list: list = [] # initialize an empty list - to allow _add_deployment and delete_deployment to work - # Track this router so a price data reload can rebuild its deployments' - # cost-map entries from the list it is serving at that moment, whether - # those deployments came from model_list or were added later from the DB. - _live_routers.add(self) + with MODEL_COST_ADOPTION_LOCK: + if model_list is not None: + # set_model_list will build indices automatically + self.set_model_list(model_list) + self.healthy_deployments: list = self._core_model_list + for m in model_list: + if "model" in m["litellm_params"]: + self._core_deployment_latency_map[m["litellm_params"]["model"]] = 0 + else: + self._core_model_list: list = [] # initialize an empty list - to allow _add_deployment and delete_deployment to work + # Track this router so a price data reload can rebuild its deployments' + # cost-map entries from the list it is serving at that moment, whether + # those deployments came from model_list or were added later from the DB. + _live_routers.add(self) if allowed_fails is not None: self._core_allowed_fails = allowed_fails diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 954815f6a4f..1a5075fd543 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,9 +1,11 @@ +import random from asyncio import Future -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import Never, final import httpx from pydantic import JsonValue +from typing_extensions import disjoint_base from litellm.rust_bridge.execution import Execution from litellm.rust_bridge.public_call import NativeCall @@ -15,6 +17,179 @@ class ProcessReservedForForking(RuntimeError): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... +def count_model_entries(model_cost: Mapping[str, object]) -> int: ... +def git_blob_id(body: bytes) -> str: ... + +@disjoint_base +class ModelCountProblem: + @final + class BelowMinimum(ModelCountProblem): + __match_args__ = ("actual", "minimum") + def __new__(cls, actual: int, minimum: int) -> ModelCountProblem.BelowMinimum: ... + @property + def actual(self) -> int: ... + @property + def minimum(self) -> int: ... + + @final + class Shrunk(ModelCountProblem): + __match_args__ = ("actual", "reference") + def __new__(cls, actual: int, reference: int) -> ModelCountProblem.Shrunk: ... + @property + def actual(self) -> int: ... + @property + def reference(self) -> int: ... + +@disjoint_base +class ModelCostMapFetch: + @final + class Adopted(ModelCostMapFetch): + __match_args__ = ("body", "revision", "etag") + def __new__(cls, body: bytes, revision: str, etag: str | None) -> ModelCostMapFetch.Adopted: ... + @property + def body(self) -> bytes: ... + @property + def revision(self) -> str: ... + @property + def etag(self) -> str | None: ... + + @final + class NotModified(ModelCostMapFetch): + __match_args__ = () + def __new__(cls) -> ModelCostMapFetch.NotModified: ... + + @final + class Failed(ModelCostMapFetch): + __match_args__ = ("reason",) + def __new__(cls, reason: str) -> ModelCostMapFetch.Failed: ... + @property + def reason(self) -> str: ... + + @final + class Retryable(ModelCostMapFetch): + __match_args__ = ("reason", "retry_after") + def __new__(cls, reason: str, retry_after: float | None) -> ModelCostMapFetch.Retryable: ... + @property + def reason(self) -> str: ... + @property + def retry_after(self) -> float | None: ... + +_ModelCostMapRetryReport = Callable[[int, int, str, float], None] + +@final +class ModelCostMapRefresh: + def __new__( + cls, + url: str, + *, + timeout: float, + max_attempts: int, + rng: random.Random, + on_retry: _ModelCostMapRetryReport, + sleep: Callable[[float], object] | None = None, + ) -> ModelCostMapRefresh: ... + @staticmethod + def at_import( + url: str, *, timeout: float, max_attempts: int, rng: random.Random, on_retry: _ModelCostMapRetryReport + ) -> ModelCostMapRefresh: ... + def attempt( + self, + ) -> ( + ModelCostMapFetch.Adopted + | ModelCostMapFetch.NotModified + | ModelCostMapFetch.Failed + | ModelCostMapFetch.Retryable + ): ... + def run( + self, first: ModelCostMapFetch.Retryable | None = None + ) -> ModelCostMapFetch.Adopted | ModelCostMapFetch.NotModified | ModelCostMapFetch.Failed: ... + def run_detached(self) -> ModelCostMapFetch.Adopted | ModelCostMapFetch.NotModified | ModelCostMapFetch.Failed: ... + def arun( + self, + ) -> Future[ModelCostMapFetch.Adopted | ModelCostMapFetch.NotModified | ModelCostMapFetch.Failed]: ... + +@final +class ModelCostMapProvenance: + @property + def source(self) -> str: ... + @property + def url(self) -> str | None: ... + @property + def is_env_forced(self) -> bool: ... + @property + def fallback_reason(self) -> str | None: ... + @property + def loaded_at(self) -> float: ... + @property + def revision(self) -> str | None: ... + @property + def etag(self) -> str | None: ... + +def install_model_cost_map_provenance( + *, + source: str, + url: str | None, + is_env_forced: bool, + fallback_reason: str | None, + revision: str | None, + etag: str | None, +) -> None: ... +def model_cost_map_provenance() -> ModelCostMapProvenance | None: ... +def find_model_cost_key(name: str) -> str | None: ... +def cost_map_omits_token_price(keys: tuple[object, ...]) -> bool: ... +def live_row_or(key: str, fallback: Mapping[str, object]) -> Mapping[str, object]: ... +def invalidate_model_cost_keys() -> None: ... +def take_fallback_generalization_rules(model_cost: dict[str, object]) -> object: ... +def stringified_number(text: str) -> int | float | None: ... +def update_dictionary(existing: Mapping[str, object], new: Mapping[str, object]) -> dict[str, object]: ... +def registration_row(base: Mapping[str, object], registration: Mapping[str, object]) -> dict[str, object]: ... +def without_overridden_price_variants( + entry: Mapping[str, object], registration: Mapping[str, object] +) -> dict[str, object]: ... +def without_synthesized_costs( + model_info: Mapping[str, object], raw_entry: Mapping[str, object], registration: Mapping[str, object] +) -> dict[str, object]: ... +def with_pricing_overrides(model_info: Mapping[str, object], overrides: Mapping[str, object]) -> dict[str, object]: ... +def model_count_problem( + fetched_map: Mapping[str, object], backup_model_count: int, min_model_count: int, max_shrink_ratio: float +) -> ModelCountProblem.BelowMinimum | ModelCountProblem.Shrunk | None: ... + +@disjoint_base +class AliasWarning: + @final + class NotAList(AliasWarning): + __match_args__ = ("model", "type_name") + def __new__(cls, model: object, type_name: str) -> AliasWarning.NotAList: ... + @property + def model(self) -> object: ... + @property + def type_name(self) -> str: ... + + @final + class ShadowsModel(AliasWarning): + __match_args__ = ("alias", "model") + def __new__(cls, alias: object, model: object) -> AliasWarning.ShadowsModel: ... + @property + def alias(self) -> object: ... + @property + def model(self) -> object: ... + + @final + class Taken(AliasWarning): + __match_args__ = ("alias", "model") + def __new__(cls, alias: object, model: object) -> AliasWarning.Taken: ... + @property + def alias(self) -> object: ... + @property + def model(self) -> object: ... + +def expand_model_aliases( + model_cost: dict[str, dict[str, object]], + entry_type: Callable[[Mapping[str, object], tuple[str, ...]], dict[str, object]], +) -> tuple[ + list[AliasWarning.NotAList | AliasWarning.ShadowsModel | AliasWarning.Taken], + BaseException | None, +]: ... @final class NativeClickHouseSpendConfig: @@ -210,8 +385,13 @@ def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... __all__ = [ + "AliasWarning", "ForkedAfterNativeRuntimeStarted", "HuggingFaceEncoding", + "ModelCostMapFetch", + "ModelCostMapProvenance", + "ModelCostMapRefresh", + "ModelCountProblem", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", @@ -230,17 +410,34 @@ __all__ = [ "atranscription", "bundled_model_cost_map", "completion", + "cost_map_omits_token_price", + "count_model_entries", "embedding", + "expand_model_aliases", + "find_model_cost_key", "gil_stats", + "git_blob_id", + "install_model_cost_map_provenance", + "invalidate_model_cost_keys", + "live_row_or", "messages", + "model_cost_map_provenance", + "model_count_problem", "ocr", "ocr_health_check_document", "ocr_passthrough_response", "process_state_started", + "registration_row", "reserve_process_for_forking", "responses", + "stringified_number", + "take_fallback_generalization_rules", "trace_encode_error", "transcription", + "update_dictionary", + "with_pricing_overrides", + "without_overridden_price_variants", + "without_synthesized_costs", ] @final diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 9b751a1ada2..e3e573e540d 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -2,7 +2,7 @@ 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) +- 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), `model_cost` (the bound `litellm.model_cost`) - 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 diff --git a/litellm/rust_bridge/host/model_cost.py b/litellm/rust_bridge/host/model_cost.py new file mode 100644 index 00000000000..98df188c5f3 --- /dev/null +++ b/litellm/rust_bridge/host/model_cost.py @@ -0,0 +1,7 @@ +from __future__ import annotations + +import litellm + + +def current() -> dict[str, object]: + return litellm.model_cost diff --git a/litellm/utils.py b/litellm/utils.py index ec7bb041fb5..50d0acfcef2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -46,7 +46,7 @@ from httpx import Proxy from httpx._utils import get_environment_proxies from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # OpenAI parser module is private from openai.types.chat.completion_create_params import ResponseFormat -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, with_config +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, with_config import litellm import litellm.litellm_core_utils @@ -108,6 +108,7 @@ from litellm.litellm_core_utils.tool_fields import ( resolve_deployment_tool_fields, tools_without_hosted_mcp, ) +from litellm.rust_bridge import _native as _model_cost_rules from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge.catalog import decision from litellm.rust_bridge.configuration import Decision @@ -3818,42 +3819,18 @@ def supports_embedding_image_input(model: str, custom_llm_provider: str | None = ####### HELPER FUNCTIONS ################ def update_dictionary(existing_dict: dict[str, object], new_dict: Mapping[str, object]) -> dict[str, object]: - for k, v in new_dict.items(): - if v is not None: - # Convert stringified numbers to appropriate numeric types - if isinstance(v, str): - existing_dict[k] = _convert_stringified_numbers(v) - elif isinstance(v, dict): - nested_dict = cast(dict[str, object], v) - existing_nested_value = existing_dict.get(k) - if isinstance(existing_nested_value, dict): - existing_nested_dict = cast(dict[str, object], existing_nested_value) - existing_dict[k] = {**existing_nested_dict, **nested_dict} - else: - existing_dict[k] = dict(nested_dict) - else: - existing_dict[k] = v - + existing_dict.update(_model_cost_rules.update_dictionary(existing_dict, new_dict)) return existing_dict _update_dictionary = update_dictionary -def _convert_stringified_numbers(value): - """Convert stringified numbers (including scientific notation) to appropriate numeric types.""" - if isinstance(value, str): - try: - # Try to convert to float first to handle scientific notation like "3e-07" - if "e" in value.lower() or "." in value: - return float(value) - # Try to convert to int for whole numbers like "8192" - else: - return int(value) - except (ValueError, TypeError): - # If conversion fails, return the original string - return value - return value +def _convert_stringified_numbers(value: object) -> object: + if not isinstance(value, str): + return value + number: Final = _model_cost_rules.stringified_number(value) + return value if number is None else number _BEDROCK_REGION_PREFIXES: Final = ( @@ -3867,8 +3844,6 @@ _BEDROCK_REGION_PREFIXES: Final = ( "ap-northeast-1.", ) -_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token") - def _builtin_cost_map() -> dict[str, Mapping[str, object]]: return cast(dict[str, Mapping[str, object]], litellm.model_cost) @@ -3893,79 +3868,12 @@ def _replace_cost_map_entry( def _without_synthesized_costs( model_info: Mapping[str, object], key: str, value: Mapping[str, object] ) -> Mapping[str, object]: - """``_get_model_info_helper`` synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` = 0 when the raw - entry has neither (#30198). Persisting those zeros would turn a name-priced entry into a free one and make - ``_is_cost_explicitly_configured`` disable budget enforcement on the next re-registration.""" builtin_cost_map: Final = _builtin_cost_map() - raw_fields: Final = frozenset(builtin_cost_map.get(str(model_info["key"]), builtin_cost_map.get(key, {}))) - return MappingProxyType( - { - name: field_value - for name, field_value in model_info.items() - if name not in _TOKEN_COST_FIELDS or name in raw_fields or name in value - } - ) + raw_entry: Final = builtin_cost_map.get(str(model_info["key"]), builtin_cost_map.get(key, {})) + return _model_cost_rules.without_synthesized_costs(model_info, raw_entry, value) -_DIRECTION_RATE_FIELDS: Final = MappingProxyType( - { - "input_cost_per_token": ("input_cost_per_token",), - "output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"), - } -) _RATE_TABLE: Final = TypeAdapter(Mapping[str, object]) -_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) - - -def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool: - is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields)) - return name in overridden_fields or (is_variant and "_cache_" not in name) - - -def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]: - return MappingProxyType( - {name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)} - ) - - -def _rate_table_with_flat_prices( - name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object] -) -> object: - try: - if name == "tiered_pricing": - return [ - {**_without_overridden_rates(tier, overridden_fields), **flat_prices} - for tier in _RATE_TIERS.validate_python(table) - ] - if name == "off_peak_pricing": - return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields)) - except ValidationError: - return table - return table - - -def _without_overridden_price_variants( - entry: Mapping[str, object], value: Mapping[str, object] -) -> Mapping[str, object]: - """A caller's flat ``input_cost_per_token`` / ``output_cost_per_token`` is that direction's price above every - token threshold, for every service tier and batch, in every ``tiered_pricing`` tier and off-peak window, so each - inherited rate that would stand in for it there, the reasoning-token rate included for output, is dropped, and - each tier carries the flat price instead.""" - flat_prices: Final = MappingProxyType( - {cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None} - ) - overridden_fields: Final = tuple( - itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices) - ) - if not overridden_fields: - return entry - return MappingProxyType( - { - name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices) - for name, field_value in entry.items() - if not _is_overridden_rate(name, overridden_fields) - } - ) @with_config(ConfigDict(extra="allow")) @@ -3982,9 +3890,7 @@ def model_info_with_pricing_overrides(model_info: ModelInfo, registration: Mappi overrides: Final = { name: value for name, value in registration.items() if value is not None and is_server_derived_pricing_key(name) } - return _PRICING_MODEL_INFO.validate_python( - update_dictionary(dict(_without_overridden_price_variants(model_info, overrides)), overrides) - ) + return _PRICING_MODEL_INFO.validate_python(_model_cost_rules.with_pricing_overrides(model_info, overrides)) @dataclass(frozen=True, slots=True) @@ -4013,7 +3919,9 @@ def _fallback_registration(key: str, provider: str, value: Mapping[str, object]) generalized: Final[Mapping[str, object]] = match_capability_generalizations(key) or {} if match is None: return _FallbackRegistration(entry=generalized, matched_builtin=False) - inherited: Final = match.entry if match.names_the_row else _without_overridden_price_variants(match.entry, value) + inherited: Final = ( + match.entry if match.names_the_row else _model_cost_rules.without_overridden_price_variants(match.entry, value) + ) own_fields: Final = {name: field_value for name, field_value in inherited.items() if name != "litellm_provider"} return _FallbackRegistration(entry={**generalized, **own_fields}, matched_builtin=True) @@ -4143,7 +4051,7 @@ def model_cost_registration_row(key: str, provider: object) -> Mapping[str, obje 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 + return _model_cost_rules.live_row_or(base.key, base.entry) def _registration_base( @@ -4174,33 +4082,11 @@ def _registration_base( return _RegistrationWrite(key=key, entry=dict(fallback.entry)) -def _base_fields_replaced_by(value: Mapping[str, object]) -> frozenset[str]: - return ( - frozenset({"explicit_cache_read_input_token_cost"}) - if value.get("cache_read_input_token_cost") is not None - and value.get("explicit_cache_read_input_token_cost") is None - else frozenset() - ) - - 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) - # ``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 - } - return _RegistrationWrite( - key=base.key, - entry={**update_dictionary(existing_model, value), **dict.fromkeys(_base_fields_replaced_by(value))}, - ) + return _RegistrationWrite(key=base.key, entry=_model_cost_rules.registration_row(base.entry, value)) def _registration_attempt( @@ -4362,14 +4248,7 @@ def cost_map_omits_token_price(*keys: object) -> bool: ``get_model_info`` substitutes 0 for a missing price, which reads exactly like a declared zero. Surfaces that report pricing use this to keep an unpriced deployment at ``None``. """ - entries: Final = tuple( - entry - for entry in (litellm.model_cost.get(key) for key in keys if isinstance(key, str)) - if isinstance(entry, dict) - ) - return len(entries) > 0 and not any( - "input_cost_per_token" in entry or "output_cost_per_token" in entry for entry in entries - ) + return _model_cost_rules.cost_map_omits_token_price(keys) def register_model( @@ -6713,9 +6592,6 @@ 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 - # 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 @@ -6728,108 +6604,21 @@ def get_model_cost_mutation_generation() -> int: 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 + """Call this whenever litellm.model_cost is modified: drops the case-insensitive key index and the LRU caches + that depend on model_cost data.""" + global _model_cost_mutation_generation _model_cost_mutation_generation += 1 + _model_cost_rules.invalidate_model_cost_keys() - # 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 - - 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)) - 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 + """The ``litellm.model_cost`` key for ``potential_key``: the key itself when present, else the key that + lowercases to the same name""" + return _model_cost_rules.find_model_cost_key(potential_key) _get_model_cost_key = get_model_cost_key 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 deleted file mode 100644 index 3ccb9b2469d..00000000000 --- a/tests/code_coverage_tests/check_get_model_cost_key_performance.py +++ /dev/null @@ -1,242 +0,0 @@ -""" -Code quality check: Ensure _get_model_cost_key only uses O(1) operations. - -Simple pattern-based check for O(n) operations in _get_model_cost_key. -""" - -import re -import os - - -def _function_has_on_operations(all_lines, func_name, visited=None): - """ - Check if a function contains O(n) operations by searching for it in the file. - Recursively checks called functions as well. - """ - if visited is None: - visited = set() - - # Prevent infinite recursion - if func_name in visited: - return False - visited.add(func_name) - - func_start = None - func_end = None - - for i, line in enumerate(all_lines): - if func_start is None and f"def {func_name}(" in line: - func_start = i - elif func_start is not None: - # Function ends when we hit next def at module level - if ( - line.strip() - and not line.startswith(" ") - and not line.startswith("\t") - and line.startswith("def ") - ): - func_end = i - break - - if func_start is None or func_end is None: - return False - - # Check function body for O(n) patterns - func_lines = all_lines[func_start:func_end] - - for line in func_lines: - # Skip comments and docstrings - line_stripped = line.strip() - if ( - line_stripped.startswith("#") - or line_stripped.startswith('"""') - or line_stripped.startswith("'''") - ): - continue - - # Check for for loops - if re.search(r"\bfor\s+\w+\s+in\s+", line): - return True - # Check for while loops - if re.search(r"\bwhile\s+", line): - return True - # Check for comprehensions - if re.search(r"\[.*\s+for\s+.*\s+in\s+", line) or re.search( - r"\{.*\s+for\s+.*\s+in\s+", line - ): - return True - - # Recursively check called functions (check all, don't skip any in recursive checks) - func_call_match = re.search(r"\b([a-z_][a-z0-9_]*)\s*\(", line) - if func_call_match: - called_func = func_call_match.group(1) - if called_func.startswith("_"): - if _function_has_on_operations(all_lines, called_func, visited): - return True - - return False - - -def check_get_model_cost_key_performance(): - """ - Check that _get_model_cost_key doesn't contain O(n) operations. - """ - utils_file = "./litellm/utils.py" - - if not os.path.exists(utils_file): - print(f"Warning: File {utils_file} does not exist.") - return [] - - with open(utils_file, "r", encoding="utf-8") as f: - lines = f.readlines() - - # Find the _get_model_cost_key function - func_start = None - func_end = None - - for i, line in enumerate(lines): - if func_start is None and "def _get_model_cost_key(" in line: - func_start = i - elif func_start is not None: - # Function ends when we hit next def at module level (no indentation) - if ( - line.strip() - and not line.startswith(" ") - and not line.startswith("\t") - and line.startswith("def ") - ): - func_end = i - break - - if func_start is None: - print("Warning: Could not find _get_model_cost_key function") - return [] - - if func_end is None: - func_end = len(lines) - - # Extract function body - func_lines = lines[func_start:func_end] - problematic_lines = [] - - # Track if we're inside a docstring - in_docstring = False - docstring_quote = None - - # Check for O(n) patterns - for i, line in enumerate(func_lines, start=func_start + 1): - line_stripped = line.strip() - - # Track docstring state (handle both single-line and multi-line docstrings) - if not in_docstring: - if line_stripped.startswith('"""') or line_stripped.startswith("'''"): - docstring_quote = '"""' if line_stripped.startswith('"""') else "'''" - # Check if it's a single-line docstring - if line_stripped.count(docstring_quote) >= 2: - in_docstring = False # Single-line, skip this line - continue - else: - in_docstring = True - continue - else: - # Inside multi-line docstring, check for closing quote - if docstring_quote is not None and docstring_quote in line: - in_docstring = False - docstring_quote = None - continue # Skip all lines inside docstring - - # Skip comments - if line_stripped.startswith("#"): - continue - - # 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)) - - # Check for while loops - if re.search(r"\bwhile\s+", line): - problematic_lines.append((i, "while loop", line_stripped)) - - # Check for comprehensions - if re.search(r"\[.*\s+for\s+.*\s+in\s+", line) or re.search( - r"\{.*\s+for\s+.*\s+in\s+", line - ): - problematic_lines.append((i, "comprehension", line_stripped)) - - # Check for problematic function calls - problematic_funcs = [ - "enumerate", - "zip", - "map", - "filter", - "sorted", - "any", - "all", - "sum", - "max", - "min", - ] - for func in problematic_funcs: - if re.search(rf"\b{func}\s*\(", line): - problematic_lines.append((i, f"call to {func}()", line_stripped)) - - # 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", - ] - - # Check for function calls (pattern: function_name(...), but not function definitions) - # Skip function definitions (def function_name(...)) - if not re.search(r"\bdef\s+", line): - func_call_match = re.search(r"\b([a-z_][a-z0-9_]*)\s*\(", line) - if func_call_match: - func_name = func_call_match.group(1) - # If it's a call to a function that might have O(n) operations, check it - if func_name not in allowed_helpers and func_name.startswith("_"): - # Check if this function has O(n) operations - if _function_has_on_operations(lines, func_name): - problematic_lines.append( - ( - i, - f"call to {func_name}() which contains O(n) operations", - line_stripped, - ) - ) - - return problematic_lines - - -def main(): - """Main function to check _get_model_cost_key performance requirements.""" - problematic_lines = check_get_model_cost_key_performance() - - if problematic_lines: - print("\nERROR: Found O(n) operations in _get_model_cost_key:") - for line_num, operation, context in problematic_lines: - print(f" Line {line_num}: {operation} - {context}") - - print( - "\nWARNING: Only O(1) lookup operations are acceptable in _get_model_cost_key." - ) - print("Any O(n) operations will cause severe CPU overhead.") - - raise Exception( - f"Found {len(problematic_lines)} O(n) operation(s) in _get_model_cost_key. " - f"This violates the performance requirement." - ) - else: - print( - "OK: No O(n) operations found in _get_model_cost_key. Performance requirement satisfied." - ) - - -if __name__ == "__main__": - main() diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index 134184c7c01..192f29f5240 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -1,5 +1,11 @@ +import json import os import textwrap +import threading +import time +from collections.abc import Callable, Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from importlib.resources import files from typing import Final import pytest @@ -224,3 +230,110 @@ assert os.waitpid(pid, 0)[1] == 0 ) assert result.returncode == 0, result.stderr + + +_IMPORT_REFRESH_CONTRACT = textwrap.dedent( + """ + import asyncio, os, time + import litellm + from litellm.rust_bridge import _native + + def native_route_error(): + async def call(): + await _native.ResponsesWebSocketConnection.connect("ws://127.0.0.1:1", {}, 0.2) + + try: + asyncio.run(call()) + except Exception as error: + return f"{type(error).__name__}: {error}" + return "" + + def adopted(deadline): + while "fork-remote-only-model" not in litellm.model_cost and time.monotonic() < deadline: + time.sleep(0.05) + return "fork-remote-only-model" in litellm.model_cost + + assert _native.process_state_started() is False, "import litellm started the native runtime" + assert "fork-remote-only-model" not in litellm.model_cost + time.sleep(float(os.environ["FORK_AFTER_SECONDS"])) + pid = os.fork() + if pid == 0: + child_adopted = adopted(time.monotonic() + 30) + error = native_route_error() + native = _native.process_state_started() and "reserved" not in error and "forked" not in error + os._exit(0 if child_adopted and native else 1) + assert os.waitpid(pid, 0)[1] == 0, "the forked child did not end up with the remote map or could not run natively" + """ +) + + +class _ScriptedCostMap(BaseHTTPRequestHandler): + """Answers request n with ``statuses[n]``, the last one repeating, each after ``delay`` seconds""" + + body: bytes = b"" + statuses: tuple[int, ...] = (200,) + delay: float = 0.0 + requests: list[str] = [] # mutable-ok: the server thread records each fetch + + def do_GET(self) -> None: + handler: Final = type(self) + handler.requests.append(self.path) + time.sleep(handler.delay) + status: Final = handler.statuses[min(len(handler.requests), len(handler.statuses)) - 1] + body: Final = handler.body if status == 200 else b"" + self.send_response(status) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + pass + + +@pytest.fixture +def cost_map_server() -> Iterator[Callable[[tuple[int, ...], float], tuple[str, type[_ScriptedCostMap]]]]: + remote: Final = json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_bytes()) + remote["fork-remote-only-model"] = {"litellm_provider": "openai", "mode": "chat"} + servers: Final[list[ThreadingHTTPServer]] = [] # mutable-ok: shut down after the test + + def serve(statuses: tuple[int, ...], delay: float) -> tuple[str, type[_ScriptedCostMap]]: + attributes: Final = {"body": json.dumps(remote).encode(), "statuses": statuses, "delay": delay, "requests": []} + handler: Final = type("Handler", (_ScriptedCostMap,), attributes) + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + servers.append(server) + return f"http://127.0.0.1:{server.server_address[1]}/model_prices.json", handler + + yield serve + for server in servers: + server.shutdown() + server.server_close() + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork only") +@pytest.mark.parametrize( + ("statuses", "delay", "fork_after", "fetches"), + [ + pytest.param((200,), 2.0, 0.0, 1, id="a_fork_mid_refresh_waits_and_the_child_inherits_the_map"), + pytest.param((404, 200), 0.0, 1.0, 2, id="a_child_of_a_parent_whose_refresh_failed_fetches_its_own"), + ], +) +def test_import_starts_no_runtime_and_a_forked_child_ends_up_with_the_remote_map( + cost_map_server: Callable[[tuple[int, ...], float], tuple[str, type[_ScriptedCostMap]]], + statuses: tuple[int, ...], + delay: float, + fork_after: float, + fetches: int, +) -> None: + url, handler = cost_map_server(statuses, delay) + env: Final = { + **{name: value for name, value in os.environ.items() if name != "LITELLM_LOCAL_MODEL_COST_MAP"}, + "LITELLM_MODEL_COST_MAP_URL": url, + "FORK_AFTER_SECONDS": str(fork_after), + "OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "YES", + } + + result: Final = run_child_interpreter(_IMPORT_REFRESH_CONTRACT, env=env, timeout=120) + + assert result.returncode == 0, result.stderr + assert handler.requests == ["/model_prices.json"] * fetches diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 914eb1f2508..94d8289d978 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -4,8 +4,10 @@ import importlib import os import shutil import tempfile -from collections.abc import Callable, Coroutine, Iterator +import threading +from collections.abc import Callable, Coroutine, Iterator, Sequence from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -480,6 +482,69 @@ def fake_provider_credentials(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv(name, value) +@dataclass +class ScriptedHttp: + """A loopback HTTP server answering request ``n`` with ``replies[n]``, the last reply repeating. An + ``httpx.TransportError`` reply closes the connection unanswered""" + + url: str + replies: Sequence[httpx.Response | httpx.TransportError] + requests: list[dict[str, str]] = field(default_factory=list) # mutable-ok: the handler thread records each request + + @property + def count(self) -> int: + return len(self.requests) + + def reply(self) -> httpx.Response | httpx.TransportError: + return self.replies[min(self.count, len(self.replies)) - 1] + + +def _scripted_handler(scripted: Callable[[], ScriptedHttp]) -> type[BaseHTTPRequestHandler]: + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self) -> None: + server: Final = scripted() + server.requests.append({"path": self.path, **{name.lower(): value for name, value in self.headers.items()}}) + reply: Final = server.reply() + self.close_connection = True + if isinstance(reply, httpx.TransportError): + return + self.send_response(reply.status_code) + for name, value in reply.headers.multi_items(): + if name.lower() not in ("content-length", "connection"): + self.send_header(name, value) + self.send_header("Content-Length", str(len(reply.content))) + self.send_header("Connection", "close") + self.end_headers() + self.wfile.write(reply.content) + + def log_message(self, format: str, *args: object) -> None: + pass + + return Handler + + +@pytest.fixture +def scripted_http() -> Iterator[Callable[[Sequence[httpx.Response | httpx.TransportError], str], ScriptedHttp]]: + servers: Final[list[ThreadingHTTPServer]] = [] # mutable-ok: every server the test started, shut down after it + + def serve( + replies: Sequence[httpx.Response | httpx.TransportError], path: str = "/model_prices.json" + ) -> ScriptedHttp: + holder: Final[list[ScriptedHttp]] = [] # mutable-ok: lets the handler class find its script + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _scripted_handler(lambda: holder[0])) + holder.append(ScriptedHttp(url=f"http://127.0.0.1:{server.server_address[1]}{path}", replies=replies)) + threading.Thread(target=server.serve_forever, daemon=True).start() + servers.append(server) + return holder[0] + + yield serve + for server in servers: + server.shutdown() + server.server_close() + + @pytest.fixture async def drained_logging_worker() -> None: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.clear_queue(), timeout=10) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 789d59bd613..00bdcbe150d 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -9,6 +9,7 @@ import threading from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import wait from functools import partial +from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias @@ -3778,13 +3779,14 @@ CatalogLoader: TypeAlias = Callable[[Mapping[str, Mapping[str, object]]], Awaita @pytest.fixture -def load_catalog_with(monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None) -> Iterator[CatalogLoader]: +def load_catalog_with( + monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None, tmp_path: Path +) -> 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)) + catalog: Final = tmp_path / "prices.json" + catalog.write_text(json.dumps({**litellm.get_model_cost_map(url=""), **rows})) 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) + loaded: Final = await refetch_model_cost_map(url=catalog.as_uri()) monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") assert isinstance(loaded, ModelCostMapReloaded) monkeypatch.setattr(litellm, "model_cost", loaded.model_cost_map) 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 d44169ae28a..b737bccc082 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 @@ -11,10 +11,8 @@ import threading from pathlib import Path 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 from litellm.litellm_core_utils.fallback_generalizations import ( @@ -26,8 +24,8 @@ from litellm.litellm_core_utils.fallback_generalizations import ( from litellm.litellm_core_utils.get_model_cost_map import ( FALLBACK_GENERALIZATIONS_KEY, GetModelCostMap, - _count_model_entries, _finalize_model_cost_map, + count_model_entries, get_model_cost_map_provenance, git_blob_id, ) @@ -57,7 +55,7 @@ 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 + assert count_model_entries(m) == 3 def test_validation_rejects_truly_shrunk_file_even_with_meta_keys(): @@ -229,17 +227,17 @@ 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(): +def test_get_model_cost_map_stamps_loaded_at(scripted_http): """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) + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) before = datetime.now(timezone.utc) - module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client) + module.get_model_cost_map(url=server.url) loaded_at = module.get_model_cost_map_loaded_at() assert loaded_at is not None @@ -261,8 +259,7 @@ from litellm.litellm_core_utils.get_model_cost_map import ( ModelCostMapReloadUnavailable, refetch_model_cost_map, ) - -_URL = "https://example.invalid/model_prices.json" +from litellm.rust_bridge._native import install_model_cost_map_provenance, model_cost_map_provenance @functools.lru_cache(maxsize=1) @@ -286,25 +283,26 @@ def _unset_local_cost_map_env(monkeypatch): 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.fixture(autouse=True) +def _restore_cost_map_provenance(): + """Loads install a new provenance, so each test puts the one it found back.""" + found = model_cost_map_provenance() + yield + if found is not None: + install_model_cost_map_provenance( + source=found.source, + url=found.url, + is_env_forced=found.is_env_forced, + fallback_reason=found.fallback_reason, + revision=found.revision, + etag=found.etag, + ) @pytest.mark.asyncio -async def test_refetch_retries_429_honoring_retry_after(): +async def test_refetch_retries_429_honoring_retry_after(scripted_http): """Two 429s with Retry-After then success: waits follow the header, not backoff.""" - client, calls = _mock_client( + server = scripted_http( [ httpx.Response(429, headers={"Retry-After": "7"}), httpx.Response(429, headers={"Retry-After": "7"}), @@ -312,117 +310,214 @@ async def test_refetch_retries_429_honoring_retry_after(): ] ) sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) assert isinstance(result, ModelCostMapReloaded) assert len(result.model_cost_map) > 100 - assert calls["count"] == 3 + assert server.count == 3 assert sleeper.waits == [7.0, 7.0] @pytest.mark.asyncio -async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff(): +async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff(scripted_http): """All 429 without Retry-After: exponential backoff waits, then a failure value.""" - client, calls = _mock_client([httpx.Response(429)]) + server = scripted_http([httpx.Response(429)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) 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 + assert result.reason == f"HTTP 429 from {server.url} (after 3 attempts)" + assert server.count == 3 + assert sleeper.waits == [2.0 + _jitter(2)[0], 4.0 + _jitter(2)[1]] + + +def _jitter(count: int) -> list[float]: + rng = random.Random(0) + return [rng.uniform(0.0, 1.0) for _ in range(count)] @pytest.mark.asyncio -async def test_refetch_caps_retry_after_wait(): +@pytest.mark.parametrize( + ("replies", "waits", "reloaded", "reason"), + [ + pytest.param( + [httpx.Response(429, headers={"Retry-After": "3"}), httpx.Response(503), None], + [3.0, 4.0 + _jitter(1)[0]], + True, + None, + id="jitter_is_drawn_only_for_a_backoff", + ), + pytest.param( + [ + httpx.Response(429, headers={"Retry-After": "-1"}), + httpx.Response(429, headers={"Retry-After": "nan"}), + None, + ], + [2.0 + _jitter(2)[0], 4.0 + _jitter(2)[1]], + True, + None, + id="negative_and_nan_retry_after_back_off", + ), + pytest.param( + [httpx.Response(429, headers={"Retry-After": "inf"}), None], + [30.0], + True, + None, + id="infinite_retry_after_is_capped", + ), + pytest.param( + [httpx.Response(503, headers={"Retry-After": " 2_0 "}), None], + [20.0], + True, + None, + id="retry_after_is_read_as_python_float_reads_it", + ), + pytest.param( + [httpx.Response(503, headers={"Retry-After": "Wed, 21 Oct 2026 07:28:00 GMT"}), None], + [2.0 + _jitter(1)[0]], + True, + None, + id="an_http_date_retry_after_backs_off", + ), + pytest.param([httpx.Response(501)], [], False, "HTTP 501 from ", id="other_5xx_is_final"), + pytest.param([httpx.Response(599)], [], False, "HTTP 599 from ", id="599_is_final"), + pytest.param([httpx.Response(400)], [], False, "HTTP 400 from ", id="400_is_final"), + pytest.param([httpx.Response(302, content=b"moved")], [], False, "invalid JSON", id="a_redirect_is_read"), + pytest.param([httpx.Response(304)], [], False, "invalid JSON", id="an_unconditional_304_is_read"), + pytest.param([httpx.Response(600, content=b"[]")], [], False, "expected a JSON object", id="600_is_read"), + ], +) +async def test_refetch_retry_policy_matches_main( + scripted_http, replies: list[httpx.Response | None], waits: list[float], reloaded: bool, reason: str | None +): + """The same status cases as ``status_rule_matches_main`` in ``model-catalog-refresh/src/status.rs``""" + server = scripted_http( + [reply if reply is not None else httpx.Response(200, content=_real_map_bytes()) for reply in replies] + ) + sleeper = _SleepRecorder() + + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) + + assert sleeper.waits == waits + assert isinstance(result, ModelCostMapReloaded) is reloaded + if reason is not None: + assert isinstance(result, ModelCostMapReloadUnavailable) + assert result.reason.startswith(reason) + + +@pytest.mark.asyncio +async def test_refetch_caps_retry_after_wait(scripted_http): """A hostile/huge Retry-After is capped so reloads never sleep unbounded.""" - client, _calls = _mock_client( + server = scripted_http( [ 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) + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) assert isinstance(result, ModelCostMapReloaded) assert sleeper.waits == [30.0] @pytest.mark.asyncio -async def test_refetch_retries_transport_errors(): +async def test_refetch_retries_transport_errors(scripted_http): """Connection failures are transient: retried like 5xx, succeeding when the network heals.""" - client, calls = _mock_client( + server = scripted_http( [ 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) + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) assert isinstance(result, ModelCostMapReloaded) - assert calls["count"] == 2 + assert server.count == 2 assert len(sleeper.waits) == 1 @pytest.mark.asyncio -async def test_refetch_non_retryable_status_fails_immediately(): +async def test_refetch_without_an_injected_sleep_waits_in_rust(scripted_http): + server = scripted_http( + [httpx.Response(503, headers={"Retry-After": "0"}), httpx.Response(200, content=_real_map_bytes())] + ) + result = await refetch_model_cost_map(url=server.url, rng=random.Random(0)) + assert isinstance(result, ModelCostMapReloaded) + assert server.count == 2 + + +@pytest.mark.asyncio +async def test_refetch_raises_what_the_injected_sleep_raises_and_stops_fetching(scripted_http): + server = scripted_http([httpx.Response(503, headers={"Retry-After": "0"})]) + + async def broken_sleep(seconds: float) -> None: + raise RuntimeError(f"no sleeping for {seconds}") + + with pytest.raises(RuntimeError, match=r"no sleeping for 0\.0"): + await refetch_model_cost_map(url=server.url, sleep=broken_sleep, rng=random.Random(0)) + assert server.count == 1 + + +@pytest.mark.asyncio +async def test_refetch_warns_that_an_injected_client_is_ignored(scripted_http): + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) + with pytest.warns(DeprecationWarning, match="client= is ignored"): + result = await refetch_model_cost_map(url=server.url, client=object()) + assert isinstance(result, ModelCostMapReloaded) + assert server.count == 1 + + +@pytest.mark.asyncio +async def test_refetch_non_retryable_status_fails_immediately(scripted_http): """A 404 is permanent: one attempt, no sleeps, failure value.""" - client, calls = _mock_client([httpx.Response(404)]) + server = scripted_http([httpx.Response(404)]) sleeper = _SleepRecorder() - result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client) + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) assert isinstance(result, ModelCostMapReloadUnavailable) - assert "404" in result.reason - assert calls["count"] == 1 + assert result.reason == f"HTTP 404 from {server.url}" + assert server.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")]) +async def test_refetch_invalid_json_fails_immediately(scripted_http): + server = scripted_http([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) + result = await refetch_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) assert isinstance(result, ModelCostMapReloadUnavailable) assert "invalid JSON" in result.reason - assert calls["count"] == 1 + assert server.count == 1 assert sleeper.waits == [] @pytest.mark.asyncio -async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(): +async def test_refetch_shrunk_map_fails_integrity_not_swapped_in(scripted_http): """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) + server = scripted_http([httpx.Response(200, content=tiny)]) + result = await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) 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): +async def test_refetch_respects_local_env_override(monkeypatch, scripted_http, flag): """LITELLM_LOCAL_MODEL_COST_MAP=True or LITELLM_OFFLINE_MODE=True short-circuits to the bundled backup, zero HTTP.""" monkeypatch.setenv(flag, "True") + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) - def _fail(request): - raise AssertionError("no HTTP request should be made when local map is forced") + result = await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) - 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 + assert server.count == 0 @pytest.mark.asyncio -async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(): +async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(scripted_http): body = _real_map_bytes() - client, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"abc123"'}, content=body)]) + server = scripted_http([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) + result = await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) assert isinstance(result, ModelCostMapReloaded) assert result.revision == git_blob_id(body) @@ -431,14 +526,42 @@ async def test_refetch_records_the_blob_id_of_the_bytes_served_and_the_fetch_eta @pytest.mark.asyncio -async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later_registrations(monkeypatch): +async def test_refetch_reads_a_file_url(tmp_path): + body = _real_map_bytes() + path = tmp_path / "model_prices.json" + path.write_bytes(body) + + result = await refetch_model_cost_map(url=path.as_uri()) + + assert isinstance(result, ModelCostMapReloaded) + assert result.revision == git_blob_id(body) + assert get_model_cost_map_source_info() == { + **get_model_cost_map_source_info(), + "source": "remote", + "url": path.as_uri(), + "etag": None, + } + + +@pytest.mark.asyncio +async def test_refetch_of_a_missing_file_is_final(tmp_path): + missing = (tmp_path / "missing.json").as_uri() + result = await refetch_model_cost_map(url=missing, sleep=_SleepRecorder()) + assert isinstance(result, ModelCostMapReloadUnavailable) + assert result.reason.startswith(f"could not fetch {missing}: could not read the catalog") + + +@pytest.mark.asyncio +async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later_registrations( + monkeypatch, scripted_http +): 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())]) + server = scripted_http([httpx.Response(200, content=json.dumps(edited).encode())]) - result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client) + result = await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) assert isinstance(result, ModelCostMapReloaded) monkeypatch.setattr(litellm, "model_cost", result.model_cost_map) @@ -448,15 +571,15 @@ async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later @pytest.mark.asyncio -async def test_loaded_catalog_fields_follow_each_fetched_map(): +async def test_loaded_catalog_fields_follow_each_fetched_map(scripted_http): 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())]) + with_field = scripted_http([httpx.Response(200, content=json.dumps(added).encode())]) + without_field = scripted_http([httpx.Response(200, content=_real_map_bytes())]) - await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=with_field) + await refetch_model_cost_map(url=with_field.url, sleep=_SleepRecorder(), rng=random.Random(0)) 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) + await refetch_model_cost_map(url=without_field.url, sleep=_SleepRecorder(), rng=random.Random(0)) assert "supports_litellm_qa_capability" in fields_after_adding assert "supports_litellm_qa_capability" not in GetModelCostMap.loaded_model_cost_map_fields() @@ -464,15 +587,15 @@ async def test_loaded_catalog_fields_follow_each_fetched_map(): @pytest.mark.asyncio -async def test_refetch_revision_follows_the_bytes_not_the_url(): +async def test_refetch_revision_follows_the_bytes_not_the_url(scripted_http): edited = json.loads(_real_map_bytes()) edited["gpt-5.4-mini"]["input_cost_per_token"] = 0.5 - client, _ = _mock_client( + server = scripted_http( [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) + first = await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) + second = await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) assert isinstance(first, ModelCostMapReloaded) and isinstance(second, ModelCostMapReloaded) assert first.revision != second.revision @@ -480,12 +603,12 @@ async def test_refetch_revision_follows_the_bytes_not_the_url(): @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) +async def test_refetch_local_override_reports_the_bundled_blob_id_without_an_etag(monkeypatch, scripted_http): + remote = scripted_http([httpx.Response(200, headers={"ETag": 'W/"remote"'}, content=_real_map_bytes())]) + await refetch_model_cost_map(url=remote.url, sleep=_SleepRecorder(), rng=random.Random(0)) monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0)) + result = await refetch_model_cost_map(url=remote.url, sleep=_SleepRecorder(), rng=random.Random(0)) assert isinstance(result, ModelCostMapReloaded) assert result.revision == _bundled_blob_id() @@ -493,19 +616,19 @@ async def test_refetch_local_override_reports_the_bundled_blob_id_without_an_eta @pytest.mark.asyncio -async def test_refetch_stamps_loaded_at_on_remote_and_local_reloads(monkeypatch): +async def test_refetch_stamps_loaded_at_on_remote_and_local_reloads(monkeypatch, scripted_http): from litellm.litellm_core_utils import get_model_cost_map as module - client, _ = _mock_client([httpx.Response(200, content=_real_map_bytes())]) + server = scripted_http([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) + await refetch_model_cost_map(url=server.url, sleep=_SleepRecorder(), rng=random.Random(0)) 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)) + await refetch_model_cost_map(url=server.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) @@ -518,52 +641,69 @@ async def test_refetch_stamps_loaded_at_on_remote_and_local_reloads(monkeypatch) from litellm.litellm_core_utils.get_model_cost_map import ( get_model_cost_map, get_model_cost_map_source_info, + load_model_cost_map_at_import, ) 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 +from unittest.mock import call, 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)) - +def test_get_model_cost_map_fetches_s3_catalog(monkeypatch: pytest.MonkeyPatch, scripted_http) -> None: 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) + body: Final = json.dumps(s3_content).encode() + server = scripted_http([httpx.Response(200, headers={"ETag": '"fake-etag"'}, content=body)], path="") + for name, value in { + "AWS_ENDPOINT_URL": server.url, + "AWS_ACCESS_KEY_ID": "test-access-key", + "AWS_SECRET_ACCESS_KEY": "test-secret-key", + "AWS_REGION_NAME": "us-east-1", + }.items(): + monkeypatch.setenv(name, value) 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" + assert [request["path"].split("?")[0] for request in server.requests] == ["/my-bucket/prices/model_prices.json"] + assert server.requests[0]["authorization"].startswith("AWS4-HMAC-SHA256 Credential=test-access-key/") + source = get_model_cost_map_source_info() + assert (source["source"], source["etag"], source["source_revision"]) == ("remote", '"fake-etag"', git_blob_id(body)) + + +def test_an_s3_failure_is_final(monkeypatch: pytest.MonkeyPatch, scripted_http) -> None: + server = scripted_http([httpx.Response(503)], path="") + for name, value in { + "AWS_ENDPOINT_URL": server.url, + "AWS_ACCESS_KEY_ID": "test-access-key", + "AWS_SECRET_ACCESS_KEY": "test-secret-key", + "AWS_REGION_NAME": "us-east-1", + }.items(): + monkeypatch.setenv(name, value) + + get_model_cost_map("s3://my-bucket/prices/model_prices.json") + + assert server.count == 1 + assert _retry_threads() == [] + assert get_model_cost_map_source_info()["fallback_reason"].startswith( + "Remote fetch failed: could not fetch s3://my-bucket/prices/model_prices.json: S3 GetObject failed" + ) + + +@pytest.mark.parametrize( + ("url", "reason"), + [ + pytest.param("s3://my-bucket/", "expected s3://bucket/key", id="s3_without_a_key"), + pytest.param("ftp://example.invalid/prices.json", 'unsupported catalog URL scheme "ftp"', id="ftp"), + pytest.param("", 'unsupported catalog URL scheme ""', id="empty"), + ], +) +def test_a_location_no_source_serves_falls_back_at_once(url: str, reason: str) -> None: + get_model_cost_map(url) + + assert _retry_threads() == [] + assert ( + get_model_cost_map_source_info()["fallback_reason"] == f"Remote fetch failed: could not fetch {url}: {reason}" + ) class _SyncSleepRecorder: @@ -586,30 +726,27 @@ 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) +def _refresh_threads(): + return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-refresh"] + + +def test_boot_load_success_does_not_start_background_retry(scripted_http): + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) sleeper = _SyncSleepRecorder() - cost_map = get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - ) - assert calls["count"] == 1 + cost_map = get_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) + assert server.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): +def _isolated_catalog_state(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)) + monkeypatch.setattr(litellm, "model_cost", dict(litellm.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)) @@ -619,31 +756,26 @@ def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_ "_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)) + +def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch, scripted_http): + import litellm + + _isolated_catalog_state(monkeypatch) remote_map = _load_root_cost_map() remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} - client, calls = _mock_client( + server = scripted_http( [ 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, - ) + cost_map = get_model_cost_map(url=server.url, max_attempts=3, sleep=sleeper, rng=random.Random(0)) - assert calls["count"] == 1 + assert server.count == 1 assert sleeper.waits == [] assert sleeper.started.wait(timeout=10) threads = _retry_threads() @@ -658,8 +790,8 @@ def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_ 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 sleeper.waits == [2.0 + _jitter(1)[0]] + assert server.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 @@ -672,65 +804,60 @@ def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_ 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) +def test_boot_load_does_not_retry_non_retryable_failure(scripted_http): + server = scripted_http([httpx.Response(404)]) sleeper = _SyncSleepRecorder() - get_model_cost_map( - url=_URL, - sleep=sleeper, - rng=random.Random(0), - client=client, - ) - assert calls["count"] == 1 + get_model_cost_map(url=server.url, sleep=sleeper, rng=random.Random(0)) + assert server.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 + assert source["fallback_reason"] == f"Remote fetch failed: HTTP 404 from {server.url}" -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) +def test_boot_load_with_one_attempt_does_not_retry_a_retryable_failure(scripted_http): + server = scripted_http([httpx.Response(503)]) + + get_model_cost_map(url=server.url, max_attempts=1) + + assert server.count == 1 + assert _retry_threads() == [] + assert get_model_cost_map_source_info()["fallback_reason"] == f"Remote fetch failed: HTTP 503 from {server.url}" + + +def test_boot_load_falls_back_when_the_url_does_not_parse(): 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 + cost_map = get_model_cost_map(url="http://[bad/model_prices.json", sleep=sleeper, rng=random.Random(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 "") + assert source["fallback_reason"].startswith( + "Remote fetch failed: could not fetch http://[bad/model_prices.json: invalid catalog URL" + ) @pytest.mark.parametrize("flag", ["LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"]) -def test_boot_load_respects_local_env_override(monkeypatch, flag): +def test_boot_load_respects_local_env_override(monkeypatch, scripted_http, 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") + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) - 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)), - ) + cost_map = get_model_cost_map(url=server.url, sleep=_SyncSleepRecorder()) assert len(cost_map) > 100 + assert server.count == 0 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(): +def test_boot_load_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(scripted_http): body = _real_map_bytes() - client, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=body)], client_cls=httpx.Client) + server = scripted_http([httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=body)]) - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=client) + get_model_cost_map(url=server.url, sleep=_SyncSleepRecorder(), rng=random.Random(0)) source = get_model_cost_map_source_info() assert source["source"] == "remote" @@ -739,14 +866,22 @@ def test_boot_load_records_the_blob_id_of_the_bytes_served_and_the_fetch_etag(): 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 +def test_boot_load_joins_a_repeated_etag_as_httpx_reads_it(scripted_http): + server = scripted_http( + [httpx.Response(200, headers=[("ETag", 'W/"a"'), ("ETag", 'W/"b"')], content=_real_map_bytes())] ) - 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) + get_model_cost_map(url=server.url) + + assert get_model_cost_map_source_info()["etag"] == 'W/"a", W/"b"' + + +def test_boot_load_fallback_to_the_backup_reports_its_blob_id_and_drops_the_remote_etag(scripted_http): + remote = scripted_http([httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=_real_map_bytes())]) + get_model_cost_map(url=remote.url, sleep=_SyncSleepRecorder(), rng=random.Random(0)) + failing = scripted_http([httpx.Response(404)]) + + get_model_cost_map(url=failing.url, sleep=_SyncSleepRecorder(), rng=random.Random(0)) source = get_model_cost_map_source_info() assert source["source"] == "local" @@ -754,17 +889,13 @@ def test_boot_load_fallback_to_the_backup_reports_its_blob_id_and_drops_the_remo 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) +def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rejected_fetch(scripted_http): + remote = scripted_http([httpx.Response(200, headers={"ETag": 'W/"boot"'}, content=_real_map_bytes())]) + get_model_cost_map(url=remote.url, sleep=_SyncSleepRecorder(), rng=random.Random(0)) 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 - ) + shrunk = scripted_http([httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)]) - get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=shrunk) + get_model_cost_map(url=shrunk.url, sleep=_SyncSleepRecorder(), rng=random.Random(0)) source = get_model_cost_map_source_info() assert source["source"] == "local" @@ -783,15 +914,15 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej ], ) def test_boot_load_skips_remote_fetch_for_cli_processes( - monkeypatch: pytest.MonkeyPatch, argv0: str, request_count: int + monkeypatch: pytest.MonkeyPatch, scripted_http, 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) + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) - cost_map = get_model_cost_map(url=_URL, client=client) + cost_map = get_model_cost_map(url=server.url) - assert calls["count"] == request_count + assert server.count == request_count assert cost_map source = get_model_cost_map_source_info() if request_count == 0: @@ -800,12 +931,95 @@ def test_boot_load_skips_remote_fetch_for_cli_processes( assert source["source"] == "remote" +def _join_refresh_threads() -> None: + for thread in _refresh_threads(): + thread.join(timeout=30) + + +def test_import_serves_the_bundled_map_then_adopts_the_remote_one_in_the_background(monkeypatch, scripted_http): + import litellm + + _isolated_catalog_state(monkeypatch) + remote_map = _load_root_cost_map() + remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} + body = json.dumps(remote_map).encode() + server = scripted_http([httpx.Response(200, headers={"ETag": 'W/"swr"'}, content=body)]) + litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}) + + served = load_model_cost_map_at_import(server.url) + at_import = get_model_cost_map_source_info() + monkeypatch.setattr(litellm, "model_cost", served) + _join_refresh_threads() + + assert "claude-remote-only-test" not in served + assert (at_import["source"], at_import["url"], at_import["fallback_reason"]) == ("local", server.url, None) + assert at_import["source_revision"] == _bundled_blob_id() + assert server.count == 1 + 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 + after = get_model_cost_map_source_info() + assert (after["source"], after["etag"], after["source_revision"]) == ("remote", 'W/"swr"', git_blob_id(body)) + + +@pytest.mark.parametrize( + ("reply", "fallback_reason"), + [ + pytest.param(httpx.Response(404), "Remote fetch failed: HTTP 404 from {url}", id="final_status"), + pytest.param( + httpx.Response(200, content=json.dumps(_make_models(60)).encode()), + "Remote data failed integrity validation", + id="shrunk", + ), + ], +) +def test_import_keeps_the_bundled_map_when_the_background_refresh_fails( + monkeypatch, scripted_http, reply, fallback_reason +): + import litellm + + _isolated_catalog_state(monkeypatch) + server = scripted_http([reply]) + + served = load_model_cost_map_at_import(server.url) + monkeypatch.setattr(litellm, "model_cost", served) + _join_refresh_threads() + + assert litellm.model_cost is served + source = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["source_revision"]) == ("local", server.url, _bundled_blob_id()) + assert source["fallback_reason"] == fallback_reason.format(url=server.url) + + +@pytest.mark.parametrize("flag", ["LITELLM_LOCAL_MODEL_COST_MAP", "LITELLM_OFFLINE_MODE"]) +def test_import_with_the_local_map_forced_starts_no_refresh(monkeypatch, scripted_http, flag): + monkeypatch.setenv(flag, "True") + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) + + load_model_cost_map_at_import(server.url) + + assert _refresh_threads() == [] + assert server.count == 0 + assert get_model_cost_map_source_info()["is_env_forced"] is True + + +def test_import_in_a_cli_process_starts_no_refresh(monkeypatch, scripted_http): + monkeypatch.setattr(sys, "argv", ["/some/venv/bin/litellm-proxy", "--version"]) + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) + + load_model_cost_map_at_import(server.url) + + assert _refresh_threads() == [] + assert server.count == 0 + + @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="session") def event_loop(): try: @@ -815,6 +1029,7 @@ def event_loop(): yield loop loop.close() + @pytest.fixture(scope="function") def setup_and_teardown(event_loop): import litellm @@ -848,6 +1063,7 @@ def setup_and_teardown(event_loop): if pending: event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "set_verbose": getattr(litellm, "set_verbose", False), @@ -862,6 +1078,7 @@ _SCALAR_DEFAULTS = { "cohere_key": getattr(litellm, "cohere_key", None), } + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckIsValidDict: """Unit tests for _check_is_valid_dict.""" @@ -886,6 +1103,7 @@ class TestCheckIsValidDict: """Non-empty dict should pass.""" assert GetModelCostMap._check_is_valid_dict({"model": {}}) is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestCheckModelCountNotReduced: """Unit tests for _check_model_count_not_reduced.""" @@ -940,6 +1158,59 @@ class TestCheckModelCountNotReduced: is True ) + +@pytest.mark.parametrize( + ("model_cost", "expected"), + [ + pytest.param({"a": {}, "sample_spec": {}, FALLBACK_GENERALIZATIONS_KEY: {}}, 1, id="sections"), + pytest.param({1: {}, None: {}, "Sample_Spec": {}}, 3, id="non_string_and_case_differing_keys"), + pytest.param({}, 0, id="empty"), + ], +) +def test_count_model_entries_matches_main(model_cost: dict[object, object], expected: int): + assert count_model_entries(model_cost) == expected + + +_BELOW_MINIMUM = ( + "LiteLLM: Fetched model cost map has only %d models (minimum=%d). " + "This may indicate a corrupted upstream file. " + "Falling back to local backup." +) +_SHRUNK = ( + "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." +) + + +@pytest.mark.parametrize( + ("models", "backup", "minimum", "ratio", "warning"), + [ + pytest.param(5, 10, 6, 0.5, (_BELOW_MINIMUM, 5, 6), id="below_minimum_wins_over_shrunk"), + pytest.param(5, 11, 0, 0.5, (_SHRUNK, 5, 11, 50.0), id="shrunk"), + pytest.param(5, 10, 0, 0.5, None, id="exactly_at_the_ratio"), + pytest.param(1, -5, -1, 0.5, None, id="negative_minimum_and_backup_pass"), + pytest.param(0, 0, 0, 2.0, None, id="empty_backup_skips_the_ratio"), + pytest.param(10, 10, 0, 1.5, (_SHRUNK, 10, 10, 150.0), id="ratio_above_one"), + pytest.param(0, 10, 0, float("nan"), None, id="nan_ratio_never_shrinks"), + ], +) +def test_model_count_check_matches_main( + models: int, backup: int, minimum: int, ratio: float, warning: tuple[object, ...] | None +): + with patch.object(get_model_cost_map_module.verbose_logger, "warning") as warn: + passed = GetModelCostMap._check_model_count_not_reduced( + fetched_map={**_make_models(models), "sample_spec": {}}, + backup_model_count=backup, + min_model_count=minimum, + max_shrink_ratio=ratio, + ) + + assert passed is (warning is None) + assert warn.call_args_list == ([] if warning is None else [call(*warning)]) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestValidateModelCostMap: """Unit tests for validate_model_cost_map (combines both checks).""" @@ -976,16 +1247,11 @@ class TestValidateModelCostMap: 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 _fetch_with_single_outcome(serve, outcome): + server = serve([outcome]) + result = get_model_cost_map(url=server.url, max_attempts=1) + return result, server.count, get_model_cost_map_source_info(), server.url def _backup_keys(): @@ -996,60 +1262,64 @@ def _backup_keys(): class TestGetModelCostMapFallback: """Tests for get_model_cost_map fallback behavior with bad upstream.""" - def test_should_fallback_to_backup_on_invalid_json(self, monkeypatch): + def test_should_fallback_to_backup_on_invalid_json(self, scripted_http): """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")) + result, fetches, source, url = _fetch_with_single_outcome( + scripted_http, 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): + def test_should_fallback_to_backup_on_network_error(self, scripted_http): """When upstream is unreachable, should fall back to local backup.""" - result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.ConnectError("Connection refused")) + result, fetches, source, url = _fetch_with_single_outcome( + scripted_http, 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): + def test_should_fallback_when_fetched_map_is_empty(self, scripted_http): """When upstream returns valid JSON but empty dict, should fall back.""" - result, fetches, source = _fetch_with_single_outcome(monkeypatch, httpx.Response(200, content=b"{}")) + result, fetches, source, url = _fetch_with_single_outcome(scripted_http, 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"[]")) + def test_should_fallback_when_fetched_json_is_not_object(self, scripted_http): + result, fetches, source, url = _fetch_with_single_outcome(scripted_http, 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" + 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): + def test_should_fallback_when_fetched_map_shrinks_dramatically(self, scripted_http): """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()) + result, fetches, source, url = _fetch_with_single_outcome( + scripted_http, 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): + def test_should_use_local_map_when_env_var_set(self, scripted_http): """LITELLM_LOCAL_MODEL_COST_MAP=True should skip remote fetch entirely.""" + server = scripted_http([httpx.Response(200, content=_real_map_bytes())]) 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() + result = get_model_cost_map(server.url) + + assert server.count == 0 + assert result.keys() == _backup_keys() - assert isinstance(result, dict) - assert len(result) > 0 @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBackupModelCostMapExists: @@ -1066,6 +1336,7 @@ class TestBackupModelCostMapExists: 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: """ @@ -1077,17 +1348,12 @@ class TestBadHostedModelCostMap: - litellm.completion() still works """ - def test_should_model_info_pass_after_bad_hosted_map(self): + def test_should_model_info_pass_after_bad_hosted_map(self, scripted_http): """ 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") + fallback_map = get_model_cost_map(scripted_http([httpx.Response(200, content=b"bad json")]).url) original = litellm.model_cost litellm.model_cost = fallback_map @@ -1099,7 +1365,7 @@ class TestBadHostedModelCostMap: finally: litellm.model_cost = original - def test_should_completion_pass_after_bad_hosted_map(self): + def test_should_completion_pass_after_bad_hosted_map(self, scripted_http): """ If the hosted map is bad, litellm.completion() should still work. @@ -1107,13 +1373,7 @@ class TestBadHostedModelCostMap: 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") + fallback_map = get_model_cost_map(scripted_http([httpx.Response(200, content=b"bad json")]).url) original = litellm.model_cost litellm.model_cost = fallback_map @@ -1130,12 +1390,14 @@ class TestBadHostedModelCostMap: 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(): """ @@ -1188,6 +1450,7 @@ def isolate_litellm_state(): 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), @@ -1208,6 +1471,7 @@ _SCALAR_DEFAULTS_local_testing = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown_local_testing(): """ @@ -1230,6 +1494,7 @@ def setup_and_teardown_local_testing(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -1241,6 +1506,7 @@ def test_get_model_cost_map(): except Exception as e: pytest.fail(f"An exception occurred: {e}") + @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", @@ -1264,3 +1530,51 @@ def test_a_bundled_backup_that_is_not_a_json_object_is_rejected_without_quoting_ GetModelCostMap.load_local_model_cost_map_with_revision() assert "only-in-the-backup" not in str(rejected.value) + + +@pytest.mark.asyncio +async def test_a_failed_import_refresh_leaves_a_map_installed_meanwhile_alone(scripted_http): + from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapRefresh, _refresh_in_background + + failing = scripted_http([httpx.Response(404)]) + load_model_cost_map_at_import(failing.url) + _join_refresh_threads() + reloaded = scripted_http([httpx.Response(200, headers={"ETag": 'W/"reload"'}, content=_real_map_bytes())]) + await refetch_model_cost_map(url=reloaded.url) + + _refresh_in_background( + ModelCostMapRefresh(failing.url, timeout=5, max_attempts=1, rng=random.Random(), on_retry=print), failing.url + ) + + source = get_model_cost_map_source_info() + assert (source["source"], source["url"], source["etag"], source["fallback_reason"]) == ( + "remote", + reloaded.url, + 'W/"reload"', + None, + ) + + +def test_a_refresh_landing_after_interpreter_exit_began_is_not_adopted(monkeypatch, scripted_http): + import litellm + from litellm.litellm_core_utils.get_model_cost_map import ( + ModelCostMapRefresh, + _BackgroundAdoption, + _refresh_in_background, + _stop_background_adoption, + ) + + _isolated_catalog_state(monkeypatch) + remote_map = _load_root_cost_map() + remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"} + server = scripted_http([httpx.Response(200, content=json.dumps(remote_map).encode())]) + refresh = ModelCostMapRefresh(server.url, timeout=5, max_attempts=1, rng=random.Random(), on_retry=print) + + _stop_background_adoption() + try: + _refresh_in_background(refresh, server.url) + finally: + _BackgroundAdoption.exiting.clear() + + assert server.count == 1 + assert "claude-remote-only-test" not in litellm.model_cost 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..e4941e558a5 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 @@ -98,7 +98,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 + client, auth_as, monkeypatch, mock_prisma, scripted_http ): import httpx @@ -112,11 +112,8 @@ def test_reload_model_cost_map_surfaces_the_blob_id_of_the_bytes_served_on_every 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)), - ) + server = scripted_http([httpx.Response(200, headers={"ETag": _SERVED_ETAG}, content=body)]) + monkeypatch.setattr(litellm, "model_cost_map_url", server.url) monkeypatch.setattr("litellm.add_known_models", lambda model_cost_map=None: None) monkeypatch.setattr("litellm.model_cost", {}, raising=False) diff --git a/tests/unit/test_model_cost_aliases.py b/tests/unit/test_model_cost_aliases.py index f9f92a85cb2..4b6aae7b6ad 100644 --- a/tests/unit/test_model_cost_aliases.py +++ b/tests/unit/test_model_cost_aliases.py @@ -5,10 +5,17 @@ The ``_expand_model_aliases`` function processes ``aliases`` lists from model entries, creating shared dict references for alias entries at load time. """ -from unittest.mock import patch +from collections.abc import Mapping +from unittest.mock import call, patch + +import pytest from litellm import verbose_logger -from litellm.litellm_core_utils.get_model_cost_map import _expand_model_aliases +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, _expand_model_aliases + +_NOT_A_LIST = "LiteLLM model alias field for '%s' is not a list (got %s) — skipping." +_SHADOWS = "LiteLLM model alias conflict: alias '%s' (from '%s') already exists as a canonical entry — skipping." +_TAKEN = "LiteLLM model alias conflict: alias '%s' (from '%s') was already claimed by another entry — skipping." # --------------------------------------------------------------------------- @@ -243,3 +250,127 @@ class TestAliasIntegration: def test_expand_on_empty_dict(self): """Expanding an empty dict returns an empty dict.""" assert _expand_model_aliases({}) == {} + + +def _served(model_cost: Mapping[str, object]) -> list[tuple[object, object, tuple[object, ...] | None]]: + return [ + (key, dict(row) if isinstance(row, dict) else row, row.names if isinstance(row, AliasedCostEntry) else None) + for key, row in model_cost.items() + ] + + +class TestExpansionMatchesMainExactly: + """Every case pins the map, the row objects and the warnings main's Python expansion produced""" + + @pytest.mark.parametrize( + ("model_cost", "served", "warnings"), + [ + pytest.param( + {"b": {"aliases": ["b1"]}, "a": {"aliases": ["a1", "a2"]}, "c": {}}, + [ + ("b", {}, ("b", "b1")), + ("a", {}, ("a", "a1", "a2")), + ("c", {}, None), + ("b1", {}, ("b", "b1")), + ("a1", {}, ("a", "a1", "a2")), + ("a2", {}, ("a", "a1", "a2")), + ], + [], + id="aliases_follow_rows_in_first_claim_order", + ), + pytest.param( + {"m": {"aliases": None, "mode": "chat"}}, + [("m", {"aliases": None, "mode": "chat"}, None)], + [], + id="null_aliases_stay_on_the_row", + ), + pytest.param( + {"m": {"aliases": "m2", "mode": "chat"}}, + [("m", {"mode": "chat"}, None)], + [call(_NOT_A_LIST, "m", "str")], + id="non_list_aliases_are_dropped_with_a_warning", + ), + pytest.param( + {"m": {"aliases": []}}, + [("m", {}, None)], + [], + id="empty_aliases_are_dropped", + ), + pytest.param( + {"m": {"aliases": ["sample_spec", "n", "m"]}, "n": {"aliases": ["n"]}, "sample_spec": {}}, + [("m", {}, None), ("n", {}, None), ("sample_spec", {}, None)], + [ + call(_SHADOWS, "sample_spec", "m"), + call(_SHADOWS, "n", "m"), + call(_SHADOWS, "m", "m"), + call(_SHADOWS, "n", "n"), + ], + id="an_alias_equal_to_any_key_is_skipped", + ), + pytest.param( + {"m": {"aliases": ["x", "x"]}, "n": {"aliases": ["x"]}}, + [("m", {}, ("m", "x")), ("n", {}, None), ("x", {}, ("m", "x"))], + [call(_TAKEN, "x", "m"), call(_TAKEN, "x", "n")], + id="the_first_claim_wins", + ), + pytest.param( + {"m": {"aliases": [1]}, "n": {"aliases": [True, 1.0]}}, + [("m", {}, ("m", 1)), ("n", {}, None), (1, {}, ("m", 1))], + [call(_TAKEN, True, "n"), call(_TAKEN, 1.0, "n")], + id="aliases_collide_under_python_equality", + ), + ], + ) + def test_expansion( + self, + model_cost: dict[str, dict[str, object]], + served: list[tuple[object, object, tuple[object, ...] | None]], + warnings: list[object], + ): + with patch.object(verbose_logger, "warning") as warn: + result = _expand_model_aliases(model_cost) + + assert result is model_cost + assert _served(result) == served + assert warn.call_args_list == warnings + + def test_each_alias_group_shares_one_copied_row(self): + row: dict[str, object] = {"aliases": ["m1", "m2"], "nested": [1]} + result = _expand_model_aliases({"m": row}) + + assert result["m"] is result["m1"] is result["m2"] + assert result["m"] is not row + assert result["m"]["nested"] is row["nested"] + + @pytest.mark.parametrize( + ("model_cost", "error", "message", "warnings"), + [ + pytest.param( + {"m": {"aliases": "x"}, "bad": None, "n": {"aliases": "y"}}, + AttributeError, + "'NoneType' object has no attribute 'get'", + [call(_NOT_A_LIST, "m", "str")], + id="a_row_that_is_not_a_dict", + ), + pytest.param( + {"m": {"aliases": "x"}, "n": {"aliases": ["ok", ["unhashable"]]}}, + TypeError, + "unhashable type: 'list'", + [call(_NOT_A_LIST, "m", "str")], + id="an_unhashable_alias", + ), + ], + ) + def test_main_raises_mid_scan_after_its_earlier_warnings_and_leaves_the_map( + self, + model_cost: dict[str, dict[str, object]], + error: type[Exception], + message: str, + warnings: list[object], + ): + before = repr(model_cost) + with patch.object(verbose_logger, "warning") as warn, pytest.raises(error, match=f"^{message}$"): + _expand_model_aliases(model_cost) + + assert warn.call_args_list == warnings + assert repr(model_cost) == before diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 931fabd85ad..ce497d8afab 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -4662,3 +4662,35 @@ def test_pricing_provider_fills_unset_rates_from_that_providers_catalog_row() -> finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() + + +def test_a_map_adopted_while_a_router_registers_its_deployments_keeps_its_rows(monkeypatch: pytest.MonkeyPatch) -> None: + """The refresh ``import litellm`` starts adopts in its own thread, often while the proxy is building its router""" + from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map + from litellm.router_python_core import PythonCore + + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + fetched: Final = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} + adoption: Final = threading.Thread(target=adopt_model_cost_map, args=(fetched,)) + register: Final = PythonCore._register_deployment_in_model_cost + + def register_then_race_an_adoption(**kwargs: object) -> None: + register(**kwargs) + if not adoption.is_alive() and adoption.ident is None: + adoption.start() + adoption.join(timeout=2) + + monkeypatch.setattr(PythonCore, "_register_deployment_in_model_cost", staticmethod(register_then_race_an_adoption)) + Router( + model_list=[ + { + "model_name": "racing-alias", + "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, + "model_info": {"id": "racing-alias-id", "max_input_tokens": 4321}, + } + ], + ) + adoption.join(timeout=10) + + assert litellm.model_cost is fetched + assert litellm.model_cost["racing-alias-id"]["max_input_tokens"] == 4321 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 32063f65765..7d4fdc3b089 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -18,6 +18,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone from functools import partial from pathlib import Path, PurePath +from types import MappingProxyType from typing import Final, SupportsIndex, cast from unittest import mock from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2547,18 +2548,17 @@ def test_register_model_with_scientific_notation(): @pytest.mark.usefixtures("isolated_openai_model_sets") -@respx.mock -def test_register_model_url_fetch_uses_single_attempt(monkeypatch): +def test_register_model_url_fetch_uses_single_attempt(monkeypatch, scripted_http): 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)) + server = scripted_http([httpx.Response(503)], path="/custom_pricing.json") - litellm.register_model(model_cost="https://example.invalid/custom_pricing.json") + litellm.register_model(model_cost=server.url) threads_after = {thread.name for thread in threading.enumerate()} - assert route.call_count == 1 + assert server.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() @@ -5177,12 +5177,7 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey 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.utils import _invalidate_model_cost_lowercase_map, get_model_cost_key class _EntryRegisteredMidWalk(dict): def get(self, field: str, default: object = None) -> object: @@ -5190,11 +5185,10 @@ 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" + assert get_model_cost_key("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() @@ -5211,6 +5205,225 @@ class _KeyWhoseLoweringLosesTheRace(str): return super().lower() +@pytest.mark.parametrize( + ("existing", "new", "expected"), + [ + pytest.param({"a": 1}, {"a": None, "b": None}, {"a": 1}, id="none_is_skipped"), + pytest.param({}, {"input_cost_per_token": "3e-07"}, {"input_cost_per_token": 3e-07}, id="scientific_string"), + pytest.param({}, {"max_tokens": "8192"}, {"max_tokens": 8192}, id="integer_string"), + pytest.param({}, {"max_tokens": " 8_192 "}, {"max_tokens": 8192}, id="padded_integer_string"), + pytest.param({}, {"max_tokens": "8192."}, {"max_tokens": 8192.0}, id="point_makes_a_float"), + pytest.param({}, {"mode": "chat"}, {"mode": "chat"}, id="text_stays"), + pytest.param({}, {"x": "inf"}, {"x": "inf"}, id="infinity_without_e_or_point_stays"), + pytest.param({}, {"x": "1e400"}, {"x": float("inf")}, id="overflowing_float"), + pytest.param({}, {"x": "9" * 30}, {"x": int("9" * 30)}, id="big_int"), + pytest.param( + {"m": {"a": 1, "b": 2}}, + {"m": {"b": 3, "c": 4}}, + {"m": {"a": 1, "b": 3, "c": 4}}, + id="dict_merges_over_dict", + ), + pytest.param({"m": 1}, {"m": {"b": 3}}, {"m": {"b": 3}}, id="dict_replaces_a_scalar"), + pytest.param({}, {"m": {"b": "3"}}, {"m": {"b": "3"}}, id="nested_strings_are_kept"), + pytest.param({"z": 1, "a": 2}, {"m": 3, "z": 4}, {"z": 4, "a": 2, "m": 3}, id="existing_order_then_new"), + ], +) +def test_update_dictionary_matches_main(existing: dict, new: dict, expected: dict) -> None: + """The same cases as ``update_dictionary_matches_main`` in ``model-catalog/src/register.rs``""" + from litellm.utils import update_dictionary + + merged = update_dictionary(existing, new) + + assert merged is existing + assert list(merged.items()) == list(expected.items()) + assert [type(value) for value in merged.values()] == [type(value) for value in expected.values()] + + +def test_update_dictionary_keeps_the_identity_of_values_it_does_not_rewrite() -> None: + from litellm.utils import update_dictionary + + tiers = [{"input_cost_per_token": 1}] + merged = update_dictionary({}, {"tiered_pricing": tiers}) + + assert merged["tiered_pricing"] is tiers + + +def test_update_dictionary_refuses_a_field_name_that_is_not_a_str() -> None: + from litellm.utils import update_dictionary + + with pytest.raises(TypeError, match="field names must be str, got int"): + update_dictionary({}, {1: "a"}) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + pytest.param("3E-07", 3e-07, id="upper_exponent"), + pytest.param("-0", 0, id="negative_zero_int"), + pytest.param("1_0.5", 10.5, id="underscore_float"), + pytest.param("1__0", "1__0", id="double_underscore_stays"), + pytest.param("\u00a012\u3000", 12, id="unicode_whitespace"), + pytest.param("\x1c12", "\x1c12", id="file_separator_is_not_whitespace"), + pytest.param("0x10", "0x10", id="hex_stays"), + pytest.param(5, 5, id="non_string_passes_through"), + ], +) +def test_convert_stringified_numbers_matches_cpython(value: object, expected: object) -> None: + from litellm.utils import _convert_stringified_numbers + + converted = _convert_stringified_numbers(value) + + assert (converted, type(converted)) == (expected, type(expected)) + + +@pytest.mark.parametrize( + ("base", "registration", "expected"), + [ + pytest.param( + {"litellm_provider": None, "max_tokens": 10}, + {"input_cost_per_token": 1}, + {"max_tokens": 10, "input_cost_per_token": 1}, + id="provider_none_is_dropped", + ), + pytest.param( + {"explicit_cache_read_input_token_cost": 1, "a": 1}, + {"cache_read_input_token_cost": 2}, + {"explicit_cache_read_input_token_cost": None, "a": 1, "cache_read_input_token_cost": 2}, + id="cache_read_clears_explicit", + ), + pytest.param( + {"explicit_cache_read_input_token_cost": 1}, + {"cache_read_input_token_cost": None}, + {"explicit_cache_read_input_token_cost": 1}, + id="cache_read_none_clears_nothing", + ), + ], +) +def test_registration_row_matches_main(base: dict, registration: dict, expected: dict) -> None: + """The same cases as ``registration_row_matches_main`` in ``model-catalog/src/register.rs``""" + from litellm.rust_bridge._native import registration_row + + assert list(registration_row(base, registration).items()) == list(expected.items()) + + +@pytest.mark.parametrize( + ("entry", "registration", "expected"), + [ + pytest.param( + { + "input_cost_per_token": 1, + "input_cost_per_token_above_200k_tokens": 2, + "input_cost_per_token_cache_hit": 3, + "output_cost_per_token_priority": 4, + }, + {"input_cost_per_token": 9}, + {"input_cost_per_token_cache_hit": 3, "output_cost_per_token_priority": 4}, + id="input_variants_dropped_cache_kept", + ), + pytest.param( + {"output_cost_per_reasoning_token": 1, "output_cost_per_token_batches": 3}, + {"output_cost_per_token": 9}, + {}, + id="output_drops_reasoning", + ), + pytest.param( + {"tiered_pricing": ({"input_cost_per_token": 1, "range": [0, 10]},)}, + {"input_cost_per_token": "9"}, + {"tiered_pricing": [{"range": [0, 10], "input_cost_per_token": "9"}]}, + id="a_tuple_of_tiers_is_read_as_pydantic_reads_it", + ), + pytest.param( + {"tiered_pricing": [1, {"input_cost_per_token": 1}]}, + {"input_cost_per_token": 9}, + {"tiered_pricing": [1, {"input_cost_per_token": 1}]}, + id="invalid_tiers_are_kept", + ), + pytest.param( + {"tiered_pricing": [{1: "x"}]}, + {"input_cost_per_token": 9}, + {"tiered_pricing": [{1: "x"}]}, + id="a_tier_with_a_non_str_key_is_kept", + ), + pytest.param( + {"off_peak_pricing": MappingProxyType({"input_cost_per_token": 1, "hours": [1, 2]})}, + {"input_cost_per_token": 9}, + {"off_peak_pricing": {"hours": [1, 2]}}, + id="an_off_peak_mapping_drops_overridden", + ), + ], +) +def test_without_overridden_price_variants_matches_main(entry: dict, registration: dict, expected: dict) -> None: + """The same cases as ``without_overridden_price_variants_matches_main`` in ``model-catalog/src/register.rs``""" + from litellm.rust_bridge._native import without_overridden_price_variants + + assert without_overridden_price_variants(entry, registration) == expected + + +def test_a_key_added_under_another_case_is_found_without_an_invalidation(monkeypatch: pytest.MonkeyPatch) -> None: + """Main's lowercase map missed a key written straight into ``litellm.model_cost`` until the next invalidation""" + from litellm.utils import get_model_cost_key + + monkeypatch.setattr(litellm, "model_cost", {"gpt-4o": {"mode": "chat"}}) + assert get_model_cost_key("GPT-4O") == "gpt-4o" + + litellm.model_cost["LIT-Direct-Write"] = {"mode": "chat"} + + assert get_model_cost_key("lit-direct-write") == "LIT-Direct-Write" + + +def test_a_popped_key_is_never_returned(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import get_model_cost_key + + monkeypatch.setattr(litellm, "model_cost", {"Old-Name": {"mode": "chat"}, "gpt-4o": {"mode": "chat"}}) + assert get_model_cost_key("old-name") == "Old-Name" + + litellm.model_cost.pop("Old-Name") + litellm.model_cost["Other"] = {"mode": "chat"} + + assert get_model_cost_key("old-name") is None + + +def test_the_later_key_wins_when_two_keys_lowercase_alike(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import _invalidate_model_cost_lowercase_map, get_model_cost_key + + monkeypatch.setattr(litellm, "model_cost", {"Model-A": {}, "MODEL-a": {}}) + _invalidate_model_cost_lowercase_map() + + assert get_model_cost_key("model-a") == "MODEL-a" + + +def test_a_rebound_model_cost_is_indexed_again(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.utils import get_model_cost_key + + monkeypatch.setattr(litellm, "model_cost", {"First-Map": {}}) + assert get_model_cost_key("first-map") == "First-Map" + + monkeypatch.setattr(litellm, "model_cost", {"Second-Map": {}}) + + assert get_model_cost_key("second-map") == "Second-Map" + assert get_model_cost_key("first-map") is None + + +@pytest.mark.parametrize( + ("model_cost", "keys", "expected"), + [ + pytest.param({"a": {"mode": "chat"}}, ("a",), True, id="unpriced_row"), + pytest.param({"a": {"input_cost_per_token": 0}}, ("a",), False, id="declared_zero"), + pytest.param({"a": {"mode": "chat"}, "b": {"output_cost_per_token": 1}}, ("a", "b"), False, id="one_priced"), + pytest.param({}, ("a",), False, id="no_row"), + pytest.param({"a": "not a dict"}, ("a", None), False, id="no_dict_row"), + ], +) +def test_cost_map_omits_token_price_matches_main( + monkeypatch: pytest.MonkeyPatch, model_cost: dict, keys: tuple[object, ...], expected: bool +) -> None: + from litellm.utils import cost_map_omits_token_price + + monkeypatch.setattr(litellm, "model_cost", model_cost) + + assert cost_map_omits_token_price(*keys) is expected + + 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 @@ -5246,7 +5459,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: @@ -5254,14 +5467,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.utils import _invalidate_model_cost_lowercase_map 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,