scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 175

Merge remote-tracking branch 'origin/main' into litellm_scim_user_teams_race
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Details
LiteLLM Rust DB / db-queries (push) Has been cancelled
Details
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details

Browse source
This commit is contained in:
Jesus Velarde 2026-11-10 13:51:44 -08:00
commit 7f605b72b1
582 changed files with 38142 additions and 5351 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -20,6 +20,7 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
from litellm.constants import (
BATCH_LINE_ITEMS_MAX_PART_BYTES,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
@ -765,7 +766,7 @@ class CheckBatchCost:
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
_file_content = await afile_content(
file_id=cast(str, raw_output_file_id),
**cast(FileContentCallOptions, credentials),
**cast(FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(cast(str, raw_output_file_id))}),
)
content_bytes: Final = await _file_content_bytes(_file_content)

66
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -49,6 +49,7 @@ from litellm.constants import (
)
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.internal_call_metadata import requests_background
@ -78,6 +79,7 @@ from litellm.proxy.litellm_pre_call_utils import (
from litellm.proxy.openai_files_endpoints.common_utils import (
BATCH_CREATE_HIDDEN_PARAM,
FILE_LIST_CONTINUATION_CHUNK_SIZE,
TeamProviderDeployment,
_is_base64_encoded_unified_file_id,
apply_unified_file_ids,
decode_model_from_file_id,
@ -395,6 +397,7 @@ class _CursorPageArgs(TypedDict, total=False):
class _RouterFileCallKwargs(TypedDict, total=False):
client: ReadOnly[AsyncOpenAI | None]
custom_llm_provider: ReadOnly[str | None]
stored_s3_bucket_name: ReadOnly[str]
_MODEL_FILE_ID_MAPPING: Final = TypeAdapter(dict[str, str])
@ -462,7 +465,7 @@ async def _delete_provider_file(
await llm_router.afile_delete(
model=model_id,
file_id=provider_file_id,
**cast(_RouterFileCallKwargs, delete_data),
**cast(_RouterFileCallKwargs, {**delete_data, **stored_s3_bucket_kwargs(provider_file_id)}),
)
@ -1156,7 +1159,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
retrieve_file: Final = cast(_ManagedFileRetrieve, litellm.afile_retrieve)
retrieve_options: Final = cast(
FileRetrieveCallOptions,
{**call_options, "max_retries": 0},
{**call_options, **stored_s3_bucket_kwargs(provider_file_id), "max_retries": 0},
)
for attempt in range(len(_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS) + 1):
if attempt > 0:
@ -1426,6 +1429,43 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict=owner,
)
async def register_provider_file(
self,
*,
provider_file_id: str,
content_type: str,
deployment: TeamProviderDeployment,
owner: UserAPIKeyAuth,
llm_router: Router | None,
) -> str | None:
unified_file_id: Final = self.get_unified_output_file_id(
provider_file_id, deployment.deployment_id, deployment.public_model_name, file_type=content_type
)
provider_file_object, _ = await self._fetch_provider_file_object(
unified_file_id=unified_file_id,
provider_file_id=provider_file_id,
model_id=deployment.deployment_id,
model_name=None,
llm_router=llm_router,
)
if provider_file_object is None:
return None
try:
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=provider_file_object,
litellm_parent_otel_span=None,
model_mappings={deployment.deployment_id: provider_file_id},
user_api_key_dict=owner,
)
except Exception as error: # noqa: BLE001 # the upload succeeded, so the caller keeps the provider id
verbose_logger.warning(
f"Failed to register provider_file_id={sanitize_for_log(provider_file_id)} as a managed file: "
f"{type(error).__name__} {sanitize_for_log(error)}"
)
return None
return unified_file_id
async def list_user_batches(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -1757,8 +1797,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if potential_file_id and "llm_output_file_id," in potential_file_id:
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
if model_id:
provider_file_id: Final = self.get_output_file_id_from_unified_file_id(potential_file_id)
data["model"] = model_id
data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id)
data["file_id"] = provider_file_id
data.update(stored_s3_bucket_kwargs(provider_file_id))
elif call_type == CallTypes.acreate_batch.value:
input_file_id = cast(Optional[str], data.get("input_file_id"))
if input_file_id:
@ -2253,10 +2295,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=")
def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str:
def get_unified_output_file_id(
self, output_file_id: str, model_id: str, model_name: Optional[str], file_type: str = "application/json"
) -> str:
deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id)
unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
"application/json",
file_type,
str(deterministic_uuid),
model_name or "",
output_file_id,
@ -2854,7 +2898,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
await llm_router.afile_content(
model=model_id,
file_id=provider_file_id,
**cast(_RouterFileCallKwargs, data),
**cast(
_RouterFileCallKwargs,
{
**data,
**(
stored_s3_bucket_kwargs(provider_file_id)
if isinstance(provider_file_id, str)
else {}
),
},
),
),
)
except Exception as e:

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql Normal file
View file

@ -0,0 +1,9 @@
SET LOCAL lock_timeout = '10s';
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundResponseSettlement" (
"response_id" TEXT NOT NULL,
"settled_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_BackgroundResponseSettlement_pkey" PRIMARY KEY ("response_id")
);

7
litellm-proxy-extras/litellm_proxy_extras/migrations/20261110000000_add_budget_model_tpm_rpm_limit/migration.sql Normal file
View file

@ -0,0 +1,7 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_tpm_limit" JSONB;
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "model_rpm_limit" JSONB;

22
litellm-proxy-extras/litellm_proxy_extras/migrations/20261110120000_add_daily_spend_guardrail_spend/migration.sql Normal file
View file

@ -0,0 +1,22 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyGlobalSpend" ADD COLUMN IF NOT EXISTS "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

14
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -19,6 +19,8 @@ model LiteLLM_BudgetTable {
rpm_limit BigInt?
tpd_limit BigInt?
model_max_budget Json?
model_tpm_limit Json?
model_rpm_limit Json?
budget_duration String?
budget_reset_at DateTime?
allowed_models String[] @default([]) // per-member model scope; empty = inherit team models
@ -881,6 +883,7 @@ model LiteLLM_DailyUserSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -917,6 +920,7 @@ model LiteLLM_DailyGlobalSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -950,6 +954,7 @@ model LiteLLM_DailyOrganizationSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -988,6 +993,7 @@ model LiteLLM_DailyEndUserSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -1025,6 +1031,7 @@ model LiteLLM_DailyAgentSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -1062,6 +1069,7 @@ model LiteLLM_DailyTeamSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -1102,6 +1110,7 @@ model LiteLLM_DailyTagSpend {
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
guardrail_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -2053,6 +2062,11 @@ model LiteLLM_BackgroundInteractionSettlement {
@@index([settled_at], map: "idx_background_interaction_settlement_settled_at")
}
model LiteLLM_BackgroundResponseSettlement {
response_id String @id
settled_at DateTime @default(now())
}
model LiteLLM_Lens {
reviews LiteLLM_LensReview[]
id String @id

33
litellm-rust/Cargo.lock generated
View file

@ -894,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
"tokio",
]
[[package]]
@ -3657,19 +3658,15 @@ 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-s3",
"rstest",
"serde_json",
"tokio",
@ -4346,12 +4343,14 @@ dependencies = [
"jsonschema",
"litellm-cost",
"litellm-error",
"litellm-tracing",
"macro_rules_attribute",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"serde_with",
"sha1 0.10.7",
"strum",
"thiserror 2.0.19",
"time",
@ -4361,9 +4360,13 @@ dependencies = [
name = "litellm-model-catalog-refresh"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"backon",
"bytes",
"litellm-error",
"litellm-http",
"litellm-model-catalog",
"litellm-s3",
"litellm-tracing",
"reqwest",
"rstest",
"thiserror 2.0.19",
@ -4402,6 +4405,7 @@ dependencies = [
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-model-catalog-refresh",
"litellm-python-compat",
"litellm-redis-client",
"litellm-redis-connect",
@ -4609,6 +4613,23 @@ dependencies = [
"serde_with",
]
[[package]]
name = "litellm-s3"
version = "0.1.0"
dependencies = [
"aws-config",
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-runtime-api",
"aws-smithy-types",
"aws-types",
"http 1.4.2",
"litellm-http",
"reqwest",
"rstest",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-secret-value"
version = "0.1.0"

3
litellm-rust/Cargo.toml
View file

@ -60,6 +60,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-s3 = { path = "crates/s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-response = { path = "crates/cache-response" }
@ -90,6 +91,8 @@ axum = { version = "0.8.9", default-features = false, features = ["http1", "toki
axum-login = "0.18.0"
tower-sessions = { version = "0.14.0", features = ["memory-store"] }
bytes = "1"
backon = { version = "1.6.0", default-features = false, features = ["std", "tokio-sleep"] }
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rt-tokio"] }
http = "1"
google-cloud-auth = { version = "1.16.0", default-features = false }
jsonwebtoken = { version = "11.1.0", default-features = false }

2
litellm-rust/crates/auth-aws/src/lib.rs
View file

@ -1,9 +1,11 @@
mod aws;
pub mod constants;
mod error;
mod provider;
mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::{CredentialFlow, Error, SigningStage};
pub use provider::AwsCredentialsProvider;
pub use signer::SigV4Signer;

135
litellm-rust/crates/auth-aws/src/provider.rs Normal file
View file

@ -0,0 +1,135 @@
use std::sync::Arc;
use aws_credential_types::Credentials;
use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
use crate::{AwsAuthConfig, AwsAuthService};
type Environment = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
#[derive(Clone)]
pub struct AwsCredentialsProvider {
auth: AwsAuthService,
config: Box<AwsAuthConfig>,
environment: Environment,
}
impl AwsCredentialsProvider {
/// Explicit keys in `config` as given, otherwise LiteLLM's `AWS_*` settings read through `environment`
pub fn configured(
auth: AwsAuthService,
config: AwsAuthConfig,
environment: impl Fn(&str) -> Option<String> + Send + Sync + 'static,
) -> Self {
Self {
auth,
config: Box::new(config),
environment: Arc::new(environment),
}
}
async fn credentials(&self) -> Result<Credentials, crate::Error> {
match (&self.config.access_key_id, &self.config.secret_access_key) {
(Some(access_key_id), Some(secret_access_key)) => Ok(Credentials::new(
access_key_id.expose(),
secret_access_key.expose(),
self.config
.session_token
.as_ref()
.map(|token| token.expose().to_owned()),
None,
"litellm-configured",
)),
_ => {
self.auth
.resolve_credentials(
AwsAuthConfig::clone(&self.config),
self.environment.as_ref(),
)
.await
}
}
}
}
impl ProvideCredentials for AwsCredentialsProvider {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
self.credentials()
.await
.map_err(CredentialsError::provider_error)
})
}
}
impl std::fmt::Debug for AwsCredentialsProvider {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("AwsCredentialsProvider")
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
session_token: session_token.map(Into::into),
region_name: Some("us-east-1".to_string()),
..Default::default()
}
}
fn ambient_token(name: &str) -> Option<String> {
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
}
fn environment_keys(name: &str) -> Option<String> {
match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
}
}
#[rstest]
#[case::explicit_keys_ignore_an_ambient_session_token(
explicit(None), ambient_token, ("key", "secret", None)
)]
#[case::explicit_keys_keep_their_session_token(
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
)]
#[case::environment_keys_resolve_with_their_session_token(
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
)]
#[tokio::test]
async fn configured_credentials_read_the_environment_they_are_given(
#[case] config: AwsAuthConfig,
#[case] environment: fn(&str) -> Option<String>,
#[case] expected: (&str, &str, Option<&str>),
) {
let credentials =
AwsCredentialsProvider::configured(AwsAuthService::default(), config, environment)
.provide_credentials()
.await
.unwrap();
assert_eq!(
(
credentials.access_key_id(),
credentials.secret_access_key(),
credentials.session_token(),
),
expected
);
}
}

10
litellm-rust/crates/cache-s3/Cargo.toml
View file

@ -9,14 +9,10 @@ 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-s3.workspace = true
aws-sdk-s3.workspace = true
aws-smithy-types = "1.6.0"
futures-util.workspace = true
http.workspace = true
reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]

61
litellm-rust/crates/cache-s3/src/auth.rs
View file

@ -1,61 +0,0 @@
use aws_credential_types::{
Credentials as AwsCredentials,
provider::{ProvideCredentials, error::CredentialsError, future},
};
use litellm_auth_aws::{AwsAuthConfig, AwsAuthService};
#[derive(Clone)]
pub struct S3Credentials {
auth: AwsAuthService,
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl S3Credentials {
pub fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self {
auth: AwsAuthService::default(),
config,
env,
}
}
}
impl ProvideCredentials for S3Credentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::new(async {
if let (Some(access_key_id), Some(secret_access_key)) = (
self.config.access_key_id.clone(),
self.config.secret_access_key.clone(),
) {
return Ok(AwsCredentials::new(
access_key_id.expose(),
secret_access_key.expose(),
self.config
.session_token
.as_ref()
.map(|value| value.expose().to_owned()),
None,
"litellm-s3-cache",
));
}
self.auth
.resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(CredentialsError::provider_error)
})
}
}
impl std::fmt::Debug for S3Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("S3Credentials").finish_non_exhaustive()
}
}

27
litellm-rust/crates/cache-s3/src/cache.rs
View file

@ -5,20 +5,18 @@ use std::{
};
use aws_sdk_s3::{
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
config::{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_auth_aws::AwsCredentialsProvider;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::{auth::S3Credentials, transport::ReqwestHttpClient};
pub struct S3Endpoint {
pub url: String,
}
@ -28,7 +26,7 @@ pub struct S3CacheConfig {
pub key_prefix: String,
pub region: String,
pub endpoint: Option<S3Endpoint>,
pub auth: AwsAuthConfig,
pub credentials: AwsCredentialsProvider,
}
pub struct S3Cache<C: CacheCodec> {
@ -49,17 +47,14 @@ impl<C: CacheCodec> S3Cache<C> {
runtime: Handle,
) -> Self {
let endpoint_url: Option<String> = 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 = litellm_s3::config(
config.region.clone(),
endpoint_url.as_deref(),
http,
config.credentials,
)
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
Self {
client: aws_sdk_s3::Client::from_conf(builder.build()),
codec,

3
litellm-rust/crates/cache-s3/src/lib.rs
View file

@ -1,6 +1,3 @@
mod auth;
mod cache;
mod transport;
pub use auth::S3Credentials;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

57
litellm-rust/crates/cache-s3/tests/auth.rs
View file

@ -1,57 +0,0 @@
use aws_credential_types::provider::ProvideCredentials;
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::S3Credentials;
use rstest::rstest;
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
session_token: session_token.map(Into::into),
region_name: Some("us-east-1".to_string()),
..Default::default()
}
}
fn ambient_token(name: &str) -> Option<String> {
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
}
fn environment_keys(name: &str) -> Option<String> {
match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
}
}
#[rstest]
#[case::explicit_keys_ignore_an_ambient_session_token(
explicit(None), ambient_token, ("key", "secret", None)
)]
#[case::explicit_keys_keep_their_session_token(
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
)]
#[case::environment_keys_resolve_with_their_session_token(
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
)]
#[tokio::test]
async fn credentials_resolve(
#[case] config: AwsAuthConfig,
#[case] env: fn(&str) -> Option<String>,
#[case] expected: (&str, &str, Option<&str>),
) {
let credentials = S3Credentials::with_env(config, env)
.provide_credentials()
.await
.unwrap();
assert_eq!(
(
credentials.access_key_id(),
credentials.secret_access_key(),
credentials.session_token(),
),
expected
);
}

18
litellm-rust/crates/cache-s3/tests/support/mod.rs
View file

@ -5,7 +5,7 @@ use std::{
sync::{Arc, Mutex},
};
use litellm_auth_aws::AwsAuthConfig;
use litellm_auth_aws::{AwsAuthConfig, AwsAuthService, AwsCredentialsProvider};
use litellm_cache::JsonCodec;
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::Value;
@ -22,12 +22,16 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
endpoint: Some(S3Endpoint {
url: endpoint.to_string(),
}),
auth: AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
credentials: AwsCredentialsProvider::configured(
AwsAuthService::default(),
AwsAuthConfig {
access_key_id: Some("key".into()),
secret_access_key: Some("secret".into()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|_| None,
),
}
}

7
litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

@ -27,6 +27,10 @@
"api_key",
"additional_args"
],
"messages_input": [
"body",
"kwargs"
],
"post_call": [
"logger",
"original_response",
@ -126,7 +130,8 @@
"chunks",
"start",
"end",
"first_chunk"
"first_chunk",
"stream_teardown"
],
"stream_failure": [
"logger",

215
litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -14,7 +14,7 @@ use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDateTime, PyDict, PyList},
types::{PyDateTime, PyDict, PyList, PyString},
};
use serde_json::Value;
@ -22,7 +22,7 @@ use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call,
python::Streaming,
python::{Logging, Streaming},
setup,
};
@ -137,6 +137,40 @@ impl LegacyLogging {
Some(params)
}
// Compared as Python compares them, so a caller's `0` is kept over a typed float's `0.0`.
fn caller_value<'py>(
&self,
py: Python<'py>,
name: &str,
sent: &Value,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match self.call.lookup(py, name)? {
Some(value) if from_py::<Value>(&value).is_ok() && value.eq(to_py(py, sent)?)? => {
Ok(Some(value))
}
_ => Ok(None),
}
}
// Python's Messages handler logs the caller's own optional params, not their typed copies.
fn as_called(&self, py: Python<'_>, context: &RequestContext) -> PyResult<RequestContext> {
let mut context = context.clone();
if self.operation != LoggingOperation::Messages {
return Ok(context);
}
for (name, sent) in context
.optional_params
.as_object_mut()
.into_iter()
.flatten()
{
if let Some(value) = self.caller_value(py, name, sent)? {
*sent = from_py(&value)?;
}
}
Ok(context)
}
/// Python times the provider exchange inside the async `AsyncHTTPHandler.post`; its sync
/// client is untimed.
fn close_api_window(&mut self, py: Python<'_>, end: f64) -> PyResult<()> {
@ -158,13 +192,20 @@ impl LegacyLogging {
}
}
fn input_description(&self) -> &'static str {
match self.operation {
fn logged_input<'py>(
&self,
py: Python<'py>,
body: &Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyAny>> {
let description = match self.operation {
LoggingOperation::Completion => "Chat completions",
LoggingOperation::Responses => "Responses",
LoggingOperation::Messages => "Messages",
LoggingOperation::Messages => {
return Logging::MessagesInput.call(py, (body, self.call.kwargs().bind(py)));
}
LoggingOperation::Ocr => "OCR document processing",
}
};
Ok(PyString::new(py, description).into_any())
}
/// The key `pre_call` logs. Python's Messages handler logs an empty string, never the key.
@ -288,7 +329,12 @@ impl LegacyLogging {
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
fn stream_success(
&self,
py: Python<'_>,
stream: &DeliveredStream,
teardown: bool,
) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billed = Streaming::Success.call(
@ -305,6 +351,7 @@ impl LegacyLogging {
&self.start,
&self.end,
&stream.first_chunk,
teardown,
),
);
match billed {
@ -456,6 +503,7 @@ impl LegacyLogging {
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
let context = &self.as_called(py, context)?;
let logger = self.logger()?;
logger.update_from_kwargs(
py,
@ -468,12 +516,17 @@ impl LegacyLogging {
.into_bound(py)
.cast_into::<PyDict>()?;
for (name, sent) in wire.body.as_object().into_iter().flatten() {
if let Some(value) = self.call.lookup(py, name)?
&& from_py::<Value>(&value).is_ok_and(|caller| caller == *sent)
{
if let Some(value) = self.caller_value(py, name, sent)? {
body.set_item(name, value)?;
}
}
if self.operation == LoggingOperation::Messages {
// Python's Messages handler merges its transformed request body into the call details.
logger
.object(py)
.getattr("model_call_details")?
.call_method1("update", (&body,))?;
}
let headers = PyDict::new(py);
for (name, value) in &wire.headers {
headers.set_item(name, value)?;
@ -483,9 +536,10 @@ impl LegacyLogging {
headers: headers.clone().unbind(),
context: context.clone(),
});
let input = self.logged_input(py, &body)?;
self.logger()?.pre_call(
py,
self.input_description(),
&input,
self.logged_api_key(context),
&body,
&headers,
@ -539,7 +593,11 @@ impl LegacyLogging {
.getattr("litellm_params")?
.cast_into::<PyDict>()?
.copy()?;
params.set_item("custom_llm_provider", &facts.provider.provider)?;
if self.operation != LoggingOperation::Messages {
// Python's Messages handler hands the provider to the logger beside its params, never
// in them, and the proxy's logger carries over to the next fallback hop.
params.set_item("custom_llm_provider", &facts.provider.provider)?;
}
crate::python::Logging::Update.call(
py,
(
@ -562,8 +620,17 @@ impl LegacyLogging {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
details.set_item("cache_hit", self.cache_key.is_some())?;
if self.cache_key.is_some() {
details.set_item("cache_hit", true)?;
} else {
details.call_method1("pop", ("cache_hit", py.None()))?;
}
details.set_item("cache_key", self.cache_key.as_deref())?;
if let (LoggingOperation::Messages, ResultSource::Provider, Some(request)) =
(self.operation, &facts.source, &self.request)
{
details.call_method1("update", (request.body.bind(py),))?;
}
Ok(HookStep::Ready(()))
}
@ -606,7 +673,7 @@ impl LegacyLogging {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
Some(stream) => self.stream_success(py, stream, false)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
@ -676,8 +743,28 @@ impl LegacyLogging {
}
}
impl LegacyLogging {
// A caller that goes away mid-stream cancels the call before it ends: Python's Messages
// stream iterator still bills the chunks it read, straight away rather than deferred.
fn bill_abandoned_stream(&mut self, py: Python<'_>) -> PyResult<()> {
if !self.asynchronous || self.end.is_some() {
return Ok(());
}
let Some(stream) = self.stream.take() else {
return Ok(());
};
self.end = Some(datetime(py, epoch_seconds())?);
self.stream_success(py, &stream, true)
}
}
impl PythonOwned for LegacyLogging {
fn close(&mut self, py: Python<'_>) {
if self.logger.is_some()
&& let Err(error) = self.bill_abandoned_stream(py)
{
error.write_unraisable(py, None);
}
if let Some(logger) = self.logger.take()
&& let Err(error) = logger.restore_context(py)
{
@ -1128,10 +1215,10 @@ logger._provider_reported_cost = None
};
let previous = litellm_host::interceptors::ExecutionFacts {
provider: facts.provider.clone(),
cost: litellm_host::interceptors::Cost::Reported {
amount: serde_json::Number::from_f64(9.0).unwrap(),
cost: litellm_host::interceptors::Cost::CacheHit,
source: litellm_host::interceptors::ResultSource::Cache {
key: "earlier-hit".into(),
},
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &previous).unwrap();
local(&locals, "logger")
@ -1149,13 +1236,59 @@ logger._provider_reported_cost = None
c"
assert logger.model_call_details.get('response_cost') is None, logger.model_call_details
assert logger._provider_reported_cost == reported, logger._provider_reported_cost
assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details
assert logger.model_call_details.get('cache_hit') is (True if cache_hit else None), logger.model_call_details
assert logger.update['custom_llm_provider'] == 'edenai', logger.update
",
);
});
}
#[rstest]
#[case::messages(crate::LoggingOperation::Messages, false)]
#[case::ocr(crate::LoggingOperation::Ocr, true)]
fn result_ready_puts_the_provider_in_the_logged_params_except_for_messages(
#[case] operation: crate::LoggingOperation,
#[case] in_params: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(
py,
&locals,
c"
logger.litellm_params = {}
logger.optional_params = {}
logger.model_call_details = {}
logger._provider_reported_cost = None
",
);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
operation,
..legacy_call(py, &locals, false)
};
let facts = litellm_host::interceptors::ExecutionFacts {
provider: litellm_host::interceptors::ProviderIdentity {
model: "native-test".into(),
provider: "anthropic".into(),
},
cost: litellm_host::interceptors::Cost::Deferred,
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &facts).unwrap();
locals.set_item("in_params", in_params).unwrap();
run(
py,
&locals,
c"
assert logger.update['custom_llm_provider'] == 'anthropic', logger.update
assert ('custom_llm_provider' in logger.update['litellm_params']) is in_params, logger.update
",
);
});
}
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";
@ -1290,7 +1423,10 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update
#[rstest]
#[case::completion(crate::LoggingOperation::Completion, "Chat completions")]
#[case::responses(crate::LoggingOperation::Responses, "Responses")]
#[case::messages(crate::LoggingOperation::Messages, "Messages")]
#[case::messages(
crate::LoggingOperation::Messages,
r#"[{'role': 'user', 'content': '{"pages": [1]}'}]"#
)]
#[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")]
fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases(
#[case] operation: crate::LoggingOperation,
@ -1331,7 +1467,7 @@ prepared = {'pages': replacement}
&locals,
c"
assert logger.pre['complete_input_dict']['pages'] is replacement
assert logger.pre_input == description
assert str(logger.pre_input) == description
assert original == [0]
",
);
@ -2076,6 +2212,45 @@ assert chunks[1] is last
});
}
#[rstest]
#[case::cancelled_mid_stream(None, &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_success(Some(true), &["stamp_stream_head", "stream_opened", "stream_success", "restore"])]
#[case::after_failure(Some(false), &["stamp_stream_head", "stream_opened", "stream_failure", "restore"])]
fn closing_an_unfinished_stream_bills_its_delivered_chunks_once(
#[case] ended: Option<bool>,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"first = b'first'\nresponse = None\nfailure = ValueError('provider')\nhead = {}",
);
let mut logging = LegacyLogging {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
logging
.on_stream_open(py, &local(&locals, "head").unbind())
.unwrap();
logging
.on_stream_chunk(py, &local(&locals, "first").unbind())
.unwrap();
match ended {
Some(true) => drop(succeed(py, &locals, &mut logging)),
Some(false) => drop(fail(py, &locals, &mut logging)),
None => {}
}
litellm_host_python::PythonOwned::close(&mut logging, py);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::synchronous(false, &["failure_handler"])]
#[case::asynchronous(true, &[])]

