scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

Merge remote-tracking branch 'origin/main' into litellm_count_tokens_forward_context_management
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:
mateo-berri 2026-11-10 08:48:44 -08:00
commit 3fd430fe23
852 changed files with 54112 additions and 12708 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution:
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
team=await self._live_team_row(team_id, job.created_at),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
async def _live_team_row(
self, team_id: str | None, created_before: datetime | None = None
) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
if team is not None and created_before is not None and team.created_at > created_before:
return None
return team
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:

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

@ -78,6 +78,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,
@ -1426,6 +1427,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,
@ -2253,10 +2291,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,

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;

7
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
@ -2053,6 +2055,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

36
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",
@ -4423,6 +4427,7 @@ dependencies = [
"pyo3-async-runtimes",
"reqwest",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"serde_with",
@ -4558,6 +4563,8 @@ name = "litellm-redis-types"
version = "0.1.0"
dependencies = [
"litellm-secret-value",
"rstest",
"url",
]
[[package]]
@ -4606,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,
),
}
}

18
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",
@ -47,9 +51,18 @@
"defers_async_logging": [
"logger"
],
"claim_async_success_log": [
"logger"
],
"mark_async_success_enqueued": [
"logger"
],
"defer_success": [
"logger",
"pending"
"pending",
"response",
"start",
"end"
],
"sync_success_for_async_call": [
"logger",
@ -117,7 +130,8 @@
"chunks",
"start",
"end",
"first_chunk"
"first_chunk",
"stream_teardown"
],
"stream_failure": [
"logger",

346
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.
@ -251,30 +292,49 @@ impl LegacyLogging {
if !self.asynchronous {
return pending().sync(py);
}
if !self.internal
&& self
if self.internal {
return Ok(());
}
let defers = logger.defers_async_logging(py);
if logger.claim_async_success_log(py)? {
let enqueues = self
.call
.kwargs()
.bind(py)
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if logger.defers_async_logging(py) {
.is_none_or(|value| value.is_none());
if defers {
let pending = Py::new(
py,
PendingLogging {
pending: Some(pending()),
enqueues,
},
)?;
logger.defer_success(py, pending.bind(py).as_any())?;
} else {
return logger.defer_success(
py,
pending.bind(py).as_any(),
&self.response,
&self.start,
&self.end,
);
}
if enqueues {
pending().asynchronous(py)?;
}
}
if defers {
return Ok(());
}
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(
@ -291,6 +351,7 @@ impl LegacyLogging {
&self.start,
&self.end,
&stream.first_chunk,
teardown,
),
);
match billed {
@ -442,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,
@ -454,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)?;
@ -469,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,
@ -525,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,
(
@ -548,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(()))
}
@ -592,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(()))
@ -662,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)
{
@ -1114,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")
@ -1135,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";
@ -1276,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,
@ -1317,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]
",
);
@ -1873,10 +2023,15 @@ mod terminal_tests {
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &[])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
#[case::async_deferred_with_fallbacks(
true,
c"logger._defer_async_logging = True\nkwargs = {'fallbacks': ['other']}",
&[]
)]
fn success_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
@ -1902,12 +2057,88 @@ mod terminal_tests {
&locals,
c"
assert all(value is response for name, value in logger.calls if name.endswith('_handler'))
assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False)
assert hasattr(logger, 'enqueue_deferred_logging') == getattr(logger, '_defer_async_logging', False)
assert getattr(logger, 'deferred_success_result', (response,))[0] is response
",
);
});
}
#[rstest]
#[case::accepted(c"", &["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"])]
#[case::accepted_with_fallbacks(c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn a_deferred_success_runs_sync_callbacks_only_once_released(
#[case] script: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"response = object()\nlogger._defer_async_logging = True",
);
run(py, &locals, script);
let mut logging = logged(py, &locals, true);
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
",
);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
});
}
#[rstest]
#[case::deferred(
c"logger._defer_async_logging = True",
c"logger.enqueue_deferred_logging()",
&["sync_success_for_async_call", "async_success_handler", "mark_async_success_enqueued", "enqueued"]
)]
#[case::deferred_and_blocked(c"logger._defer_async_logging = True", c"", &[])]
#[case::immediate(
c"",
c"",
&["async_success_handler", "mark_async_success_enqueued", "enqueued", "sync_success_for_async_call", "sync_success_for_async_call"]
)]
fn a_fallback_attempt_on_the_same_logger_leaves_the_first_success_claimed(
#[case] script: &CStr,
#[case] release: &CStr,
#[case] expected: &[&str],
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"refused = object()\nresponse = refused");
run(py, &locals, script);
let mut first = logged(py, &locals, true);
succeed(py, &locals, &mut first);
run(py, &locals, c"response = object()");
let mut fallback = logged(py, &locals, true);
succeed(py, &locals, &mut fallback);
run(py, &locals, release);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
.extract()
.unwrap();
assert_eq!(names, expected);
run(
py,
&locals,
c"assert all(value is refused for name, value in logger.calls if name == 'async_success_handler')",
);
});
}
#[rstest]
fn dropped_observations_preserve_deferred_success_and_response_identity() {
Python::initialize();
@ -1933,10 +2164,10 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
));
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
assert logger.names() == ['sync_success_for_async_call'], logger.calls
logger._native_pending_logging.release(True)
logger._native_pending_logging.release(True)
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
assert logger.names() == [], logger.calls
logger.enqueue_deferred_logging()
logger.enqueue_deferred_logging()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
assert logger.calls[1][1] is response
");
@ -1981,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, &[])]
@ -2009,7 +2279,7 @@ assert chunks[1] is last
}
#[rstest::rstest]
fn internal_async_calls_skip_the_async_success_fan_out() {
fn internal_async_calls_skip_success_logging() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"response = object()");
@ -2018,11 +2288,7 @@ assert chunks[1] is last
..logged(py, &locals, true)
};
succeed(py, &locals, &mut logging);
run(
py,
&locals,
c"assert logger.names() == ['sync_success_for_async_call'], logger.calls",
);
run(py, &locals, c"assert logger.calls == [], logger.calls");
});
}

