From edc99fdefd26f3c9911429a7389fc7da9b252af2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 4 Nov 2026 22:51:58 +0900 Subject: [PATCH] refactor(model-cost): split S3 into storage-s3 and keep catalog rules in model-catalog The S3 client config, credentials and HTTP adapter move from cache-s3 to a new litellm-storage-s3 crate, which cache-s3 and model-catalog-refresh both build on The model name fold, the litellm_provider group field, the token price check and Provenance now live in model-catalog, so python-bridge only adapts Python objects. refresh and resume return Attempt, dropping Refreshed and the bridge's second mapping, and the unused RetryPolicy is deleted --- litellm-rust/Cargo.lock | 26 +++++++--- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/cache-s3/Cargo.toml | 9 +--- litellm-rust/crates/cache-s3/src/cache.rs | 30 +---------- litellm-rust/crates/cache-s3/src/lib.rs | 6 +-- .../crates/model-catalog-refresh/AGENTS.md | 5 ++ .../crates/model-catalog-refresh/Cargo.toml | 3 +- .../crates/model-catalog-refresh/src/lib.rs | 4 -- .../model-catalog-refresh/src/policy.rs | 20 ------- .../model-catalog-refresh/src/refresh.rs | 52 +++++-------------- .../crates/model-catalog-refresh/src/s3.rs | 6 +-- .../model-catalog-refresh/src/source.rs | 4 ++ .../model-catalog-refresh/src/source_fetch.rs | 2 +- litellm-rust/crates/model-catalog/AGENTS.md | 6 +++ .../crates/model-catalog/src/catalog.rs | 11 ++-- litellm-rust/crates/model-catalog/src/lib.rs | 2 + .../src/provenance.rs} | 0 .../crates/model-catalog/src/register.rs | 32 +++++++++++- .../crates/model-catalog/src/store/AGENTS.md | 2 +- .../python-bridge/src/model_registry/fetch.rs | 33 +++--------- .../python-bridge/src/model_registry/keys.rs | 7 ++- .../python-bridge/src/model_registry/live.rs | 32 +++++------- .../src/model_registry/provenance.rs | 2 +- litellm-rust/crates/storage-s3/AGENTS.md | 5 ++ litellm-rust/crates/storage-s3/Cargo.toml | 21 ++++++++ .../{cache-s3 => storage-s3}/src/auth.rs | 2 +- litellm-rust/crates/storage-s3/src/config.rs | 27 ++++++++++ litellm-rust/crates/storage-s3/src/lib.rs | 7 +++ .../{cache-s3 => storage-s3}/src/transport.rs | 0 .../{cache-s3 => storage-s3}/tests/auth.rs | 2 +- 30 files changed, 187 insertions(+), 172 deletions(-) create mode 100644 litellm-rust/crates/model-catalog-refresh/AGENTS.md delete mode 100644 litellm-rust/crates/model-catalog-refresh/src/policy.rs rename litellm-rust/crates/{model-catalog-refresh/src/installed.rs => model-catalog/src/provenance.rs} (100%) create mode 100644 litellm-rust/crates/storage-s3/AGENTS.md create mode 100644 litellm-rust/crates/storage-s3/Cargo.toml rename litellm-rust/crates/{cache-s3 => storage-s3}/src/auth.rs (97%) create mode 100644 litellm-rust/crates/storage-s3/src/config.rs create mode 100644 litellm-rust/crates/storage-s3/src/lib.rs rename litellm-rust/crates/{cache-s3 => storage-s3}/src/transport.rs (100%) rename litellm-rust/crates/{cache-s3 => storage-s3}/tests/auth.rs (97%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9607fb0afb4..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", @@ -4269,11 +4264,10 @@ version = "0.1.0" dependencies = [ "http 1.4.2", "litellm-auth-aws", - "litellm-cache-s3", "litellm-error", "litellm-http", - "litellm-model-catalog", "litellm-python-compat", + "litellm-storage-s3", "reqwest", "rstest", "tempfile", @@ -4624,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 c1d582a2286..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, } @@ -31,27 +26,6 @@ pub struct S3CacheConfig { pub auth: AwsAuthConfig, } -/// 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, - } -} - pub struct S3Cache { client: aws_sdk_s3::Client, codec: C, diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs index 9b413509f72..43006532441 100644 --- a/litellm-rust/crates/cache-s3/src/lib.rs +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -1,7 +1,3 @@ -mod auth; mod cache; -mod transport; -pub use auth::S3Credentials; -pub use aws_sdk_s3; -pub use cache::{S3Cache, S3CacheConfig, S3Endpoint, s3_config}; +pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; 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 a93df0094e7..9703fc36f56 100644 --- a/litellm-rust/crates/model-catalog-refresh/Cargo.toml +++ b/litellm-rust/crates/model-catalog-refresh/Cargo.toml @@ -8,10 +8,9 @@ repository.workspace = true [dependencies] http.workspace = true litellm-auth-aws.workspace = true -litellm-cache-s3.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"] } diff --git a/litellm-rust/crates/model-catalog-refresh/src/lib.rs b/litellm-rust/crates/model-catalog-refresh/src/lib.rs index d835a13add0..a072d07496b 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/lib.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/lib.rs @@ -2,8 +2,6 @@ mod error; mod fetch; mod file; mod http; -mod installed; -mod policy; mod refresh; mod retry; mod s3; @@ -15,8 +13,6 @@ pub use error::Error; pub use fetch::*; pub use file::*; pub use http::*; -pub use installed::*; -pub use policy::*; pub use refresh::*; pub use retry::*; pub use s3::*; 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 index 2b825ce6239..52af9e3e902 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/refresh.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/refresh.rs @@ -2,21 +2,6 @@ use crate::retry::{AfterRetryable, after_retryable}; use crate::status::{Attempt, classify}; use crate::{Fetch, FetchRequest}; -/// How a refresh ended -#[derive(Clone, Debug, PartialEq)] -pub enum Refreshed { - Adopted { - body: Vec, - /// `git_blob_id` of `body` - revision: String, - etag: Option, - }, - NotModified, - Failed { - reason: String, - }, -} - /// A retryable failure about to be waited out #[derive(Clone, Debug, PartialEq)] pub struct Retrying<'a> { @@ -37,15 +22,15 @@ pub trait Pause: Send { fn sleep(&mut self, seconds: f64) -> impl Future> + Send; } -/// Up to `max_attempts` attempts, as main's `_fetch_remote_model_cost_map_with_retry` +/// 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 { +) -> Result { if max_attempts < 1 { - return Ok(Refreshed::Failed { + return Ok(Attempt::Failed { reason: "model cost map fetch failed".to_owned(), }); } @@ -58,7 +43,7 @@ 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` +/// Carries on after `outcome`, the result of attempt number `attempt`. Never `Attempt::Retryable` pub async fn resume( fetch: &impl Fetch, request: &FetchRequest, @@ -66,28 +51,20 @@ pub async fn resume( max_attempts: i64, outcome: Attempt, pause: &mut P, -) -> Result { +) -> Result { let (mut attempt_number, mut outcome) = (attempt, outcome); loop { let (reason, retry_after) = match outcome { - Attempt::Fetched { body, etag } => { - return Ok(Refreshed::Adopted { - revision: litellm_model_catalog::revision(&body), - body, - etag, - }); - } - Attempt::NotModified => return Ok(Refreshed::NotModified), - Attempt::Failed { reason } => return Ok(Refreshed::Failed { reason }), 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(Refreshed::Failed { + return Ok(Attempt::Failed { reason: format!("{reason} (after {max_attempts} attempts)"), }); } @@ -216,15 +193,14 @@ mod tests { #[rstest] #[tokio::test] - async fn retries_then_adopts_with_the_revision_of_the_body() { + 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, - Refreshed::Adopted { + Attempt::Fetched { body: b"hello\n".to_vec(), - revision: "ce013625030ba8dba906f756967f9e9ca394464a".to_owned(), etag: Some("\"v2\"".to_owned()), } ); @@ -247,7 +223,7 @@ mod tests { let refreshed = refresh(&fetch, &request(), 3, &mut pause).await.unwrap(); assert_eq!( refreshed, - Refreshed::Failed { + Attempt::Failed { reason: format!("HTTP 429 from {URL} (after 3 attempts)"), } ); @@ -263,7 +239,7 @@ mod tests { let refreshed = refresh(&fetch, &request(), 3, &mut pause).await.unwrap(); assert_eq!( refreshed, - Refreshed::Failed { + Attempt::Failed { reason: format!("HTTP 404 from {URL}"), } ); @@ -281,7 +257,7 @@ mod tests { .unwrap(); assert_eq!( refreshed, - Refreshed::Failed { + Attempt::Failed { reason: "model cost map fetch failed".to_owned(), } ); @@ -300,7 +276,7 @@ mod tests { let refreshed = resume(&fetch, &request(), 1, 3, first, &mut pause) .await .unwrap(); - assert!(matches!(refreshed, Refreshed::Adopted { .. })); + assert!(matches!(refreshed, Attempt::Fetched { .. })); assert_eq!( pause.retried, vec![(1, 3, "HTTP 503 from x".to_owned(), 2.0)] @@ -345,7 +321,7 @@ mod tests { let refreshed = refresh(&fetch, &request, 3, &mut Recording::default()) .await .unwrap(); - assert_eq!(refreshed, Refreshed::NotModified); + 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/s3.rs b/litellm-rust/crates/model-catalog-refresh/src/s3.rs index 0123c37e575..ff6407e97ab 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/s3.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/s3.rs @@ -1,4 +1,4 @@ -use litellm_cache_s3::aws_sdk_s3; +use litellm_storage_s3::aws_sdk_s3; use crate::source::S3Object; use crate::{Error, Fetch, FetchRequest, Reply}; @@ -12,7 +12,7 @@ pub struct S3Fetch { } impl S3Fetch { - /// `config` is built by `litellm_cache_s3::s3_config`, and retries are left to the refresh + /// `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 { @@ -79,7 +79,7 @@ mod tests { ) .mount(&server) .await; - let config = litellm_cache_s3::s3_config( + let config = litellm_storage_s3::s3_config( "us-east-1".to_owned(), Some(server.uri()), litellm_http::Client::for_test(), diff --git a/litellm-rust/crates/model-catalog-refresh/src/source.rs b/litellm-rust/crates/model-catalog-refresh/src/source.rs index d56e936e55f..b6a082dbed7 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/source.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/source.rs @@ -4,6 +4,10 @@ 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 { diff --git a/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs b/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs index 5d17164e751..ba59bc90b11 100644 --- a/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs +++ b/litellm-rust/crates/model-catalog-refresh/src/source_fetch.rs @@ -32,7 +32,7 @@ impl SourceFetch { region_name: Some(region.clone()), ..AwsAuthConfig::default() }; - let config = litellm_cache_s3::s3_config( + let config = litellm_storage_s3::s3_config( region, s3_endpoint(clients.env), clients.http, diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index 8e67a9a2843..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 @@ -48,6 +52,8 @@ Every price in a row lowers to one `Price { meter, channel, when }` and its `Uni `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/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index b4c82abd3f9..0d0015ae16c 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -81,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 {} @@ -105,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); } } diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index 4b8aa5986bf..6df27f94411 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -44,6 +44,7 @@ mod integrity; mod lower; mod model; mod model_info; +mod provenance; mod register; mod resolve; mod store; @@ -58,6 +59,7 @@ pub use fallback::*; 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-refresh/src/installed.rs b/litellm-rust/crates/model-catalog/src/provenance.rs similarity index 100% rename from litellm-rust/crates/model-catalog-refresh/src/installed.rs rename to litellm-rust/crates/model-catalog/src/provenance.rs diff --git a/litellm-rust/crates/model-catalog/src/register.rs b/litellm-rust/crates/model-catalog/src/register.rs index a6c93dd8ea7..6b2694ffd90 100644 --- a/litellm-rust/crates/model-catalog/src/register.rs +++ b/litellm-rust/crates/model-catalog/src/register.rs @@ -10,7 +10,15 @@ 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"; -const LITELLM_PROVIDER: &str = "litellm_provider"; +/// 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)) +} /// How the merge rule sees one field value #[derive(Clone, Debug, PartialEq)] @@ -467,4 +475,26 @@ mod tests { 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/src/model_registry/fetch.rs b/litellm-rust/crates/python-bridge/src/model_registry/fetch.rs index c98eea3f715..07995365430 100644 --- a/litellm-rust/crates/python-bridge/src/model_registry/fetch.rs +++ b/litellm-rust/crates/python-bridge/src/model_registry/fetch.rs @@ -2,8 +2,8 @@ use std::time::Duration; use litellm_http::ClientVariant; use litellm_model_catalog_refresh::{ - Attempt, FetchRequest, Pause, Refreshed, Retrying, SourceClients, SourceFetch, attempt, - refresh, refused, resume, + Attempt, FetchRequest, Pause, Retrying, SourceClients, SourceFetch, attempt, refresh, refused, + resume, }; use pyo3::intern; use pyo3::prelude::*; @@ -22,7 +22,7 @@ pub(crate) enum ModelCostMapFetch { Failed { reason: String, }, - /// Only from `fetch_model_cost_map_once`, to resume with `refresh_model_cost_map` + /// Only from `attempt`, to resume with `run` Retryable { reason: String, retry_after: Option, @@ -123,7 +123,7 @@ impl PyRefresh { let refreshed = litellm_host_python::run_sync_value(py, async move { run.refresh(first, &mut pause).await })?; - Python::attach(|py| Ok(from_refreshed(py, refreshed))) + 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 @@ -132,7 +132,7 @@ impl PyRefresh { 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_refreshed(py, refreshed))) + Python::attach(|py| Ok(from_attempt(py, refreshed))) } /// `run` as an awaitable, where `sleep` returns an awaitable @@ -141,7 +141,7 @@ impl PyRefresh { 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_refreshed(py, refreshed))) + Ok(Python::attach(|py| from_attempt(py, refreshed))) }) } } @@ -189,11 +189,10 @@ struct Run { } impl Run { - async fn refresh(self, first: Option, pause: &mut PythonPause) -> PyResult { + async fn refresh(self, first: Option, pause: &mut PythonPause) -> PyResult { let fetch = match self.source { Ok(fetch) => fetch, - Err(Attempt::Failed { reason }) => return Ok(Refreshed::Failed { reason }), - Err(_) => unreachable!("a source that cannot be built is a final failure"), + Err(failed) => return Ok(failed), }; match first { Some(first) => resume(&fetch, &self.request, 1, self.max_attempts, first, pause).await, @@ -248,22 +247,6 @@ fn to_attempt(py: Python<'_>, fetch: &ModelCostMapFetch) -> Attempt { } } -fn from_refreshed(py: Python<'_>, refreshed: Refreshed) -> ModelCostMapFetch { - match refreshed { - Refreshed::Adopted { - body, - revision, - etag, - } => ModelCostMapFetch::Adopted { - body: PyBytes::new(py, &body).unbind(), - revision, - etag, - }, - Refreshed::NotModified => ModelCostMapFetch::NotModified {}, - Refreshed::Failed { reason } => ModelCostMapFetch::Failed { reason }, - } -} - /// 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 { diff --git a/litellm-rust/crates/python-bridge/src/model_registry/keys.rs b/litellm-rust/crates/python-bridge/src/model_registry/keys.rs index d4d331dd149..03cb69556cb 100644 --- a/litellm-rust/crates/python-bridge/src/model_registry/keys.rs +++ b/litellm-rust/crates/python-bridge/src/model_registry/keys.rs @@ -1,8 +1,7 @@ use std::borrow::Cow; use std::sync::LazyLock; -use litellm_model_catalog::{Rows, Store}; -use pyo3::intern; +use litellm_model_catalog::{LITELLM_PROVIDER, Rows, Store, fold_model_name}; use pyo3::prelude::*; use pyo3::types::{PyDict, PyString}; @@ -17,7 +16,7 @@ impl Rows for LiveModelCost { type Stamp = Option<(usize, usize)>; fn fold(name: &str) -> String { - name.to_lowercase() + fold_model_name(name) } fn stamp(&self) -> Self::Stamp { @@ -63,7 +62,7 @@ impl Rows for LiveModelCost { fn provider(row: &Bound<'_, PyAny>) -> Option { let row = row.cast::().ok()?; - let provider = row.get_item(intern!(row.py(), "litellm_provider")).ok()??; + let provider = row.get_item(LITELLM_PROVIDER).ok()??; provider.extract().ok() } diff --git a/litellm-rust/crates/python-bridge/src/model_registry/live.rs b/litellm-rust/crates/python-bridge/src/model_registry/live.rs index ff50fc102ae..c19b4f90b08 100644 --- a/litellm-rust/crates/python-bridge/src/model_registry/live.rs +++ b/litellm-rust/crates/python-bridge/src/model_registry/live.rs @@ -4,7 +4,6 @@ use pyo3::sync::PyOnceLock; use pyo3::types::{PyDict, PyString, PyTuple}; const MODULE: &str = "litellm.rust_bridge.host.model_cost"; -const TOKEN_PRICES: [&str; 2] = ["input_cost_per_token", "output_cost_per_token"]; /// The bound `litellm.model_cost` pub(super) fn current(py: Python<'_>) -> PyResult> { @@ -16,27 +15,22 @@ pub(super) fn current(py: Python<'_>) -> PyResult> { .map_err(PyErr::from) } -/// Whether some `litellm.model_cost` rows under the `str` keys are dicts and none of them carries a per-token price. -/// `get_model_info` fills a missing price with `0`, so this tells an unpriced row from a declared zero +/// `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 mut found_row = false; - for key in keys.iter().filter(|key| key.is_instance_of::()) { - let Some(row) = model_cost.get_item(&key)? else { - continue; - }; - let Ok(row) = row.cast_into::() else { - continue; - }; - for price in TOKEN_PRICES { - if row.contains(price)? { - return Ok(false); - } - } - found_row = true; - } - Ok(found_row) + 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` diff --git a/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs b/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs index cf27cec0fdf..607428017b3 100644 --- a/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs +++ b/litellm-rust/crates/python-bridge/src/model_registry/provenance.rs @@ -1,7 +1,7 @@ use std::sync::{LazyLock, PoisonError, RwLock}; use std::time::{SystemTime, UNIX_EPOCH}; -use litellm_model_catalog_refresh::{Origin, Provenance}; +use litellm_model_catalog::{Origin, Provenance}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; 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 {