4
litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs
View file

@ -26,7 +26,7 @@ pub trait LegacyCallbacks {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
@ -157,7 +157,7 @@ impl LegacyCallbacks for PythonLogger {
fn pre_call(
&self,
py: Python<'_>,
input: &str,
input: &Bound<'_, PyAny>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,

2
litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -42,6 +42,8 @@ pub(crate) enum Logging {
Update,
#[strum(serialize = "pre_call")]
PreCall,
#[strum(serialize = "messages_input")]
MessagesInput,
#[strum(serialize = "post_call")]
PostCall,
#[strum(serialize = "record_response")]

4
litellm-rust/crates/callbacks-legacy-python/src/test_support.rs
View file

@ -62,6 +62,7 @@ FAKES = {
custom_llm_provider=provider,
),
'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args),
'messages_input': lambda body, kwargs: [{'role': 'user', 'content': json.dumps(body)}],
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
original_response, api_key, additional_args
),
@ -97,7 +98,7 @@ FAKES = {
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
'stream_opened': lambda logger, head: logger.record('stream_opened', head),
'stamp_stream_head': lambda head, logger, kwargs, start, end: logger.record('stamp_stream_head', head),
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk, stream_teardown: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),
@ -134,6 +135,7 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.model_call_details = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None

5
litellm-rust/crates/host-python/src/binding.rs
View file

@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned {
) -> PyResult<Py<PyAny>>;
/// What the stream carries at hand-off, as the caller's stream receives it.
/// The unary response after the deployment transform, before success hooks run.
fn deliver_response(&mut self, _py: Python<'_>, _response: &Bound<'_, PyAny>) -> PyResult<()> {
Ok(())
}
fn encode_stream_head(
&mut self,
py: Python<'_>,

12
litellm-rust/crates/host-python/src/driver.rs
View file

@ -399,8 +399,16 @@ where
Ok(step) => step,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
match self.settle(step).await {
Ok(response) => self.succeeded(response).await,
let response = match self.settle(step).await {
Ok(response) => response,
Err(error) => return self.failure(error, FailureOrigin::Host).await,
};
let delivered = {
let mut parts = self.parts();
Python::attach(|py| parts.binding.deliver_response(py, response.bind(py)))
};
match delivered {
Ok(()) => self.succeeded(response).await,
Err(error) => self.failure(error, FailureOrigin::Host).await,
}
}

2
litellm-rust/crates/host-python/src/lib.rs
View file

@ -36,7 +36,7 @@ pub use marshal::{
pub use owned::PythonOwned;
pub use runtime::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, enter_native, poll_async_value,
ready_future, reserve_process_for_forking, run_async, run_async_value, run_sync,
ready_future, reserve_process_for_forking, run_async, run_async_value, run_isolated, run_sync,
run_sync_value, runtime_handle, runtime_started,
};
pub use services::PythonHostCalls;

29
litellm-rust/crates/host-python/src/runtime.rs
View file

@ -68,6 +68,18 @@ fn runtime() -> PyResult<&'static Runtime> {
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
/// Runs `future` on a current-thread runtime built and shut down inside the call, its blocking pool joined before
/// it returns. No thread outlives the call, so unlike every other door this one never claims the process in the fork
/// gate: a process that only used it can still fork. The future must not hold resources tied to another runtime
pub fn run_isolated<F: Future>(future: F) -> std::io::Result<F::Output> {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let output = runtime.block_on(future);
runtime.shutdown_timeout(Duration::from_secs(5));
Ok(output)
}
/// The native runtime, for work that must outlive the call that schedules it.
pub fn runtime_handle() -> PyResult<Handle> {
Ok(runtime()?.handle().clone())
@ -818,4 +830,21 @@ asyncio.run(exercise())
.expect("async route contract should hold");
});
}
#[rstest]
fn an_isolated_run_joins_blocking_work_it_left_running_before_returning() {
let finished = Arc::new(AtomicBool::new(false));
let flag = Arc::clone(&finished);
let output = run_isolated(async move {
drop(tokio::task::spawn_blocking(move || {
thread::sleep(Duration::from_millis(50));
flag.store(true, Ordering::SeqCst);
}));
7
})
.unwrap();
assert_eq!(output, 7);
assert!(finished.load(Ordering::SeqCst));
}
}

9
litellm-rust/crates/host-python/tests/isolated_run.rs Normal file
View file

@ -0,0 +1,9 @@
use litellm_host_python::{reserve_process_for_forking, run_isolated, runtime_started};
#[test]
fn an_isolated_run_leaves_the_process_free_to_reserve_for_forking() {
run_isolated(async { tokio::task::spawn_blocking(|| ()).await.unwrap() }).unwrap();
assert!(!runtime_started());
assert!(reserve_process_for_forking().is_ok());
}

104
litellm-rust/crates/http/src/error.rs
View file

@ -16,6 +16,15 @@ pub enum TlsSource {
#[error(transparent)]
pub struct Error(Arc<Repr>);
// A streamed body the connection ended early, after the response head arrived; `expected` is
// the declared `content-length`, `None` for a chunked body.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BodyCut {
pub received: u64,
pub expected: Option<u64>,
pub reset: bool,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Repr {
#[error("could not read {}", path.display())]
@ -56,6 +65,12 @@ pub(crate) enum Repr {
Connect(#[source] reqwest::Error),
#[error("upstream network error")]
Network(#[source] reqwest::Error),
#[error("upstream network error")]
BodyCut {
#[source]
source: reqwest::Error,
cut: BodyCut,
},
#[error("upstream request timed out")]
Timeout(#[source] reqwest::Error),
/// No bytes arrived within the read timeout, which the caller enforces per read; `body` says
@ -201,6 +216,37 @@ impl Error {
matches!(&*self.0, Repr::ReadTimeout { body: true, .. })
}
pub fn body_read(error: reqwest::Error, received: u64, expected: Option<u64>) -> Self {
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(&error);
let kind = std::iter::from_fn(|| {
let current = source?;
source = current.source();
Some(current)
})
.find_map(|cause| cause.downcast_ref::<io::Error>().map(io::Error::kind));
let reset = match kind {
Some(io::ErrorKind::UnexpectedEof) => false,
Some(io::ErrorKind::ConnectionReset) => true,
_ => return error.into(),
};
Repr::BodyCut {
source: error.without_url(),
cut: BodyCut {
received,
expected,
reset,
},
}
.into()
}
pub fn body_cut(&self) -> Option<BodyCut> {
match &*self.0 {
Repr::BodyCut { cut, .. } => Some(*cut),
_ => None,
}
}
pub fn computed_header(name: impl Into<String>) -> Self {
Repr::ComputedHeader(name.into()).into()
}
@ -226,6 +272,7 @@ impl Error {
pub fn dispatched(&self) -> bool {
match &*self.0 {
Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -285,6 +332,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -318,6 +366,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -355,6 +404,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -394,6 +444,7 @@ impl Error {
| Repr::UnsignableHeader(_)
| Repr::Signature(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -437,6 +488,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Status { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
@ -473,6 +525,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::WebSocketUrl(_)
@ -513,6 +566,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::WebSocketUrl(_)
| Repr::WebSocketScheme(_)
| Repr::WebSocketSchemeRewrite(_)
@ -548,6 +602,7 @@ impl Error {
| Repr::Signature(_)
| Repr::Connect(_)
| Repr::Network(_)
| Repr::BodyCut { .. }
| Repr::Timeout(_)
| Repr::ReadTimeout { .. }
| Repr::Status { .. }
@ -742,4 +797,53 @@ mod tests {
assert_eq!(error.upstream_response(), Some((429, body.as_str())));
assert_eq!(error.status(), Some(429));
}
#[rstest]
#[case::chunked("transfer-encoding: chunked\r\n\r\n5\r\nhello\r\n", None)]
#[case::content_length("content-length: 100\r\n\r\nhello", Some(100))]
#[tokio::test]
async fn a_connection_that_ends_mid_body_is_a_body_cut(
#[case] tail: &'static str,
#[case] expected: Option<u64>,
) {
use std::io::{Read as _, Write as _};
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
let address = listener.local_addr().expect("address");
let server = std::thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("accept");
let mut request = [0_u8; 1024];
let _ = stream.read(&mut request).expect("request");
stream
.write_all(format!("HTTP/1.1 200 OK\r\n{tail}").as_bytes())
.expect("response");
});
let mut response = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}/"))
.send()
.await
.expect("head");
server.join().expect("server");
let mut received = 0_u64;
let error = loop {
match response.chunk().await {
Ok(Some(chunk)) => received += chunk.len() as u64,
Ok(None) => panic!("the body must end with an error"),
Err(error) => break crate::Error::body_read(error, received, expected),
}
};
assert_eq!(
error.body_cut(),
Some(super::BodyCut {
received: 5,
expected,
reset: false,
})
);
assert!(error.dispatched());
assert_eq!(error.to_string(), "upstream network error");
}
}

2
litellm-rust/crates/http/src/lib.rs
View file

@ -18,7 +18,7 @@ pub mod websocket;
pub use client::Client;
pub use config::{ClientIdentity, HttpClientConfig, Resolution, Verify};
pub use error::{Error, TlsSource};
pub use error::{BodyCut, Error, TlsSource};
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};

70
litellm-rust/crates/http/src/pool.rs
View file

@ -55,36 +55,13 @@ impl HttpClientPool {
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<Client, Error> {
let effective = match variant {
ClientVariant::Media => HttpClientConfig {
client_certificate: None,
proxies: EnvironmentProxies::default(),
..config.clone()
},
ClientVariant::UnpinnedMedia => HttpClientConfig {
client_certificate: None,
..config.clone()
},
ClientVariant::WebSocket => HttpClientConfig {
http2: false,
verify: match &config.verify {
Verify::BuiltInRoots => Verify::NativeRoots,
verify => verify.clone(),
},
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
};
let key = (effective, variant);
let key = (effective(config, variant), variant);
if let Some(pooled) = self.lock().get(&key)
&& pooled.built_at.elapsed() < self.ttl
{
return Ok(Client::new(pooled.client.clone()));
}
let client = self
.apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
.build()
.map_err(Repr::Client)?;
let client = self.build(&key.0, variant)?;
self.lock().insert(
key,
PooledClient {
@ -95,6 +72,26 @@ impl HttpClientPool {
Ok(Client::new(client))
}
/// A client nobody else holds, for a caller that drives it on a runtime of its own and drops both together
pub fn unpooled(
&self,
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<Client, Error> {
self.build(&effective(config, variant), variant)
.map(Client::new)
}
fn build(
&self,
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<reqwest::Client, Error> {
self.apply(variant, reqwest::ClientBuilder::try_from(config)?)
.build()
.map_err(|error| Error::from(Repr::Client(error)))
}
fn lock(&self) -> MutexGuard<'_, Clients> {
self.clients.lock().unwrap_or_else(PoisonError::into_inner)
}
@ -116,6 +113,29 @@ impl HttpClientPool {
}
}
fn effective(config: &HttpClientConfig, variant: ClientVariant) -> HttpClientConfig {
match variant {
ClientVariant::Media => HttpClientConfig {
client_certificate: None,
proxies: EnvironmentProxies::default(),
..config.clone()
},
ClientVariant::UnpinnedMedia => HttpClientConfig {
client_certificate: None,
..config.clone()
},
ClientVariant::WebSocket => HttpClientConfig {
http2: false,
verify: match &config.verify {
Verify::BuiltInRoots => Verify::NativeRoots,
verify => verify.clone(),
},
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
}
}
#[cfg(test)]
mod tests {
use std::{

83
litellm-rust/crates/inference-messages/src/common_utils.rs
View file

@ -1,5 +1,5 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers as shared_string_headers;
use litellm_llms::ErrorDetail;
use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG;
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
@ -19,8 +19,6 @@ use serde_json::{Map, Value};
use super::Error;
const HEADER_CONTEXT: &str = "messages";
#[derive(Clone, Copy)]
pub(crate) struct MessagesProvider {
provider: LlmProviders,
@ -75,7 +73,52 @@ pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<M
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(name, value)| match value {
Value::String(value) => Ok((name, value)),
value => Err(Error::invalid_request(ErrorDetail::invalid(
"extra_headers",
NonStringHeader {
actual: python_type_name(&value),
name,
},
))),
})
.collect()
}
/// A header value the caller set to something other than a string, named by its Python type
/// since the caller wrote it in Python.
#[derive(Debug)]
struct NonStringHeader {
name: String,
actual: &'static str,
}
impl std::fmt::Display for NonStringHeader {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"header {} must be a string, got {}",
self.name, self.actual
)
}
}
impl std::error::Error for NonStringHeader {}
fn python_type_name(value: &Value) -> &'static str {
match value {
Value::Null => "NoneType",
Value::Bool(_) => "bool",
Value::Number(number) if number.is_f64() => "float",
Value::Number(_) => "int",
Value::String(_) => "str",
Value::Array(_) => "list",
Value::Object(_) => "dict",
}
}
#[cfg(test)]
@ -85,7 +128,6 @@ mod tests {
use rstest::rstest;
use super::{messages_provider, string_headers};
use crate::Error;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
@ -113,17 +155,28 @@ mod tests {
assert!(messages_provider(provider, model).is_none());
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({"x-count": 3}).as_object().unwrap().clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
#[rstest]
#[case::int(json!(12345), "int")]
#[case::float(json!(1.5), "float")]
#[case::bool(json!(true), "bool")]
#[case::null(json!(null), "NoneType")]
#[case::list(json!([1]), "list")]
#[case::dict(json!({"a": 1}), "dict")]
fn non_string_header_is_a_bad_request_naming_the_header_and_its_type(
#[case] value: serde_json::Value,
#[case] python_type: &str,
) {
let headers = json!({"x-kept": "v", "x-count": value})
.as_object()
.unwrap()
.clone();
let error = string_headers(Some(headers)).expect_err("non-string header refused");
assert!(error.is_request(), "{error:?}");
assert_eq!(error.upstream_response(), None);
let rejection = error.rejected_request().expect("a rejected request");
assert_eq!(
err,
Error::Headers(litellm_http::request::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
})
litellm_error::Report::new(rejection).to_string(),
format!("invalid extra_headers: header x-count must be a string, got {python_type}")
);
}
}

64
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -26,10 +26,10 @@ use litellm_llms_types::formats::messages::{
};
use litellm_secrets::source::Secrets;
use litellm_tracing::ByteChunk;
use serde_json::Value;
use serde_json::{Map, Value};
use super::{
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, merge_messages_header_sources, string_headers},
constants::MESSAGES_TIMEOUT_SECS,
route::Messages,
@ -85,7 +85,10 @@ pub(super) async fn execute(
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
MessagesCallResponse::Complete(response) => match response {
MessagesBody::Message(message) => provider.config().reported_cost(message),
MessagesBody::Unchecked(_) => None,
},
MessagesCallResponse::Stream { .. } => None,
};
context
@ -130,9 +133,10 @@ async fn provider_call(
)?;
let transform_context =
MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params);
let filtered = requested_optional_params(sanitized, provider, &transform_context)?;
let trimmed =
without_additional_drop_params(filtered, &shaping.settings.additional_drop_params)?;
let without_listed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let trimmed = requested_optional_params(without_listed, provider, &transform_context)?;
let requested_params = serde_json::to_value(&trimmed.params).map_err(serialize_failure)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
@ -173,7 +177,7 @@ async fn provider_call(
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
optional_params: requested_params,
secret_fields: Vec::new(),
api_key: api_key.map(SecretValue::new),
};
@ -346,8 +350,7 @@ async fn call_provider(
// body it cannot parse, and the public mapping classifies it by that status.
return Err(litellm_http::Error::upstream(status, text).into());
}
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
decode_response(config, &identity.model, &text).map(MessagesCallResponse::Complete)
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -383,13 +386,20 @@ fn timed_chunks(
response: reqwest::Response,
timeout: Duration,
) -> impl futures_util::Stream<Item = Result<Bytes, litellm_http::Error>> + Send + 'static {
futures_util::stream::try_unfold(response, move |mut response| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(litellm_http::Error::from)?;
Ok(chunk.map(|chunk| (chunk, response)))
})
let expected = response.content_length();
futures_util::stream::try_unfold(
(response, 0_u64),
move |(mut response, received)| async move {
let chunk = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)?
.map_err(|error| litellm_http::Error::body_read(error, received, expected))?;
Ok(chunk.map(|chunk| {
let received = received + chunk.len() as u64;
(chunk, (response, received))
}))
},
)
}
/// The provider's error, its body read under `timeout`: a stream has no overall deadline once
@ -411,15 +421,23 @@ fn decode_response(
config: &dyn BaseMessagesConfig,
model: &str,
text: &str,
) -> Result<MessagesResponse, Error> {
let response = serde_json::from_str(text).map_err(|err| {
Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
))
})?;
) -> Result<MessagesBody, Error> {
let response = match serde_json::from_str::<MessagesResponse>(text) {
Ok(response) => response,
Err(err) => {
let Ok(mut body) = serde_json::from_str::<Map<String, Value>>(text) else {
return Err(Error::invalid_response(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
err,
)));
};
body.entry("stop_sequence").or_insert(Value::Null);
return Ok(MessagesBody::Unchecked(body));
}
};
config
.transform_anthropic_messages_response(model, response)
.map(|message| MessagesBody::Message(Box::new(message)))
.map_err(Error::from)
.map_err(Error::after_send)
}

4
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -17,8 +17,8 @@ use crate::common_utils::messages_provider;
pub use common_utils::merge_messages_header_sources;
pub use litellm_inference::Error;
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
messages_body,
MessagesBody, MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping,
litellm_params, messages_body,
};
#[derive(Clone)]

7
litellm-rust/crates/inference-messages/src/route.rs
View file

@ -5,11 +5,10 @@ use litellm_host::{
call::{HostedCompletion, HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_llms_types::formats::messages::MessagesResponse;
use super::{Error, MessagesCall};
use super::{Error, MessagesBody, MessagesCall};
pub type MessagesOutput = HostedCompletion<Box<MessagesResponse>>;
pub type MessagesOutput = HostedCompletion<MessagesBody>;
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
pub struct MessagesStreamHead {
@ -19,7 +18,7 @@ pub struct MessagesStreamHead {
pub struct Messages;
impl Protocol for Messages {
type Response = Box<MessagesResponse>;
type Response = MessagesBody;
type Error = Error;
type Request = MessagesCall;
type HostCall = Infallible;

11
litellm-rust/crates/inference-messages/src/types.rs
View file

@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error {
}
pub type MessagesCallResponse =
CallOutput<Box<MessagesResponse>, super::route::MessagesStreamHead, Bytes, Error>;
CallOutput<MessagesBody, super::route::MessagesStreamHead, Bytes, Error>;
// Python's `transform_anthropic_messages_response` returns any JSON object the provider sends
// with a 200, message or not, so one that is not a message passes through as sent.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessagesBody {
Message(Box<MessagesResponse>),
Unchecked(Map<String, Value>),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MessagesShaping {

9
litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs
View file

@ -1,7 +1,9 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesShaping,
};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use rstest::{fixture, rstest};
@ -74,7 +76,10 @@ async fn run(call: MessagesCall) -> Result<MessagesCallResponse, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.unwrap() {
MessagesCallResponse::Complete(message) => *message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("expected a message"),
},
MessagesCallResponse::Stream { .. } => panic!("expected a completed message"),
}
}

72
litellm-rust/crates/inference-messages/tests/caching.rs
View file

@ -50,42 +50,47 @@ async fn messages_cache_identity_includes_provider_native_parameters(
#[case] original: Value,
#[case] changed: Value,
) {
use litellm_inference_messages::route::Messages;
use litellm_inference_messages::{MessagesBody, route::Messages};
use litellm_llms_types::formats::messages::MessagesResponse;
let calls = AtomicUsize::new(0);
for (value, expected_call) in [(original.clone(), 0), (changed, 1), (original, 0)] {
let response =
execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(Box::new(serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
})).unwrap()))
let response = execute_unary::<Messages, _, _>(
CacheRequest::from_wire(
ProviderIdentity {
model: "test".into(),
provider: "anthropic".into(),
},
)
.await
.unwrap();
Some(&WireRequest {
url: "https://example.test/v1/messages".into(),
headers: vec![],
body: json!({
"model":"test", "messages":[{"role":"user","content":"hello"}],
"max_tokens":32, (field):value
}),
}),
),
Some(cache.clone()),
Some(CacheOptions::new(CacheScope::Shared)),
&(),
None,
|| async {
let call = calls.fetch_add(1, Ordering::SeqCst);
Ok(MessagesBody::Message(Box::new(
serde_json::from_value::<MessagesResponse>(json!({
"id":call.to_string(), "type":"message", "role":"assistant", "model":"test",
"content":[{"type":"text","text":format!("answer {call}")}],
"stop_reason":"end_turn", "stop_sequence":null
}))
.unwrap(),
)))
},
)
.await
.unwrap();
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected_call.to_string());
assert_eq!(
response.content[0]["text"],
@ -253,7 +258,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
#[case] cost: Value,
#[case] expected: Option<f64>,
) {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
use litellm_inference_messages::{MessagesBody, MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost});
@ -277,6 +282,9 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_
panic!("expected Messages responses");
};
assert_eq!(first, second);
let MessagesBody::Message(first) = first else {
panic!("expected a message");
};
assert_eq!(first.extra.get("cost"), Some(&cost));
assert_eq!(support::received(&upstream).await.len(), 1);
let facts = hooks.facts.lock().unwrap();

9
litellm-rust/crates/inference-messages/tests/live/support.rs
View file

@ -4,7 +4,9 @@ use futures_util::{FutureExt, TryStreamExt};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
};
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
use litellm_inference_messages::{
Error, MessagesBody, MessagesCall, MessagesCallResponse, MessagesRoute,
};
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
use litellm_llms_types::formats::messages::MessagesResponse;
use serde_json::{Value, json};
@ -111,7 +113,10 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
match response {
MessagesCallResponse::Complete(message) => message,
MessagesCallResponse::Complete(message) => match message {
MessagesBody::Message(message) => message,
MessagesBody::Unchecked(_) => panic!("a complete request returned no message"),
},
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
}
}

92
litellm-rust/crates/inference-messages/tests/messages/host.rs
View file

@ -6,9 +6,11 @@ use litellm_host::{
interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest},
lifecycle::CallEvent,
};
use litellm_inference_messages::{MessagesCallResponse, route::Messages};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, route::Messages};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities;
use litellm_llms::base_llm::messages::context::{
MessagesModelCapabilities as AnthropicModelCapabilities, SupportedEffortTiers,
};
use rstest::rstest;
use super::*;
@ -338,7 +340,10 @@ async fn response_mode_follows_the_intercepted_request(
}
MessagesCallResponse::Complete(message) => {
assert_eq!(expected_stream, Some(false));
assert_eq!(*message, serde_json::from_value(message_body()).unwrap());
assert_eq!(
message,
MessagesBody::Message(serde_json::from_value(message_body()).unwrap())
);
}
}
Ok::<_, Error>(())
@ -419,7 +424,7 @@ async fn no_raw_response_is_emitted_for_a_stream_or_a_failure(
assert!(host.raw_responses().is_empty());
}
/// Python logs `optional_params` as what it is about to send, so a dropped param must
/// Python logs `optional_params` after its requested-param gate, so a dropped param must
/// not resurface in callbacks.
#[rstest]
#[tokio::test]
@ -450,3 +455,82 @@ async fn the_request_context_carries_the_shaped_params_without_model_or_messages
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[tokio::test]
async fn the_request_context_leaves_out_a_listed_param_the_model_refuses(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["temperature".into()],
..MessagesSettings::default()
},
capabilities: AnthropicModelCapabilities {
supports_sampling_params: false,
..AnthropicModelCapabilities::default()
},
},
..with_fields(call, json!({"temperature": 0.2}))
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(optional_params, json!({"max_tokens": 16}));
}
#[rstest]
#[case::reasoning_effort_on_an_adaptive_model(
AnthropicModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
supports_output_config: true,
effort_tiers: SupportedEffortTiers { low: true, ..SupportedEffortTiers::default() },
..AnthropicModelCapabilities::default()
},
json!({"reasoning_effort": "low"})
)]
#[case::adaptive_payload_on_a_legacy_model(
AnthropicModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
},
json!({"max_tokens": 3000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}})
)]
#[tokio::test]
async fn the_request_context_carries_the_requested_reasoning_not_its_translation(
call: MessagesCall,
#[case] capabilities: AnthropicModelCapabilities,
#[case] fields: Value,
) {
let upstream = upstream([message_response()]).await;
let host = RecordingHost::passthrough(authenticated(
MessagesCall {
shaping: MessagesShaping {
capabilities,
..MessagesShaping::default()
},
..with_fields(call, fields.clone())
},
upstream.uri(),
));
run_through(&host).await.expect("messages call succeeds");
let [optional_params] = <[Value; 1]>::try_from(host.optional_params.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
assert_eq!(
optional_params,
Value::Object(
object(json!({"max_tokens": 16}))
.into_iter()
.chain(object(fields))
.collect()
)
);
}

7
litellm-rust/crates/inference-messages/tests/messages/main.rs
View file

@ -5,7 +5,7 @@ use std::{
use litellm_http::{HttpSettings, Resolution};
use litellm_inference_messages::{
Error, MessagesCall, MessagesSettings, MessagesShaping,
Error, MessagesBody, MessagesCall, MessagesSettings, MessagesShaping,
route::{Messages, MessagesMachine, MessagesOutput},
};
use litellm_inference_testing::RecordingSecrets;
@ -119,7 +119,10 @@ async fn run(call: MessagesCall) -> Result<MessagesOutput, Error> {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.expect("messages call succeeds") {
MessagesOutput::Complete(message) => *message,
MessagesOutput::Complete(message) => match message {
MessagesBody::Message(message) => *message,
MessagesBody::Unchecked(_) => panic!("a call returned no message"),
},
MessagesOutput::StreamEnded | MessagesOutput::Detached => {
panic!("a non-streaming call returned a stream")
}

24
litellm-rust/crates/inference-messages/tests/messages/request.rs
View file

@ -530,6 +530,30 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[tokio::test]
async fn additional_drop_params_clear_a_param_the_model_does_not_support(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities: sampling_removed(),
settings: MessagesSettings {
additional_drop_params: vec!["top_k".into()],
..MessagesSettings::default()
},
},
..with_fields(call, json!({"top_k": 5}))
})
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("top_k"), None);
assert_eq!(sent["max_tokens"], 16);
}
#[rstest]
#[case::adaptive_thinking(json!({"type": "adaptive"}), json!({"type": "adaptive", "display": "summarized"}))]
#[case::disabled_thinking(json!({"type": "disabled"}), json!({"type": "disabled"}))]