33
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>,
@ -59,7 +59,16 @@ pub trait LegacyCallbacks {
fn defers_async_logging(&self, py: Python<'_>) -> bool;
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()>;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool>;
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()>;
fn sync_success_for_async_call(
&self,
@ -148,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>,
@ -203,8 +212,21 @@ impl LegacyCallbacks for PythonLogger {
.unwrap_or(false)
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending))?;
fn claim_async_success_log(&self, py: Python<'_>) -> PyResult<bool> {
Logging::ClaimAsyncSuccess
.call(py, (self.object(py),))?
.extract()
}
fn defer_success(
&self,
py: Python<'_>,
pending: &Bound<'_, PyAny>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
Logging::DeferSuccess.call(py, (self.object(py), pending, response, start, end))?;
Ok(())
}
@ -252,6 +274,7 @@ impl LegacyCallbacks for PythonLogger {
) -> PyResult<()> {
let coroutine =
Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?;
Logging::MarkAsyncSuccessEnqueued.call(py, (self.object(py),))?;
let enqueue = Logging::Enqueue.call(py, (&coroutine,));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")

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

@ -22,21 +22,35 @@ impl PendingSuccess {
self.logger
.enqueue_success(py, &self.response, &self.start, &self.end)
}
pub(crate) fn sync_for_async_call(&self, py: Python<'_>) -> PyResult<()> {
self.logger
.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
}
#[pyclass]
pub(crate) struct PendingLogging {
pub(crate) pending: Option<PendingSuccess>,
pub(crate) enqueues: bool,
}
#[pymethods]
impl PendingLogging {
fn release(slf: &Bound<'_, Self>, py: Python<'_>, success: bool) -> PyResult<()> {
let pending = slf.borrow_mut().pending.take();
if let Some(pending) = pending
&& success
{
match pending.asynchronous(py) {
fn __call__(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<()> {
let (pending, enqueues) = {
let mut slot = slf.borrow_mut();
(slot.pending.take(), slot.enqueues)
};
if let Some(pending) = pending {
let released = pending.sync_for_async_call(py).and_then(|()| {
if enqueues {
pending.asynchronous(py)
} else {
Ok(())
}
});
match released {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(pending.logger.object(py)));
}
@ -87,6 +101,7 @@ mod tests {
start: py.None(),
end: Some(py.None()),
}),
enqueues: true,
},
)
.unwrap();
@ -108,7 +123,7 @@ observed = []
def on_enqueue(coroutine):
observed.append(marker.get())
pending.release(True)
pending()
logger.on_enqueue = on_enqueue
",
@ -118,28 +133,12 @@ logger.on_enqueue = on_enqueue
&locals,
c"
marker.set('release')
pending.release(True)
pending.release(True)
pending()
pending()
assert observed == ['release'], observed
assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
",
);
});
}
#[test]
fn a_blocked_release_drops_the_success_for_good() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"");
run(
py,
&locals,
c"
pending.release(False)
pending.release(True)
assert logger.calls == [], logger.calls
assert logger.calls[1][1] is response
",
);
});
@ -168,7 +167,7 @@ logger.on_enqueue = on_enqueue
locals
.set_item("failure", py.eval(failure, None, Some(&locals)).unwrap())
.unwrap();
let released = local(&locals, "pending").call_method1("release", (true,));
let released = local(&locals, "pending").call0();
match released {
Ok(_) => assert!(!propagates),
Err(error) => {
@ -181,8 +180,8 @@ logger.on_enqueue = on_enqueue
py,
&locals,
c"
pending.release(True)
assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls
pending()
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'mark_async_success_enqueued', 'enqueued', 'closed'], logger.calls
assert unraisable_from(logger) == ([] if propagates else [failure])
",
);

6
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")]
@ -50,6 +52,8 @@ pub(crate) enum Logging {
RecordApiWindow,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "claim_async_success_log")]
ClaimAsyncSuccess,
#[strum(serialize = "defer_success")]
DeferSuccess,
#[strum(serialize = "sync_success_for_async_call")]
@ -58,6 +62,8 @@ pub(crate) enum Logging {
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "mark_async_success_enqueued")]
MarkAsyncSuccessEnqueued,
#[strum(serialize = "enqueue_logging")]
Enqueue,
#[strum(serialize = "failure_handler")]

15
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
),
@ -72,7 +73,11 @@ FAKES = {
llm_api_duration_ms=(end - start) * 1000
),
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
'claim_async_success_log': lambda logger: logger.claim_async_success_log(),
'defer_success': lambda logger, pending, response, start, end: (
setattr(logger, 'deferred_success_result', (response, start, end)),
setattr(logger, 'enqueue_deferred_logging', pending),
),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
response, start, end
),
@ -81,6 +86,7 @@ FAKES = {
)(error, ''.join(traceback.format_exception(error)), start, end),
'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)),
'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end),
'mark_async_success_enqueued': lambda logger: logger.record('mark_async_success_enqueued', None),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
@ -92,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),
@ -129,6 +135,7 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.model_call_details = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
@ -156,6 +163,10 @@ class StubLogger:
self.record('async_success_handler', response)
return StubCoroutine(self)
def claim_async_success_log(self):
claimed, self.claimed = getattr(self, 'claimed', False), True
return not claimed
def handle_sync_success_callbacks_for_async_calls(self, response, start, end):
self.record('sync_success_for_async_call', response)

9
litellm-rust/crates/circuit-breaker/AGENTS.md
View file

@ -1,13 +1,14 @@
- `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer
- `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary
- Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition
- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition
- The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric
- A disabled breaker is `Option<Breaker>::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job
## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`)
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps
- Divergences, each on purpose
## The only Redis breaker
- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone
- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees
- Divergences from the old Python class, each on purpose
- monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker
- a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common
- `_failure_count` is dropped, it only fed a log line

15
litellm-rust/crates/circuit-breaker/src/breaker.rs
View file