48
litellm-rust/crates/inference-messages/tests/messages/response.rs
View file

@ -2,7 +2,7 @@ use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use litellm_inference_messages::{MessagesCallResponse, messages_body};
use litellm_inference_messages::{MessagesBody, MessagesCallResponse, messages_body};
use litellm_inference_testing::{
RecordingSecrets, http_config, no_secrets, provider_http, resources,
};
@ -48,6 +48,9 @@ async fn calls_defer_execution_until_polled(
let MessagesCallResponse::Complete(response) = future.await.unwrap() else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -240,10 +243,42 @@ async fn an_unparseable_success_body_is_reported_as_the_upstream_answer(call: Me
assert_eq!(error.upstream_response(), Some((200, "not json")));
}
#[rstest]
#[case::error_object(
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}),
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"},
"stop_sequence": null}),
)]
#[case::message_missing_id(
json!({"type": "message", "content": [], "stop_sequence": "###"}),
json!({"type": "message", "content": [], "stop_sequence": "###"}),
)]
#[tokio::test]
async fn a_success_body_that_is_not_a_message_passes_through(
call: MessagesCall,
#[case] body: Value,
#[case] expected: Value,
) {
let upstream = upstream([json_response(body)]).await;
let response = run(MessagesCall {
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
})
.await
.expect("any JSON object is returned");
let MessagesOutput::Complete(response) = response else {
panic!("expected a completed response");
};
assert_eq!(serde_json::to_value(response).unwrap(), expected);
}
#[rstest]
#[tokio::test]
async fn an_unreadable_success_body_is_an_invalid_response(call: MessagesCall) {
let upstream = upstream([json_response(json!({"unexpected": true}))]).await;
let upstream = upstream([json_response(json!(["unexpected"]))]).await;
let error = run(MessagesCall {
api_key: Some("sk".into()),
@ -320,6 +355,9 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
let MessagesCallResponse::Complete(message) = response else {
panic!("a non-streaming request returns a message");
};
let MessagesBody::Message(message) = message else {
panic!("expected a message");
};
assert_eq!(message.id, "msg_1");
let sent = only_request(&upstream).await;
assert_eq!(sent.header("x-api-key"), Some("sk-ant"));
@ -412,6 +450,9 @@ async fn route_uses_injected_dependencies_and_optional_cache(
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
@ -484,6 +525,9 @@ async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall)
else {
panic!("expected a completed message");
};
let MessagesBody::Message(response) = response else {
panic!("expected a message");
};
assert_eq!(response.id, expected["id"].as_str().unwrap());
}
assert_eq!(received(&upstream).await.len(), 2);

8
litellm-rust/crates/inference/src/error.rs
View file

@ -164,6 +164,14 @@ impl Error {
}
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Llm(error) => error.body_cut(),
Self::Http(error) => error.body_cut(),
_ => None,
}
}
/// `Some(connect)` when the call timed out, where `connect` says the connection was never
/// established; `None` for every other error.
pub fn timeout(&self) -> Option<bool> {

7
litellm-rust/crates/llms/src/error.rs
View file

@ -103,6 +103,13 @@ impl Error {
matches!(self, Self::Http(error) if error.is_body_read_timeout())
}
pub fn body_cut(&self) -> Option<litellm_http::BodyCut> {
match self {
Self::Http(error) => error.body_cut(),
_ => None,
}
}
pub fn is_timeout(&self) -> Option<bool> {
match self {
Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()),

9
litellm-rust/crates/model-catalog-refresh/AGENTS.md Normal file
View file

@ -0,0 +1,9 @@
# Model catalog refresh
- `fetch` fetches one document over http(s) or `s3://` and never parses it
- Only 429, 500, 502, 503, 504 and transport errors are retried. S3 errors are final
- `Failed.retryable` means only the attempt limit ended the fetch
- A server's Retry-After overrides the backoff, capped at `max_delay`
- `Start::After(failed)` resumes the backoff schedule where another caller's first attempt left off
- Error reasons use the URL with userinfo, query and fragment stripped, since presigned URLs carry credentials
- Bodies over `MAX_BODY_BYTES` are rejected

13
litellm-rust/crates/model-catalog-refresh/Cargo.toml
View file

@ -6,14 +6,19 @@ license.workspace = true
repository.workspace = true
[dependencies]
backon.workspace = true
bytes.workspace = true
litellm-error.workspace = true
litellm-http.workspace = true
litellm-model-catalog.workspace = true
litellm-s3.workspace = true
litellm-tracing.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["time"] }
url.workspace = true
[dev-dependencies]
litellm-error.workspace = true
aws-credential-types = "1.3.0"
litellm-http = { workspace = true, features = ["test-support"] }
reqwest.workspace = true
rstest.workspace = true
tokio.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }

153
litellm-rust/crates/model-catalog-refresh/src/attempt.rs Normal file
View file

@ -0,0 +1,153 @@
use std::time::Duration;
use bytes::Bytes;
use litellm_http::Client;
use litellm_s3::S3Uri;
use litellm_s3::sdk::config::retry::RetryConfig;
use litellm_s3::sdk::config::timeout::TimeoutConfig;
use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER};
use url::Url;
use crate::{Error, RETRYABLE_STATUSES, S3Access};
pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Fetched {
pub body: Vec<u8>,
pub etag: Option<String>,
}
#[derive(Debug)]
pub(crate) enum Cause {
Status(u16),
Transport(Error),
}
#[derive(Debug)]
pub(crate) struct Rejected {
pub(crate) cause: Cause,
pub(crate) retryable: bool,
pub(crate) retry_after: Option<Duration>,
}
impl Rejected {
pub(crate) fn fatal(error: Error) -> Self {
Self {
cause: Cause::Transport(error),
retryable: false,
retry_after: None,
}
}
}
/// Any status below 400 is the document
pub(crate) async fn http(
client: &Client,
url: Url,
timeout: Duration,
) -> Result<Fetched, Rejected> {
let transport = |error: litellm_http::Error| Rejected {
cause: Cause::Transport(Error::from(error)),
retryable: true,
retry_after: None,
};
let mut response = tokio::time::timeout(timeout, client.get(url).send())
.await
.map_err(litellm_http::Error::read_timeout)
.and_then(|sent| sent.map_err(litellm_http::Error::from))
.map_err(transport)?;
let status = response.status().as_u16();
let etag = etag(response.headers());
let retry_after = response
.headers()
.get(RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse().ok())
.map(Duration::from_secs);
if status >= 400 {
return Err(Rejected {
cause: Cause::Status(status),
retryable: RETRYABLE_STATUSES.contains(&status),
retry_after,
});
}
if response
.content_length()
.is_some_and(|length| length > MAX_BODY_BYTES)
{
return Err(Rejected::fatal(Error::body_too_large(MAX_BODY_BYTES)));
}
let mut body = Vec::new();
while let Some(chunk) = tokio::time::timeout(timeout, response.chunk())
.await
.map_err(litellm_http::Error::body_read_timeout)
.and_then(|read| read.map_err(litellm_http::Error::from))
.map_err(transport)?
{
append(&mut body, &chunk).map_err(Rejected::fatal)?;
}
Ok(Fetched { body, etag })
}
/// S3 failures are final: the SDK already classifies its own errors and a missing or denied object stays so
pub(crate) async fn s3(
client: &Client,
s3: &S3Access,
object: &S3Uri,
timeout: Duration,
) -> Result<Fetched, Rejected> {
let fetched = async {
let config = match s3 {
S3Access::Environment { region } => {
litellm_s3::environment_config(client.clone(), region.clone()).await
}
S3Access::Config(config) => config.to_builder(),
}
.retry_config(RetryConfig::disabled())
.timeout_config(
TimeoutConfig::builder()
.operation_attempt_timeout(timeout)
.build(),
)
.build();
let output = litellm_s3::sdk::Client::from_conf(config)
.get_object()
.bucket(&object.bucket)
.key(&object.key)
.send()
.await?;
if output
.content_length()
.is_some_and(|length| length.unsigned_abs() > MAX_BODY_BYTES)
{
return Err(Error::body_too_large(MAX_BODY_BYTES));
}
let etag = output.e_tag().map(str::to_owned);
let mut stream = output.body;
let mut body = Vec::new();
while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? {
append(&mut body, &chunk)?;
}
Ok::<_, Error>(Fetched { body, etag })
};
fetched.await.map_err(Rejected::fatal)
}
fn append(body: &mut Vec<u8>, chunk: &Bytes) -> Result<(), Error> {
if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES {
return Err(Error::body_too_large(MAX_BODY_BYTES));
}
body.extend_from_slice(chunk);
Ok(())
}
fn etag(headers: &HeaderMap) -> Option<String> {
let values: Vec<&str> = headers
.get_all(ETAG)
.iter()
.filter_map(|value| value.to_str().ok())
.collect();
(!values.is_empty()).then(|| values.join(", "))
}

70
litellm-rust/crates/model-catalog-refresh/src/error.rs
View file