@ -1,5 +1,5 @@
use std::{
sync::Arc,
sync::{Arc, Weak},
time::{Duration, Instant},
};
@ -128,6 +128,9 @@ impl Machine for Config {
#[derive(Clone)]
pub struct Breaker(Arc<Shared<Config>>);
#[derive(Clone)]
pub struct WeakBreaker(Weak<Shared<Config>>);
#[derive(Debug, PartialEq, Eq)]
pub struct Rejected;
@ -165,6 +168,16 @@ impl Breaker {
pub fn state(&self) -> State {
self.0.state()
}
pub fn downgrade(&self) -> WeakBreaker {
WeakBreaker(Arc::downgrade(&self.0))
}
}
impl WeakBreaker {
pub fn upgrade(&self) -> Option<Breaker> {
self.0.upgrade().map(Breaker)
}
}
impl Admission {

4
litellm-rust/crates/circuit-breaker/src/lib.rs
View file

@ -1,5 +1,7 @@
mod breaker;
mod fsm;
pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State};
pub use breaker::{
Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker,
};
pub use fsm::{Machine, Observer, Shared, Transition};

5
litellm-rust/crates/gateway/src/redis/mod.rs
View file

@ -4,7 +4,7 @@ use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_config::{CacheSettings, Config, RedisSettings, Value};
use litellm_core_utils::settings::Lookup;
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{HostServices, SharedEnvironment, TracingEvents};
use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents};
use litellm_redis_resolve::Expected;
use litellm_redis_types::Resolution;
@ -80,8 +80,9 @@ pub(crate) async fn connect(
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let events: Arc<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
let hosts = RedisHosts::default();
let client = |params: Params| -> Result<RedisClient, Error> {
Ok(litellm_redis_connect::connect(
Ok(hosts.client(
params,
HostServices {
events: Some(events.clone()),

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;

12
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 {
@ -144,6 +153,7 @@ mod tests {
supports_reasoning: true,
supports_adaptive_thinking: true,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: true,
supports_sampling_params: false,

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> {

1
litellm-rust/crates/llms/src/anthropic/common_utils.rs
View file

@ -1913,6 +1913,7 @@ mod tests {
supports_reasoning: false,
supports_adaptive_thinking: false,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: false,
supports_sampling_params: true,

37
litellm-rust/crates/llms/src/anthropic/messages/thinking.rs
View file

@ -2,8 +2,8 @@ use litellm_llms_types::{
formats::{
chat_completions::ReasoningEffort,
messages::{
AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig,
ThinkingConfig, ThinkingDisplay,
AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams,
MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay,
},
},
recognized::Recognized,
@ -172,7 +172,11 @@ fn translate_reasoning_effort(
) else {
return Ok(MessagesRequest {
params: MessagesOptionalParams {
thinking: None,
thinking: (context.capabilities.thinking_on_by_default
&& !context.capabilities.thinking_always_on)
.then(|| {
Recognized::Known(ThinkingConfig::Disabled(DisabledThinking::default()))
}),
output_config: None,
..request.params
},
@ -748,6 +752,33 @@ mod tests {
);
}
#[rstest]
#[case::default_thinking(true, false, Some(serde_json::json!({"type": "disabled"}))) ]
#[case::legacy(false, false, None)]
#[case::always_on(true, true, None)]
fn none_disables_only_optional_default_thinking(
#[case] thinking_on_by_default: bool,
#[case] thinking_always_on: bool,
#[case] thinking: Option<Value>,
) {
let capabilities = MessagesModelCapabilities {
thinking_on_by_default,
thinking_always_on,
..opus_4_7()
};
let expected = match thinking {
Some(thinking) => serde_json::json!({"max_tokens": 1024, "thinking": thinking}),
None => serde_json::json!({"max_tokens": 1024}),
};
assert_eq!(
translate(
capabilities,
serde_json::json!({"max_tokens": 1024, "reasoning_effort": "none"})
),
Ok(request(expected))
);
}
#[rstest]
#[case::adaptive_model(opus_4_7())]
#[case::effort_capable_model(opus_4_5())]

3
litellm-rust/crates/llms/src/base_llm/messages/context.rs
View file

@ -32,6 +32,8 @@ pub struct MessagesModelCapabilities {
#[serde(default)]
pub thinking_always_on: bool,
#[serde(default)]
pub thinking_on_by_default: bool,
#[serde(default)]
pub supports_legacy_thinking: bool,
#[serde(default)]
pub supports_output_config: bool,
@ -55,6 +57,7 @@ impl Default for MessagesModelCapabilities {
supports_reasoning: false,
supports_adaptive_thinking: false,
thinking_always_on: false,
thinking_on_by_default: false,
supports_legacy_thinking: false,
supports_output_config: false,
supports_sampling_params: true,

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

109
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,
@ -142,6 +148,7 @@ impl Catalog {
let mut sections = Vec::new();
let mut entries = IndexMap::with_capacity(document.len());
let mut alias_lists = Vec::new();
let mut pointer_rows = Vec::new();
for (index, (name, value)) in document.into_iter().enumerate() {
let fields = match value {
Value::Object(fields) if !is_section(&name) => fields,
@ -157,16 +164,34 @@ impl Catalog {
continue;
}
};
let (model, aliases, row_problems) = read_row(&fields);
let row = read_row(&fields);
let readable = row.problems.is_empty();
problems.extend(
row_problems
row.problems
.into_iter()
.map(|problem| problem.prepend(&name)),
);
if let Some(aliases) = aliases {
if let Some(target) = row.alias_of {
if readable {
pointer_rows.push((name.clone(), target));
}
sections.push(Section {
index,
key: name,
value: Value::Object(fields),
});
continue;
}
if let Some(aliases) = row.aliases {
alias_lists.push((name.clone(), aliases));
}
entries.insert(name, ModelEntry { fields, model });
entries.insert(
name,
ModelEntry {
fields,
model: row.model,
},
);
}
if entries.is_empty() {
return Err(Error::new("", Repr::Empty));
@ -190,6 +215,21 @@ impl Catalog {
}
}
}
for (alias, target) in pointer_rows {
if !entries.contains_key(&target) {
problems.push(Error::from_key(&alias, Repr::AliasTargetMissing { target }));
} else if let Some(taken_by) = aliases.get(&alias) {
problems.push(Error::from_key(
&alias,
Repr::AliasTaken {
alias: alias.clone(),
taken_by: taken_by.clone(),
},
));
} else {
aliases.insert(alias, target);
}
}
let fallback_generalizations = sections
.iter()
@ -242,8 +282,8 @@ impl Catalog {
}
pub fn row_errors(model: &str, fields: &Fields) -> Vec<Error> {
let (_, _, errors) = read_row(fields);
errors
read_row(fields)
.problems
.into_iter()
.map(|error| error.prepend(model))
.collect()
@ -255,7 +295,13 @@ impl Catalog {
(0..self.rows().entries.len() + self.sections.len()).map_while(move |index| match sections
.next_if(|section| section.index == index)
{
Some(section) => Some((section.key.as_str(), Written::Section(&section.value))),
Some(section) => Some((
section.key.as_str(),
match &section.value {
Value::Object(fields) if !is_section(&section.key) => Written::Entry(fields),
value => Written::Section(value),
},
)),
None => entries
.next()
.map(|(name, entry)| (name.as_str(), Written::Entry(&entry.fields))),
@ -378,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()
}
@ -406,16 +466,41 @@ impl Catalog {
}
}
fn read_row(fields: &Fields) -> (Option<Model>, Option<Vec<String>>, Vec<Error>) {
struct RowRead {
model: Option<Model>,
aliases: Option<Vec<String>>,
alias_of: Option<String>,
problems: Vec<Error>,
}
fn read_row(fields: &Fields) -> RowRead {
let wire = EntryWire::read(fields);
if let Some(target) = wire.attributes.alias_of.clone() {
let other_fields = fields
.keys()
.filter(|key| *key != "alias_of")
.map(|key| Error::from_key(key, Repr::AliasRowField));
return RowRead {
model: None,
aliases: None,
alias_of: Some(target),
problems: wire.problems.into_iter().chain(other_fields).collect(),
};
}
let aliases = wire.attributes.aliases.clone();
let timezone_error = unknown_timezone(&wire);
match Model::try_from(wire) {
Ok(model) => (Some(model), aliases, timezone_error.into_iter().collect()),
let (model, problems) = match Model::try_from(wire) {
Ok(model) => (Some(model), timezone_error.into_iter().collect()),
Err(mut errors) => {
errors.extend(timezone_error);
(None, aliases, errors)
(None, errors)
}
};
RowRead {
model,
aliases,
alias_of: None,
problems,
}
}

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}")
}

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

@ -67,6 +67,12 @@ pub(crate) enum Repr {
#[error("alias {alias:?} already belongs to {taken_by}")]
#[strum_discriminants(strum(serialize = "alias_taken"))]
AliasTaken { alias: String, taken_by: String },
#[error("alias_of {target:?} is not a model name")]
#[strum_discriminants(strum(serialize = "alias_target_missing"))]
AliasTargetMissing { target: String },
#[error("an alias_of row carries no other field")]
#[strum_discriminants(strum(serialize = "alias_row_field"))]
AliasRowField,
#[error("{first} and {second} price the same meter under the same conditions")]
#[strum_discriminants(strum(serialize = "duplicate_price"))]
DuplicatePrice { first: String, second: String },
@ -141,6 +147,8 @@ impl Error {
| Repr::UnknownTimezone { .. }
| Repr::AliasShadowsModel { .. }
| Repr::AliasTaken { .. }
| Repr::AliasTargetMissing { .. }
| Repr::AliasRowField
| Repr::DuplicatePrice { .. }
| Repr::SuffixInsideBlock { .. }
| Repr::TierRange
@ -159,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,
@ -327,6 +403,8 @@ mod tests {
#[case::unknown_timezone(Repr::UnknownTimezone { zone: "invalid".into() }, Rejects::Row)]
#[case::alias_shadows_model(Repr::AliasShadowsModel { alias: "model".into() }, Rejects::Row)]
#[case::alias_taken(Repr::AliasTaken { alias: "alias".into(), taken_by: "owner".into() }, Rejects::Row)]
#[case::alias_target_missing(Repr::AliasTargetMissing { target: "owner".into() }, Rejects::Row)]
#[case::alias_row_field(Repr::AliasRowField, Rejects::Row)]
#[case::duplicate_price(Repr::DuplicatePrice { first: "a".into(), second: "b".into() }, Rejects::Row)]
#[case::suffix_inside_block(Repr::SuffixInsideBlock { block: "tiered_pricing" }, Rejects::Row)]
#[case::tier_range(Repr::TierRange, Rejects::Row)]

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::*;

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<_>>()
);
}
}

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

@ -104,9 +104,18 @@ pub fn registry_json_schema() -> Schema {
"title".into(),
"LiteLLM model prices and context window registry".into(),
);
definitions.insert(
"aliasRow".into(),
json!({
"type": "object",
"required": ["alias_of"],
"properties": {"alias_of": {"type": "string"}},
"additionalProperties": false,
}),
);
root.insert(
"additionalProperties".into(),
json!({"$ref": "#/$defs/modelEntry"}),
json!({"anyOf": [{"$ref": "#/$defs/modelEntry"}, {"$ref": "#/$defs/aliasRow"}]}),
);
root.insert("$defs".into(), Value::Object(definitions));
Schema::from(root.clone())
@ -153,6 +162,7 @@ fn decorate_model_entry(schema: &mut Value) {
let rates = json!({FLAT_KEY_PATTERN: {"type": "number", "minimum": 0}});
let object = schema.as_object_mut().unwrap();
object.insert("required".into(), json!(["litellm_provider"]));
object.insert("not".into(), json!({"required": ["alias_of"]}));
object.insert("additionalProperties".into(), Value::Bool(true));
object.insert(
"patternProperties".into(),

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

2
litellm-rust/crates/model-catalog/src/wire/attributes.rs
View file

@ -47,6 +47,8 @@ macro_rules! attributes {
}
attributes! {
/// Model this row is another name for; such a row carries no other key.
alias_of: String,
/// Other names this entry is served under.
aliases: Vec<String>,
audio_transcription_config: String,

79
litellm-rust/crates/model-catalog/tests/catalog.rs
View file

@ -139,6 +139,18 @@ fn a_bad_entry_is_reported_by_model_and_key_and_kept_raw() {
json!({"First": {"litellm_provider": "t", "aliases": ["shared"]}, "Second": {"litellm_provider": "t", "aliases": ["shared"]}}),
r#"Second: alias "shared" already belongs to First"#
)]
#[case::alias_of_target_missing(
json!({"Real": {"litellm_provider": "t"}, "ghost-latest": {"alias_of": "Ghost"}}),
r#"ghost-latest: alias_of "Ghost" is not a model name"#
)]
#[case::alias_of_row_with_other_fields(
json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real", "mode": "chat"}}),
"real-latest.mode: an alias_of row carries no other field"
)]
#[case::alias_of_claimed_by_a_list(
json!({"Real": {"litellm_provider": "t"}, "real-latest": {"alias_of": "Real"}, "Other": {"litellm_provider": "t", "aliases": ["real-latest"]}}),
r#"real-latest: alias "real-latest" already belongs to Other"#
)]
#[case::aliases_not_a_list(
json!({"a": {"litellm_provider": "t", "aliases": "bad"}}),
"a.aliases: invalid value: invalid type: string \"bad\", expected a sequence"
@ -399,3 +411,70 @@ fn a_bad_value_keeps_its_json_error_as_the_source_and_names_its_key_and_field()
assert_eq!(error.key().as_deref(), Some("open/router"));
assert_eq!(error.field().as_deref(), Some("input_cost_per_token"));
}
#[rstest]
fn an_alias_of_row_is_served_as_its_target_and_written_as_itself() {
let catalog = catalog(json!({
"Real": {"litellm_provider": "t", "input_cost_per_token": 1e-6},
"real-latest": {"alias_of": "Real"},
"After": {"litellm_provider": "t"}
}));
let found = catalog.lookup("REAL-LATEST").unwrap();
assert_eq!(
(found.matched_key, found.canonical_key),
("real-latest", "Real")
);
assert_eq!(
found.entry.field("input_cost_per_token"),
Some(&json!(1e-6))
);
assert_eq!(catalog.model_count(), 2);
assert_eq!(
catalog.served().map(|(key, _)| key).collect::<Vec<_>>(),
["Real", "After", "real-latest"]
);
assert_eq!(
written(&catalog),
r#"{"Real":{"litellm_provider":"t","input_cost_per_token":1e-6},"real-latest":{"alias_of":"Real"},"After":{"litellm_provider":"t"}}"#
);
}
#[rstest]
fn a_dangling_alias_of_row_is_rejected_and_dropped_from_the_written_document() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t"},
"ghost-latest": {"alias_of": "Ghost"}
}));
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert!(catalog.lookup("ghost-latest").is_none());
assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#);
}
#[rstest]
fn an_alias_of_row_with_other_fields_is_never_served() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t"},
"real-latest": {"alias_of": "Real", "mode": "chat"}
}));
assert!(parsed.value.lookup("real-latest").is_none());
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert!(catalog.lookup("real-latest").is_none());
assert_eq!(written(&catalog), r#"{"Real":{"litellm_provider":"t"}}"#);
}
#[rstest]
fn a_pointer_row_that_repeats_a_listed_alias_is_dropped_and_the_model_is_kept() {
let parsed = loaded(json!({
"Real": {"litellm_provider": "t", "aliases": ["real-latest"]},
"real-latest": {"alias_of": "Real"}
}));
let (catalog, removed) = parsed.value.without_rejected(&parsed.errors);
assert_eq!(removed.len(), 1);
assert_eq!(catalog.lookup("real-latest").unwrap().canonical_key, "Real");
assert_eq!(
written(&catalog),
r#"{"Real":{"litellm_provider":"t","aliases":["real-latest"]}}"#
);
}

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());
}

4
litellm-rust/crates/model-catalog/tests/schema.rs
View file

@ -35,6 +35,7 @@ fn generated_registry_schema_validates_checked_in_catalog() {
#[case(json!({"example": {"litellm_provider": "test"}}))]
#[case(json!({"example": {"litellm_provider": "test", "future_field": true}}))]
#[case(json!({"sample_spec": {"litellm_provider": "placeholder"}}))]
#[case(json!({"example": {"litellm_provider": "test"}, "example-latest": {"alias_of": "example"}}))]
fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value) {
assert!(registry_validator().is_valid(&document));
}
@ -54,6 +55,9 @@ fn generated_registry_schema_keeps_reader_compatibility(#[case] document: Value)
#[case::empty_windows(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": []}}))]
#[case::invalid_weekday(json!({"litellm_provider": "test", "off_peak_pricing": {"windows": [{"hours_utc": "00:00-01:00", "weekdays": [0]}]}}))]
#[case::invalid_aliases(json!({"litellm_provider": "test", "aliases": "wrong"}))]
#[case::alias_of_on_a_model(json!({"litellm_provider": "test", "alias_of": "other"}))]
#[case::alias_row_with_other_fields(json!({"alias_of": "other", "mode": "chat"}))]
#[case::alias_of_not_a_string(json!({"alias_of": 4}))]
#[case::non_object_model(json!(4))]
fn generated_registry_schema_rejects_invalid_entries(#[case] entry: Value) {
assert!(!registry_validator().is_valid(&json!({"example": entry})));

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

@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"]
tiktoken = ["litellm-token-counter/tiktoken"]
[dependencies]
rustls.workspace = true
litellm-secret-value.workspace = true
fancy-regex.workspace = true
litellm-tracing.workspace = true
@ -59,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(

44
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,11 +37,17 @@ 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::client::NativeRedisClient;
use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states};
#[pymodule_export]
use crate::redis::client::native_redis_breaker_states;
use crate::redis::client::NativeRedisClient;
#[pymodule_export]
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
#[pymodule_export]
@ -55,7 +62,9 @@ mod _native {
#[pymodule_export]
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
#[pymodule_export]
use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead};
use crate::redis::store::spend_counters::{
NativeSpendCounters, PyPostCallSettlement, PySpendRead,
};
#[pymodule_export]
use crate::redis::store::{
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
@ -90,6 +99,7 @@ mod _native {
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let _ = rustls::crypto::ring::default_provider().install_default();
let py = module.py();
let dict = module.dict();
dict.set_item(
@ -136,6 +146,8 @@ mod tests {
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeRedisClient",
"NativeRedisBreaker",
"NativeAdmission",
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
@ -152,6 +164,7 @@ mod tests {
"NativeSensitiveRoutePins",
"NativeRedisStoreError",
"NativeSpendCounters",
"PostCallSettlement",
"SpendRead",
"LocalSizes",
"native_redis_breaker_states",
@ -160,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",
@ -184,4 +208,18 @@ mod tests {
assert_eq!(public_names, expected);
});
}
#[rstest::rstest]
fn module_import_installs_ring_as_the_process_crypto_provider() {
Python::initialize();
Python::attach(|py| {
native_module(py);
});
let installed = rustls::crypto::CryptoProvider::get_default()
.expect("importing _native should install a process crypto provider");
assert_eq!(
format!("{installed:?}"),
format!("{:?}", rustls::crypto::ring::default_provider())
);
}
}

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 }
}

134
litellm-rust/crates/python-bridge/src/redis/breaker.rs Normal file
View file

@ -0,0 +1,134 @@
use std::{
sync::{Arc, Mutex, PoisonError},
time::{Duration, Instant},
};
use litellm_circuit_breaker::{Admission, Breaker, Config, Failure};
use litellm_redis_client::RedisEvents;
use litellm_redis_types::Endpoint;
use pyo3::{exceptions::PyValueError, prelude::*};
use crate::redis::{
client::{RedisSpec, resolved_endpoint},
events::{ServiceLoggingEvents, state_label},
hosts,
};
/// `litellm.redis.store_settings.RedisBreakerSettings`, filled from `litellm.constants`.
#[derive(FromPyObject)]
pub(crate) struct BreakerSettings {
failure_threshold: u32,
recovery_timeout: f64,
timeout_min_duration: f64,
}
impl BreakerSettings {
pub(crate) fn config(&self) -> PyResult<Config> {
let seconds = |seconds: f64| {
Duration::try_from_secs_f64(seconds)
.map_err(|error| PyValueError::new_err(error.to_string()))
};
Ok(Config {
failure_threshold: self.failure_threshold,
recovery_timeout: seconds(self.recovery_timeout)?,
timeout_min_duration: seconds(self.timeout_min_duration)?,
})
}
}
/// The one health signal for a Redis, shared by its `RedisCache` and its native client, so a Redis the
/// cache found wedged is fast-failed by every store too.
#[pyclass(frozen, name = "NativeRedisBreaker")]
pub(crate) struct NativeRedisBreaker(pub(crate) Breaker);
impl NativeRedisBreaker {
fn build(
py: Python<'_>,
settings: &BreakerSettings,
endpoint: Option<Endpoint>,
) -> PyResult<Self> {
let config = settings.config()?;
let runtime = litellm_host_python::runtime_handle()?;
let events: Arc<dyn RedisEvents> =
Arc::new(ServiceLoggingEvents::transitions_only(py, runtime)?);
Ok(Self(py.detach(|| {
hosts().breaker(endpoint, config, Some(events))
})))
}
}
#[pymethods]
impl NativeRedisBreaker {
#[new]
fn new(py: Python<'_>, settings: BreakerSettings) -> PyResult<Self> {
Self::build(py, &settings, None)
}
/// The breaker of the Redis server `redis_kwargs` names, which every cache and client on that server shares
/// whatever its db, credentials or namespace. The first caller's settings win.
#[staticmethod]
fn shared(
py: Python<'_>,
redis_kwargs: RedisSpec,
settings: BreakerSettings,
) -> PyResult<Self> {
let endpoint = resolved_endpoint(py, &redis_kwargs)?;
Self::build(py, &settings, endpoint)
}
fn admit(&self) -> Option<NativeAdmission> {
self.0
.admit(Instant::now())
.ok()
.map(|admission| NativeAdmission(Mutex::new(Some(admission))))
}
fn failed(&self, timeout: bool) {
self.0.failed(failure(timeout), Instant::now());
}
#[getter]
fn state(&self) -> &'static str {
state_label(&self.0.state())
}
}
/// One admitted call. Settling it twice records nothing more, and dropping an unsettled probe reopens the breaker.
#[pyclass(frozen, name = "NativeAdmission")]
pub(crate) struct NativeAdmission(Mutex<Option<Admission>>);
impl NativeAdmission {
fn take(&self) -> Option<Admission> {
self.0.lock().unwrap_or_else(PoisonError::into_inner).take()
}
}
#[pymethods]
impl NativeAdmission {
fn succeeded(&self) {
if let Some(admission) = self.take() {
admission.succeeded();
}
}
fn failed(&self, timeout: bool) {
if let Some(admission) = self.take() {
admission.failed(failure(timeout), Instant::now());
}
}
}
fn failure(timeout: bool) -> Failure {
match timeout {
true => Failure::Timeout,
false => Failure::Hard,
}
}
#[pyfunction]
pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> {
py.detach(|| hosts().breakers())
.iter()
.map(|breaker| state_label(&breaker.state()))
.collect()
}

131
litellm-rust/crates/python-bridge/src/redis/client/mod.rs
View file

@ -1,28 +1,23 @@
mod credentials;
use std::{
collections::BTreeMap,
sync::{Arc, LazyLock, Mutex, PoisonError},
time::Duration,
};
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use litellm_circuit_breaker::Config as BreakerConfig;
use litellm_core_utils::settings::ProcessEnvironment;
use litellm_host_python::PythonContext;
use litellm_python_compat::number::{self, SecondsError};
use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient};
use litellm_redis_client::{RedisClient, RedisEvents};
use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment};
use litellm_redis_resolve::{Kwargs, PythonException};
use litellm_redis_types::{ProviderHandle, Resolution};
use litellm_redis_types::{Endpoint, ProviderHandle, Resolution};
use pyo3::{
exceptions::{PyOverflowError, PyTypeError, PyValueError},
prelude::*,
sync::MutexExt,
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple},
};
use self::credentials::PythonCredentialSource;
use crate::redis::events::{ServiceLoggingEvents, state_label};
use crate::redis::{breaker::BreakerSettings, events::ServiceLoggingEvents, hosts};
pyo3::import_exception!(redis.exceptions, RedisClusterException);
@ -67,42 +62,13 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
}
}
/// `litellm.redis.store_settings.RedisBreakerSettings`, which `RedisCache` fills from `litellm.constants`.
#[derive(FromPyObject)]
pub(crate) struct RedisBreaker {
#[pyo3(attribute)]
failure_threshold: u32,
#[pyo3(attribute)]
recovery_timeout: f64,
#[pyo3(attribute)]
timeout_min_duration: f64,
}
impl RedisBreaker {
fn config(&self) -> PyResult<BreakerConfig> {
let seconds = |seconds: f64| {
Duration::try_from_secs_f64(seconds)
.map_err(|error| PyValueError::new_err(error.to_string()))
};
Ok(BreakerConfig {
failure_threshold: self.failure_threshold,
recovery_timeout: seconds(self.recovery_timeout)?,
timeout_min_duration: seconds(self.timeout_min_duration)?,
})
}
}
/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built.
#[derive(FromPyObject)]
pub(crate) struct StoreSettings {
#[pyo3(attribute)]
redis_kwargs: RedisSpec,
#[pyo3(attribute)]
namespace: Option<String>,
#[pyo3(attribute)]
default_ttl: f64,
#[pyo3(attribute)]
breaker: Option<RedisBreaker>,
breaker: Option<BreakerSettings>,
}
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
@ -115,22 +81,6 @@ fn is_plain(value: &Bound<'_, PyAny>) -> bool {
|| value.is_instance_of::<PyDict>()
}
/// A client the process built, held weakly so a discarded `Router` drops its client and provider.
struct Shared {
resolved: Resolution,
client: WeakRedisClient,
}
static CLIENTS: LazyLock<Mutex<Vec<Shared>>> = LazyLock::new(|| Mutex::new(Vec::new()));
fn live_clients(py: Python<'_>) -> std::sync::MutexGuard<'static, Vec<Shared>> {
let mut clients = CLIENTS
.lock_py_attached(py)
.unwrap_or_else(PoisonError::into_inner);
clients.retain(|shared| shared.client.upgrade().is_some());
clients
}
/// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool
/// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on
/// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error.
@ -141,17 +91,8 @@ fn shared_client(
breaker: Option<BreakerConfig>,
) -> PyResult<Option<RedisClient>> {
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
let RedisSpec {
kwargs,
credential_provider,
} = spec;
// The registry entry's client keeps the provider alive, so no live entry shares its address
let handle = credential_provider
.as_ref()
.map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64));
let logger = crate::logger::capture(py);
let environ = py.import("os")?.getattr("environ")?;
let resolved = match logger.scope(|| resolve(&kwargs, handle, &environ))? {
let resolved = match logger.scope(|| resolve_spec(py, &spec))? {
Ok(resolution) => {
logger.scope(|| {
for ignored in &resolution.ignored {
@ -175,7 +116,8 @@ fn shared_client(
events: Some(events),
http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?,
environment,
credential_provider: credential_provider
credential_provider: spec
.credential_provider
.as_ref()
.map(|provider| -> PyResult<Arc<dyn CredentialSource>> {
Ok(Arc::new(PythonCredentialSource::new(
@ -185,30 +127,14 @@ fn shared_client(
})
.transpose()?,
};
let mut clients = live_clients(py);
if let Some(client) = clients
.iter()
.filter(|shared| shared.resolved == resolved)
.find_map(|shared| shared.client.upgrade())
{
return Ok(Some(client));
}
let connected = {
let connected = py.detach(|| {
let _runtime = runtime.enter();
litellm_redis_connect::connect(resolved.clone(), host, breaker)
};
let client = match connected {
Ok(client) => client,
Err(error) => {
drop(clients);
return Ok(unconnected(py, &error));
}
};
clients.push(Shared {
resolved,
client: client.downgrade(),
hosts().client(resolved, host, breaker)
});
Ok(Some(client))
match connected {
Ok(client) => Ok(Some(client)),
Err(error) => Ok(unconnected(py, &error)),
}
}
/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store
@ -232,7 +158,7 @@ impl NativeRedisClient {
breaker,
} = settings;
let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?;
let breaker = breaker.as_ref().map(RedisBreaker::config).transpose()?;
let breaker = breaker.as_ref().map(BreakerSettings::config).transpose()?;
Ok(
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
client,
@ -247,16 +173,23 @@ impl NativeRedisClient {
}
}
/// The breaker state of every native Redis client this process built, so the Prometheus state
/// gauge counts them next to the Python `RedisCache` breakers.
#[pyfunction]
pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> {
live_clients(py)
.iter()
.filter_map(|shared| shared.client.upgrade())
.filter_map(|client| client.breaker_state())
.map(|state| state_label(&state))
.collect()
/// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it.
pub(crate) fn resolved_endpoint(py: Python<'_>, spec: &RedisSpec) -> PyResult<Option<Endpoint>> {
Ok(resolve_spec(py, spec)?
.ok()
.and_then(|resolution| resolution.target.endpoint()))
}
// A live registry entry's client keeps its provider alive, so no live entry shares its address
fn resolve_spec(
py: Python<'_>,
spec: &RedisSpec,
) -> PyResult<Result<Resolution, litellm_redis_resolve::Error>> {
let handle = spec
.credential_provider
.as_ref()
.map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64));
resolve(&spec.kwargs, handle, &py.import("os")?.getattr("environ")?)
}
/// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth

32
litellm-rust/crates/python-bridge/src/redis/events.rs
View file

@ -54,21 +54,33 @@ pub(crate) struct ServiceLoggingEvents {
}
impl ServiceLoggingEvents {
const MODULE: &str = "litellm.rust_bridge.host.redis_events";
pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
const MODULE: &str = "litellm.rust_bridge.host.redis_events";
Self::from_module(py, &py.import(MODULE)?, runtime)
Self::from_module(
py,
&py.import(Self::MODULE)?,
runtime,
EventLoop::running(py),
)
}
/// For a breaker, whose transitions never run on the caller's loop, so building it never looks one up.
pub(crate) fn transitions_only(py: Python<'_>, runtime: Handle) -> PyResult<Self> {
Self::from_module(py, &py.import(Self::MODULE)?, runtime, None)
}
fn from_module(
py: Python<'_>,
module: &Bound<'_, PyModule>,
runtime: Handle,
event_loop: Option<EventLoop>,
) -> PyResult<Self> {
Ok(Self {
record_call: Arc::new(module.getattr("record_redis_call")?.unbind()),
record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()),
context: PythonContext::capture(py)?,
event_loop: EventLoop::running(py),
event_loop,
runtime,
})
}
@ -226,14 +238,14 @@ def record_redis_breaker(state):
let module =
PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name)
.unwrap();
let events = ServiceLoggingEvents::from_module(py, &module, Handle::current()).unwrap();
(
ServiceLoggingEvents {
event_loop: client_loop.map(|client_loop| stand_in_loop(py, client_loop)),
..events
},
module.unbind(),
let events = ServiceLoggingEvents::from_module(
py,
&module,
Handle::current(),
client_loop.map(|client_loop| stand_in_loop(py, client_loop)),
)
.unwrap();
(events, module.unbind())
})
}

12
litellm-rust/crates/python-bridge/src/redis/mod.rs
View file

@ -1,3 +1,15 @@
use std::sync::LazyLock;
use litellm_redis_connect::RedisHosts;
pub(crate) mod breaker;
pub(crate) mod client;
mod events;
pub(crate) mod store;
static HOSTS: LazyLock<RedisHosts> = LazyLock::new(RedisHosts::default);
/// The Redis servers this process talks to, shared by every `RedisCache` and native store in it.
pub(crate) fn hosts() -> &'static RedisHosts {
&HOSTS
}

14
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs
View file

@ -64,9 +64,7 @@ impl<'py> IntoPyObject<'py> for LocalValue {
/// `litellm.redis.stores.rate_limits.RateLimitScope`.
#[derive(FromPyObject)]
struct PyScope {
#[pyo3(attribute)]
descriptor_key: String,
#[pyo3(attribute)]
descriptor_value: String,
}
@ -82,9 +80,7 @@ impl From<PyScope> for Scope {
/// `RateLimitCounter`.
#[derive(FromPyObject)]
struct PyCounter {
#[pyo3(attribute)]
scope: PyScope,
#[pyo3(attribute)]
kind: Spelled<Kind>,
}
@ -100,9 +96,7 @@ impl From<PyCounter> for Counter {
/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names.
#[derive(FromPyObject)]
struct PyEntry {
#[pyo3(attribute)]
suffix: String,
#[pyo3(attribute)]
scope: PyScope,
}
@ -118,11 +112,8 @@ impl TryFrom<PyEntry> for Entry {
/// `CounterIncrement`.
#[derive(FromPyObject)]
struct PyIncrement {
#[pyo3(attribute)]
counter: PyCounter,
#[pyo3(attribute)]
amount: PyNumber,
#[pyo3(attribute)]
ttl: Option<u64>,
}
@ -139,15 +130,10 @@ impl From<PyIncrement> for Increment {
/// `AtomicCounter`.
#[derive(FromPyObject)]
struct PyAtomicCounter {
#[pyo3(attribute)]
counter: PyCounter,
#[pyo3(attribute)]
limit: i64,
#[pyo3(attribute)]
increment: i64,
#[pyo3(attribute)]
ttl: i64,
#[pyo3(attribute)]
window_size: i64,
}

21
litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs
View file

@ -1,13 +1,10 @@
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Send,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use pyo3::{prelude::*, types::PyTuple};
use super::{
NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies,
@ -173,11 +170,13 @@ impl NativeRouterCooldowns {
self.0.flush_local_counters().map_err(kind_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, cooldown) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyData(cooldown))?;
}
Ok(snapshot)
fn snapshot(&self) -> PyResult<HashMap<String, PyData<Cooldown>>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, cooldown)| (key, PyData(cooldown)))
.collect())
}
}

22
litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs
View file

@ -1,7 +1,7 @@
use std::{sync::Arc, time::Duration};
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
use pyo3::{prelude::*, types::PyDict};
use pyo3::prelude::*;
use super::{
PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote,
@ -11,13 +11,9 @@ use crate::redis::client::NativeRedisClient;
/// `litellm.redis.stores.router_usage.UsageCounter`.
#[derive(FromPyObject)]
struct PyCounter {
#[pyo3(attribute)]
deployment_id: String,
#[pyo3(attribute)]
model: Option<String>,
#[pyo3(attribute)]
minute: String,
#[pyo3(attribute)]
kind: Spelled<Kind>,
}
@ -174,11 +170,13 @@ impl NativeRouterUsage {
self.0.flush_local_counters().map_err(kind_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, count) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyNumber(count))?;
}
Ok(snapshot)
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, count)| (key, PyNumber(count)))
.collect())
}
}

45
litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs
View file

@ -1,13 +1,13 @@
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters},
};
use pyo3::{prelude::*, types::PyDict};
use pyo3::prelude::*;
use super::{
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl, unwrap_all,
};
@ -29,6 +29,18 @@ impl PySpendRead {
}
}
/// Whether every increment one `increment_many_post_call` declared has its reply.
#[pyclass(frozen, name = "PostCallSettlement")]
pub(crate) struct PyPostCallSettlement(PostCallSettlement);
#[pymethods]
impl PyPostCallSettlement {
#[getter]
fn settled(&self) -> bool {
self.0.is_settled()
}
}
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
increments
.into_iter()
@ -73,6 +85,12 @@ impl NativeSpendCounters {
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
}
/// The least TTL a Redis write gives a committed spend twin.
fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> {
self.0.set_committed_ttl(duration(ttl_seconds)?);
Ok(())
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
fn read<'py>(
&self,
@ -125,14 +143,15 @@ impl NativeSpendCounters {
})
}
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
/// Declares the increments on the request's post-call batch; `None` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> bool {
) -> Option<PyPostCallSettlement> {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
.map(PyPostCallSettlement)
}
fn set_max<'py>(
@ -246,11 +265,13 @@ impl NativeSpendCounters {
self.0.flush_local().map_err(kind_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, spend) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyNumber(spend))?;
}
Ok(snapshot)
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
Ok(self
.0
.snapshot()
.map_err(kind_error)?
.into_iter()
.map(|(key, spend)| (key, PyNumber(spend)))
.collect())
}
}