@ -1,6 +1,9 @@
use std::io;
use litellm_s3::sdk::config::http::HttpResponse;
use litellm_s3::sdk::error::SdkError;
use litellm_s3::sdk::operation::get_object::GetObjectError;
use litellm_s3::sdk::primitives::ByteStreamError;
use thiserror::Error;
use tokio::time::error::Elapsed;
#[derive(Debug, Error)]
#[error(transparent)]
@ -10,18 +13,48 @@ pub struct Error(Box<Repr>);
enum Repr {
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error("could not read the catalog")]
Read(#[source] io::Error),
#[error("invalid catalog URL")]
InvalidUrl(#[source] url::ParseError),
#[error("unsupported catalog URL scheme {scheme}")]
UnsupportedScheme { scheme: String },
#[error("invalid S3 URL")]
S3Uri(#[source] litellm_s3::Error),
#[error("S3 GetObject failed")]
S3GetObject(#[source] Box<SdkError<GetObjectError, HttpResponse>>),
#[error("S3 GetObject body read failed")]
S3Body(#[source] ByteStreamError),
#[error("S3 GetObject timed out")]
S3Timeout(#[source] Elapsed),
#[error("response body is over the {limit} byte limit")]
BodyTooLarge { limit: u64 },
}
impl Error {
pub fn unsupported_scheme(scheme: String) -> Self {
Self(Box::new(Repr::UnsupportedScheme { scheme }))
}
pub(crate) fn body_too_large(limit: u64) -> Self {
Self(Box::new(Repr::BodyTooLarge { limit }))
}
}
impl From<SdkError<GetObjectError, HttpResponse>> for Error {
fn from(cause: SdkError<GetObjectError, HttpResponse>) -> Self {
Self(Box::new(Repr::S3GetObject(Box::new(cause))))
}
}
impl From<ByteStreamError> for Error {
fn from(cause: ByteStreamError) -> Self {
Self(Box::new(Repr::S3Body(cause)))
}
}
impl From<Elapsed> for Error {
fn from(cause: Elapsed) -> Self {
Self(Box::new(Repr::S3Timeout(cause)))
}
}
impl From<litellm_http::Error> for Error {
@ -30,9 +63,9 @@ impl From<litellm_http::Error> for Error {
}
}
impl From<io::Error> for Error {
fn from(cause: io::Error) -> Self {
Self(Box::new(Repr::Read(cause)))
impl From<litellm_s3::Error> for Error {
fn from(cause: litellm_s3::Error) -> Self {
Self(Box::new(Repr::S3Uri(cause)))
}
}
@ -44,29 +77,18 @@ impl From<url::ParseError> for Error {
#[cfg(test)]
mod tests {
use std::io;
use litellm_error::{Report, find};
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 an_s3_failure_keeps_its_typed_cause_once() {
let error = Error::from(SdkError::<GetObjectError, HttpResponse>::timeout_error(
io::Error::from(io::ErrorKind::ConnectionReset),
));
let cause = find::<io::Error>(&error).expect("typed I/O cause");
assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset);
assert_eq!(

122
litellm-rust/crates/model-catalog-refresh/src/fetch.rs Normal file
View file

@ -0,0 +1,122 @@
use std::borrow::Cow;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;
use backon::Retryable;
use litellm_error::Report;
use litellm_http::Client;
use url::Url;
use crate::attempt::{self, Cause, Fetched, Rejected};
use crate::{RetryPolicy, Source};
/// How an `s3://` fetch reaches S3
#[derive(Clone, Debug)]
pub enum S3Access {
/// What an AWS SDK finds in the environment and shared config files, loaded on each `s3://` attempt, with
/// `region` winning over the region it finds
Environment { region: Option<String> },
/// A config the caller built with its region, endpoint and credentials
Config(Box<litellm_s3::sdk::Config>),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Failed {
pub reason: String,
/// Another attempt could succeed. Only the attempt limit ended the fetch
pub retryable: bool,
pub retry_after: Option<Duration>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Start {
Fresh,
/// Continues a fetch whose first attempt failed retryably elsewhere, waiting before its next attempt
After(Failed),
}
pub async fn fetch(
client: &Client,
s3: &S3Access,
url: &str,
policy: RetryPolicy,
start: Start,
) -> Result<Fetched, Failed> {
let shown = redacted(url);
let source =
Source::parse(url).map_err(|error| Failed::of(&shown, &Rejected::fatal(error), 1))?;
let mut backoff = policy.backoff();
let made = AtomicU32::new(0);
if let Start::After(failed) = start {
let Some(wait) = backoff.next() else {
return Err(failed);
};
tokio::time::sleep(policy.delay(wait, failed.retry_after)).await;
made.store(1, Ordering::Relaxed);
}
let attempt = || {
made.fetch_add(1, Ordering::Relaxed);
async {
match &source {
Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await,
Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await,
}
}
};
attempt
.retry(backoff)
.sleep(tokio::time::sleep)
.when(|rejected| rejected.retryable)
.adjust(|rejected, wait| wait.map(|wait| policy.delay(wait, rejected.retry_after)))
.notify(|rejected, wait| {
litellm_tracing::warn!(
"LiteLLM: model cost map fetch attempt {}/{} failed ({}); retrying in {:.1}s",
made.load(Ordering::Relaxed),
policy.max_attempts,
reason(&shown, &rejected.cause),
wait.as_secs_f64(),
);
})
.await
.map_err(|rejected| Failed::of(&shown, &rejected, made.load(Ordering::Relaxed)))
}
impl Failed {
fn of(url: &str, rejected: &Rejected, attempts: u32) -> Self {
let reason = reason(url, &rejected.cause);
Self {
reason: match rejected.retryable && attempts > 1 {
true => format!("{reason} (after {attempts} attempts)"),
false => reason,
},
retryable: rejected.retryable,
retry_after: rejected.retry_after,
}
}
}
fn reason(url: &str, cause: &Cause) -> String {
match cause {
Cause::Status(status) => format!("HTTP {status} from {url}"),
Cause::Transport(error) => format!("error fetching {url}: {}", Report::new(error)),
}
}
/// Drops userinfo, query and fragment, where presigned URLs and tokens carry credentials
fn redacted(url: &str) -> Cow<'_, str> {
let Ok(mut parsed) = Url::parse(url) else {
return Cow::Borrowed(url);
};
if parsed.username().is_empty()
&& parsed.password().is_none()
&& parsed.query().is_none()
&& parsed.fragment().is_none()
{
return Cow::Borrowed(url);
}
let _ = parsed.set_username("");
let _ = parsed.set_password(None);
parsed.set_query(None);
parsed.set_fragment(None);
Cow::Owned(parsed.into())
}

6
litellm-rust/crates/model-catalog-refresh/src/lib.rs
View file

@ -1,9 +1,11 @@
mod attempt;
mod error;
mod fetch;
mod policy;
mod source;
mod step;
pub use attempt::Fetched;
pub use error::Error;
pub use fetch::*;
pub use policy::*;
pub use source::*;
pub use step::*;

74
litellm-rust/crates/model-catalog-refresh/src/policy.rs
View file

@ -1,20 +1,76 @@
use std::time::Duration;
use litellm_model_catalog::IntegrityLimits;
use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder};
/// 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)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
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,
pub min_delay: Duration,
pub max_delay: Duration,
}
pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
/// 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";
impl RetryPolicy {
pub(crate) fn backoff(&self) -> ExponentialBackoff {
ExponentialBuilder::new()
.with_min_delay(self.min_delay)
.with_max_delay(self.max_delay)
.with_max_times(self.max_attempts.saturating_sub(1) as usize)
.with_jitter()
.build()
}
pub(crate) fn delay(&self, backoff: Duration, retry_after: Option<Duration>) -> Duration {
retry_after.map_or(backoff, |server| server.min(self.max_delay))
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
const POLICY: RetryPolicy = RetryPolicy {
max_attempts: 3,
timeout: Duration::from_secs(5),
min_delay: Duration::from_secs(1),
max_delay: Duration::from_secs(30),
};
#[rstest]
#[case::backoff_without_retry_after(None, 4)]
#[case::retry_after_wins(Some(7), 7)]
#[case::retry_after_zero(Some(0), 0)]
#[case::retry_after_capped(Some(120), 30)]
fn the_server_retry_after_overrides_the_backoff_up_to_max_delay(
#[case] retry_after: Option<u64>,
#[case] seconds: u64,
) {
assert_eq!(
POLICY.delay(Duration::from_secs(4), retry_after.map(Duration::from_secs)),
Duration::from_secs(seconds)
);
}
#[rstest]
#[case::one_attempt(1, 0)]
#[case::three_attempts(3, 2)]
#[case::zero_attempts(0, 0)]
fn the_backoff_allows_one_retry_fewer_than_the_attempts(
#[case] max_attempts: u32,
#[case] retries: usize,
) {
let policy = RetryPolicy {
max_attempts,
..POLICY
};
assert_eq!(policy.backoff().count(), retries);
}
}

58
litellm-rust/crates/model-catalog-refresh/src/source.rs
View file

@ -1,17 +1,59 @@
use std::path::PathBuf;
use litellm_s3::S3Uri;
use url::Url;
use crate::Error;
/// Where a fetched catalog comes from
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Source {
Http(Url),
File(PathBuf),
S3(S3Object),
S3(S3Uri),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct S3Object {
pub bucket: String,
pub key: String,
impl Source {
pub fn parse(url: &str) -> Result<Self, Error> {
let parsed = Url::parse(url)?;
match parsed.scheme() {
"http" | "https" => Ok(Self::Http(parsed)),
"s3" => Ok(Self::S3(S3Uri::try_from(url)?)),
scheme => Err(Error::unsupported_scheme(scheme.to_owned())),
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::https("https://example.com/cost.json")]
#[case::http("http://127.0.0.1:8000/cost.json")]
fn http_urls_are_fetched_over_http(#[case] url: &str) {
assert_eq!(
Source::parse(url).unwrap(),
Source::Http(Url::parse(url).unwrap())
);
}
#[rstest]
fn s3_urls_name_a_bucket_and_key() {
assert_eq!(
Source::parse("s3://bucket/prices/cost.json").unwrap(),
Source::S3(S3Uri {
bucket: "bucket".into(),
key: "prices/cost.json".into(),
})
);
}
#[rstest]
#[case::s3_without_a_key("s3://bucket/")]
#[case::unsupported("ftp://example.com/cost.json")]
#[case::file("file:///tmp/cost.json")]
#[case::not_a_url("not a url")]
fn other_urls_are_refused(#[case] url: &str) {
assert!(Source::parse(url).is_err());
}
}

66
litellm-rust/crates/model-catalog-refresh/src/step.rs
View file

@ -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<String>,
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<String>,
retry_after: Option<Duration>,
body: Option<Vec<u8>>,
},
/// 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<Catalog>,
installed: Installed,
errors: Vec<litellm_model_catalog::Error>,
},
/// 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<String>,
/// `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 },
}

541
litellm-rust/crates/model-catalog-refresh/tests/fetch.rs Normal file
View file

@ -0,0 +1,541 @@
use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aws_credential_types::Credentials;
use aws_credential_types::provider::error::CredentialsError;
use aws_credential_types::provider::{ProvideCredentials, future};
use litellm_http::Client;
use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch};
use rstest::rstest;
const POLICY: RetryPolicy = RetryPolicy {
max_attempts: 3,
timeout: Duration::from_secs(5),
min_delay: Duration::from_millis(10),
max_delay: Duration::from_millis(50),
};
const OK: &str =
"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}";
const UNAVAILABLE: &str = "HTTP/1.1 503 Service Unavailable\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
const NOT_FOUND: &str = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
const SLOW_BODY: &[u8] = b"{ }";
#[derive(Clone, Copy)]
enum Answer {
Raw(&'static str),
Stall,
StallBody,
SlowBody,
Oversized,
}
/// A scripted upstream: the n-th request gets the n-th answer, repeating the last, and every request head is kept
struct Upstream {
base: String,
heads: Arc<Mutex<Vec<String>>>,
}
impl Upstream {
fn serve(answers: Vec<Answer>) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let base = format!("http://{}", listener.local_addr().unwrap());
let heads = Arc::new(Mutex::new(Vec::new()));
let recorded = Arc::clone(&heads);
std::thread::spawn(move || {
for (index, stream) in listener.incoming().enumerate() {
let answer = answers[index.min(answers.len() - 1)];
let recorded = Arc::clone(&recorded);
std::thread::spawn(move || {
let mut stream = stream.unwrap();
let mut reader = BufReader::new(stream.try_clone().unwrap());
let mut head = String::new();
loop {
let mut line = String::new();
if reader.read_line(&mut line).unwrap() <= 2 {
break;
}
head.push_str(&line);
}
recorded.lock().unwrap().push(head);
match answer {
Answer::Raw(raw) => stream.write_all(raw.as_bytes()).unwrap(),
Answer::Stall => std::thread::sleep(Duration::from_secs(5)),
Answer::Oversized => {
let chunk = vec![b' '; 1024 * 1024];
let written = stream
.write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n")
.and_then(|()| {
(0..=100).try_for_each(|_| stream.write_all(&chunk))
});
drop(written);
}
Answer::StallBody => {
stream
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{")
.unwrap();
std::thread::sleep(Duration::from_secs(5));
}
Answer::SlowBody => {
stream.set_nodelay(true).unwrap();
stream
.write_all(
b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 8\r\nConnection: close\r\n\r\n",
)
.unwrap();
for byte in SLOW_BODY {
std::thread::sleep(Duration::from_millis(60));
stream.write_all(&[*byte]).unwrap();
}
}
}
});
}
});
Self { base, heads }
}
fn url(&self) -> String {
format!("{}/cost.json", self.base)
}
fn heads(&self) -> Vec<String> {
self.heads.lock().unwrap().clone()
}
}
fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> {
head.lines().find_map(|line| {
let (key, value) = line.split_once(':')?;
key.eq_ignore_ascii_case(name).then(|| value.trim())
})
}
#[derive(Clone, Debug)]
struct NoCredentials;
impl ProvideCredentials for NoCredentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
{
future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no credentials")))
}
}
fn keys() -> Credentials {
Credentials::new(
"AKIDMODELCOST",
"model-cost-secret",
None,
None,
"model-cost-test",
)
}
async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result<Fetched, Failed> {
fetch(
&Client::for_test(),
&S3Access::Environment { region: None },
url,
policy,
start,
)
.await
}
#[rstest]
#[case::every_etag_joined(
"HTTP/1.1 200 OK\r\nETag: \"a\"\r\nETag: \"b\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
Some("\"a\", \"b\"")
)]
#[case::no_etag(
"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
None
)]
#[tokio::test]
async fn a_200_returns_the_body_and_its_etags(
#[case] response: &'static str,
#[case] etag: Option<&str>,
) {
let upstream = Upstream::serve(vec![Answer::Raw(response)]);
assert_eq!(
fetched(&upstream.url(), POLICY, Start::Fresh).await,
Ok(Fetched {
body: b"{}".to_vec(),
etag: etag.map(str::to_owned),
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
let url = upstream.url();
let fetched = tokio::spawn(async move {
fetch(
&Client::for_test(),
&S3Access::Environment { region: None },
&url,
POLICY,
Start::Fresh,
)
.await
})
.await
.unwrap();
assert!(fetched.is_ok());
}
#[rstest]
#[tokio::test]
async fn a_redirect_is_not_followed_and_its_body_is_the_document() {
let upstream = Upstream::serve(vec![
Answer::Raw(
"HTTP/1.1 301 Moved Permanently\r\nLocation: /elsewhere.json\r\nContent-Length: 7\r\nConnection: close\r\n\r\n{\"a\":1}",
),
Answer::Raw(OK),
]);
assert_eq!(
fetched(&upstream.url(), POLICY, Start::Fresh).await,
Ok(Fetched {
body: b"{\"a\":1}".to_vec(),
etag: None,
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_retry_after_shorter_than_the_backoff_is_honored() {
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]);
let policy = RetryPolicy {
min_delay: Duration::from_secs(10),
max_delay: Duration::from_secs(10),
..POLICY
};
let fetched = tokio::time::timeout(
Duration::from_secs(2),
fetched(&upstream.url(), policy, Start::Fresh),
)
.await
.expect("Retry-After: 0 replaces the 10s backoff");
assert!(fetched.is_ok(), "{fetched:?}");
assert_eq!(upstream.heads().len(), 2);
}
#[rstest]
#[tokio::test]
async fn retryable_failures_stop_at_the_attempt_limit() {
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
let url = upstream.url();
assert_eq!(
fetched(&url, POLICY, Start::Fresh).await,
Err(Failed {
reason: format!("HTTP 503 from {url} (after 3 attempts)"),
retryable: true,
retry_after: Some(Duration::ZERO),
})
);
assert_eq!(upstream.heads().len(), 3);
}
#[rstest]
#[tokio::test]
async fn a_404_fails_on_the_first_attempt() {
let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
let url = upstream.url();
assert_eq!(
fetched(&url, POLICY, Start::Fresh).await,
Err(Failed {
reason: format!("HTTP 404 from {url}"),
retryable: false,
retry_after: None,
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_stalled_response_times_out_and_is_retried() {
let upstream = Upstream::serve(vec![Answer::Stall]);
let url = upstream.url();
let policy = RetryPolicy {
max_attempts: 2,
timeout: Duration::from_millis(100),
..POLICY
};
let failed = fetched(&url, policy, Start::Fresh).await.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: "))
&& failed.reason.ends_with(" (after 2 attempts)"),
"{}",
failed.reason
);
assert!(failed.retryable);
assert_eq!(upstream.heads().len(), 2);
}
#[rstest]
#[tokio::test]
async fn a_resumed_fetch_makes_only_the_attempts_left() {
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
let url = upstream.url();
let first = Failed {
reason: "first".into(),
retryable: true,
retry_after: Some(Duration::ZERO),
};
let failed = fetched(&url, POLICY, Start::After(first))
.await
.unwrap_err();
assert_eq!(
failed.reason,
format!("HTTP 503 from {url} (after 3 attempts)")
);
assert_eq!(upstream.heads().len(), 2);
}
#[rstest]
#[tokio::test]
async fn a_resumed_fetch_with_no_attempts_left_returns_the_failure_it_resumed() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
let first = Failed {
reason: "first".into(),
retryable: true,
retry_after: None,
};
let policy = RetryPolicy {
max_attempts: 1,
..POLICY
};
assert_eq!(
fetched(&upstream.url(), policy, Start::After(first.clone())).await,
Err(first)
);
assert!(upstream.heads().is_empty());
}
#[rstest]
#[tokio::test]
async fn a_failure_reason_drops_url_credentials() {
let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
let address = upstream.base.trim_start_matches("http://");
let url = format!("http://user:hunter2@{address}/cost.json?X-Amz-Signature=sig-secret#frag");
let failed = fetched(&url, POLICY, Start::Fresh).await.unwrap_err();
assert_eq!(failed.reason, format!("HTTP 404 from {}", upstream.url()));
}
#[rstest]
#[case::unsupported_scheme("ftp://example.com/cost.json")]
#[case::not_a_url("not a url")]
#[case::s3_without_a_key("s3://prices")]
#[tokio::test]
async fn an_unfetchable_url_fails_without_a_request(#[case] url: &str) {
let failed = fetched(url, POLICY, Start::Fresh).await.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
}
fn s3(upstream: &Upstream, credentials: impl ProvideCredentials + 'static) -> S3Access {
S3Access::Config(Box::new(
litellm_s3::config(
"us-west-2".into(),
Some(&upstream.base),
Client::for_test(),
credentials,
)
.build(),
))
}
async fn fetched_from_s3(
url: &str,
upstream: &Upstream,
credentials: impl ProvideCredentials + 'static,
) -> Result<Fetched, Failed> {
fetch(
&Client::for_test(),
&s3(upstream, credentials),
url,
POLICY,
Start::Fresh,
)
.await
}
#[rstest]
#[tokio::test]
async fn an_s3_fetch_sends_a_signed_get_object_to_the_endpoint() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
assert_eq!(
fetched_from_s3("s3://prices/maps/cost.json", &upstream, keys()).await,
Ok(Fetched {
body: b"{}".to_vec(),
etag: Some("\"v1\"".into()),
})
);
let heads = upstream.heads();
let [head] = heads.as_slice() else {
panic!("expected one request, got {heads:?}");
};
assert!(head.starts_with("GET /prices/maps/cost.json"), "{head}");
let authorization = header(head, "authorization").unwrap_or_default();
assert!(
authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIDMODELCOST/")
&& authorization.contains("/us-west-2/s3/aws4_request"),
"{authorization}"
);
}
#[rstest]
#[tokio::test]
async fn an_s3_body_that_stalls_after_the_headers_times_out() {
let upstream = Upstream::serve(vec![Answer::StallBody]);
let url = "s3://prices/cost.json";
let policy = RetryPolicy {
timeout: Duration::from_millis(200),
..POLICY
};
let failed = tokio::time::timeout(
Duration::from_secs(2),
fetch(
&Client::for_test(),
&s3(&upstream, keys()),
url,
policy,
Start::Fresh,
),
)
.await
.expect("the attempt timeout bounds the body read")
.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[case::http(|upstream: &Upstream| (S3Access::Environment { region: None }, upstream.url()))]
#[case::s3(|upstream: &Upstream| (s3(upstream, keys()), "s3://prices/cost.json".to_owned()))]
#[tokio::test]
async fn a_body_that_keeps_arriving_finishes_after_the_timeout(
#[case] source: fn(&Upstream) -> (S3Access, String),
) {
let upstream = Upstream::serve(vec![Answer::SlowBody]);
let (access, url) = source(&upstream);
let policy = RetryPolicy {
timeout: Duration::from_millis(300),
..POLICY
};
assert_eq!(
fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await,
Ok(Fetched {
body: SLOW_BODY.to_vec(),
etag: Some("\"v1\"".into()),
})
);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[case::missing_key(NOT_FOUND)]
#[case::denied("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")]
#[tokio::test]
async fn an_s3_error_fails_without_a_retry(#[case] response: &'static str) {
let upstream = Upstream::serve(vec![Answer::Raw(response), Answer::Raw(OK)]);
let url = "s3://prices/cost.json";
let failed = fetched_from_s3(url, &upstream, keys()).await.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn an_s3_fetch_without_credentials_fails_without_a_request() {
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
let url = "s3://prices/cost.json";
let failed = fetched_from_s3(url, &upstream, NoCredentials)
.await
.unwrap_err();
assert!(
failed
.reason
.starts_with(&format!("error fetching {url}: ")),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert!(upstream.heads().is_empty());
}
const DECLARED_OVERSIZED: &str =
"HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n";
#[rstest]
#[case::http_declared(Answer::Raw(DECLARED_OVERSIZED), false)]
#[case::http_streamed(Answer::Oversized, false)]
#[case::s3_declared(Answer::Raw(DECLARED_OVERSIZED), true)]
#[case::s3_streamed(Answer::Oversized, true)]
#[tokio::test]
async fn a_body_over_100_mib_fails_without_a_retry(#[case] answer: Answer, #[case] s3: bool) {
let upstream = Upstream::serve(vec![answer, Answer::Raw(OK)]);
let failed = match s3 {
true => fetched_from_s3("s3://prices/cost.json", &upstream, keys()).await,
false => fetched(&upstream.url(), POLICY, Start::Fresh).await,
}
.unwrap_err();
assert!(
failed
.reason
.contains("response body is over the 104857600 byte limit"),
"{}",
failed.reason
);
assert!(!failed.retryable);
assert_eq!(upstream.heads().len(), 1);
}
#[rstest]
#[tokio::test]
async fn a_body_of_exactly_100_mib_is_fetched() {
let upstream = Upstream::serve(vec![Answer::Raw(Box::leak(
format!(
"HTTP/1.1 200 OK\r\nContent-Length: 104857600\r\nConnection: close\r\n\r\n{}",
" ".repeat(100 * 1024 * 1024)
)
.into_boxed_str(),
))]);
let fetched = fetched(&upstream.url(), POLICY, Start::Fresh)
.await
.unwrap();
assert_eq!(fetched.body.len(), 100 * 1024 * 1024);
}

9
litellm-rust/crates/model-catalog/AGENTS.md
View file

@ -46,7 +46,14 @@ 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
`cost_map.rs` holds a host's inputs, the installed document and the `register_model` overlay in `register.rs`. It serves the document in the key order and alias sharing Python's `litellm.model_cost` has, without the overlay, since where a registration lands is a lookup in the published table. It never holds the live table, which belongs to the host
- `CostMap` is immutable behind an `Arc`. A new document is a new `CostMap`
- `Registrations<V>` is generic over the host's value type, because Python registrations can hold `datetime` and other values JSON cannot. The host keeps its own objects (`Py<PyDict>` in python-bridge) and does the merging
- Entry order is latest-registration order. `landed_on` records which registered keys landed on each row key
- `write_pass` stops at the first `place` that returns `None` and reports `Pass::Swapped`. Retrying is the host's job
`resolve/` and `model_info.rs` hold types only. The functions over them land with their first caller
Catalog errors use one opaque `Error` with typed causes and JSON Pointer accessors (`path`, `key`, `field`). `Error::rejects` alone decides whether an error fails the document, rejects its row or neither, and `ErrorKind` is the stable name hosts report it by

2
litellm-rust/crates/model-catalog/Cargo.toml
View file

@ -11,6 +11,7 @@ schema = ["dep:schemars", "litellm-cost/schema"]
[dependencies]
litellm-cost = { workspace = true, features = ["serde"] }
litellm-error.workspace = true
litellm-tracing.workspace = true
indexmap = { version = "2.14.0", features = ["serde"] }
jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] }
@ -19,6 +20,7 @@ schemars = { workspace = true, optional = true }
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_with.workspace = true
sha1.workspace = true
strum.workspace = true
thiserror.workspace = true
time.workspace = true

22
litellm-rust/crates/model-catalog/src/catalog.rs
View file

@ -12,7 +12,8 @@ use crate::store::{Rows, Store};
use crate::wire::{EntryWire, Fields};
const SAMPLE_SPEC: &str = "sample_spec";
const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations";
/// The top-level key of the rules block, which is never a model row
pub const FALLBACK_GENERALIZATIONS: &str = "fallback_generalizations";
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct IntegrityLimits {
@ -111,6 +112,11 @@ impl Rows for CatalogRows {
}
impl IntegrityLimits {
pub const UNCHECKED: Self = Self {
min_model_count: 0,
min_reference_ratio: 0.0,
};
pub fn new(min_model_count: usize, min_reference_ratio: f64) -> Option<Self> {
(0.0..=1.0).contains(&min_reference_ratio).then_some(Self {
min_model_count,
@ -418,6 +424,20 @@ impl Catalog {
})
}
/// The row under an exact model or alias key, an alias reading its owner's row
pub fn get(&self, key: &str) -> Option<&Fields> {
self.rows().fields(key)
}
/// A top-level value that is not a model row, such as `sample_spec`
pub fn section(&self, key: &str) -> Option<&Value> {
self.sections
.iter()
.filter(|section| is_section(&section.key))
.find(|section| section.key == key)
.map(|section| &section.value)
}
pub fn model_count(&self) -> usize {
self.rows().entries.len()
}

130
litellm-rust/crates/model-catalog/src/cost_map.rs Normal file
View file

@ -0,0 +1,130 @@
use std::sync::{Arc, LazyLock};
use indexmap::IndexMap;
use serde_json::Value;
use sha1::{Digest, Sha1};
use crate::bundled::BUNDLED_REGISTRY;
use crate::catalog::{Catalog, FALLBACK_GENERALIZATIONS, IntegrityLimits};
use crate::error::Error;
/// Which document is installed: the `git_blob_id` of its bytes and the ETag it was served with
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Provenance {
pub revision: String,
pub etag: Option<String>,
}
/// An installed cost map document. A host builds its live table from it
#[derive(Clone, Debug)]
pub struct CostMap {
base: Arc<Base>,
provenance: Provenance,
}
#[derive(Debug)]
struct Base {
catalog: Catalog,
/// Each owner's names, the owner first, for owners with at least one alias
groups: IndexMap<String, Vec<String>>,
}
impl Base {
fn new(catalog: Catalog) -> Self {
let groups = catalog.aliases().iter().fold(
IndexMap::<String, Vec<String>>::new(),
|mut groups, (alias, owner)| {
groups
.entry(owner.clone())
.or_insert_with(|| vec![owner.clone()])
.push(alias.clone());
groups
},
);
Self { catalog, groups }
}
}
static BUNDLED: LazyLock<CostMap> = LazyLock::new(|| CostMap {
base: Arc::new(Base::new(
Catalog::parse(BUNDLED_REGISTRY)
.expect("the bundled registry parses")
.value,
)),
provenance: Provenance {
revision: git_blob_id(BUNDLED_REGISTRY),
etag: None,
},
});
impl CostMap {
pub fn bundled() -> Self {
BUNDLED.clone()
}
/// Reads a fetched document. One that does not parse, or that has fewer models than `limits` allow against the
/// bundled one, is refused. Every alias problem it has is logged
pub fn read(body: &[u8], etag: Option<String>, limits: IntegrityLimits) -> Result<Self, Error> {
let loaded = Catalog::parse(body)?;
loaded
.value
.check_size(BUNDLED.base.catalog.model_count(), limits)?;
for problem in loaded.errors.iter().filter(|error| error.is_alias()) {
litellm_tracing::warn!(
"LiteLLM: model cost map alias problem: {}",
litellm_error::Report::new(problem)
);
}
Ok(Self {
base: Arc::new(Base::new(loaded.value)),
provenance: Provenance {
revision: git_blob_id(body),
etag,
},
})
}
pub fn catalog(&self) -> &Catalog {
&self.base.catalog
}
pub fn provenance(&self) -> &Provenance {
&self.provenance
}
/// Every name the row under `key` is served under, the owner first. Empty for a row without aliases
pub fn names(&self, key: &str) -> &[String] {
let owner = self
.catalog()
.aliases()
.get(key)
.map_or(key, String::as_str);
self.base
.groups
.get(owner)
.map(Vec::as_slice)
.unwrap_or_default()
}
/// Each owner with its names, the owner first, in the order the document declared the aliases
pub fn alias_groups(&self) -> impl Iterator<Item = (&str, &[String])> {
self.base
.groups
.iter()
.map(|(owner, names)| (owner.as_str(), names.as_slice()))
}
/// The `fallback_generalizations` block as the document wrote it
pub fn fallback_generalizations(&self) -> Option<&Value> {
self.catalog().section(FALLBACK_GENERALIZATIONS)
}
}
/// `git hash-object` of the bytes, which main reports as `source_revision`
pub fn git_blob_id(body: &[u8]) -> String {
let digest = Sha1::new()
.chain_update(format!("blob {}\0", body.len()))
.chain_update(body)
.finalize();
format!("{digest:x}")
}

68
litellm-rust/crates/model-catalog/src/error.rs
View file

@ -167,6 +167,74 @@ impl Error {
}
}
impl Error {
/// The document parsed but holds fewer models than the integrity limits trust
pub fn is_integrity(&self) -> bool {
match self.repr() {
Repr::Empty | Repr::BelowMinimum { .. } | Repr::Shrunk { .. } => true,
Repr::Json(_)
| Repr::NotAnObject
| Repr::Invalid(_)
| Repr::UnknownBlockKey
| Repr::UnknownTimezone { .. }
| Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
| Repr::TierReversed
| Repr::TierGap { .. }
| Repr::MixedTierBasis
| Repr::NotPricedByResults
| Repr::BandsFromTwoSources
| Repr::NotAnOffPeakRate
| Repr::NoOffPeakPrices
| Repr::NoOffPeakHours
| Repr::EmptyList
| Repr::InvalidDate { .. }
| Repr::MissingProvider
| Repr::UnreadField
| Repr::NoFallbackRules => false,
}
}
/// A problem with a row's aliases, which the served map skips
pub fn is_alias(&self) -> bool {
match self.repr() {
Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField => true,
Repr::Invalid(_) => matches!(self.field().as_deref(), Some("aliases" | "alias_of")),
Repr::Json(_)
| Repr::Empty
| Repr::BelowMinimum { .. }
| Repr::Shrunk { .. }
| Repr::NotAnObject
| Repr::UnknownBlockKey
| Repr::UnknownTimezone { .. }
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
| Repr::TierReversed
| Repr::TierGap { .. }
| Repr::MixedTierBasis
| Repr::NotPricedByResults
| Repr::BandsFromTwoSources
| Repr::NotAnOffPeakRate
| Repr::NoOffPeakPrices
| Repr::NoOffPeakHours
| Repr::EmptyList
| Repr::InvalidDate { .. }
| Repr::MissingProvider
| Repr::UnreadField
| Repr::NoFallbackRules => false,
}
}
}
#[derive(Debug)]
pub struct Loaded<T> {
pub value: T,

32
litellm-rust/crates/model-catalog/src/layers.rs
View file

@ -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<Catalog>,
/// 1, the last installed fetch, rejected rows left out
fetched: Option<Arc<Catalog>>,
/// 2, `register(persist=False)`, cleared by every install and refilled by the router replay
transient: IndexMap<String, Fields>,
/// 3, `register(persist=True)`, applied over every install
persisted: Registrations,
}

7
litellm-rust/crates/model-catalog/src/lib.rs
View file

@ -17,6 +17,9 @@
//! failing the document.
//! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known
//! families. They are data in the file, not entries.
//! - `CostMap`: an installed document with its provenance, which a host builds its live table from.
//! `Registrations` is the ordered `register_model` overlay over the host's own values, and
//! `write_pass` writes registrations into a host's table through `RegistrationTable`.
//! - Validation: `Catalog::check_size` checks a catalog against `IntegrityLimits` (a remote copy
//! that lost most of its models is refused), and `Catalog::row_errors` validates one row.
//!
@ -37,9 +40,9 @@ macro_rules_attribute::attribute_alias! {
mod bundled;
mod catalog;
mod cost_map;
mod error;
mod fallback;
mod layers;
mod lower;
mod model;
mod model_info;
@ -51,9 +54,9 @@ pub mod wire;
pub use bundled::*;
pub use catalog::*;
pub use cost_map::*;
pub use error::*;
pub use fallback::*;
pub use layers::*;
pub use model::*;
pub use model_info::*;
pub use register::*;

1
litellm-rust/crates/model-catalog/src/model_info.rs
View file

@ -28,6 +28,7 @@ pub struct ModelInfoFields {
pub cache_read_input_audio_token_cost_flex: Option<f64>,
pub cache_read_input_audio_token_cost_priority: Option<f64>,
pub cache_read_input_audio_token_cost_batches: Option<f64>,
pub cache_read_input_image_token_cost_batches: Option<f64>,
pub explicit_cache_read_input_token_cost: Option<f64>,
pub prompt_cache_min_tokens: Option<i64>,
pub cache_read_input_token_cost_above_200k_tokens: Option<f64>,

266
litellm-rust/crates/model-catalog/src/register.rs
View file

@ -1,26 +1,252 @@
use std::collections::{BTreeSet, HashMap};
use std::collections::BTreeMap;
use indexmap::IndexMap;
use crate::wire::Fields;
/// Where one `register` call writes and which row it inherits from
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RegistrationTarget {
/// Always the caller's key, never a case variant of it
pub write: String,
/// The first candidate that is an exact key whose row passes `provider_matches`
pub inherit_from: Option<String>,
/// `register_model(persist_across_reloads=True)` calls, replayed above every installed document
#[derive(Clone, Debug, PartialEq)]
pub struct Registrations<V> {
/// Each registered key's merged value, ordered by its latest registration
entries: IndexMap<String, V>,
/// The registered keys whose write landed on each row key, in landing order
landed_on: IndexMap<String, Vec<String>>,
}
/// 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<String>);
/// 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<String, Fields>,
registered_fields: HashMap<String, BTreeSet<String>>,
impl<V> Default for Registrations<V> {
fn default() -> Self {
Self {
entries: IndexMap::new(),
landed_on: IndexMap::new(),
}
}
}
/// One registration written to a host's table: the key it was registered under and the row key it landed on
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Landed {
pub registered: String,
pub written: String,
}
impl<V> Registrations<V> {
/// Keeps `value` as the key's registration, already merged with `get(key)` by the host, and makes it the latest
pub fn register(&mut self, key: &str, value: V) {
self.entries.shift_remove(key);
self.entries.insert(key.to_owned(), value);
}
pub fn get(&self, key: &str) -> Option<&V> {
self.entries.get(key)
}
pub fn record(&mut self, landed: &[Landed]) {
for Landed {
registered,
written,
} in landed
{
let names = self.landed_on.entry(written.clone()).or_default();
if !names.contains(registered) {
names.push(registered.clone());
}
}
}
/// Every registration, oldest first
pub fn entries(&self) -> impl Iterator<Item = (&str, &V)> {
self.entries
.iter()
.map(|(key, value)| (key.as_str(), value))
}
/// The registrations that apply to the row under `key`, oldest first: `key`'s own and every one that landed on it
pub fn over<'a>(&'a self, key: &str) -> Vec<&'a V> {
let landed = self.landed_on.get(key).into_iter().flatten();
let by_order: BTreeMap<usize, &V> = landed
.map(String::as_str)
.chain([key])
.filter_map(|name| self.entries.get_full(name))
.map(|(order, _, value)| (order, value))
.collect();
by_order.into_values().collect()
}
pub fn map<W>(&self, mut convert: impl FnMut(&V) -> W) -> Registrations<W> {
Registrations {
entries: self
.entries
.iter()
.map(|(key, value)| (key.clone(), convert(value)))
.collect(),
landed_on: self.landed_on.clone(),
}
}
}
/// A host's live table that registrations are written into
pub trait RegistrationTable {
type Registration;
type Error;
/// Writes `registration` where it lands, `None` when this table is no longer the one the host publishes
fn place(
&mut self,
key: &str,
registration: &Self::Registration,
) -> Result<Option<Placed>, Self::Error>;
}
/// Where a registration was written and whether a built-in row matched it
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Placed {
pub key: String,
pub matched_builtin: bool,
}
#[derive(Debug, PartialEq)]
pub enum Pass<'a, R> {
/// Every entry was written. `unmatched` are those no built-in row matched
Finished {
landed: Vec<Landed>,
unmatched: Vec<(&'a str, &'a R)>,
},
/// The table stopped being the published one before every entry was written
Swapped { landed: Vec<Landed> },
}
impl<R> Pass<'_, R> {
pub fn landed(&self) -> &[Landed] {
match self {
Pass::Finished { landed, .. } | Pass::Swapped { landed } => landed,
}
}
}
/// Writes `entries` into `table` in order
pub fn write_pass<'a, T: RegistrationTable>(
table: &mut T,
entries: &'a [(String, T::Registration)],
) -> Result<Pass<'a, T::Registration>, T::Error> {
let mut landed = Vec::with_capacity(entries.len());
let mut unmatched = Vec::new();
for (key, registration) in entries {
let Some(written) = table.place(key, registration)? else {
return Ok(Pass::Swapped { landed });
};
if !written.matched_builtin {
unmatched.push((key.as_str(), registration));
}
landed.push(Landed {
registered: key.clone(),
written: written.key,
});
}
Ok(Pass::Finished { landed, unmatched })
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn landed(registered: &str, written: &str) -> Landed {
Landed {
registered: registered.to_owned(),
written: written.to_owned(),
}
}
#[rstest]
fn reregistering_a_key_makes_it_the_latest() {
let mut registrations = Registrations::default();
registrations.register("a", 1);
registrations.register("b", 2);
registrations.register("a", 3);
assert_eq!(
registrations.entries().collect::<Vec<_>>(),
[("b", &2), ("a", &3)]
);
}
#[rstest]
fn over_applies_names_landed_on_the_row_in_registration_order() {
let mut registrations = Registrations::default();
registrations.register("openai/o3-pro", 1);
registrations.register("o3-pro", 2);
registrations.register("other", 3);
registrations.record(&[
landed("openai/o3-pro", "o3-pro"),
landed("o3-pro", "o3-pro"),
]);
registrations.register("openai/o3-pro", 4);
assert_eq!(registrations.over("o3-pro"), [&2, &4]);
}
#[rstest]
fn over_applies_a_keys_own_registration_that_landed_on_another_row() {
let mut registrations = Registrations::default();
registrations.register("openai/o3-pro", 1);
registrations.record(&[landed("openai/o3-pro", "o3-pro")]);
assert_eq!(
(
registrations.over("openai/o3-pro"),
registrations.over("o3-pro")
),
(vec![&1], vec![&1])
);
}
/// A table published until `swaps_after` writes, whose rows match a built-in only when the key starts with `b`
struct Table {
writes: Vec<(String, i32)>,
swaps_after: usize,
}
impl RegistrationTable for Table {
type Registration = i32;
type Error = ();
fn place(&mut self, key: &str, registration: &i32) -> Result<Option<Placed>, ()> {
if self.writes.len() >= self.swaps_after {
return Ok(None);
}
let written = key.to_uppercase();
self.writes.push((written.clone(), *registration));
Ok(Some(Placed {
matched_builtin: written.starts_with('B'),
key: written,
}))
}
}
#[rstest]
#[case::published_throughout(
usize::MAX,
Pass::Finished { landed: vec![landed("a", "A"), landed("b", "B")], unmatched: vec![("a", &1)] },
&[("A", 1), ("B", 2)],
)]
#[case::swapped_after_the_first_write(
1,
Pass::Swapped { landed: vec![landed("a", "A")] },
&[("A", 1)],
)]
fn a_pass_writes_entries_in_order_until_the_table_is_swapped(
#[case] swaps_after: usize,
#[case] expected: Pass<'static, i32>,
#[case] writes: &[(&str, i32)],
) {
let entries = [("a".to_owned(), 1), ("b".to_owned(), 2)];
let mut table = Table {
writes: Vec::new(),
swaps_after,
};
assert_eq!(write_pass(&mut table, &entries), Ok(expected));
assert_eq!(
table.writes,
writes
.iter()
.map(|(key, value)| ((*key).to_owned(), *value))
.collect::<Vec<_>>()
);
}
}

2
litellm-rust/crates/model-catalog/src/store/AGENTS.md
View file

@ -24,7 +24,7 @@
| Source | Stamp | Blind spot |
|---|---|---|
| `Catalog`'s `CatalogRows` | `()`, every change goes through `Store::edit`, which drops the index | none |
| python-bridge's bound `litellm.model_cost` | the dict object and its `len` | a delete plus an add between two reads, or a `litellm_provider` rewrite on an existing key |
| python-bridge's bound `litellm.model_cost` | the dict object, its `len` and its last key | deletes plus adds between two reads that put the old last key back last, or a `litellm_provider` rewrite on an existing key |
A source with a blind spot calls `Store::invalidate` when it knows keys changed. Until then a key added in the blind spot is not found, and a key moved to another provider stays listed under the old one. Removed keys are still never returned

53
litellm-rust/crates/model-catalog/tests/cost_map.rs Normal file
View file

@ -0,0 +1,53 @@
use litellm_model_catalog::{
BUNDLED_REGISTRY, CostMap, ErrorKind, IntegrityLimits, Provenance, git_blob_id,
};
use rstest::rstest;
use serde_json::json;
fn any_size() -> IntegrityLimits {
IntegrityLimits::new(0, 0.0).unwrap()
}
#[rstest]
#[case::not_json(b"{".as_slice(), false, ErrorKind::Json)]
#[case::empty(b"{}".as_slice(), true, ErrorKind::Empty)]
fn a_refused_document_says_whether_it_failed_the_integrity_limits(
#[case] body: &[u8],
#[case] integrity: bool,
#[case] kind: ErrorKind,
) {
let error = CostMap::read(body, None, any_size()).unwrap_err();
assert_eq!((error.is_integrity(), error.kind()), (integrity, kind));
}
#[rstest]
fn provenance_is_the_read_bytes_and_etag() {
let body = br#"{"m":{"litellm_provider":"p"}}"#;
let cost_map = CostMap::read(body, Some("\"v1\"".into()), any_size()).unwrap();
assert_eq!(
cost_map.provenance(),
&Provenance {
revision: "eaa747f210a1527d5ea7138be0f768c47578ca68".to_owned(),
etag: Some("\"v1\"".into()),
}
);
assert_eq!(
CostMap::bundled().provenance().revision,
git_blob_id(BUNDLED_REGISTRY)
);
}
#[rstest]
fn every_name_of_an_aliased_row_lists_its_owner_first(
#[values("owner", "short", "owner-latest")] key: &str,
) {
let body = json!({
"owner": {"litellm_provider": "p", "aliases": ["short"]},
"owner-latest": {"alias_of": "owner"},
"plain": {"litellm_provider": "p"},
})
.to_string();
let cost_map = CostMap::read(body.as_bytes(), None, any_size()).unwrap();
assert_eq!(cost_map.names(key), ["owner", "short", "owner-latest"]);
assert!(cost_map.names("plain").is_empty());
}

1
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -60,6 +60,7 @@ litellm-router-types.workspace = true
litellm-router-core.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-model-catalog-refresh.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
indexmap = "2.14.0"
pyo3 = { workspace = true, features = ["indexmap"] }

2
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -6,6 +6,8 @@ pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandleError {
#[error("this inference route does not support configured caches")]

2
litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md vendored
View file

@ -4,6 +4,6 @@ This directory lets Rust inference use a selected Python cache. `service.rs` imp
Keep the shared inference protocol wrapper and adapter selection in the parent module. Receive the configured cache and prepared arguments from the parent module. Do not discover global configuration, choose native backends, or move inference to Python. Core remains independent of Python objects and cache implementation details
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Await asynchronous cache operations through the existing host driver in the caller's task. Do not create another asyncio task or event loop. The one exception is the asynchronous unary messages store: Python's `_async_set_cache` writes the returned response from `create_cache_write_task` after the router has stamped its `_hidden_params`, so `host.rs` defers that store until the deployment transform has returned the response the caller receives and hands that response to `litellm.rust_bridge.host.cache.store_response`, which schedules the same write task. Lookups, sync stores and stream stores stay in the caller's task. Cancellation must prevent subsequent provider requests and cache writes. Preserve ordinary cache failure handling without swallowing cancellation or other Python base exceptions
Traverse retained Python references for GC and release pending replies when the call closes. Regression tests must require native inference with Python fallback disabled and assert observable hits, provider request counts, cache-key headers, task identity and cancellation

27
litellm-rust/crates/python-bridge/src/cache/python/host.rs vendored
View file

@ -9,7 +9,7 @@ use pyo3::{
use serde_json::Value;
use super::service::CacheCall;
use crate::cache::HandleError;
use crate::cache::{CACHE_HOST_MODULE, HandleError};
enum Pending {
Lookup(Reply<Result<Option<Value>, Error>>),
@ -21,6 +21,7 @@ pub(crate) struct PythonCache {
arguments: Option<Py<PyDict>>,
pending: Option<Pending>,
asynchronous: bool,
deferred: bool,
}
impl PythonCache {
@ -30,6 +31,7 @@ impl PythonCache {
arguments: None,
pending: None,
asynchronous,
deferred: false,
}
}
@ -64,7 +66,14 @@ impl PythonCache {
None,
)
}
CacheCall::Store { value, reply } => {
CacheCall::Store {
unary: true, reply, ..
} if self.asynchronous => {
self.deferred = true;
reply.send(Ok(()));
return Ok(None);
}
CacheCall::Store { value, reply, .. } => {
self.pending = Some(Pending::Store(reply));
(
if self.asynchronous {
@ -89,6 +98,19 @@ impl PythonCache {
self.resume(py, result)
}
pub fn store_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
if !std::mem::take(&mut self.deferred) {
return Ok(());
}
let (Some(cache), Some(arguments)) = (self.cache.as_ref(), self.arguments.as_ref()) else {
return Err(HandleError::Unconfigured.into());
};
py.import(CACHE_HOST_MODULE)?
.getattr("store_response")?
.call((cache, response), Some(&arguments.bind(py).copy()?))?;
Ok(())
}
pub fn resume(
&mut self,
py: Python<'_>,
@ -124,6 +146,7 @@ impl PythonCache {
pub fn close(&mut self) {
self.pending = None;
self.deferred = false;
self.cache = None;
self.arguments = None;
}

8
litellm-rust/crates/python-bridge/src/cache/python/service.rs vendored
View file

@ -45,6 +45,7 @@ pub(crate) enum CacheCall {
},
Store {
value: Value,
unary: bool,
reply: Reply<Result<(), Error>>,
},
}
@ -101,8 +102,13 @@ where
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
Box::pin(async move {
let value = to_python(value)?;
let unary = value.get(STREAM_EVENTS_KEY).is_none();
self.services
.call(|reply| CacheCall::Store { value, reply })
.call(|reply| CacheCall::Store {
value,
unary,
reply,
})
.await
.map_err(Error::unavailable)?
})

3
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -1,11 +1,10 @@
use super::{HandleError, python};
use super::{CACHE_HOST_MODULE, HandleError, python};
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::Protocol,
};
use pyo3::{prelude::*, types::PyDict};
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);

5
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -29,6 +29,8 @@ const MESSAGES_UNSUPPORTED_PARAM: &str = "messages_unsupported_param";
/// The phase label (`Connection` or `Read`) of a timed-out call, for Python's timeout text, or
/// `Body` for a streamed body read that timed out after the head arrived.
const MESSAGES_TIMEOUT: &str = "messages_timeout";
// `(received, expected, reset)` of a streamed body the connection ended early.
const MESSAGES_BODY_CUT: &str = "messages_body_cut";
/// No credential reached the provider, which Python raises as `AuthenticationError`.
const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential";
/// Python refuses the request only once its logger knows the resolved model and provider (the
@ -134,6 +136,9 @@ pub(crate) fn route_failure(
};
value.setattr(MESSAGES_TIMEOUT, phase)?;
}
if let Some(cut) = error.body_cut() {
value.setattr(MESSAGES_BODY_CUT, (cut.received, cut.expected, cut.reset))?;
}
if error.is_missing_credential() {
value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?;
}

30
litellm-rust/crates/python-bridge/src/http.rs
View file

@ -108,13 +108,35 @@ pub(crate) fn call_config(
Some(snapshot) => configured(&snapshot)?,
None => HttpSettingsLayer::default(),
},
])
.without_missing_files(&|path: &Path| path.exists());
let resolution = Resolution::from(&settings);
]);
Ok(resolved(py, settings))
}
/// The client config for a value of `litellm.rust_bridge.host.settings.http_settings()` the caller handed in, under
/// the process environment. `None` reads the environment alone
pub(crate) fn settings_config(
py: Python<'_>,
http_settings: Option<&Bound<'_, PyAny>>,
) -> PyResult<HttpClientConfig> {
let configured = match http_settings {
Some(value) => configured(&PythonSettings::Http.snapshot(value.clone()))?,
None => HttpSettingsLayer::default(),
};
let settings = HttpSettings::from_layers([
HttpSettingsLayer::from_environment(&ProcessEnvironment),
configured,
]);
Ok(resolved(py, settings))
}
/// The client config `settings` resolve to, warning once per process about each setting the client cannot apply
fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig {
let resolution =
Resolution::from(&settings.without_missing_files(&|path: &Path| path.exists()));
for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) {
crate::logger::capture(py).scope(|| litellm_tracing::warn!("{unsupported}"));
}
Ok(resolution.config)
resolution.config
}
pub(crate) fn provider_client(

18
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -9,6 +9,7 @@ mod execution;
mod http;
mod logger;
mod marshal;
mod model_cost;
mod model_registry;
mod preflight;
mod python_settings;
@ -36,6 +37,12 @@ mod _native {
#[pymodule_export]
use crate::logger::NativeDiagnosticProcessor;
#[pymodule_export]
use crate::model_cost::{
NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched, NativeInstalled,
NativeModelCostMap, NativeModelCostMapState, NativeProvenance, NativeRegistered,
NativeRejected, NativeTableSwapped,
};
#[pymodule_export]
use crate::model_registry::bundled_model_cost_map;
#[pymodule_export]
use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states};
@ -166,6 +173,17 @@ mod tests {
"NativeDatabase",
"trace_encode_error",
"bundled_model_cost_map",
"NativeModelCostMap",
"NativeModelCostMapState",
"NativeInstalled",
"NativeRejected",
"NativeRegistered",
"NativeTableSwapped",
"NativeProvenance",
"NativeCostMapSource",
"NativeFetchPolicy",
"NativeFetched",
"NativeFetchFailed",
"TokenCounter",
"Tokenizer",
"RouterCoreHandle",

16
litellm-rust/crates/python-bridge/src/model_cost/AGENTS.md Normal file
View file

@ -0,0 +1,16 @@
# Model cost bridge
- `NativeModelCostMap` holds inputs only, the installed `CostMap`, kept `Registrations` and the name index. Python's `litellm.model_cost` is the only copy of the live rows
- Every call that reads or writes rows takes the dict as an argument. Never store rows or build a second table here
- `LiveDict` holds a reference to that same dict and reads it under the GIL on every query, so a Python write is seen on the next read
- It needs the GIL, so Rust code running without the GIL cannot use it
- Its stamp is the dict's address, length and last key. [model-catalog/src/store/AGENTS.md](../../../model-catalog/src/store/AGENTS.md) says what that misses
- `host.rs` is the only code that calls Python, by name through `MODULE`. Calls happen with the GIL held and no `NativeModelCostMap` lock held
- `resolve` reaches `find_key` on this same object, so holding a lock across a host call deadlocks
- Test `a_host_lookup_that_reaches_find_key_on_the_same_map_finishes` guards this
- Rows and registrations cross as Python objects, never JSON, because router `model_info` holds `datetime`
- `table::build` serves the document as main's `_expand_model_aliases` did, so all names of a model share one row object made by `host.entry`
- `install` touches no state until `CostMap::read` succeeds. `NativeRejected` means nothing changed
- `register` keeps a registration only after its pass finished, so an exception in a host call keeps nothing
- `rows=None` writes each entry into whichever table is published at that moment and never reports a swap
- `fetch_blocking` runs on a throwaway runtime with an unpooled client so an importing process can still fork. `fetch` is for workers after the fork

239
litellm-rust/crates/python-bridge/src/model_cost/fetch.rs Normal file
View file

@ -0,0 +1,239 @@
use std::time::Duration;
use litellm_http::{ClientVariant, HttpClientConfig};
use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
/// A cost map URL and how to connect to it. An `s3://` URL reaches S3 as an AWS SDK finds it in the environment,
/// with LiteLLM's `AWS_REGION_NAME` winning over the region it finds
#[pyclass(
frozen,
name = "NativeCostMapSource",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeCostMapSource {
url: String,
http: HttpClientConfig,
}
#[pymethods]
impl NativeCostMapSource {
/// `http_settings` is `litellm.rust_bridge.host.settings.http_settings()`, or `None` for the environment alone
#[new]
#[pyo3(signature = (url, *, http_settings))]
fn new(py: Python<'_>, url: String, http_settings: Option<Bound<'_, PyAny>>) -> PyResult<Self> {
Ok(Self {
url,
http: crate::http::settings_config(py, http_settings.as_ref())?,
})
}
/// Fetches on the shared runtime with a pooled client
fn fetch<'py>(
&self,
py: Python<'py>,
policy: PyRef<'_, NativeFetchPolicy>,
) -> PyResult<Bound<'py, PyAny>> {
let client = crate::http::pool()
.client(&self.http, ClientVariant::Provider)
.map_err(crate::http::client_error)?;
let url = self.url.clone();
let policy = policy.0;
crate::execution::run_async_value(py, async move {
Ok(Outcome(
litellm_model_catalog_refresh::fetch(&client, &s3(), &url, policy, Start::Fresh)
.await,
))
})
}
/// Fetches on a throwaway runtime with a client of its own. `resume_from` continues a fetch whose first attempt
/// failed retryably
#[pyo3(signature = (policy, resume_from = None))]
fn fetch_blocking(
&self,
py: Python<'_>,
policy: PyRef<'_, NativeFetchPolicy>,
resume_from: Option<PyRef<'_, NativeFetchFailed>>,
) -> PyResult<Py<PyAny>> {
let start = resume_from.map_or(Start::Fresh, |failed| Start::After(failed.0.clone()));
let (config, url, policy) = (self.http.clone(), self.url.clone(), policy.0);
let logger = crate::logger::capture(py);
let fetched = litellm_host_python::release_gil(py, move || {
logger.scope(|| {
let client = crate::http::pool()
.unpooled(&config, ClientVariant::Provider)
.map_err(crate::http::client_error)?;
litellm_host_python::run_isolated(litellm_model_catalog_refresh::fetch(
&client,
&s3(),
&url,
policy,
start,
))
.map_err(PyErr::from)
})
})?;
Outcome(fetched).into_pyobject(py).map(Bound::unbind)
}
}
fn s3() -> S3Access {
S3Access::Environment {
region: std::env::var("AWS_REGION_NAME").ok(),
}
}
#[pyclass(
frozen,
name = "NativeFetchPolicy",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeFetchPolicy(RetryPolicy);
#[pymethods]
impl NativeFetchPolicy {
#[new]
/// Retries wait an exponential backoff with jitter from `min_delay_seconds` up to `max_wait_seconds`, or the
/// server's `Retry-After` up to that cap
#[pyo3(signature = (*, max_attempts, max_wait_seconds, timeout_seconds, min_delay_seconds))]
fn new(
max_attempts: u32,
max_wait_seconds: f64,
timeout_seconds: f64,
min_delay_seconds: f64,
) -> PyResult<Self> {
if max_attempts == 0 {
return Err(PyValueError::new_err("max_attempts must be at least 1"));
}
let seconds = |name: &str, value: f64| {
Duration::try_from_secs_f64(value).map_err(|_| {
PyValueError::new_err(format!("{name} must be a non-negative number of seconds"))
})
};
let min_delay = seconds("min_delay_seconds", min_delay_seconds)?;
let max_delay = seconds("max_wait_seconds", max_wait_seconds)?;
if min_delay > max_delay {
return Err(PyValueError::new_err(
"min_delay_seconds must not exceed max_wait_seconds",
));
}
Ok(Self(RetryPolicy {
max_attempts,
timeout: seconds("timeout_seconds", timeout_seconds)?,
min_delay,
max_delay,
}))
}
}
#[pyclass(frozen, name = "NativeFetched", module = "litellm.rust_bridge._native")]
pub(crate) struct NativeFetched {
#[pyo3(get)]
body: Py<PyBytes>,
#[pyo3(get)]
etag: Option<String>,
}
#[pyclass(
frozen,
skip_from_py_object,
name = "NativeFetchFailed",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeFetchFailed(Failed);
#[pymethods]
impl NativeFetchFailed {
#[getter]
fn reason(&self) -> &str {
&self.0.reason
}
/// Another attempt could succeed. Only the attempt limit ended the fetch
#[getter]
fn retryable(&self) -> bool {
self.0.retryable
}
#[getter]
fn retry_after(&self) -> Option<f64> {
self.0.retry_after.map(|wait| wait.as_secs_f64())
}
}
struct Outcome(Result<Fetched, Failed>);
impl<'py> IntoPyObject<'py> for Outcome {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
match self.0 {
Ok(fetched) => Bound::new(
py,
NativeFetched {
body: PyBytes::new(py, &fetched.body).unbind(),
etag: fetched.etag,
},
)
.map(Bound::into_any),
Err(failed) => Bound::new(py, NativeFetchFailed(failed)).map(Bound::into_any),
}
}
}
#[cfg(test)]
mod tests {
use std::io::{BufRead, BufReader, Write};
use std::net::TcpListener;
use rstest::rstest;
use super::*;
#[rstest]
#[case::equal_bounds(2.0, 2.0, true)]
#[case::min_below_max(2.0, 30.0, true)]
#[case::min_above_max(31.0, 30.0, false)]
fn a_policy_accepts_a_min_delay_only_up_to_the_max_wait(
#[case] min_delay_seconds: f64,
#[case] max_wait_seconds: f64,
#[case] accepted: bool,
) {
Python::initialize();
let policy = NativeFetchPolicy::new(3, max_wait_seconds, 5.0, min_delay_seconds);
assert_eq!(policy.is_ok(), accepted);
}
#[rstest]
fn a_blocking_fetch_hands_python_the_document_without_starting_the_shared_runtime() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let url = format!("http://{}/cost.json", listener.local_addr().unwrap());
std::thread::spawn(move || {
let mut stream = listener.incoming().next().unwrap().unwrap();
let mut reader = BufReader::new(stream.try_clone().unwrap());
let mut line = String::new();
while reader.read_line(&mut line).unwrap() > 2 {
line.clear();
}
stream
.write_all(b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}")
.unwrap();
});
Python::initialize();
Python::attach(|py| {
let source = NativeCostMapSource::new(py, url, None).unwrap();
let policy =
Bound::new(py, NativeFetchPolicy::new(3, 30.0, 5.0, 2.0).unwrap()).unwrap();
let outcome = source.fetch_blocking(py, policy.borrow(), None).unwrap();
let fetched = outcome.bind(py).cast::<NativeFetched>().unwrap().get();
assert_eq!(fetched.body.bind(py).as_bytes(), b"{}");
assert_eq!(fetched.etag.as_deref(), Some("\"v1\""));
});
assert!(!litellm_host_python::runtime_started());
}
}

81
litellm-rust/crates/python-bridge/src/model_cost/host.rs Normal file
View file

@ -0,0 +1,81 @@
//! Calls into `litellm.rust_bridge.host.model_cost`
use pyo3::call::PyCallArgs;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};
pub(crate) const MODULE: &str = "litellm.rust_bridge.host.model_cost";
/// Where a registration lands in the published table and the row it starts from
#[derive(FromPyObject)]
pub(crate) struct Landing<'py> {
#[pyo3(attribute)]
pub(crate) key: String,
#[pyo3(attribute)]
pub(crate) base: Bound<'py, PyAny>,
#[pyo3(attribute)]
pub(crate) matched_builtin: bool,
}
fn call<'py>(
py: Python<'py>,
name: &str,
args: impl PyCallArgs<'py>,
) -> PyResult<Bound<'py, PyAny>> {
py.import(MODULE)?.getattr(name)?.call1(args)
}
/// The row object every name in `names` shares, the owner first
pub(crate) fn entry<'py>(
fields: Bound<'py, PyAny>,
names: &[String],
) -> PyResult<Bound<'py, PyAny>> {
let py = fields.py();
call(py, "entry", (fields, PyTuple::new(py, names)?))
}
/// Where `registration` lands in `rows`, `None` when `rows` is no longer the published table. `rows` of `None` is the
/// published table
pub(crate) fn resolve<'py>(
rows: Option<&Bound<'py, PyDict>>,
key: &str,
registration: &Bound<'py, PyAny>,
custom_llm_provider: Option<&str>,
) -> PyResult<Option<Landing<'py>>> {
call(
registration.py(),
"resolve",
(rows, key, registration, custom_llm_provider),
)?
.extract()
}
/// The row `registration` writes over `current`
pub(crate) fn merged<'py>(
current: Option<&Bound<'py, PyAny>>,
landing: &Landing<'py>,
registration: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> {
call(
registration.py(),
"merged",
(current, &landing.base, registration),
)
}
/// The registration kept for a key: `registration` merged into the one kept before it
pub(crate) fn kept<'py>(
previous: Option<&Bound<'py, PyDict>>,
registration: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyDict>> {
Ok(call(registration.py(), "kept", (previous, registration))?.cast_into::<PyDict>()?)
}
pub(crate) fn published(py: Python<'_>) -> PyResult<Bound<'_, PyDict>> {
Ok(call(py, "published", ())?.cast_into::<PyDict>()?)
}
/// Writes `written` into `rows`, the published table, through the host, which tracks the table's version
pub(crate) fn write(rows: &Bound<'_, PyDict>, written: &Bound<'_, PyDict>) -> PyResult<()> {
call(rows.py(), "write", (rows, written)).map(drop)
}

176
litellm-rust/crates/python-bridge/src/model_cost/live.rs Normal file
View file

@ -0,0 +1,176 @@
use std::borrow::Cow;
use litellm_model_catalog::Rows;
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::types::PyDict;
/// The dict Python last handed in as `litellm.model_cost`, read under the GIL on every query
pub(crate) struct LiveDict(Py<PyDict>);
/// One row exactly as the dict holds it
pub(crate) struct LiveRow(Py<PyAny>);
impl Clone for LiveRow {
fn clone(&self) -> Self {
Python::attach(|py| Self(self.0.clone_ref(py)))
}
}
impl LiveDict {
pub(crate) fn new(dict: &Bound<'_, PyDict>) -> Self {
Self(dict.clone().unbind())
}
pub(crate) fn is(&self, dict: &Bound<'_, PyDict>) -> bool {
self.0.bind(dict.py()).is(dict)
}
}
impl Rows for LiveDict {
type Row = LiveRow;
/// The dict's address, length and last key
type Stamp = (usize, usize, Option<String>);
fn fold(name: &str) -> String {
name.to_lowercase()
}
fn stamp(&self) -> Self::Stamp {
Python::attach(|py| {
let rows = self.0.bind(py);
let last = rows
.call_method0(intern!(py, "__reversed__"))
.and_then(|keys| keys.try_iter())
.ok()
.and_then(|mut keys| keys.next())
.and_then(Result::ok)
.and_then(|key| key.extract::<String>().ok());
(self.0.as_ptr() as usize, rows.len(), last)
})
}
fn contains(&self, key: &str) -> bool {
Python::attach(|py| self.0.bind(py).contains(key).unwrap_or(false))
}
fn row(&self, key: &str) -> Option<Cow<'_, LiveRow>> {
Python::attach(|py| {
let row = self.0.bind(py).get_item(key).ok().flatten()?;
Some(Cow::Owned(LiveRow(row.unbind())))
})
}
/// Groups each key by its row's `litellm_provider` when that is a string
fn scan(&self, visit: &mut dyn FnMut(&str, Option<&str>)) {
Python::attach(|py| {
for (key, row) in self.0.bind(py).iter() {
let Ok(key) = key.extract::<String>() else {
continue;
};
let provider = row
.cast::<PyDict>()
.ok()
.and_then(|row| row.get_item("litellm_provider").ok().flatten())
.and_then(|provider| provider.extract::<String>().ok());
visit(&key, provider.as_deref());
}
});
}
}
#[cfg(test)]
mod tests {
use litellm_model_catalog::Store;
use pyo3::ffi::c_str;
use rstest::rstest;
use super::*;
fn dict<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
py.eval(source, None, None)
.unwrap()
.cast_into::<PyDict>()
.unwrap()
}
#[rstest]
fn a_key_added_to_the_dict_is_found_on_the_next_query() {
Python::initialize();
Python::attach(|py| {
let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}"));
let store = Store::new(LiveDict::new(&rows));
assert_eq!(store.find_key("GPT-4O"), Some("gpt-4o".to_owned()));
rows.set_item("Claude-X", PyDict::new(py)).unwrap();
assert_eq!(store.find_key("claude-x"), Some("Claude-X".to_owned()));
});
}
#[rstest]
fn a_removed_key_is_never_returned() {
Python::initialize();
Python::attach(|py| {
let rows = dict(
py,
c_str!("{'gpt-4o': {'litellm_provider': 'openai'}, 'o3': {}}"),
);
let store = Store::new(LiveDict::new(&rows));
assert_eq!(store.keys_in("openai"), ["gpt-4o"]);
rows.del_item("gpt-4o").unwrap();
rows.set_item("o4", PyDict::new(py)).unwrap();
assert_eq!(store.find_key("GPT-4O"), None);
assert!(store.keys_in("openai").is_empty());
});
}
#[rstest]
fn a_key_swapped_in_for_a_removed_one_is_found() {
Python::initialize();
Python::attach(|py| {
let rows = dict(py, c_str!("{'gpt-4o': {}, 'Lit-Gone': {}}"));
let store = Store::new(LiveDict::new(&rows));
assert_eq!(store.find_key("lit-gone"), Some("Lit-Gone".to_owned()));
rows.del_item("Lit-Gone").unwrap();
rows.set_item("Lit-New", PyDict::new(py)).unwrap();
assert_eq!(store.find_key("lit-new"), Some("Lit-New".to_owned()));
assert_eq!(store.find_key("lit-gone"), None);
});
}
#[rstest]
fn two_keys_folding_alike_resolve_to_the_later_one() {
Python::initialize();
Python::attach(|py| {
let rows = dict(py, c_str!("{'Model': {}, 'MODEL': {}}"));
assert_eq!(
Store::new(LiveDict::new(&rows)).find_key("model"),
Some("MODEL".to_owned())
);
});
}
#[rstest]
fn rows_are_the_objects_the_dict_holds() {
Python::initialize();
Python::attach(|py| {
let rows = dict(py, c_str!("{'gpt-4o': {'litellm_provider': 'openai'}}"));
let store = Store::new(LiveDict::new(&rows));
let row = store.get("gpt-4o").unwrap().into_owned();
assert!(row.0.bind(py).is(rows.get_item("gpt-4o").unwrap().unwrap()));
});
}
#[rstest]
fn providers_group_only_string_providers() {
Python::initialize();
Python::attach(|py| {
let rows = dict(
py,
c_str!(
"{'a': {'litellm_provider': 'p'}, 'b': {'litellm_provider': None}, 'c': 1, 1: {'litellm_provider': 'p'}}"
),
);
assert_eq!(Store::new(LiveDict::new(&rows)).keys_in("p"), ["a"]);
});
}
}