46
litellm-rust/crates/python-bridge/src/router/host.rs
View file

@ -357,7 +357,11 @@ impl HostCallResponse {
Self::Select(ids) | Self::Filtered(Ok((ids, _))),
HostCall::Select { candidates, .. },
) = (self, call)
&& let Some(id) = ids.iter().find(|id| !candidates.contains(id))
&& let Some(id) = {
let known: std::collections::HashSet<&str> =
candidates.iter().map(String::as_str).collect();
ids.iter().find(|id| !known.contains(id.as_str()))
}
{
return Err(DriveError::UnknownCandidate { id: id.clone() });
}
@ -688,14 +692,22 @@ fn shares_ids(candidates: &[Arc<Deployment>]) -> bool {
/// The candidates a host answer names, in its order. An id names every candidate that carries
/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them.
fn answered_by_id(candidates: &[Arc<Deployment>], ids: &[String]) -> Vec<Arc<Deployment>> {
let mut taken = vec![false; candidates.len()];
let mut positions: std::collections::HashMap<&str, Vec<usize>> =
std::collections::HashMap::new();
for (position, deployment) in candidates.iter().enumerate() {
positions
.entry(deployment.id.as_str())
.or_default()
.push(position);
}
let mut kept = Vec::new();
for id in ids {
for (position, deployment) in candidates.iter().enumerate() {
if !taken[position] && deployment.id == *id {
taken[position] = true;
kept.push(deployment.clone());
}
if let Some(found) = positions.remove(id.as_str()) {
kept.extend(
found
.into_iter()
.map(|position| candidates[position].clone()),
);
}
}
kept
@ -1540,6 +1552,26 @@ mod tests {
);
}
#[rstest]
fn an_answer_by_id_keeps_its_order_and_every_candidate_sharing_an_id() {
let candidates: Vec<Arc<Deployment>> = ["a", "b", "a", "c"]
.into_iter()
.map(|id| {
let spec: litellm_router_core::DeploymentSpec =
serde_json::from_value(deployment("g", id)).unwrap();
Arc::new(Deployment::from_spec(spec).unwrap())
})
.collect();
let ids = ["c", "a", "x", "a"].map(String::from);
let kept = answered_by_id(&candidates, &ids);
let positions: Vec<usize> = kept
.iter()
.map(|d| candidates.iter().position(|c| Arc::ptr_eq(c, d)).unwrap())
.collect();
assert_eq!(positions, [3, 0, 2]);
}
#[rstest]
fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() {
let mut second = deployment("primary", "a1");

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(

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