497
litellm-rust/crates/python-bridge/src/model_cost/mod.rs Normal file
View file

@ -0,0 +1,497 @@
mod fetch;
mod host;
mod live;
mod table;
use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
use litellm_host_python::to_py;
use litellm_model_catalog::{
CostMap, IntegrityLimits, Landed, Pass, Placed, Provenance, RegistrationTable, Registrations,
Store, write_pass,
};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
pub(crate) use fetch::{NativeCostMapSource, NativeFetchFailed, NativeFetchPolicy, NativeFetched};
use live::LiveDict;
/// The installed document and the `register_model` calls kept above it, each as the dict Python merged
struct Inputs {
cost_map: CostMap,
registrations: Registrations<Py<PyDict>>,
}
impl Inputs {
fn clone_ref(&self, py: Python<'_>) -> Self {
Self {
cost_map: self.cost_map.clone(),
registrations: self.registrations.map(|kept| kept.clone_ref(py)),
}
}
}
/// The cost map's inputs and the name index over `litellm.model_cost`
#[pyclass(
frozen,
name = "NativeModelCostMap",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeModelCostMap {
inputs: RwLock<Inputs>,
limits: IntegrityLimits,
index: RwLock<Option<Arc<Store<LiveDict>>>>,
}
#[pyclass(
frozen,
name = "NativeInstalled",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeInstalled {
/// The new live table, the document without registrations
#[pyo3(get)]
rows: Py<PyDict>,
/// How many keys the document itself serves, aliases included and registrations left out
#[pyo3(get)]
model_count: usize,
/// The document's `fallback_generalizations` block as written, `None` when it has none
#[pyo3(get)]
fallback_generalizations: Py<PyAny>,
}
#[pyclass(
frozen,
name = "NativeRejected",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeRejected {
#[pyo3(get)]
reason: String,
/// The document parsed but has fewer models than the integrity limits allow
#[pyo3(get)]
integrity: bool,
}
/// A finished registration pass, with the registrations no built-in row matched, by registered key
#[pyclass(
frozen,
name = "NativeRegistered",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeRegistered {
#[pyo3(get)]
unmatched: Py<PyDict>,
}
/// A registration pass that stopped because the table it was handed is no longer the published one
#[pyclass(
frozen,
name = "NativeTableSwapped",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeTableSwapped;
/// The inputs at one moment, for a caller that has to put them back, such as test isolation
#[pyclass(
frozen,
name = "NativeModelCostMapState",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeModelCostMapState(Inputs);
#[pyclass(
frozen,
name = "NativeProvenance",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeProvenance {
#[pyo3(get)]
revision: String,
#[pyo3(get)]
etag: Option<String>,
}
#[pymethods]
impl NativeModelCostMap {
/// Starts from the bundled document. A fetched document with fewer than `min_model_count` models, or fewer than
/// `max_shrink_ratio` times the bundled count, is rejected
#[staticmethod]
fn bundled(min_model_count: i64, max_shrink_ratio: f64) -> PyResult<Self> {
let min_model_count = usize::try_from(min_model_count).map_err(|_| {
PyValueError::new_err(format!(
"min_model_count must be 0 or more, got {min_model_count}"
))
})?;
let limits = IntegrityLimits::new(min_model_count, max_shrink_ratio).ok_or_else(|| {
PyValueError::new_err(format!(
"max_shrink_ratio must be between 0 and 1, got {max_shrink_ratio}"
))
})?;
Ok(Self {
inputs: RwLock::new(Inputs {
cost_map: CostMap::bundled(),
registrations: Registrations::default(),
}),
limits,
index: RwLock::new(None),
})
}
/// The live table built from a fetched document, or why it was rejected, in which case nothing changed
#[pyo3(signature = (body, etag = None))]
fn install(&self, py: Python<'_>, body: &[u8], etag: Option<String>) -> PyResult<Py<PyAny>> {
match crate::logger::capture(py).scope(|| CostMap::read(body, etag, self.limits)) {
Ok(cost_map) => self.install_document(py, cost_map),
Err(error) => rejected(py, &error),
}
}
/// A caller's document as a table, read without size limits and installed nowhere
#[staticmethod]
fn finalize(py: Python<'_>, body: &[u8]) -> PyResult<Py<PyAny>> {
match crate::logger::capture(py)
.scope(|| CostMap::read(body, None, IntegrityLimits::UNCHECKED))
{
Ok(cost_map) => Ok(Py::new(py, built(py, &cost_map)?)?.into_any()),
Err(error) => rejected(py, &error),
}
}
/// The live table built from the bundled document, which is never size checked
fn install_bundled(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
self.install_document(py, CostMap::bundled())
}
/// Writes each entry into `rows`, or into the published table when `rows` is `None`
#[pyo3(signature = (rows, entries, persist, custom_llm_provider = None))]
fn register(
&self,
py: Python<'_>,
rows: Option<&Bound<'_, PyDict>>,
entries: &Bound<'_, PyDict>,
persist: bool,
custom_llm_provider: Option<&str>,
) -> PyResult<Py<PyAny>> {
let entries: Vec<(String, Bound<'_, PyAny>)> = entries
.iter()
.map(|(key, registration)| Ok((key.str()?.to_string(), registration)))
.collect::<PyResult<_>>()?;
let kept = match persist {
true => self.merged_with_kept(&entries)?,
false => Vec::new(),
};
let pass = self.write_pass(rows, &entries, custom_llm_provider)?;
if persist {
self.keep(kept, pass.landed());
}
into_py(py, pass)
}
/// Writes every kept registration into `rows` again, oldest first
fn reapply(&self, rows: &Bound<'_, PyDict>) -> PyResult<Py<PyAny>> {
let py = rows.py();
let entries: Vec<(String, Bound<'_, PyAny>)> = self
.read_inputs()
.registrations
.entries()
.map(|(key, kept)| (key.to_owned(), kept.bind(py).clone().into_any()))
.collect();
let pass = self.write_pass(Some(rows), &entries, None)?;
self.write_inputs().registrations.record(pass.landed());
into_py(py, pass)
}
/// The installed row under `key` with the kept registrations that apply to it, without anything the router wrote
fn registered_row<'py>(
&self,
py: Python<'py>,
key: &str,
) -> PyResult<Option<Bound<'py, PyDict>>> {
let (cost_map, applied) = {
let inputs = self.read_inputs();
let applied: Vec<Bound<'py, PyDict>> = inputs
.registrations
.over(key)
.into_iter()
.map(|kept| kept.bind(py).clone())
.collect();
(inputs.cost_map.clone(), applied)
};
table::registered_row(py, &cost_map, key, &applied)
}
fn registered_keys(&self) -> Vec<String> {
let inputs = self.read_inputs();
table::registered_keys(
&inputs.cost_map,
inputs.registrations.entries().map(|(key, _)| key),
)
}
/// The installed document as a table, without any registration
fn loaded<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let cost_map = self.read_inputs().cost_map.clone();
table::build(py, &cost_map)
}
/// The key itself when `rows` has it, otherwise the key that differs from `name` only in case
fn find_key(&self, rows: &Bound<'_, PyDict>, name: &str) -> Option<String> {
self.index_over(rows).find_key(name)
}
/// Keys whose row's `litellm_provider` is exactly `provider`, in row order
fn keys_in(&self, rows: &Bound<'_, PyDict>, provider: &str) -> Vec<String> {
self.index_over(rows).keys_in(provider)
}
/// Drops the name index, for a caller that changed keys in a way the dict's length does not show
fn invalidate(&self) {
if let Some(index) = self
.index
.read()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
{
index.invalidate();
}
}
fn snapshot(&self, py: Python<'_>) -> NativeModelCostMapState {
NativeModelCostMapState(self.read_inputs().clone_ref(py))
}
/// Puts back the document, registrations and provenance a `snapshot` saw
fn restore(&self, py: Python<'_>, state: PyRef<'_, NativeModelCostMapState>) {
*self.write_inputs() = state.0.clone_ref(py);
self.invalidate();
}
#[getter]
fn provenance(&self) -> NativeProvenance {
let Provenance { revision, etag } = self.read_inputs().cost_map.provenance().clone();
NativeProvenance { revision, etag }
}
}
impl NativeModelCostMap {
fn read_inputs(&self) -> RwLockReadGuard<'_, Inputs> {
self.inputs.read().unwrap_or_else(PoisonError::into_inner)
}
fn write_inputs(&self) -> RwLockWriteGuard<'_, Inputs> {
self.inputs.write().unwrap_or_else(PoisonError::into_inner)
}
fn install_document(&self, py: Python<'_>, cost_map: CostMap) -> PyResult<Py<PyAny>> {
let installed = built(py, &cost_map)?;
self.write_inputs().cost_map = cost_map;
*self.index.write().unwrap_or_else(PoisonError::into_inner) =
Some(Arc::new(Store::new(LiveDict::new(installed.rows.bind(py)))));
Ok(Py::new(py, installed)?.into_any())
}
/// Each entry merged into the registration kept for its key
fn merged_with_kept<'py>(
&self,
entries: &[(String, Bound<'py, PyAny>)],
) -> PyResult<Vec<(String, Bound<'py, PyDict>)>> {
entries
.iter()
.map(|(key, registration)| {
let previous = self
.read_inputs()
.registrations
.get(key)
.map(|kept| kept.bind(registration.py()).clone());
Ok((key.clone(), host::kept(previous.as_ref(), registration)?))
})
.collect()
}
fn keep(&self, kept: Vec<(String, Bound<'_, PyDict>)>, landed: &[Landed]) {
let mut inputs = self.write_inputs();
for (key, registration) in kept {
inputs.registrations.register(&key, registration.unbind());
}
inputs.registrations.record(landed);
}
fn write_pass<'a, 'py>(
&self,
rows: Option<&Bound<'py, PyDict>>,
entries: &'a [(String, Bound<'py, PyAny>)],
custom_llm_provider: Option<&str>,
) -> PyResult<Pass<'a, Bound<'py, PyAny>>> {
let mut table = PublishedTable {
map: self,
rows,
custom_llm_provider,
};
write_pass(&mut table, entries)
}
/// The index over `rows`, rebuilt when `rows` is not the dict it was built over
fn index_over(&self, rows: &Bound<'_, PyDict>) -> Arc<Store<LiveDict>> {
let current = self
.index
.read()
.unwrap_or_else(PoisonError::into_inner)
.clone();
if let Some(index) = current.filter(|index| index.rows().is(rows)) {
return index;
}
let index = Arc::new(Store::new(LiveDict::new(rows)));
*self.index.write().unwrap_or_else(PoisonError::into_inner) = Some(Arc::clone(&index));
index
}
}
/// The dict Python publishes as `litellm.model_cost`, written through the host
struct PublishedTable<'a, 'py> {
map: &'a NativeModelCostMap,
rows: Option<&'a Bound<'py, PyDict>>,
custom_llm_provider: Option<&'a str>,
}
impl<'py> RegistrationTable for PublishedTable<'_, 'py> {
type Registration = Bound<'py, PyAny>;
type Error = PyErr;
fn place(&mut self, key: &str, registration: &Self::Registration) -> PyResult<Option<Placed>> {
let Some(landing) = host::resolve(self.rows, key, registration, self.custom_llm_provider)?
else {
return Ok(None);
};
let rows = match self.rows {
Some(rows) => rows.clone(),
None => host::published(registration.py())?,
};
table::write(&rows, &landing, registration)?;
self.map.invalidate();
Ok(Some(Placed {
key: landing.key,
matched_builtin: landing.matched_builtin,
}))
}
}
fn built(py: Python<'_>, cost_map: &CostMap) -> PyResult<NativeInstalled> {
let rows = table::build(py, cost_map)?;
let fallback_generalizations = match cost_map.fallback_generalizations() {
Some(block) => to_py(py, block)?,
None => py.None(),
};
Ok(NativeInstalled {
model_count: rows.len(),
rows: rows.unbind(),
fallback_generalizations,
})
}
fn rejected(py: Python<'_>, error: &litellm_model_catalog::Error) -> PyResult<Py<PyAny>> {
let rejected = NativeRejected {
reason: litellm_error::Report::new(error).to_string(),
integrity: error.is_integrity(),
};
Ok(Py::new(py, rejected)?.into_any())
}
fn into_py<'py>(py: Python<'py>, pass: Pass<'_, Bound<'py, PyAny>>) -> PyResult<Py<PyAny>> {
match pass {
Pass::Swapped { .. } => Ok(Py::new(py, NativeTableSwapped)?.into_any()),
Pass::Finished { unmatched, .. } => {
let table = PyDict::new(py);
for (key, registration) in unmatched {
table.set_item(key, registration)?;
}
let registered = NativeRegistered {
unmatched: table.unbind(),
};
Ok(Py::new(py, registered)?.into_any())
}
}
}
#[cfg(test)]
mod tests {
use std::sync::mpsc;
use std::time::Duration;
use pyo3::ffi::c_str;
use rstest::rstest;
use super::*;
const HOST: &std::ffi::CStr = c_str!(
r#"
from types import SimpleNamespace
class Row(dict):
names = ()
found = []
def resolve(rows, key, registration, custom_llm_provider):
found.append(cost_map.find_key(rows, key))
return SimpleNamespace(key=key, base={}, matched_builtin=True)
def merged(current, base, registration):
return Row({**(current or {}), **base, **registration})
def write(rows, written):
rows.update(written)
"#
);
#[rstest]
fn a_host_lookup_that_reaches_find_key_on_the_same_map_finishes() {
Python::initialize();
let (sender, receiver) = mpsc::channel();
std::thread::spawn(move || {
let outcome = Python::attach(|py| -> PyResult<(Vec<Option<String>>, f64)> {
let cost_map = Py::new(py, NativeModelCostMap::bundled(0, 0.0)?)?;
let module_name = std::ffi::CString::new(host::MODULE).unwrap();
let host =
PyModule::from_code(py, HOST, c_str!("host_model_cost.py"), &module_name)?;
host.setattr("cost_map", &cost_map)?;
let modules = py.import("sys")?.getattr("modules")?;
let packages = ["litellm", "litellm.rust_bridge", "litellm.rust_bridge.host"];
for package in packages {
modules.set_item(package, PyModule::new(py, package)?)?;
}
modules.set_item(host::MODULE, &host)?;
let rows = PyDict::new(py);
rows.set_item("Lit-Reentry", PyDict::new(py))?;
let entry = PyDict::new(py);
entry.set_item("input_cost_per_token", 1.5)?;
let entries = PyDict::new(py);
entries.set_item("lit-reentry", entry)?;
let registered = cost_map
.bind(py)
.call_method1("register", (&rows, &entries, false));
for name in [host::MODULE].into_iter().chain(packages) {
modules.del_item(name)?;
}
registered?;
let price = rows
.get_item("lit-reentry")?
.expect("the registration landed")
.get_item("input_cost_per_token")?
.extract()?;
Ok((host.getattr("found")?.extract()?, price))
});
sender
.send(outcome.map_err(|error| error.to_string()))
.unwrap();
});
let (found, price) = receiver
.recv_timeout(Duration::from_secs(30))
.expect("register finished instead of deadlocking on its own lock")
.unwrap();
assert_eq!(found, [Some("Lit-Reentry".to_owned())]);
assert!((price - 1.5).abs() < f64::EPSILON);
}
}

166
litellm-rust/crates/python-bridge/src/model_cost/table.rs Normal file
View file

@ -0,0 +1,166 @@
//! Builds `litellm.model_cost` from an installed document
use std::borrow::Cow;
use std::collections::HashMap;
use indexmap::IndexSet;
use litellm_host_python::to_py;
use litellm_model_catalog::wire::Fields;
use litellm_model_catalog::{CostMap, FALLBACK_GENERALIZATIONS, Written};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::Value;
use super::host::{self, Landing};
const ALIASES: &str = "aliases";
enum Served<'a> {
Model,
AliasOf(&'a str),
/// A top-level value that is not a model row, such as `sample_spec`
Section(&'a Value),
}
/// Every served key in dict order. An `alias_of` row repeats among the aliases and keeps its document position
fn served(cost_map: &CostMap) -> impl Iterator<Item = (&str, Served<'_>)> {
let catalog = cost_map.catalog();
let documented = catalog
.written()
.filter(|(key, _)| *key != FALLBACK_GENERALIZATIONS)
.filter_map(|(key, written)| {
let served = match written {
Written::Section(value) => Served::Section(value),
Written::Entry(_) => match catalog.aliases().get(key) {
Some(owner) => Served::AliasOf(owner),
None => catalog.get(key).map(|_| Served::Model)?,
},
};
Some((key, served))
});
let listed = cost_map.alias_groups().flat_map(|(owner, names)| {
names
.iter()
.skip(1)
.map(move |alias| (alias.as_str(), Served::AliasOf(owner)))
});
documented.chain(listed)
}
/// The installed document as a new dict
pub(crate) fn build<'py>(py: Python<'py>, cost_map: &CostMap) -> PyResult<Bound<'py, PyDict>> {
let table = PyDict::new(py);
let mut rows = HashMap::new();
for (key, served) in served(cost_map) {
let value = match served {
Served::Model => row(py, cost_map, key, &mut rows)?,
Served::AliasOf(owner) => row(py, cost_map, owner, &mut rows)?,
Served::Section(value) => to_py(py, value)?.into_bound(py),
};
table.set_item(key, value)?;
}
Ok(table)
}
fn row<'a, 'py>(
py: Python<'py>,
cost_map: &'a CostMap,
owner: &'a str,
rows: &mut HashMap<&'a str, Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
if let Some(row) = rows.get(owner) {
return Ok(row.clone());
}
let fields = cost_map
.catalog()
.get(owner)
.expect("a served model or alias owner is a catalog row");
let fields = to_py(py, without_alias_list(fields).as_ref())?.into_bound(py);
let row = match cost_map.names(owner) {
[] => fields,
names => host::entry(fields, names)?,
};
rows.insert(owner, row.clone());
Ok(row)
}
fn without_alias_list(fields: &Fields) -> Cow<'_, Fields> {
if fields.get(ALIASES).is_none_or(Value::is_null) {
return Cow::Borrowed(fields);
}
Cow::Owned(
fields
.iter()
.filter(|(name, _)| *name != ALIASES)
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
)
}
/// The installed row under `key` with `applied` written over it in order, `None` when neither has one
pub(crate) fn registered_row<'py>(
py: Python<'py>,
cost_map: &CostMap,
key: &str,
applied: &[Bound<'py, PyDict>],
) -> PyResult<Option<Bound<'py, PyDict>>> {
let catalog = cost_map.catalog();
let installed = catalog
.get(key)
.or_else(|| match catalog.section(key) {
Some(Value::Object(fields)) if key != FALLBACK_GENERALIZATIONS => Some(fields),
_ => None,
})
.map(without_alias_list);
let row = match (installed, applied) {
(None, []) => return Ok(None),
(Some(fields), _) => to_py(py, fields.as_ref())?
.into_bound(py)
.cast_into::<PyDict>()?,
(None, _) => PyDict::new(py),
};
for registration in applied {
row.update(registration.as_mapping())?;
}
Ok(Some(row))
}
/// The keys `registered_row` answers for, installed keys first
pub(crate) fn registered_keys<'a>(
cost_map: &'a CostMap,
registered: impl Iterator<Item = &'a str>,
) -> Vec<String> {
let served: IndexSet<&str> = served(cost_map).map(|(key, _)| key).collect();
let only_registered: Vec<&str> = registered.filter(|key| !served.contains(key)).collect();
served
.into_iter()
.chain(only_registered)
.map(str::to_owned)
.collect()
}
/// Writes `registration` over the row under `landing.key`, and under every name that shares that row object
pub(crate) fn write<'py>(
rows: &Bound<'py, PyDict>,
landing: &Landing<'py>,
registration: &Bound<'py, PyAny>,
) -> PyResult<()> {
let current = rows.get_item(&landing.key)?;
let row = host::merged(current.as_ref(), landing, registration)?;
let names: Vec<String> = row.getattr("names")?.extract()?;
let sharing: Vec<&String> = names
.iter()
.filter(|name| {
let served = rows.get_item(name.as_str()).ok().flatten();
served
.zip(current.as_ref())
.is_some_and(|(served, current)| served.is(current))
})
.collect();
let written = PyDict::new(rows.py());
written.set_item(&landing.key, &row)?;
sharing
.into_iter()
.try_for_each(|name| written.set_item(name, &row))?;
host::write(rows, &written)
}

2
litellm-rust/crates/python-bridge/src/python_settings.rs
View file

@ -56,7 +56,7 @@ impl PythonSettings {
}
}
#[cfg(test)]
/// A snapshot of a value Python handed in, read with the same field rules as the accessor's
pub(crate) fn snapshot(self, value: Bound<'_, PyAny>) -> Snapshot<'_> {
Snapshot { group: self, value }
}

8
litellm-rust/crates/python-bridge/src/routes/messages/host.rs
View file

@ -239,9 +239,13 @@ impl PythonBinding for MessagesPythonHost {
fn encode_response(
&mut self,
py: Python<'_>,
response: Box<litellm_llms_types::formats::messages::MessagesResponse>,
response: litellm_inference_messages::MessagesBody,
) -> PyResult<Py<PyAny>> {
public_response(py, ROUTE_HOST_MODULE, response.as_ref())
public_response(py, ROUTE_HOST_MODULE, &response)
}
fn deliver_response(&mut self, py: Python<'_>, response: &Bound<'_, PyAny>) -> PyResult<()> {
self.cache.store_response(py, response)
}
fn encode_stream_head(

9
litellm-rust/crates/router-core/src/call.rs
View file

@ -1037,8 +1037,13 @@ impl Core {
self.log_retry(target, at, 0, &failed.error);
}
let named = failure.failed_deployment_id.as_ref();
if let (Some(status), Some(id)) =
(failure.status, named.or(failed.deployment_id.as_ref()))
let failed_on = named.or(failed.deployment_id.as_ref());
if let Some(id) = failed_on
&& !inputs.failed_ids.contains(id)
{
inputs.failed_ids.push(id.clone());
}
if let (Some(status), Some(id)) = (failure.status, failed_on)
&& !should_retry_status(status)
&& !inputs.skipped_ids.contains(id)
{

11
litellm-rust/crates/router-core/src/call/stream.rs
View file

@ -236,10 +236,11 @@ impl Router {
.unwrap_or((0, 0));
let (budget, policy_applies) =
stream_retry_budget(&snapshot, &request, &origin, failure.kind);
let inputs = SelectInputs {
let mut inputs = SelectInputs {
target_order: target.order,
requested: Some(target.group.clone()),
fallback_hop: at.level > 0,
failed_ids: origin.deployment_id.iter().cloned().collect(),
..SelectInputs::default()
};
let mut last = failure;
@ -290,6 +291,14 @@ impl Router {
// `_stream_fallback_error_for_raised`: a status the pre-stream retry rule
// would not retry ends the retries with that failure.
Error::Attempt(failure) => {
if let Some(id) = failure
.failed_deployment_id
.clone()
.or(failed.deployment_id)
&& !inputs.failed_ids.contains(&id)
{
inputs.failed_ids.push(id);
}
let stop = failure.status.is_some_and(|s| !should_retry_status(s));
last = failure;
if stop {

5
litellm-rust/crates/router-core/src/select.rs
View file

@ -108,6 +108,9 @@ impl Candidates {
pub(crate) struct SelectInputs {
pub target_order: Option<i64>,
pub skipped_ids: Vec<String>,
/// Every deployment an attempt of this request failed on, whatever the status: the callback
/// filters read it as `_retry_skipped_deployment_ids`, so an affinity pin can leave one.
pub failed_ids: Vec<String>,
/// A weighted failover hop's `_excluded_deployment_ids`: never picked.
pub excluded_ids: Arc<[String]>,
/// Stop before the plugins, order and skip filters: the group less cooldowns and blocked
@ -470,7 +473,7 @@ impl Core {
StrategyScope::Default,
);
ctx.excluded_ids = &inputs.excluded_ids;
ctx.skipped_ids = &inputs.skipped_ids;
ctx.skipped_ids = &inputs.failed_ids;
ctx.target_order = inputs.target_order.or(request_order);
let filtered = self.host.callback_filters(&ctx, healthy).await?;
healthy = filtered.kept;

99
litellm-rust/crates/router-core/tests/call_loop.rs
View file

@ -5,7 +5,7 @@ use std::{collections::VecDeque, sync::Arc, time::Duration};
use async_trait::async_trait;
use litellm_router_core::{
Attempt, AttemptError, Deployment, Error, FailureKind, FailureReport, Host,
Attempt, AttemptError, CallbackFiltered, Deployment, Error, FailureKind, FailureReport, Host,
InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig,
RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage,
SuccessReport,
@ -37,6 +37,9 @@ struct ScriptedHost {
pick_failures: Mutex<VecDeque<AttemptError>>,
/// The next attempt's failure, whatever group it serves.
first_failure: Mutex<Option<AttemptError>>,
/// The callback filters pin to this deployment until an attempt of the request failed on it,
/// as `encrypted_content_affinity` does.
pinned: Option<&'static str>,
}
impl ScriptedHost {
@ -112,6 +115,24 @@ impl Host for ScriptedHost {
}
}
async fn callback_filters(
&self,
ctx: &RoutingContext<'_>,
candidates: Vec<Arc<Deployment>>,
) -> Result<CallbackFiltered, Error> {
let kept = match self.pinned {
Some(pin) if !ctx.skipped_ids.iter().any(|id| id == pin) => {
candidates.into_iter().filter(|d| d.id == pin).collect()
}
Some(pin) => candidates.into_iter().filter(|d| d.id != pin).collect(),
None => candidates,
};
Ok(CallbackFiltered {
kept,
preferred: None,
})
}
async fn allow_fallback(
&self,
_request: &RouterRequest,
@ -223,6 +244,32 @@ async fn a_rate_limit_retries_in_the_same_group(config: RouterConfig) {
);
}
#[rstest]
#[case::server_error(server_error())]
#[case::rate_limit(failure(FailureKind::RateLimit, 429))]
#[case::connection_error(AttemptError::new(FailureKind::ApiConnection, None))]
#[tokio::test]
async fn a_retry_leaves_a_pinned_deployment_that_failed_the_request(
mut config: RouterConfig,
#[case] origin_failure: AttemptError,
) {
config.callback_filters = true;
let host = Arc::new(ScriptedHost {
first_failure: Mutex::new(Some(origin_failure)),
pinned: Some("a1"),
..ScriptedHost::default()
});
let router = router(config, json!({"num_retries": 3}), host.clone());
let response = router
.call(Op::from("acompletion"), request())
.await
.unwrap();
assert_eq!(
(host.selected(), response.attempted_retries),
(vec!["a1".to_owned(), "a2".to_owned()], 1)
);
}
#[rstest]
#[tokio::test]
async fn a_bad_request_falls_back_without_retrying(config: RouterConfig) {
@ -945,6 +992,56 @@ async fn a_stream_retry_is_reported_before_its_attempt(config: RouterConfig) {
assert_eq!((response.attempted_retries, host.retries()), (1, 1));
}
fn pinned_to_a1() -> Arc<ScriptedHost> {
Arc::new(ScriptedHost {
pinned: Some("a1"),
..ScriptedHost::default()
})
}
#[rstest]
#[tokio::test]
async fn a_stream_retry_leaves_the_pinned_deployment_the_stream_failed_on(
mut config: RouterConfig,
) {
config.callback_filters = true;
let host = pinned_to_a1();
let router = router(config, json!({}), host.clone());
let origin = StreamOrigin {
deployment_id: Some("a1".to_owned()),
..stream_origin()
};
router
.retry_stream(
Op::from("acompletion"),
stream_request(),
origin,
server_error(),
)
.await
.unwrap();
assert_eq!(host.selected(), vec!["a2".to_owned()]);
}
#[rstest]
#[tokio::test]
async fn a_stream_retry_leaves_a_pinned_deployment_a_retry_failed_on(mut config: RouterConfig) {
config.callback_filters = true;
let host = pinned_to_a1();
*host.first_failure.lock() = Some(server_error());
let router = router(config, json!({}), host.clone());
router
.retry_stream(
Op::from("acompletion"),
stream_request(),
stream_origin(),
server_error(),
)
.await
.unwrap();
assert_eq!(host.selected(), vec!["a1".to_owned(), "a2".to_owned()]);
}
fn fallback_list(value: serde_json::Value) -> Fallbacks {
serde_json::from_value(value).unwrap()
}

3
litellm-rust/crates/s3/AGENTS.md Normal file
View file

@ -0,0 +1,3 @@
`litellm-s3` builds `aws-sdk-s3` clients that send through `litellm_http::Client`, and owns the S3 concepts the SDK lacks, such as the `s3://bucket/key` URI. It reads no environment and picks no credentials: the host passes the region, the endpoint and a credentials provider, usually `litellm_auth_aws::AwsCredentialsProvider`
Operation policy such as retries, timeouts and checksum settings stays with the caller, which adjusts the returned config builder

21
litellm-rust/crates/s3/Cargo.toml Normal file
View file

@ -0,0 +1,21 @@
[package]
name = "litellm-s3"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-http.workspace = true
aws-sdk-s3.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
aws-types = "1.4.0"
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
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true

7
litellm-rust/crates/s3/src/error.rs Normal file
View file

@ -0,0 +1,7 @@
use thiserror::Error;
#[derive(Debug, Error)]
pub enum Error {
#[error("expected s3://bucket/key")]
InvalidUri,
}

60
litellm-rust/crates/s3/src/lib.rs Normal file
View file

@ -0,0 +1,60 @@
mod error;
mod transport;
mod uri;
pub use aws_credential_types::provider::ProvideCredentials;
pub use aws_sdk_s3 as sdk;
pub use error::Error;
pub use uri::S3Uri;
use aws_config::meta::region::RegionProviderChain;
use aws_sdk_s3::config::{BehaviorVersion, Builder, Region};
use aws_types::SdkConfig;
use aws_types::service_config::ServiceConfigKey;
use crate::transport::ReqwestHttpClient;
/// A custom endpoint is addressed path style, as boto3 addresses one
pub fn config(
region: String,
endpoint_url: Option<&str>,
http: litellm_http::Client,
credentials: impl ProvideCredentials + 'static,
) -> Builder {
let builder = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(region))
.http_client(ReqwestHttpClient(http))
.credentials_provider(credentials);
match endpoint_url {
Some(url) => builder.endpoint_url(url).force_path_style(true),
None => builder,
}
}
/// The region, endpoint and credentials an AWS SDK finds in the environment and shared config files, as boto3's
/// default `s3` client found them. `region` wins over that chain, and `us-east-1` is used when nothing names one
pub async fn environment_config(http: litellm_http::Client, region: Option<String>) -> Builder {
let region = RegionProviderChain::first_try(region.map(Region::new))
.or_default_provider()
.or_else(Region::from_static("us-east-1"));
let shared = aws_config::defaults(BehaviorVersion::latest())
.region(region)
.http_client(ReqwestHttpClient(http))
.load()
.await;
Builder::from(&shared).force_path_style(has_custom_endpoint(&shared))
}
fn has_custom_endpoint(shared: &SdkConfig) -> bool {
let s3_endpoint = ServiceConfigKey::builder()
.service_id("S3")
.env("AWS_ENDPOINT_URL")
.profile("endpoint_url")
.build()
.ok();
shared.endpoint_url().is_some()
|| s3_endpoint
.and_then(|key| shared.service_config()?.load_config(key))
.is_some()
}

0
litellm-rust/crates/cache-s3/src/transport.rs → litellm-rust/crates/s3/src/transport.rs
View file

76
litellm-rust/crates/s3/src/uri.rs Normal file
View file

@ -0,0 +1,76 @@
use crate::Error;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct S3Uri {
pub bucket: String,
pub key: String,
}
/// Splits the way Python's `urllib.parse.urlsplit` does, so the key keeps dot segments and characters a URL parser
/// would normalize or percent-encode. S3 keys are opaque strings
impl TryFrom<&str> for S3Uri {
type Error = Error;
fn try_from(uri: &str) -> Result<Self, Error> {
let rest = match uri.split_once("://") {
Some((scheme, rest)) if scheme.eq_ignore_ascii_case("s3") => rest,
_ => return Err(Error::InvalidUri),
};
let (bucket, path) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len()));
let key = path
.split(['?', '#'])
.next()
.unwrap_or_default()
.trim_start_matches('/');
if bucket.is_empty() || key.is_empty() {
return Err(Error::InvalidUri);
}
Ok(Self {
bucket: bucket.to_owned(),
key: key.to_owned(),
})
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::key("s3://bucket/cost.json", "bucket", "cost.json")]
#[case::nested_key("s3://bucket/prices/cost.json", "bucket", "prices/cost.json")]
#[case::dotted_bucket("s3://my.bucket/cost.json", "my.bucket", "cost.json")]
#[case::dot_segments("s3://bucket/archives/../cost.json", "bucket", "archives/../cost.json")]
#[case::current_dir_segment("s3://bucket/./cost.json", "bucket", "./cost.json")]
#[case::space_kept_raw("s3://bucket/my prices.json", "bucket", "my prices.json")]
#[case::query_dropped("s3://bucket/cost.json?versionId=1", "bucket", "cost.json")]
#[case::fragment_dropped("s3://bucket/cost.json#top", "bucket", "cost.json")]
#[case::leading_slashes_trimmed("s3://bucket//cost.json", "bucket", "cost.json")]
#[case::uppercase_scheme("S3://bucket/cost.json", "bucket", "cost.json")]
fn an_s3_url_names_its_bucket_and_key(
#[case] url: &str,
#[case] bucket: &str,
#[case] key: &str,
) {
assert_eq!(
S3Uri::try_from(url).unwrap(),
S3Uri {
bucket: bucket.into(),
key: key.into(),
}
);
}
#[rstest]
#[case::no_key("s3://bucket/")]
#[case::no_path("s3://bucket")]
#[case::query_only("s3://bucket?cost.json")]
#[case::no_bucket("s3:///cost.json")]
#[case::other_scheme("https://bucket/cost.json")]
#[case::no_scheme("bucket/cost.json")]
fn a_url_without_a_bucket_and_key_is_refused(#[case] url: &str) {
assert!(matches!(S3Uri::try_from(url), Err(Error::InvalidUri)));
}
}

1
litellm/accounting/projection.py
View file

@ -133,6 +133,7 @@ def _daily_counters(row: Mapping[str, object], metadata: Mapping[str, object]) -
("prompt_caching_savings_spend", _number(savings.get("prompt_caching"))),
("gateway_injected_caching_savings_spend", _number(savings.get("gateway_injected_caching"))),
("autorouter_savings_spend", 0.0 if internal else _number(savings.get("autorouter"))),
("guardrail_spend", _number(metadata.get("accounting_guardrail_spend"))),
("total_response_time_ms", _number(duration) if timed else 0),
("timed_requests", int(timed)),
)

2
litellm/accounting/snapshot.py
View file

@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_share_of_spend
from litellm.litellm_core_utils.safe_json_dumps import safe_json_structure
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens
@ -73,6 +74,7 @@ def build_accounting_payload(
**{key: _accounting_metadata_field(key, value) for key, value in metadata.items() if key in _METADATA_FIELDS},
"accounting_entities": entities,
"accounting_savings": savings._asdict(),
"accounting_guardrail_spend": guardrail_share_of_spend(metadata, payload["spend"]),
"accounting_model_access_groups": model_access_groups,
"accounting_windows": windows,
}

64
litellm/accounting/worker.py
View file

@ -46,9 +46,6 @@ def _progressed(result: DeliveryBatchResult) -> bool:
async def flush_accounting_deliveries(
writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging"
) -> None:
acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False)
if acquired is False:
return
applied: tuple[Contribution, ...] = () # rebind-ok: keep committed contributions for cache invalidation
async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None:
@ -57,36 +54,39 @@ async def flush_accounting_deliveries(
transaction, deliveries, prisma_client=prisma_client, writer=writer, proxy_logging_obj=proxy_logging_obj
)
database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction
LedgerDatabase, prisma_client
)
repository: Final = AccountingDeliveryRepository(database)
started: Final = time.monotonic()
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500"))
async def run_batch(*, retry_only: bool) -> DeliveryBatchResult:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
if result.applied:
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
if result.parked:
verbose_proxy_logger.error(
"Accounting delivery parked %d requests after isolated failures", len(result.parked)
)
await _alert(
proxy_logging_obj,
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
)
return result
fresh_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline:
break
retry_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
break
acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False)
if acquired is False:
return
try:
database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction
LedgerDatabase, prisma_client
)
repository: Final = AccountingDeliveryRepository(database)
started: Final = time.monotonic()
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500"))
async def run_batch(*, retry_only: bool) -> DeliveryBatchResult:
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
if result.applied:
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
if result.parked:
verbose_proxy_logger.error(
"Accounting delivery parked %d requests after isolated failures", len(result.parked)
)
await _alert(
proxy_logging_obj,
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
)
return result
fresh_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline:
break
retry_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
break
retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1"))
cleanup_deadline: Final = time.monotonic() + 5
for _ in range(batch_size):

3
litellm/batches/batch_utils.py
View file

@ -11,6 +11,7 @@ from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils import batch_utils as core_batch_utils
from litellm.litellm_core_utils.cloud_storage_security import stored_s3_bucket_kwargs
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
@ -551,7 +552,7 @@ async def fetch_batch_managed_file_content(
"file_id": provider_output_file_id,
"custom_llm_provider": custom_llm_provider,
**cast( # cast-ok: preserve dynamic provider credentials without validation
FileContentCallOptions, credentials
FileContentCallOptions, {**credentials, **stored_s3_bucket_kwargs(provider_output_file_id)}
),
}

11
litellm/caching/redis_batch.py
View file

@ -16,7 +16,7 @@ import logging
import time
import weakref
from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence
from contextvars import ContextVar, Token
from contextvars import ContextVar, Token, copy_context
from dataclasses import dataclass, field
from datetime import timedelta
from enum import Enum, auto
@ -528,6 +528,7 @@ class RequestRedisBatches:
"_deadline",
"_deadline_flush",
"_native",
"_native_context",
"_native_factory",
"_native_flush_hooks",
"_native_misses",
@ -551,6 +552,7 @@ class RequestRedisBatches:
self._native_factory: Final = native_factory
self._native_send_factory: Final = native_send_factory
self._native: NativeRequestRedis | None = None
self._native_context: Final = copy_context()
self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry
self._native_misses: frozenset[str] = frozenset()
self._deadline: asyncio.TimerHandle | None = None
@ -580,9 +582,12 @@ class RequestRedisBatches:
return batch
def native(self) -> NativeRequestRedis:
"""The request's native batches, built on first use so requests that touch no native store pay nothing."""
"""The request's native batches, built on first use so requests that touch no native store pay nothing.
They keep a copy of the context they are built in, and the cycle collector cannot see through native
objects, so they are built in the context this scope was opened from, which cannot reach the scope."""
if self._native is None:
self._native = self._native_factory(self.post_call_deadline)
self._native = self._native_context.run(self._native_factory, self.post_call_deadline)
return self._native
def native_post_call(self) -> NativeRequestRedis:

3
litellm/chat_completions/dispatch.py
View file

@ -16,6 +16,7 @@ from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_acall_hook,
native_call,
native_call_hook,
optional_sequence,
@ -131,7 +132,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_ACOMPLETION,
native=native_call_hook,
native=native_acall_hook,
)

4
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -36,7 +36,7 @@ from litellm.completion_extras.litellm_responses_transformation.history import (
)
from litellm.integrations.anthropic_cache_control_hook import routed_model_info, supports_openai_prompt_cache_breakpoint
from litellm.litellm_core_utils.allowed_tools_tool_choice import (
parse_chat_allowed_tools_tool_choice,
parse_allowed_tools_tool_choice,
responses_allowed_tools_tool_choice,
)
from litellm.litellm_core_utils.core_helpers import (
@ -900,7 +900,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if isinstance(choice_type, str) and choice_type in ("auto", "none", "required"):
return choice_type
if choice_type == "allowed_tools":
parsed: Final = parse_chat_allowed_tools_tool_choice(
parsed: Final = parse_allowed_tools_tool_choice(
cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic
object, tool_choice
)

2
litellm/constants.py
View file

@ -1706,6 +1706,7 @@ ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags"
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_usage_based_routing_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
RESPONSES_WEBSOCKET_CALL_ID_METADATA_KEY: Final = "responses_websocket_call_id"
BILLED_RETRY_ATTEMPT_METADATA_KEY: Final = "billed_retry_attempt"
@ -1968,6 +1969,7 @@ BATCH_LINE_ITEMS_MAX_PART_BYTES: Final = get_env_int_in_range(
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress")
TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"})
BACKGROUND_RESPONSE_SETTLEMENT_TIMEOUT_SECONDS: Final = 2.0
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60
STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))

132
litellm/cost_calculator.py
View file

@ -49,7 +49,6 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
BilledTokenRates,
CostCalculatorUtils,
calculate_cache_writing_cost,
calculate_cost_component,
generic_cost_per_character,
generic_cost_per_token,
get_batch_cost_rates,
@ -3177,6 +3176,21 @@ def default_video_cost_calculator(
return 0.0
def _optional_batch_rate(
model_info: ModelInfo,
key: Literal[
"input_cost_per_audio_token_batches",
"cache_read_input_audio_token_cost_batches",
"input_cost_per_image_token_batches",
"input_cost_per_video_token_batches",
],
prompt_tokens: int,
) -> float | None:
context_key: Final = get_modality_token_cost_key(model_info, key, prompt_tokens)
context_rate: Final = get_cost_per_unit(model_info, context_key, None)
return model_info.get(key) if context_rate is None else context_rate
def _batch_rate(
model_info: ModelInfo,
key: Literal[
@ -3189,9 +3203,7 @@ def _batch_rate(
*,
prompt_tokens: int = 0,
) -> float:
context_key: Final = get_modality_token_cost_key(model_info, key, prompt_tokens)
context_rate: Final = get_cost_per_unit(model_info, context_key, None)
rate: Final = model_info.get(key) if context_rate is None else context_rate
rate: Final = _optional_batch_rate(model_info, key, prompt_tokens)
return fallback if rate is None else rate
@ -3275,30 +3287,14 @@ def batch_cost_calculator(
model_info,
batch_rates.input,
batch_rates.cache_read,
batch_rates.cache_read_image,
batch_rates.cache_creation,
batch_rates.cache_creation_rates_by_ttl,
)
elif input_cost_per_token is not None:
details: Final = parse_prompt_tokens_details(usage)
cache_read_tokens: Final = details["cache_hit_tokens"]
cache_creation_tokens: Final = details["cache_creation_tokens"]
# Subtract cached tokens from prompt_tokens before calculating cost
# Fixes issue where cached tokens are being charged again
base_input_tokens: Final = get_billable_input_tokens(usage) - cache_creation_tokens
total_prompt_cost = (
base_input_tokens * (input_cost_per_token) / 2
) # batch cost is usually half of the regular token cost
# Add cache read cost if applicable
cache_read_cost_key: Final = get_service_tier_cost_key("cache_read_input_token_cost", None)
total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2
configured_cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost")
cache_creation_cost: Final = (
input_cost_per_token if configured_cache_creation_cost is None else configured_cache_creation_cost
total_prompt_cost = _half_price_prompt_cost(
usage, model_info, input_cost_per_token, batch_rates.cache_read_image
)
total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2
text_rate: Final = _batch_or_half(batch_rates.output, output_cost_per_token)
image_rate: Final = _batch_or_half(
model_info.get("output_cost_per_image_token_batches"),
@ -3324,15 +3320,54 @@ def _standard_rate(model_info: ModelInfo, key: str, prompt_tokens: int) -> float
return get_cost_per_unit(model_info, get_modality_token_cost_key(model_info, key, prompt_tokens), None)
def _cached_audio_batch_fallback(
model_info: ModelInfo, batch_cache_read_rate: float | None, cache_read_rate: float | None, prompt_tokens: int
) -> float:
def _cached_modality_rate(
model_info: ModelInfo,
standard_key: Literal["cache_read_input_audio_token_cost", "cache_read_input_image_token_cost"],
batch_modality_rate: float | None,
batch_cache_read_rate: float | None,
cache_read_rate: float | None,
prompt_tokens: int,
) -> float | None:
if batch_modality_rate is not None:
return batch_modality_rate
if batch_cache_read_rate is not None:
return batch_cache_read_rate
standard_audio_rate: Final = _standard_rate(model_info, "cache_read_input_audio_token_cost", prompt_tokens)
if standard_audio_rate is not None:
return standard_audio_rate
return 0.0 if cache_read_rate is None else cache_read_rate
standard_modality_rate: Final = _standard_rate(model_info, standard_key, prompt_tokens)
return cache_read_rate if standard_modality_rate is None else standard_modality_rate
def _half_price_prompt_cost(
usage: Usage, model_info: ModelInfo, input_cost_per_token: float, batch_cache_read_image_rate: float | None
) -> float:
details: Final = parse_prompt_tokens_details(usage)
cache_creation_tokens: Final = details["cache_creation_tokens"]
base_input_tokens: Final = get_billable_input_tokens(usage) - cache_creation_tokens
cache_read_cost_key: Final = get_service_tier_cost_key("cache_read_input_token_cost", None)
half_cache_read_rate: Final = (get_cost_per_unit(model_info, cache_read_cost_key, 0.0) or 0.0) / 2
cached_audio_rate: Final = _batch_or_half(
_optional_batch_rate(model_info, "cache_read_input_audio_token_cost_batches", usage.prompt_tokens),
_standard_rate(model_info, "cache_read_input_audio_token_cost", usage.prompt_tokens),
default=half_cache_read_rate,
)
cached_image_rate: Final = _batch_or_half(
batch_cache_read_image_rate,
_standard_rate(model_info, "cache_read_input_image_token_cost", usage.prompt_tokens),
default=half_cache_read_rate,
)
cached_audio_tokens: Final = details["cache_hit_audio_tokens"]
cached_image_tokens: Final = details["cached_image_tokens"]
cached_text_tokens: Final = max(details["cache_hit_tokens"] - cached_audio_tokens - cached_image_tokens, 0)
configured_cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost")
cache_creation_cost: Final = (
input_cost_per_token if configured_cache_creation_cost is None else configured_cache_creation_cost
)
return (
base_input_tokens * input_cost_per_token / 2
+ cached_text_tokens * half_cache_read_rate
+ cached_audio_tokens * cached_audio_rate
+ cached_image_tokens * cached_image_rate
+ cache_creation_tokens * cache_creation_cost / 2
)
def _batch_prompt_cost(
@ -3340,6 +3375,7 @@ def _batch_prompt_cost(
model_info: ModelInfo,
input_rate: float,
batch_cache_read_rate: float | None,
batch_cache_read_image_rate: float | None,
cache_creation_rate: float | None,
cache_creation_rates_by_ttl: Mapping[str, float],
) -> float:
@ -3349,8 +3385,30 @@ def _batch_prompt_cost(
if batch_cache_read_rate is None
else batch_cache_read_rate
)
cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0
cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0
cached_audio_rate: Final = _cached_modality_rate(
model_info,
"cache_read_input_audio_token_cost",
_optional_batch_rate(model_info, "cache_read_input_audio_token_cost_batches", usage.prompt_tokens),
batch_cache_read_rate,
cache_read_rate,
usage.prompt_tokens,
)
cached_image_rate: Final = _cached_modality_rate(
model_info,
"cache_read_input_image_token_cost",
batch_cache_read_image_rate,
batch_cache_read_rate,
cache_read_rate,
usage.prompt_tokens,
)
cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cached_audio_rate is not None else 0
cached_image_tokens: Final = details["cached_image_tokens"] if cached_image_rate is not None else 0
cached_text_tokens: Final = (
max(details["cache_hit_tokens"] - details["cache_hit_audio_tokens"] - details["cached_image_tokens"], 0)
if cache_read_rate is not None
else 0
)
cached_tokens: Final = cached_text_tokens + cached_audio_tokens + cached_image_tokens
bills_cache_writes: Final = cache_creation_rate is not None or bool(cache_creation_rates_by_ttl)
written_tokens: Final = details["cache_creation_tokens"] if bills_cache_writes else 0
audio_tokens, image_tokens, video_tokens = (
@ -3367,13 +3425,9 @@ def _batch_prompt_cost(
* _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate, prompt_tokens=usage.prompt_tokens)
+ image_tokens * _batch_rate(model_info, "input_cost_per_image_token_batches", input_rate)
+ video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate)
+ (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0)
+ cached_audio_tokens
* _batch_rate(
model_info,
"cache_read_input_audio_token_cost_batches",
_cached_audio_batch_fallback(model_info, batch_cache_read_rate, cache_read_rate, usage.prompt_tokens),
)
+ cached_text_tokens * (cache_read_rate or 0.0)
+ cached_audio_tokens * (cached_audio_rate or 0.0)
+ cached_image_tokens * (cached_image_rate or 0.0)
+ (
calculate_cache_writing_cost(
written_tokens,

Some files were not shown because too many files have changed in this diff Show more