scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge origin/main into litellm_rust_qdrant_semantic_cache

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-22 00:19:34 +00:00
commit e3ca80b19a
143 changed files with 14806 additions and 2004 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test-rust.yml vendored
View file

@ -130,7 +130,7 @@ jobs:
- name: Test secret manager feature combinations
run: |
cargo test -p litellm-auth-gcp --locked --no-default-features
for features in '' aws google azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark; do
for features in '' aws google hashicorp azure cyberark aws,google aws,azure google,azure aws,google,azure aws,google,cyberark aws,google,azure,cyberark aws,google,hashicorp,azure,cyberark; do
cargo test -p litellm-secrets --locked --no-default-features --features "$features"
done

42
litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql Normal file
View file

@ -0,0 +1,42 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterUserSession" (
"user_id" TEXT NOT NULL,
"api_key" TEXT NOT NULL,
"session_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"router_type" TEXT NOT NULL,
"first_turn_at" TIMESTAMP(3) NOT NULL,
"last_turn_at" TIMESTAMP(3) NOT NULL,
"last_model" TEXT NOT NULL,
"models" JSONB NOT NULL DEFAULT '{}',
"turns" INTEGER NOT NULL DEFAULT 0,
"unordered_turns" INTEGER NOT NULL DEFAULT 0,
"covered_turns" INTEGER NOT NULL DEFAULT 0,
"cache_hits" INTEGER NOT NULL DEFAULT 0,
"same_model_turns" INTEGER NOT NULL DEFAULT 0,
"same_model_hits" INTEGER NOT NULL DEFAULT 0,
"first_visit_turns" INTEGER NOT NULL DEFAULT 0,
"first_visit_hits" INTEGER NOT NULL DEFAULT 0,
"return_turns" INTEGER NOT NULL DEFAULT 0,
"return_hits" INTEGER NOT NULL DEFAULT 0,
"return_expired_misses" INTEGER NOT NULL DEFAULT 0,
"return_within_ttl_misses" INTEGER NOT NULL DEFAULT 0,
"ttl_5m_turns" INTEGER NOT NULL DEFAULT 0,
"ttl_1h_turns" INTEGER NOT NULL DEFAULT 0,
"total_tokens" BIGINT NOT NULL DEFAULT 0,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
"savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}',
"classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
"classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0,
"tier_turns" JSONB NOT NULL DEFAULT '{}',
"baseline_models" JSONB NOT NULL DEFAULT '{}',
CONSTRAINT "LiteLLM_AutoRouterUserSession_pkey" PRIMARY KEY ("user_id", "api_key", "session_id", "router_name")
);
CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_last_turn" ON "LiteLLM_AutoRouterUserSession"("last_turn_at");
CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_user_last_turn" ON "LiteLLM_AutoRouterUserSession"("user_id", "last_turn_at");

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

@ -1623,6 +1623,47 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
model LiteLLM_AutoRouterUserSession {
user_id String
api_key String
session_id String
router_name String
router_type String
first_turn_at DateTime
last_turn_at DateTime
last_model String
models Json @default("{}")
turns Int @default(0)
unordered_turns Int @default(0)
covered_turns Int @default(0)
cache_hits Int @default(0)
same_model_turns Int @default(0)
same_model_hits Int @default(0)
first_visit_turns Int @default(0)
first_visit_hits Int @default(0)
return_turns Int @default(0)
return_hits Int @default(0)
return_expired_misses Int @default(0)
return_within_ttl_misses Int @default(0)
ttl_5m_turns Int @default(0)
ttl_1h_turns Int @default(0)
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
baseline_models Json @default("{}")
@@id([user_id, api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_user_session_last_turn")
@@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn")
}
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
// either direction. forward duplicates the requests the keys did not route through the
// router through it, answering whether they should adopt it; reverse duplicates the

299
litellm-rust/Cargo.lock generated
View file

@ -40,6 +40,12 @@ dependencies = [
"cc",
]
[[package]]
name = "allocator-api2"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
[[package]]
name = "android_system_properties"
version = "0.1.6"
@ -230,6 +236,7 @@ dependencies = [
"aws-credential-types",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
@ -238,7 +245,9 @@ dependencies = [
"bytes",
"bytes-utils",
"fastrand",
"http 0.2.12",
"http 1.4.2",
"http-body 0.4.6",
"http-body 1.1.0",
"percent-encoding",
"pin-project-lite",
@ -272,6 +281,43 @@ dependencies = [
"tracing",
]
[[package]]
name = "aws-sdk-s3"
version = "1.146.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2cd651b4400d4011b8927b83a9552bf90ff11e6e5da0b9f0a7583247aceec971"
dependencies = [
"arc-swap",
"aws-credential-types",
"aws-runtime",
"aws-sigv4",
"aws-smithy-async",
"aws-smithy-checksums",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-json",
"aws-smithy-observability",
"aws-smithy-runtime",
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml 0.62.1",
"aws-types",
"bytes",
"fastrand",
"hex",
"hmac",
"http 0.2.12",
"http 1.4.2",
"http-body 1.1.0",
"lru",
"percent-encoding",
"regex-lite",
"sha2 0.11.0",
"tracing",
"url",
]
[[package]]
name = "aws-sdk-secretsmanager"
version = "1.117.0"
@ -316,7 +362,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"aws-smithy-xml",
"aws-smithy-xml 0.61.1",
"aws-types",
"fastrand",
"http 0.2.12",
@ -332,6 +378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70"
dependencies = [
"aws-credential-types",
"aws-smithy-eventstream",
"aws-smithy-http",
"aws-smithy-runtime-api",
"aws-smithy-types",
@ -359,10 +406,31 @@ dependencies = [
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
name = "aws-smithy-checksums"
version = "0.65.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
checksum = "b67ecd999972b58e67cab052f5129906c08c25883bd0788ceefc55ef97d61307"
dependencies = [
"aws-smithy-http",
"aws-smithy-types",
"bytes",
"crc-fast",
"hex",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"md-5",
"pin-project-lite",
"sha1 0.11.0",
"sha2 0.11.0",
"tracing",
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "80c2051c2f1016fb8e6548dd07b8bc2ac9c3fe583721444b92f515e856d31609"
dependencies = [
"aws-smithy-types",
"bytes",
@ -375,6 +443,7 @@ version = "0.64.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-runtime-api",
"aws-smithy-types",
"bytes",
@ -554,6 +623,18 @@ dependencies = [
"xmlparser",
]
[[package]]
name = "aws-smithy-xml"
version = "0.62.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b932c8d6dc127fc980eecd78f8694ae9b9551b69a93a7def2a199c1c0033daf"
dependencies = [
"aws-smithy-runtime-api",
"aws-smithy-schema",
"aws-smithy-types",
"xmlparser",
]
[[package]]
name = "aws-types"
version = "1.6.0"
@ -1023,6 +1104,16 @@ dependencies = [
"libc",
]
[[package]]
name = "crc-fast"
version = "1.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5"
dependencies = [
"digest 0.10.7",
"spin",
]
[[package]]
name = "crc16"
version = "0.4.0"
@ -1968,6 +2059,8 @@ version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
dependencies = [
"allocator-api2",
"equivalent",
"foldhash",
]
@ -2711,6 +2804,21 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-redis-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"r2d2",
"redis",
"redis-test",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-cache-response"
version = "0.1.0"
@ -2727,6 +2835,37 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-s3"
version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-types",
"aws-types",
"litellm-auth-aws",
"litellm-cache",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
"uuid",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2889,6 +3028,7 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
@ -2897,7 +3037,10 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-qdrant-semantic",
"litellm-cache-redis",
"litellm-cache-redis-semantic",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-cache-valkey-semantic",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2909,11 +3052,13 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"qdrant-client",
"redis",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"tokio",
"tokio-tungstenite",
"url",
@ -2933,6 +3078,7 @@ dependencies = [
"litellm-secrets-azure",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-hashicorp",
"litellm-secrets-types",
"moka",
"reqwest 0.12.28",
@ -3030,6 +3176,26 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-secrets-hashicorp"
version = "0.1.0"
dependencies = [
"litellm-core-utils",
"litellm-secrets-types",
"moka",
"rstest",
"rustify",
"rustify_derive",
"serde",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
"vaultrs",
"veil",
"wiremock",
]
[[package]]
name = "litellm-secrets-types"
version = "0.1.0"
@ -3121,6 +3287,15 @@ version = "0.4.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
[[package]]
name = "lru"
version = "0.18.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "lru-slab"
version = "0.1.2"
@ -3149,6 +3324,16 @@ version = "0.8.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
[[package]]
name = "md-5"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98"
dependencies = [
"cfg-if",
"digest 0.11.3",
]
[[package]]
name = "memchr"
version = "2.8.3"
@ -4295,6 +4480,40 @@ dependencies = [
"semver",
]
[[package]]
name = "rustify"
version = "0.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4800ce4c1cc2fec12c559dae2ddbf0e17fcee7569b796e6d75898efef443368b"
dependencies = [
"anyhow",
"async-trait",
"bytes",
"http 1.4.2",
"reqwest 0.13.5",
"rustify_derive",
"serde",
"serde_json",
"serde_urlencoded",
"thiserror 1.0.69",
"tracing",
"url",
]
[[package]]
name = "rustify_derive"
version = "0.5.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "78ea7fda74240f7410d0198b603a8a2f662acc7d76b6667a49f9b162cd8d9b4f"
dependencies = [
"proc-macro2",
"quote",
"regex",
"serde_urlencoded",
"syn 1.0.109",
"synstructure 0.12.6",
]
[[package]]
name = "rustix"
version = "1.1.5"
@ -4649,6 +4868,17 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1"
version = "0.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214"
dependencies = [
"cfg-if",
"cpufeatures 0.3.0",
"digest 0.11.3",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
@ -4765,6 +4995,12 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "spin"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3"
[[package]]
name = "spm_precompiled"
version = "0.1.4"
@ -4847,6 +5083,17 @@ version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
[[package]]
name = "syn"
version = "1.0.109"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "syn"
version = "2.0.119"
@ -4878,6 +5125,18 @@ dependencies = [
"futures-core",
]
[[package]]
name = "synstructure"
version = "0.12.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f"
dependencies = [
"proc-macro2",
"quote",
"syn 1.0.109",
"unicode-xid",
]
[[package]]
name = "synstructure"
version = "0.13.2"
@ -5297,6 +5556,7 @@ version = "0.1.44"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
dependencies = [
"log",
"pin-project-lite",
"tracing-attributes",
"tracing-core",
@ -5379,7 +5639,7 @@ dependencies = [
"rand 0.8.7",
"rustls 0.23.42",
"rustls-pki-types",
"sha1",
"sha1 0.10.7",
"thiserror 1.0.69",
"utf-8",
]
@ -5481,6 +5741,12 @@ version = "1.13.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
[[package]]
name = "unicode-xid"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
[[package]]
name = "unicode_categories"
version = "0.1.1"
@ -5540,6 +5806,25 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
[[package]]
name = "vaultrs"
version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30ffcc0e81025065dda612ec1e26a3d81bb16ef3062354873d17a35965d68522"
dependencies = [
"async-trait",
"derive_builder",
"http 1.4.2",
"reqwest 0.13.5",
"rustify",
"rustify_derive",
"serde",
"serde_json",
"thiserror 2.0.19",
"tracing",
"url",
]
[[package]]
name = "vcpkg"
version = "0.2.15"
@ -5989,7 +6274,7 @@ dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure",
"synstructure 0.13.2",
]
[[package]]
@ -6030,7 +6315,7 @@ dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure",
"synstructure 0.13.2",
]
[[package]]

6
litellm-rust/Cargo.toml
View file

@ -22,6 +22,7 @@ litellm-secrets = { path = "crates/secrets" }
litellm-secrets-types = { path = "crates/secrets-types" }
litellm-secrets-aws = { path = "crates/secrets-aws" }
litellm-secrets-google = { path = "crates/secrets-google" }
litellm-secrets-hashicorp = { path = "crates/secrets-hashicorp" }
litellm-secrets-azure = { path = "crates/secrets-azure" }
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
litellm-http = { path = "crates/http" }
@ -32,8 +33,10 @@ litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
litellm-token-counter = { path = "crates/token-counter" }
@ -58,6 +61,9 @@ uuid = { version = "1", features = ["v4"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustify = "=0.7.0"
rustify_derive = "=0.5.5"
vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }

7
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs
View file

@ -101,10 +101,13 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
}
fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
if context.messages.is_empty() {
let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else {
return Err(Error::MissingPrompt);
};
if messages.is_empty() {
return Err(Error::MissingPrompt);
}
Ok(prompt_from_messages(&context.messages))
Ok(prompt_from_messages(messages))
}
async fn set(

7
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
View file

@ -3,9 +3,7 @@ mod support;
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext, SemanticCacheScope,
};
use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
use litellm_cache_qdrant_semantic::{
Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
};
@ -60,8 +58,7 @@ fn config(quantization: Quantization) -> QdrantSemanticConfig {
fn context(prompt: &str) -> SemanticCacheContext {
SemanticCacheContext {
messages: vec![json!({"role": "user", "content": prompt})],
scope: SemanticCacheScope::default(),
messages: Some(json!([{"role": "user", "content": prompt}])),
..Default::default()
}
}

21
litellm-rust/crates/cache-redis-semantic/Cargo.toml Normal file
View file

@ -0,0 +1,21 @@
[package]
name = "litellm-cache-redis-semantic"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
r2d2 = "0.8.10"
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"
serde_json.workspace = true
tokio.workspace = true

618
litellm-rust/crates/cache-redis-semantic/src/cache.rs Normal file
View file

@ -0,0 +1,618 @@
use std::{
future::Future,
sync::{Arc, OnceLock},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
};
use litellm_cache_redis::{
RedisTopology,
connection::{ConnectionRef, Connections},
};
use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::prompt::prompt_from_context;
const CACHE_KEY_FIELD: &str = "litellm_cache_key";
const VECTOR_FIELD: &str = "prompt_vector";
pub trait Embedder: Send + Sync + 'static {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error>;
fn async_embed(
&self,
prompt: &str,
metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
#[derive(Clone, Debug)]
pub struct RedisSemanticConfig {
pub index_name: String,
pub similarity_threshold: f32,
}
struct Inner {
index_name: String,
distance_threshold: f64,
resolved_index: OnceLock<String>,
codec: ResponseCacheCodec,
clock: fn() -> f64,
}
impl Inner {
fn new(config: RedisSemanticConfig) -> Self {
Self {
index_name: config.index_name,
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
resolved_index: OnceLock::new(),
codec: ResponseCacheCodec,
clock: timestamp,
}
}
fn ensure_index(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
if let Some(name) = self.resolved_index.get() {
return Ok(name.clone());
}
let name = match index_compatible(connection, &self.index_name, dims)? {
Some(true) => self.index_name.clone(),
Some(false) => self.isolated_index(connection, dims)?,
None => match create_index(connection, &self.index_name, dims) {
Ok(()) => self.index_name.clone(),
Err(_) => match index_compatible(connection, &self.index_name, dims)? {
Some(true) => self.index_name.clone(),
Some(false) => self.isolated_index(connection, dims)?,
None => return Err(Error::Unavailable),
},
},
};
let _ = self.resolved_index.set(name.clone());
Ok(name)
}
fn isolated_index(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
let name = format!("{}_isolated", self.index_name);
match index_compatible(connection, &name, dims)? {
Some(true) => Ok(name),
Some(false) => {
redis::cmd("FT.DROPINDEX")
.arg(&name)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
create_index(connection, &name, dims)?;
Ok(name)
}
None => {
create_index(connection, &name, dims)?;
Ok(name)
}
}
}
fn store(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
value: &CacheEntry,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> Result<(), Error> {
let index = self.ensure_index(connection, vector.len())?;
let entry_id = entry_id(prompt, tag);
let hash_key = format!("{index}:{entry_id}");
let response = self.codec.encode(value)?;
redis::cmd("HSET")
.arg(&hash_key)
.arg("entry_id")
.arg(&entry_id)
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(response)
.arg(VECTOR_FIELD)
.arg(vector_buffer(vector))
.arg("inserted_at")
.arg(format!("{}", (self.clock)()))
.arg("updated_at")
.arg(format!("{}", (self.clock)()))
.arg(CACHE_KEY_FIELD)
.arg(tag)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
if let Some(ttl) = ttl {
redis::cmd("EXPIRE")
.arg(&hash_key)
.arg(ttl_seconds(ttl))
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
}
fn lookup(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
vector: &[f32],
) -> Result<Option<CacheEntry>, Error> {
let index = self.ensure_index(connection, vector.len())?;
let query = format!(
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
escape_tag(tag)
);
let result = redis::cmd("FT.SEARCH")
.arg(&index)
.arg(query)
.arg("RETURN")
.arg(8)
.arg("entry_id")
.arg("prompt")
.arg("response")
.arg("inserted_at")
.arg("updated_at")
.arg("metadata")
.arg(CACHE_KEY_FIELD)
.arg("vector_distance")
.arg("SORTBY")
.arg("vector_distance")
.arg("ASC")
.arg("DIALECT")
.arg(2)
.arg("LIMIT")
.arg(0)
.arg(1)
.arg("PARAMS")
.arg(2)
.arg("vector")
.arg(vector_buffer(vector))
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(None);
};
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
return Ok(None);
}
if number_field(fields, "vector_distance")
.is_none_or(|distance| distance > self.distance_threshold)
{
return Ok(None);
}
let Some(response) = bytes_field(fields, "response") else {
return Ok(None);
};
self.codec.decode(&response).map(Some)
}
}
pub struct RedisSemanticCache<E: Embedder, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
inner: Arc<Inner>,
}
impl<E: Embedder> RedisSemanticCache<E> {
pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
inner: Arc::new(Inner::new(config)),
})
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
inner: Arc::new(Inner::new(config)),
}
}
pub fn with_clock(self, clock: fn() -> f64) -> Self {
Self {
inner: Arc::new(Inner {
index_name: self.inner.index_name.clone(),
distance_threshold: self.inner.distance_threshold,
resolved_index: OnceLock::new(),
codec: self.inner.codec,
clock,
}),
..self
}
}
pub fn embedder(&self) -> &E {
&self.embedder
}
pub fn index_name(&self) -> &str {
&self.inner.index_name
}
pub fn similarity_threshold(&self) -> f32 {
(1.0 - self.inner.distance_threshold) as f32
}
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
context.scope.as_deref().unwrap_or(key)
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
for RedisSemanticCache<E, C>
{
type Value = CacheEntry;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(());
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections.execute(|connection| {
self.inner
.store(connection, &tag, &value, &prompt, &vector, context.ttl)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections
.execute(|connection| self.inner.lookup(connection, &tag, &vector))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, &context).to_string();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.store(connection, &tag, &value, &prompt, &vector, context.ttl)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, context).to_string();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.lookup(connection, &tag, &vector)
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(connection) {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
},
})
})
.await
{
Ok(result) => Ok(result),
Err(error) => Ok(CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
}),
}
}
}
fn timestamp() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or_default()
}
fn entry_id(prompt: &str, tag: &str) -> String {
let mut digest = Sha256::new();
digest.update(prompt.as_bytes());
digest.update(CACHE_KEY_FIELD.as_bytes());
digest.update(tag.as_bytes());
format!("{:x}", digest.finalize())
}
fn vector_buffer(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|component| component.to_le_bytes())
.collect()
}
fn escape_tag(value: &str) -> String {
value
.chars()
.flat_map(|ch| {
if matches!(
ch,
',' | '.'
| '<'
| '>'
| '{'
| '}'
| '['
| ']'
| '\\'
| '"'
| '\''
| ':'
| ';'
| '!'
| '@'
| '#'
| '$'
| '%'
| '^'
| '&'
| '*'
| '('
| ')'
| '-'
| '+'
| '='
| '~'
| '|'
| '/'
| ' '
| '?'
) {
vec!['\\', ch]
} else {
vec![ch]
}
})
.collect()
}
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
redis::cmd("FT.CREATE")
.arg(name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(name)
.arg("SCORE")
.arg(1.0)
.arg("SCHEMA")
.arg("prompt")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("response")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("inserted_at")
.arg("NUMERIC")
.arg("updated_at")
.arg("NUMERIC")
.arg(VECTOR_FIELD)
.arg("VECTOR")
.arg("FLAT")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dims)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.arg(CACHE_KEY_FIELD)
.arg("TAG")
.arg("SEPARATOR")
.arg(",")
.query::<()>(connection)
.map_err(|_| Error::Unavailable)
}
fn index_compatible(
connection: &mut ConnectionRef<'_>,
name: &str,
dims: usize,
) -> Result<Option<bool>, Error> {
let info = match redis::cmd("FT.INFO")
.arg(name)
.query::<redis::Value>(connection)
{
Ok(info) => info,
Err(error) if unknown_index(&error) => return Ok(None),
Err(_) => return Err(Error::Unavailable),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
let message = error.to_string().to_lowercase();
message.contains("unknown") && message.contains("index")
}
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
let redis::Value::Array(entries) = info else {
return false;
};
let attributes = entries
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
.map(|pair| &pair[1]);
let Some(redis::Value::Array(attributes)) = attributes else {
return false;
};
let fields = attributes
.iter()
.map(|attribute| {
let redis::Value::Array(attribute) = attribute else {
return (None, None, None, None, None);
};
let mut name = None;
let mut field_type = None;
let mut dim = None;
let mut data_type = None;
let mut distance_metric = None;
for pair in attribute.as_chunks::<2>().0 {
match string_value(&pair[0]).as_deref() {
Some("identifier") => name = string_value(&pair[1]),
Some("type") => field_type = string_value(&pair[1]),
Some("dim") => dim = number_value(&pair[1]),
Some("data_type") => data_type = string_value(&pair[1]),
Some("distance_metric") => distance_metric = string_value(&pair[1]),
_ => {}
}
}
(name, field_type, dim, data_type, distance_metric)
})
.collect::<Vec<_>>();
let has_field = |name: &str, field_type: &str| {
fields
.iter()
.any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type))
};
has_field("prompt", "TEXT")
&& has_field("response", "TEXT")
&& has_field("inserted_at", "NUMERIC")
&& has_field("updated_at", "NUMERIC")
&& has_field(CACHE_KEY_FIELD, "TAG")
&& fields.iter().any(|(n, t, d, data, metric)| {
n.as_deref() == Some(VECTOR_FIELD)
&& t.as_deref() == Some("VECTOR")
&& *d == Some(dims as f64)
&& data
.as_deref()
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
&& metric
.as_deref()
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
})
}
fn string_value(value: &redis::Value) -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
redis::Value::SimpleString(text) => Some(text.clone()),
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
_ => None,
}
}
fn number_value(value: &redis::Value) -> Option<f64> {
match value {
redis::Value::Int(number) => Some(*number as f64),
redis::Value::Double(number) => Some(*number),
_ => string_value(value).and_then(|text| text.parse().ok()),
}
}
fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
let redis::Value::Array(items) = result else {
return None;
};
let [count, _document_id, fields, ..] = items.as_slice() else {
return None;
};
if !matches!(count, redis::Value::Int(count) if *count > 0) {
return None;
}
match fields {
redis::Value::Array(fields) => Some(fields.as_slice()),
_ => None,
}
}
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
fields
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
.map(|pair| &pair[1])
}
fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
field_value(fields, name).and_then(string_value)
}
fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
match field_value(fields, name)? {
redis::Value::BulkString(bytes) => Some(bytes.clone()),
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
_ => None,
}
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}

5
litellm-rust/crates/cache-redis-semantic/src/lib.rs Normal file
View file

@ -0,0 +1,5 @@
mod cache;
mod prompt;
pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig};
pub use prompt::prompt_from_context;

97
litellm-rust/crates/cache-redis-semantic/src/prompt.rs Normal file
View file

@ -0,0 +1,97 @@
use litellm_cache::SemanticCacheContext;
use serde_json::Value;
pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
if let Some(messages) = context.messages.as_ref().and_then(Value::as_array)
&& !messages.is_empty()
{
return Some(messages_text(messages));
}
let input = context.input.as_ref()?;
let mut parts = Vec::new();
collect_input_text(input, &mut parts);
let prompt = parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some(prompt)
}
fn messages_text(messages: &[Value]) -> String {
let mut text = String::new();
for message in messages {
let Some(message) = message.as_object() else {
continue;
};
match message.get("content") {
Some(Value::String(content)) => text.push_str(content),
Some(Value::Array(parts)) => {
for part in parts {
if let Some(text_content) = part.get("text").and_then(Value::as_str) {
text.push_str(text_content);
}
}
}
_ => {}
}
text.push_str(&search_results_text(message.get("search_results")));
}
text
}
fn search_results_text(search_results: Option<&Value>) -> String {
let Some(Value::Array(results)) = search_results else {
return String::new();
};
let mut text = String::new();
for result in results {
let Some(result) = result.as_object() else {
continue;
};
for key in ["source", "title"] {
if let Some(value) = result.get(key).and_then(Value::as_str) {
text.push_str(value);
}
}
if let Some(Value::Array(content)) = result.get("content") {
for block in content {
if let Some(value) = block.get("text").and_then(Value::as_str) {
text.push_str(value);
}
}
}
if let Some(citations) = result.get("citations") {
text.push_str(&citations.to_string());
}
}
text
}
fn collect_input_text(value: &Value, parts: &mut Vec<String>) {
match value {
Value::String(text) => {
let trimmed = text.trim();
if !trimmed.is_empty() {
parts.push(trimmed.to_string());
}
}
Value::Array(items) => {
for item in items {
collect_input_text(item, parts);
}
}
Value::Object(map) => {
if let Some(content) = map.get("content").filter(|content| !content.is_null()) {
collect_input_text(content, parts);
return;
}
for key in ["text", "output", "input_text", "output_text"] {
if let Some(Value::String(text)) = map.get(key) {
let trimmed = text.trim();
if !trimmed.is_empty() {
parts.push(trimmed.to_string());
return;
}
}
}
}
_ => {}
}
}

1003
litellm-rust/crates/cache-redis-semantic/tests/cache.rs Normal file
View file

File diff suppressed because it is too large Load diff

67
litellm-rust/crates/cache-redis/src/cache.rs
View file

@ -14,7 +14,7 @@ use crate::topology::RedisTopology;
mod connection;
mod operations;
pub(crate) use connection::ConnectionRef;
pub use connection::ConnectionRef;
use connection::{ClusterConnectionManager, ConnectionManager};
pub use operations::{
@ -40,7 +40,8 @@ const CLAIM_SCRIPT: &str = concat!(
);
const CLAIM_ATTEMPTS: usize = 8;
enum Connections<C> {
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
@ -50,7 +51,7 @@ impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn execute<T>(
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
@ -73,6 +74,29 @@ where
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
match topology {
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
ClusterConnectionManager::open(url, startup_nodes)?,
)?)),
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
@ -94,12 +118,7 @@ impl<S: CacheCodec> RedisCache<S> {
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = match topology {
RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?),
RedisTopology::Cluster { startup_nodes } => {
Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?)
}
};
let connections = Connections::open(url, topology)?;
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
@ -127,7 +146,7 @@ where
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::Fixed(Mutex::new(connection))),
connections: Arc::new(Connections::fixed(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
@ -203,16 +222,6 @@ where
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
@ -271,7 +280,7 @@ where
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
@ -285,7 +294,7 @@ where
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
@ -311,7 +320,7 @@ where
if entries.is_empty() {
return Ok(());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
@ -330,7 +339,7 @@ where
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match connection.ping() {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
@ -391,7 +400,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -418,7 +427,7 @@ where
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
@ -438,7 +447,7 @@ where
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
@ -470,7 +479,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
@ -581,7 +590,7 @@ where
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await

8
litellm-rust/crates/cache-redis/src/cache/connection.rs vendored
View file

@ -12,7 +12,7 @@ use redis::{
use super::REDIS_TIMEOUT;
use crate::topology::RedisNode;
pub(super) struct PooledConnection<C> {
pub struct PooledConnection<C> {
pub(super) connection: C,
pub(super) failed: bool,
}
@ -20,7 +20,7 @@ pub(super) struct PooledConnection<C> {
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub(super) struct ConnectionManager(redis::Client);
pub struct ConnectionManager(redis::Client);
impl ConnectionManager {
pub(super) fn open(url: &str) -> Result<Self, Error> {
@ -54,7 +54,7 @@ impl r2d2::ManageConnection for ConnectionManager {
}
}
pub(super) struct ClusterConnectionManager(ClusterClient);
pub struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
@ -117,7 +117,7 @@ impl r2d2::ManageConnection for ClusterConnectionManager {
}
}
pub(crate) enum ConnectionRef<'a> {
pub enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}

28
litellm-rust/crates/cache-redis/src/cache/operations.rs vendored
View file

@ -144,7 +144,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del(keys).map_err(|_| Error::Unavailable)
})
.await
@ -172,7 +172,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -188,7 +188,7 @@ where
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
Connections::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().map_err(|_| Error::Unavailable)
})
.await
@ -196,7 +196,7 @@ where
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
let key = self.namespaced_key(key);
let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
@ -208,7 +208,7 @@ where
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
@ -231,7 +231,7 @@ where
}
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut sadd = redis::cmd("SADD");
sadd.arg(&key).arg(values);
let mut expire = redis::cmd("EXPIRE");
@ -253,7 +253,7 @@ where
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
@ -279,7 +279,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, values)| {
@ -304,7 +304,7 @@ where
) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let multiple = count.is_some();
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
@ -333,7 +333,7 @@ where
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, count)| {
@ -365,7 +365,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
@ -426,7 +426,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut commands = Vec::with_capacity(operations.len() * 2);
let mut increments = Vec::with_capacity(operations.len());
for (key, amount, ttl) in operations {
@ -460,7 +460,7 @@ where
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment_with_floor(connection, key, amount, ttl)
})
.await
@ -474,7 +474,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)

4
litellm-rust/crates/cache-redis/src/lib.rs
View file

@ -1,6 +1,10 @@
mod cache;
mod topology;
pub mod connection {
pub use crate::cache::{ConnectionRef, Connections};
}
pub use cache::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};

2
litellm-rust/crates/cache-response/README.md
View file

@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees

15
litellm-rust/crates/cache-response/src/response.rs
View file

@ -1,15 +1,14 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error,
ExactCacheContext, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
#[derive(Clone)]
pub struct ResponseCacheRequest<C: CacheContext = ExactCacheContext> {
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: C,
@ -64,6 +63,10 @@ where
&self.backend
}
pub fn backend_arc(&self) -> &Arc<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&B::Context::default())
}
@ -267,9 +270,9 @@ where
Ok(())
}
fn partial_hits<C: CacheContext>(
requests: &[ResponseCacheRequest<C>],
readable: Vec<(usize, &ResponseCacheRequest<C>)>,
fn partial_hits(
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {

2
litellm-rust/crates/cache-response/tests/response.rs
View file

@ -89,7 +89,7 @@ fn semantic_context_reaches_backend_for_store_and_lookup() {
});
let cache = ResponseCache::new(backend.clone());
let context = SemanticCacheContext {
messages: vec![json!({"role": "user", "content": "hello"})],
messages: Some(json!([{"role": "user", "content": "hello"}])),
..Default::default()
};
let request = request().with_context(context.clone());

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

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-s3"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
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-types = "1.6.0"
aws-types = "1.6.0"
tokio.workspace = true
[dev-dependencies]
wiremock = "0.6.5"
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

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

@ -0,0 +1,101 @@
use aws_credential_types::{
Credentials as AwsCredentials,
provider::{ProvideCredentials, error::CredentialsError, future},
};
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
#[derive(Clone)]
pub(crate) struct Credentials {
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl Credentials {
pub(crate) fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self { config, env }
}
}
impl ProvideCredentials for Credentials {
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,
secret_access_key,
self.config.session_token.clone(),
None,
"litellm-s3-cache",
));
}
resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
})
}
}
impl std::fmt::Debug for Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Credentials").finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn explicit_keys_ignore_an_ambient_session_token() {
let provider = Credentials::with_env(
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
);
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "key");
assert_eq!(credentials.secret_access_key(), "secret");
assert_eq!(credentials.session_token(), None);
}
#[tokio::test]
async fn explicit_keys_keep_their_session_token() {
let provider = Credentials::new(AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: Some("t".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.session_token(), Some("t"));
}
#[tokio::test]
async fn environment_keys_resolve_with_their_session_token() {
let provider = Credentials::with_env(AwsAuthConfig::default(), |name| 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,
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "env-key");
assert_eq!(credentials.secret_access_key(), "env-secret");
assert_eq!(credentials.session_token(), Some("env-token"));
}
}

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

@ -0,0 +1,220 @@
use std::{
future::Future,
sync::Arc,
time::{Duration, SystemTime},
};
use aws_sdk_s3::{
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
error::SdkError,
primitives::ByteStream,
};
use aws_smithy_types::{DateTime, date_time::Format};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::auth::Credentials;
pub struct S3Endpoint {
pub url: String,
}
pub struct S3CacheConfig {
pub bucket: String,
pub key_prefix: String,
pub region: String,
pub endpoint: Option<S3Endpoint>,
pub auth: AwsAuthConfig,
}
pub struct S3Cache<C: CacheCodec> {
client: aws_sdk_s3::Client,
codec: C,
runtime: Handle,
bucket: Arc<str>,
key_prefix: Arc<str>,
region: Arc<str>,
endpoint: Option<Arc<str>>,
}
impl<C: CacheCodec> S3Cache<C> {
pub fn new(config: S3CacheConfig, codec: 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()))
.credentials_provider(Credentials::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,
};
Self {
client: aws_sdk_s3::Client::from_conf(builder.build()),
codec,
runtime,
bucket: config.bucket.into(),
key_prefix: config.key_prefix.into(),
region: config.region.into(),
endpoint: endpoint_url.map(Into::into),
}
}
pub fn bucket(&self) -> &str {
&self.bucket
}
pub fn key_prefix(&self) -> &str {
&self.key_prefix
}
pub fn region(&self) -> &str {
&self.region
}
pub fn endpoint(&self) -> Option<&str> {
self.endpoint.as_deref()
}
pub fn to_s3_key(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key.replace(':', "/"))
}
fn block_on<F: Future>(&self, future: F) -> F::Output {
if Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
async fn put(
&self,
key: &str,
value: C::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let s3_key = self.to_s3_key(key);
let body = self.codec.encode(&value)?;
let request = self
.client
.put_object()
.bucket(self.bucket.as_ref())
.key(&s3_key)
.body(ByteStream::from(body))
.content_type("application/json")
.content_language("en")
.content_disposition(format!("inline; filename=\"{s3_key}.json\""));
let request = match context.ttl {
Some(ttl) => {
let seconds = ttl.as_secs_f64();
request
.cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}"))
.expires(DateTime::from(SystemTime::now() + ttl))
}
None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"),
};
request.send().await.map_err(|_| Error::Unavailable)?;
Ok(())
}
async fn get(&self, key: &str) -> Result<Option<C::Value>, Error> {
let output = match self
.client
.get_object()
.bucket(self.bucket.as_ref())
.key(self.to_s3_key(key))
.send()
.await
{
Ok(output) => output,
Err(error) => {
if let SdkError::ServiceError(service) = &error {
let status = error
.raw_response()
.map(|response| response.status().as_u16());
let not_found = service.err().is_no_such_key()
|| service.err().meta().code() == Some("AccessDenied")
|| status == Some(404)
|| status == Some(403);
if not_found {
return Ok(None);
}
}
return Err(Error::Unavailable);
}
};
if let Some(expires) = output.expires_string()
&& let Ok(expires) = DateTime::from_str(expires, Format::HttpDate)
&& expires < DateTime::from(SystemTime::now())
{
return Ok(None);
}
let bytes = output
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?
.into_bytes();
self.codec.decode(&bytes).map(Some)
}
}
impl<C: CacheCodec> BaseCache for S3Cache<C> {
type Value = C::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
self.block_on(self.put(key, value, context))
}
fn get_cache(&self, key: &str, _context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.block_on(self.get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
self.put(key, value, &context).await
}
async fn async_get_cache(
&self,
key: &str,
_context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key).await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<C: CacheCodec> BatchCache for S3Cache<C> {}
impl<C: CacheCodec> FlushCache for S3Cache<C> {
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
}

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

@ -0,0 +1,4 @@
mod auth;
mod cache;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

278
litellm-rust/crates/cache-s3/tests/cache.rs Normal file
View file

@ -0,0 +1,278 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
fn config(endpoint: String) -> S3CacheConfig {
S3CacheConfig {
bucket: "cache-bucket".to_string(),
key_prefix: "team/".to_string(),
region: "us-east-1".to_string(),
endpoint: Some(S3Endpoint { url: endpoint }),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
}
}
fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
S3Cache::new(
config(endpoint.to_string()),
JsonCodec::<Value>::new(),
Handle::current(),
)
}
async fn mock_server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
.mount(&server)
.await;
server
}
fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
use aws_smithy_types::{DateTime, date_time::Format};
headers
.get(name)
.and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
.map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn set_writes_python_metadata_with_and_without_ttl() {
let server = mock_server().await;
let cache = cache(&server.uri());
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(90)),
};
cache
.set_cache("alpha:beta", json!({"answer": 1}), &context)
.unwrap();
cache
.set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
.unwrap();
let requests = server.received_requests().await.unwrap();
let ttl_request = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/alpha/beta")
.expect("ttl write should hit the converted S3 key");
assert_eq!(
ttl_request.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=90, s-maxage=90"
);
assert_eq!(
ttl_request.headers["content-type"].to_str().unwrap(),
"application/json"
);
assert_eq!(
ttl_request.headers["content-language"].to_str().unwrap(),
"en"
);
assert_eq!(
ttl_request.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/alpha/beta.json\""
);
let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires");
let remaining = expires.duration_since(SystemTime::now()).unwrap();
assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91));
assert_eq!(
serde_json::from_slice::<Value>(&ttl_request.body).unwrap(),
json!({"answer": 1})
);
let plain = requests
.iter()
.find(|request| request.url.path() == "/cache-bucket/team/plain")
.expect("no-ttl write should hit the converted S3 key");
assert_eq!(
plain.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=31536000, s-maxage=31536000"
);
assert!(plain.headers.get("expires").is_none());
assert_eq!(
plain.headers["content-disposition"].to_str().unwrap(),
"inline; filename=\"team/plain.json\""
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_hit_miss_expired_and_invalid_entries() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/missing"))
.respond_with(
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/denied"))
.respond_with(
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/expired"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
.set_body_json(json!({"answer": 4})),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/malformed"))
.respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
assert_eq!(
cache.get_cache("hit", &context).unwrap(),
Some(json!({"answer": 3}))
);
assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
assert_eq!(
cache.get_cache("malformed", &context),
Err(Error::InvalidEntry)
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
let server = mock_server().await;
for (key, status, body) in [
("first", 200, "{\"answer\": 1}"),
("invalid", 200, "garbage"),
] {
Mock::given(method("GET"))
.and(path(format!("/cache-bucket/team/{key}")))
.respond_with(ResponseTemplate::new(status).set_body_string(body))
.mount(&server)
.await;
}
Mock::given(method("GET"))
.and(path("/cache-bucket/team/miss"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
let keys = vec![
"first".to_string(),
"miss".to_string(),
"invalid".to_string(),
];
let entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"answer": 1})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unsupported_and_noop_capabilities_match_python() {
let server = mock_server().await;
let cache = cache(&server.uri());
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
cache.flush_cache().unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(45)),
}),
Some(Duration::from_secs(45))
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn key_conversion_prefixes_and_splits_colons() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap();
let _guard = runtime.enter();
let cache = S3Cache::new(
S3CacheConfig {
key_prefix: "team/".to_string(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(cache.bucket(), "cache-bucket");
assert_eq!(cache.key_prefix(), "team/");
assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
assert_eq!(cache.to_s3_key("plain"), "team/plain");
let unprefixed = S3Cache::new(
S3CacheConfig {
key_prefix: String::new(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_methods_block_inside_and_outside_the_runtime() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
.mount(&server)
.await;
let uri = server.uri();
let cache = tokio::task::spawn_blocking(move || {
let cache = cache(&uri);
let context = ExactCacheContext::default();
cache
.set_cache("key", json!({"answer": 9}), &context)
.unwrap();
cache.get_cache("key", &context).unwrap()
})
.await
.unwrap();
assert_eq!(cache, Some(json!({"answer": 9})));
}

20
litellm-rust/crates/cache-valkey-semantic/Cargo.toml Normal file
View file

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
redis-test = "1.0.4"
rstest.workspace = true

1153
litellm-rust/crates/cache-valkey-semantic/src/lib.rs Normal file
View file

File diff suppressed because it is too large Load diff

50
litellm-rust/crates/cache/src/base_cache.rs vendored
View file

@ -32,6 +32,28 @@ impl CacheContext for ExactCacheContext {
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct SemanticCacheContext {
pub input: Option<serde_json::Value>,
pub messages: Option<serde_json::Value>,
pub metadata: Option<serde_json::Value>,
pub scope: Option<String>,
pub ttl: Option<Duration>,
}
impl CacheContext for SemanticCacheContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self {
ttl,
..self.clone()
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum CacheConnectionStatus {
@ -105,3 +127,31 @@ pub trait BaseCache: Send + Sync {
fn test_connection(&self) -> impl Future<Output = Result<CacheConnectionResult, Error>> + Send;
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use serde_json::json;
use super::{CacheContext, SemanticCacheContext};
#[test]
fn semantic_context_with_ttl_only_replaces_ttl() {
let context = SemanticCacheContext {
input: Some(json!({"input": "hello"})),
messages: Some(json!([{"role": "user", "content": "hello"}])),
metadata: Some(json!({"tenant": "team"})),
scope: Some("scope".into()),
ttl: Some(Duration::from_secs(10)),
};
let updated = context.with_ttl(Some(Duration::from_secs(20)));
assert_eq!(updated.ttl, Some(Duration::from_secs(20)));
assert_eq!(updated.input, context.input);
assert_eq!(updated.messages, context.messages);
assert_eq!(updated.metadata, context.metadata);
assert_eq!(updated.scope, context.scope);
}
}

4
litellm-rust/crates/cache/src/lib.rs vendored
View file

@ -5,11 +5,10 @@ mod capabilities;
mod codec;
mod dual;
mod error;
mod semantic;
pub use base_cache::{
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
ExactCacheContext,
ExactCacheContext, SemanticCacheContext,
};
pub use cache_type::CacheType;
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
@ -20,4 +19,3 @@ pub use capabilities::{
pub use codec::{CacheCodec, JsonCodec};
pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy};
pub use error::Error;
pub use semantic::{SemanticCacheContext, SemanticCacheScope};

72
litellm-rust/crates/cache/src/semantic.rs vendored
View file

@ -1,72 +0,0 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::CacheContext;
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum SemanticCacheScope {
#[default]
Key,
EndUser,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct SemanticCacheContext {
pub input: Option<String>,
pub messages: Vec<Value>,
pub metadata: Map<String, Value>,
pub scope: SemanticCacheScope,
pub ttl: Option<Duration>,
}
impl CacheContext for SemanticCacheContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self {
ttl,
..self.clone()
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn with_ttl_keeps_request_fields() {
let context = SemanticCacheContext {
input: Some("query".to_owned()),
messages: vec![serde_json::json!({"role": "user", "content": "hi"})],
metadata: Map::from_iter([("user".to_owned(), Value::from("u1"))]),
scope: SemanticCacheScope::EndUser,
ttl: None,
};
let updated = context.with_ttl(Some(Duration::from_secs(5)));
assert_eq!(updated.ttl(), Some(Duration::from_secs(5)));
assert_eq!(updated.input, context.input);
assert_eq!(updated.messages, context.messages);
assert_eq!(updated.metadata, context.metadata);
assert_eq!(updated.scope, SemanticCacheScope::EndUser);
}
#[test]
fn scope_serializes_like_python_cache_scope() {
assert_eq!(
serde_json::to_value(SemanticCacheScope::EndUser).unwrap(),
Value::from("end_user")
);
assert_eq!(
serde_json::from_value::<SemanticCacheScope>(Value::from("key")).unwrap(),
SemanticCacheScope::Key
);
}
}

21
litellm-rust/crates/cache/tests/caching.rs vendored
View file

@ -1,7 +1,8 @@
use std::{sync::Mutex, time::Duration};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext,
get_cache,
};
struct TestCache {
@ -126,6 +127,24 @@ fn associated_context_preserves_backend_specific_lookup_inputs() {
);
}
#[test]
fn semantic_context_with_ttl_preserves_lookup_inputs() {
let context = SemanticCacheContext {
input: Some(serde_json::json!("text")),
messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])),
metadata: Some(serde_json::json!({"key": "value"})),
scope: Some("scope".into()),
ttl: None,
};
let updated = context.with_ttl(Some(Duration::from_secs(30)));
assert_eq!(updated.ttl(), Some(Duration::from_secs(30)));
assert_eq!(updated.input, context.input);
assert_eq!(updated.messages, context.messages);
assert_eq!(updated.metadata, context.metadata);
assert_eq!(updated.scope, context.scope);
assert_eq!(context.with_ttl(None).ttl(), None);
}
#[tokio::test]
async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
let cache = TestCache {

2
litellm-rust/crates/framer/Cargo.toml
View file

@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
sse = ["dep:sse-stream"]
[dependencies]
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
aws-smithy-eventstream = { version = "=0.61.4", optional = true }
aws-smithy-types = { version = "1.6.1", optional = true }
bytes = "1"
futures-util.workspace = true

2
litellm-rust/crates/llms/Cargo.toml
View file

@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] }
url.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-eventstream = "=0.61.4"
aws-smithy-types = "1.6.1"
rstest.workspace = true
tokio.workspace = true

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

@ -24,13 +24,17 @@ litellm-cache.workspace = true
litellm-cache-azure-blob.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-redis-semantic.workspace = true
litellm-cache-response.workspace = true
litellm-cache-qdrant-semantic.workspace = true
qdrant-client.workspace = true
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-callbacks-legacy-python.workspace = true
litellm-core.workspace = true
litellm-core-utils.workspace = true
@ -43,9 +47,10 @@ litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
url.workspace = true
tokio = { workspace = true, features = ["sync"] }
tokio = { workspace = true, features = ["rt", "sync"] }
[dev-dependencies]
serde.workspace = true
@ -53,6 +58,7 @@ serde_with.workspace = true
criterion.workspace = true
futures-util.workspace = true
rstest.workspace = true
sha2.workspace = true
tokio-tungstenite.workspace = true
[[bench]]

21
litellm-rust/crates/python-bridge/src/cache/binding.rs vendored
View file

@ -56,12 +56,7 @@ impl ResolvedCache {
CacheBinding::Disabled => ready_none(py)?,
CacheBinding::Native(service) => {
let request = request(input)?;
let service = service.clone();
run_async(
py,
async move { service.async_lookup(&request, now()).await },
cache_error,
)?
service.async_lookup_py(py, request)?
}
CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?,
};
@ -179,12 +174,7 @@ impl ResolvedCache {
CacheBinding::Native(service) => {
let request = self::request(request)?;
let response: Value = from_py(response)?;
let service = service.clone();
run_async(
py,
async move { service.async_store(&request, response, now()).await },
cache_error,
)
service.async_store_py(py, request, response)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store(py, response, callback_kwargs)
@ -241,12 +231,7 @@ impl ResolvedCache {
));
}
let entries = requests.into_iter().zip(responses).collect();
let service = service.clone();
run_async(
py,
async move { service.async_store_batch(entries, now()).await },
cache_error,
)
service.async_store_batch_py(py, entries)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store_batch(py, callback_result, callback_kwargs)

939
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

File diff suppressed because it is too large Load diff

156
litellm-rust/crates/python-bridge/src/cache/embedder.rs vendored Normal file
View file

@ -0,0 +1,156 @@
use std::future::Future;
use litellm_cache::Error;
use litellm_host_python::to_py;
use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict};
use serde_json::Value;
tokio::task_local! {
static PREPARED_EMBEDDING: Result<Vec<f32>, Error>;
}
pub(super) fn with_prepared_embedding<F: Future>(
vector: Result<Vec<f32>, Error>,
future: F,
) -> impl Future<Output = F::Output> {
PREPARED_EMBEDDING.scope(vector, future)
}
pub(super) struct PythonEmbedder(Py<PyAny>);
impl Clone for PythonEmbedder {
fn clone(&self) -> Self {
Python::attach(|py| Self(self.0.clone_ref(py)))
}
}
impl PythonEmbedder {
pub(super) fn new(object: Py<PyAny>) -> Self {
Self(object)
}
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self(backend.clone().unbind()))
}
pub(super) fn object(&self) -> &Py<PyAny> {
&self.0
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
pub(super) fn async_embed_awaitable<'py>(
&self,
py: Python<'py>,
prompt: &str,
metadata: &Option<Value>,
) -> PyResult<Bound<'py, PyAny>> {
let metadata = to_py(py, metadata)?;
self.0
.bind(py)
.call_method1("_get_async_embedding", (prompt, metadata))
}
fn metadata_kwargs<'py>(
py: Python<'py>,
metadata: Option<&Value>,
) -> PyResult<Bound<'py, PyDict>> {
let kwargs = PyDict::new(py);
kwargs.set_item("metadata", to_py(py, &metadata)?)?;
Ok(kwargs)
}
pub(super) fn async_embedding_coroutine(
&self,
py: Python<'_>,
prompt: &str,
metadata: Option<&Value>,
) -> PyResult<Py<PyAny>> {
let kwargs = Self::metadata_kwargs(py, metadata)?;
self.0
.bind(py)
.call_method("_get_async_embedding", (prompt,), Some(&kwargs))
.map(Bound::unbind)
}
pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult<Vec<f32>> {
Ok(vector
.extract::<Vec<f64>>()?
.into_iter()
.map(|value| value as f32)
.collect())
}
}
impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
let metadata = to_py(py, &metadata)?;
self.0
.bind(py)
.call_method1("_get_embedding", (prompt, metadata))?
.extract()
})
.map_err(|_| Error::Unavailable)?;
Ok(result.into_iter().map(|value| value as f32).collect())
}
fn async_embed(
&self,
_prompt: &str,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
let seeded = PREPARED_EMBEDDING
.try_with(Clone::clone)
.unwrap_or(Err(Error::Unavailable));
std::future::ready(seeded)
}
}
impl litellm_cache_redis_semantic::Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
Python::attach(|py| {
let kwargs = Self::metadata_kwargs(py, metadata)?;
Self::extract(self.0.bind(py).call_method(
"_get_embedding",
(prompt,),
Some(&kwargs),
)?)
})
.map_err(|_| Error::Unavailable)
}
fn async_embed(
&self,
_prompt: &str,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
let seeded = PREPARED_EMBEDDING
.try_with(Clone::clone)
.unwrap_or(Err(Error::Unavailable));
std::future::ready(seeded)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn async_embed_returns_the_seeded_vector_or_unavailable() {
Python::initialize();
let object = Python::attach(|py| py.None());
let embedder = PythonEmbedder::new(object);
let scoped_embedder = embedder.clone();
let scoped = with_prepared_embedding(Ok(vec![0.25]), async move {
litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None)
.await
});
assert_eq!(scoped.await, Ok(vec![0.25]));
let unscoped =
litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await;
assert_eq!(unscoped, Err(Error::Unavailable));
}
}

260
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -31,6 +31,7 @@ struct RedisPoolGuard {
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: Option<usize>,
client_name: &'static str,
attributes: RedisPoolAttributes,
}
@ -46,12 +47,18 @@ struct AzureBlobClientGuard {
container_name: String,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
S3(S3ClientGuard),
}
#[derive(Clone, Copy)]
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -70,6 +77,8 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
max_connections: None,
};
const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL;
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
@ -78,34 +87,6 @@ pub(super) struct FacadeGuard {
}
impl ObjectGuard {
fn class_behaviors(class: &Bound<'_, PyType>) -> PyResult<Vec<(String, Py<PyAny>)>> {
let py = class.py();
let builtins = py.import("builtins")?;
let property_type = builtins.getattr("property")?;
let staticmethod_type = builtins.getattr("staticmethod")?;
let classmethod_type = builtins.getattr("classmethod")?;
class
.getattr("__dict__")?
.call_method0("items")?
.try_iter()?
.map(|item| {
let item = item?;
let (name, value): (String, Py<PyAny>) = item.extract()?;
let value_bound = value.bind(py);
let is_behavior = value_bound.is_callable()
|| value_bound.is_instance(&property_type)?
|| value_bound.is_instance(&staticmethod_type)?
|| value_bound.is_instance(&classmethod_type)?;
Ok(is_behavior.then_some((name, value)))
})
.filter_map(|result| match result {
Ok(Some(attribute)) => Some(Ok(attribute)),
Ok(None) => None,
Err(error) => Some(Err(error)),
})
.collect()
}
fn capture(
py: Python<'_>,
object: &Bound<'_, PyAny>,
@ -118,7 +99,12 @@ impl ObjectGuard {
.iter()
.map(|class| {
let class = class.cast_into::<PyType>()?;
let attributes = Self::class_behaviors(&class)?;
let attributes = class
.getattr("__dict__")?
.call_method0("items")?
.try_iter()?
.map(|item| item?.extract::<(String, Py<PyAny>)>())
.collect::<PyResult<Vec<_>>>()?;
Ok(ClassGuard {
class: class.unbind(),
attributes,
@ -171,20 +157,16 @@ impl ObjectGuard {
}
let instance = object.getattr("__dict__")?.cast_into::<PyDict>()?;
for (class, expected) in mro.iter().zip(&self.classes) {
let class = class.cast_into::<PyType>()?;
if !class.is(expected.class.bind(py)) {
return Ok(false);
}
let attributes = Self::class_behaviors(&class)?;
if attributes.len() != expected.attributes.len() {
let attributes = class.getattr("__dict__")?;
if attributes.len()? != expected.attributes.len() {
return Ok(false);
}
for ((name, value), (expected_name, expected_value)) in
attributes.iter().zip(&expected.attributes)
{
if name != expected_name
|| instance.contains(name)?
|| !value.bind(py).is(expected_value.bind(py))
for (name, value) in &expected.attributes {
if (instance.contains(name)? && !self.config_names.contains(&name.as_str()))
|| !attributes.get_item(name)?.is(value.bind(py))
{
return Ok(false);
}
@ -206,8 +188,12 @@ impl ObjectGuard {
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
fn capture(
backend: &Bound<'_, PyAny>,
client_name: &'static str,
attributes: RedisPoolAttributes,
) -> PyResult<Self> {
let pool = backend.getattr(client_name)?.getattr(attributes.pool)?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr(attributes.connection_class)?.unbind(),
@ -215,31 +201,30 @@ impl RedisPoolGuard {
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: Self::max_connections(&pool, &attributes)?,
max_connections: attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?,
client_name,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr(self.client_name)?
.getattr(self.attributes.pool)?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
.bind(py)
.is(&pool.getattr(self.attributes.connection_class)?)
&& self.max_connections == Self::max_connections(&pool, &self.attributes)?
&& self.max_connections
== self
.attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?
&& self
.connection_kwargs
.bind(py)
@ -301,12 +286,43 @@ impl AzureBlobClientGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl ConnectionGuard {
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(match (kind, cluster) {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
STANDALONE_POOL,
)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
CLUSTER_POOL,
)?),
("valkey-semantic", _) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"sync_client",
VALKEY_POOL,
)?),
("disk", _) => Self::None,
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
_ => Self::None,
})
}
@ -316,6 +332,7 @@ impl ConnectionGuard {
Self::None => Ok(true),
Self::RedisPool(guard) => guard.matches(py, backend),
Self::AzureBlob(guard) => guard.matches(py, backend),
Self::S3(guard) => guard.matches(py, backend),
}
}
@ -324,6 +341,7 @@ impl ConnectionGuard {
Self::None => Ok(()),
Self::RedisPool(guard) => guard.traverse(visit),
Self::AzureBlob(guard) => guard.traverse(visit),
Self::S3(guard) => guard.traverse(visit),
}
}
}
@ -345,23 +363,34 @@ impl FacadeGuard {
let (module, name, cache_kind) = match (kind, cluster) {
("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"),
("redis", true) => (
"litellm.caching.redis_cluster_cache",
"RedisClusterCache",
"redis",
("redis_semantic", _) => (
"litellm.caching.redis_semantic_cache",
"RedisSemanticCache",
"redis-semantic",
),
("qdrant_semantic", _) => (
"litellm.caching.qdrant_semantic_cache",
"QdrantSemanticCache",
"qdrant-semantic",
),
("redis", true) => (
"litellm.caching.redis_cluster_cache",
"RedisClusterCache",
"redis",
),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("valkey-semantic", false) => (
"litellm.caching.valkey_semantic_cache",
"ValkeySemanticCache",
"valkey-semantic",
),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -381,30 +410,15 @@ impl FacadeGuard {
if let Some(message) = config.service_mismatch(service) {
return Err(PyTypeError::new_err(message));
}
let backend_config_names = match kind {
"memory" | "redis" | "azure-blob" | "disk" | "gcs" => &[
"namespace",
"default_ttl",
"max_size_in_memory",
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
"bucket_name",
"key_prefix",
"path_service_account",
][..],
"qdrant_semantic" => &[
"qdrant_api_base",
"qdrant_api_key",
"collection_name",
"similarity_threshold",
"embedding_model",
"vector_size",
"embedding_max_input_tokens",
"embedding_timeout",
][..],
_ => unreachable!(),
};
if kind == "redis_semantic"
&& service
.embedder_object()
.is_none_or(|embedder| !backend.is(embedder.bind(py)))
{
return Err(PyTypeError::new_err(
"facade backend must be the native embedder",
));
}
Ok(Self {
outer: ObjectGuard::capture(
py,
@ -419,7 +433,37 @@ impl FacadeGuard {
"semantic_cache_scope",
],
)?,
backend: ObjectGuard::capture(py, &backend, backend_config_names)?,
backend: ObjectGuard::capture(
py,
&backend,
&[
"namespace",
"default_ttl",
"max_size_in_memory",
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
"similarity_threshold",
"distance_threshold",
"embedding_model",
"embedding_max_input_tokens",
"embedding_timeout",
"qdrant_api_base",
"qdrant_api_key",
"collection_name",
"vector_size",
"_index_name",
"_redis_url",
"similarity_threshold",
"embedding_model",
"index_name",
"embedding_max_input_tokens",
"embedding_timeout",
"bucket_name",
"key_prefix",
"path_service_account",
],
)?,
disk_store: (kind == "disk")
.then(|| DiskStoreGuard::capture(&backend))
.transpose()?,
@ -477,51 +521,3 @@ pub(super) fn resolve(
}
handle.service().map(Some)
}
#[cfg(test)]
mod tests {
use super::ObjectGuard;
use pyo3::{prelude::*, types::PyDict};
#[test]
fn class_data_shadowing_is_ignored_but_method_mutations_are_rejected() {
Python::initialize();
Python::attach(|py| {
let namespace = PyDict::new(py);
py.run(
c"class Example:\n data = 1\n def method(self):\n return 1\nobject = Example()\nobject.data = 2",
None,
Some(&namespace),
)
.unwrap();
let object = namespace.get_item("object").unwrap().unwrap();
let guard = ObjectGuard::capture(py, &object, &[]).unwrap();
assert!(guard.matches(py, &object).unwrap());
py.run(c"object.method = lambda: 2", None, Some(&namespace))
.unwrap();
assert!(!guard.matches(py, &object).unwrap());
});
}
#[test]
fn class_method_replacement_is_rejected() {
Python::initialize();
Python::attach(|py| {
let namespace = PyDict::new(py);
py.run(
c"class Example:\n def method(self):\n return 1\nobject = Example()",
None,
Some(&namespace),
)
.unwrap();
let object = namespace.get_item("object").unwrap().unwrap();
let guard = ObjectGuard::capture(py, &object, &[]).unwrap();
py.run(c"Example.method = lambda self: 2", None, Some(&namespace))
.unwrap();
assert!(!guard.matches(py, &object).unwrap());
});
}
}

144
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored
View file

@ -1,19 +1,27 @@
use std::env;
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_http::ClientVariant;
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_redis_semantic::RedisSemanticConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*, types::PyDict};
use litellm_http::ClientVariant;
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyRuntimeError, PyTypeError},
prelude::*,
types::PyDict,
};
use url::Url;
use super::{
cache_error, config::QdrantSemanticCacheConfig, facade::FacadeGuard,
native::NativeResponseCache, request::duration,
cache_error,
config::{QdrantSemanticCacheConfig, project_redis_semantic},
embedder::PythonEmbedder,
facade::FacadeGuard,
native::NativeResponseCache,
request::duration,
};
use crate::http;
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
@ -75,6 +83,40 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (bucket, *, region, endpoint_url=None, key_prefix="", access_key_id=None, secret_access_key=None, session_token=None))]
fn s3(
py: Python<'_>,
bucket: String,
region: String,
endpoint_url: Option<String>,
key_prefix: &str,
access_key_id: Option<String>,
secret_access_key: Option<String>,
session_token: Option<String>,
) -> PyResult<Self> {
let config = S3CacheConfig {
bucket,
key_prefix: key_prefix.to_string(),
region: region.clone(),
endpoint: endpoint_url.map(|url| S3Endpoint { url }),
auth: AwsAuthConfig {
access_key_id,
secret_access_key,
session_token,
region_name: Some(region),
..Default::default()
},
};
let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))]
fn gcs(
@ -140,7 +182,7 @@ impl CacheTestHandle {
|| (!parsed.path().is_empty() && parsed.path() != "/")
|| parsed.query().is_some()
|| parsed.host_str().is_none()
|| parsed.port().is_some_and(|port| port != 6333)
|| parsed.port() != Some(6333)
{
return Err(pyo3::exceptions::PyValueError::new_err(
"native Qdrant requires the default REST port so the gRPC port can be derived",
@ -156,7 +198,7 @@ impl CacheTestHandle {
grpc_url.set_query(None);
let embedding_api_key = embedding_api_key
.or_else(|| {
env::var("OPENAI_API_KEY")
std::env::var("OPENAI_API_KEY")
.ok()
.filter(|value| !value.is_empty())
})
@ -166,8 +208,8 @@ impl CacheTestHandle {
)
})?;
let embedding_api_base = embedding_api_base.unwrap_or_else(|| {
env::var("OPENAI_BASE_URL")
.or_else(|_| env::var("OPENAI_API_BASE"))
std::env::var("OPENAI_BASE_URL")
.or_else(|_| std::env::var("OPENAI_API_BASE"))
.unwrap_or_else(|_| "https://api.openai.com/v1".to_owned())
});
let quantization = match quantization {
@ -194,10 +236,10 @@ impl CacheTestHandle {
},
quantization,
};
let http_config = http::call_config(py, &PyDict::new(py), true)?;
let client = http::pool()
let http_config = crate::http::call_config(py, &PyDict::new(py), true)?;
let client = crate::http::pool()
.client(&http_config, ClientVariant::Provider)
.map_err(http::client_error)?;
.map_err(crate::http::client_error)?;
let service = run_sync_value(py, async move {
let handle = tokio::runtime::Handle::current();
NativeResponseCache::qdrant_semantic(config, client, handle)
@ -211,6 +253,29 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (url, similarity_threshold, index_name, embedder))]
fn valkey_semantic(
url: String,
similarity_threshold: f64,
index_name: String,
embedder: &Bound<'_, PyAny>,
) -> PyResult<Self> {
let python_embedder = PythonEmbedder::from_backend(embedder)?;
let service = NativeResponseCache::valkey_semantic(
&url,
similarity_threshold,
index_name,
python_embedder,
)
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
@ -226,6 +291,36 @@ impl CacheTestHandle {
})
}
#[staticmethod]
fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult<Self> {
let class = py
.import("litellm.caching.redis_semantic_cache")?
.getattr("RedisSemanticCache")?;
if !backend.get_type().is(&class) {
return Err(PyTypeError::new_err(
"native redis-semantic handles require the built-in RedisSemanticCache",
));
}
let config = project_redis_semantic(&backend)?;
let embedder = PythonEmbedder::new(backend.unbind());
let service = release_gil(py, move || {
NativeResponseCache::redis_semantic(
&config.redis_url,
embedder,
RedisSemanticConfig {
index_name: config.index_name,
similarity_threshold: config.similarity_threshold as f32,
},
)
})
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()
@ -234,11 +329,17 @@ impl CacheTestHandle {
fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
let service = self.service()?;
let guard = FacadeGuard::capture(py, facade, &service)?;
let service = service.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let service = service
.with_scope(
facade
.getattr("semantic_cache_scope")?
.extract::<String>()?,
)
.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let handle = Py::new(
py,
Self {
@ -251,6 +352,7 @@ impl CacheTestHandle {
}
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.service.traverse(&visit)?;
if let Some(guard) = &self.guard {
guard.traverse(visit)?;
}

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

@ -1,12 +1,15 @@
mod binding;
mod callback;
mod config;
mod embedder;
mod facade;
mod future;
mod handle;
mod native;
mod request;
mod resolver;
mod semantic;
mod semantic_step;
use litellm_cache::Error;
use pyo3::{

879
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

File diff suppressed because it is too large Load diff

73
litellm-rust/crates/python-bridge/src/cache/request.rs vendored
View file

@ -1,11 +1,11 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{ExactCacheContext, SemanticCacheContext, SemanticCacheScope};
use litellm_cache::ExactCacheContext;
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
use litellm_host_python::from_py;
use pyo3::{exceptions::PyValueError, prelude::*};
use serde::Deserialize;
use serde_json::{Map, Value};
use serde_json::Value;
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
@ -14,51 +14,58 @@ struct RequestInput {
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
messages: Option<Vec<Value>>,
input: Option<String>,
metadata: Option<Map<String, Value>>,
scope: Option<SemanticCacheScope>,
messages: Option<Value>,
input: Option<Value>,
metadata: Option<Value>,
litellm_metadata: Option<Value>,
litellm_params: Option<Value>,
scope: Option<String>,
}
pub(super) fn request(
value: &Bound<'_, PyAny>,
) -> PyResult<ResponseCacheRequest<SemanticCacheContext>> {
#[derive(Clone)]
pub(super) struct NativeRequest {
pub(super) key: CacheKeyInput,
pub(super) controls: CacheControls,
pub(super) ttl: Option<Duration>,
pub(super) max_age: Option<Duration>,
pub(super) messages: Option<Value>,
pub(super) input: Option<Value>,
pub(super) metadata: Option<Value>,
pub(super) litellm_metadata: Option<Value>,
pub(super) litellm_params: Option<Value>,
pub(super) scope: Option<String>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
let input: RequestInput = from_py(value)?;
request_input(input)
}
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest<SemanticCacheContext>> {
let mut request: ResponseCacheRequest<SemanticCacheContext> =
ResponseCacheRequest::new(input.key);
if let Some(controls) = input.controls {
request.controls = controls;
}
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
request.context.messages = input.messages.unwrap_or_default();
request.context.input = input.input;
request.context.metadata = input.metadata.unwrap_or_default();
request.context.scope = input.scope.unwrap_or_default();
request.max_age = input.max_age_seconds.map(duration).transpose()?;
Ok(request)
fn request_input(input: RequestInput) -> PyResult<NativeRequest> {
let controls = input.controls.unwrap_or_else(|| {
ResponseCacheRequest::<ExactCacheContext>::new(input.key.clone()).controls
});
Ok(NativeRequest {
key: input.key,
controls,
ttl: input.ttl_seconds.map(duration).transpose()?,
max_age: input.max_age_seconds.map(duration).transpose()?,
messages: input.messages,
input: input.input,
metadata: input.metadata,
litellm_metadata: input.litellm_metadata,
litellm_params: input.litellm_params,
scope: input.scope,
})
}
pub(super) fn requests(
value: &Bound<'_, PyAny>,
) -> PyResult<Vec<ResponseCacheRequest<SemanticCacheContext>>> {
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<NativeRequest>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)
.collect()
}
pub(super) fn exact(
request: &ResponseCacheRequest<SemanticCacheContext>,
) -> ResponseCacheRequest<ExactCacheContext> {
request.clone().with_context(ExactCacheContext {
ttl: request.context.ttl,
})
}
pub(super) fn duration(seconds: f64) -> PyResult<Duration> {
Duration::try_from_secs_f64(seconds)
.map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative"))

175
litellm-rust/crates/python-bridge/src/cache/semantic.rs vendored Normal file
View file

@ -0,0 +1,175 @@
use std::collections::VecDeque;
use litellm_cache::Error;
use litellm_cache_redis_semantic::prompt_from_context;
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyException, PyRuntimeError},
prelude::*,
};
use serde_json::Value;
use super::{
cache_error,
embedder::{PythonEmbedder, with_prepared_embedding},
native::NativeResponseCache,
request::{NativeRequest, now},
};
pub(super) enum SemanticOperation {
Lookup(NativeRequest),
Store(NativeRequest, Value),
StoreBatch(VecDeque<(NativeRequest, Value)>),
}
enum Phase {
Start,
AwaitingEmbedding,
AwaitingBackend,
}
pub(super) struct SemanticBody {
service: NativeResponseCache,
operation: SemanticOperation,
pending: Option<(NativeRequest, Option<Value>)>,
phase: Phase,
}
impl SemanticBody {
pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self {
Self {
service,
operation,
pending: None,
phase: Phase::Start,
}
}
fn backend_step(
&mut self,
py: Python<'_>,
seed: Result<Vec<f32>, Error>,
) -> PyResult<ExecutionStep> {
self.phase = Phase::AwaitingBackend;
let (request, response) = self.pending.take().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution resumed without a pending operation")
})?;
let service = self.service.clone();
let future = async move {
match response {
None => service.async_lookup(&request, now()).await,
Some(response) => service
.async_store(&request, response, now())
.await
.map(|_| None),
}
};
let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?;
Ok(ExecutionStep::Await(awaitable.unbind()))
}
}
impl ExecutionBody for SemanticBody {
fn resume(&mut self, mut result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| {
loop {
match self.phase {
Phase::Start => {
if result.is_some() {
return Err(PyRuntimeError::new_err(
"semantic execution received a result before starting",
));
}
if self.pending.is_none() {
match &mut self.operation {
SemanticOperation::Lookup(request) => {
self.pending = Some((request.clone(), None));
}
SemanticOperation::Store(request, response) => {
let response = std::mem::replace(response, Value::Null);
self.pending = Some((request.clone(), Some(response)));
}
SemanticOperation::StoreBatch(queue) => {
let Some((request, response)) = queue.pop_front() else {
return Ok(ExecutionStep::Return(py.None()));
};
self.pending = Some((request, Some(response)));
}
}
}
let (request, _) = self.pending.as_ref().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution has no pending operation")
})?;
let semantic = NativeResponseCache::semantic_request(request);
let Some(prompt) = prompt_from_context(&semantic.context) else {
return self.backend_step(py, Err(Error::Unavailable));
};
let embedder = self.service.semantic_embedder().ok_or_else(|| {
PyRuntimeError::new_err(
"semantic execution requires a redis-semantic backend",
)
})?;
let coroutine = embedder.async_embedding_coroutine(
py,
&prompt,
semantic.context.metadata.as_ref(),
)?;
self.phase = Phase::AwaitingEmbedding;
return Ok(ExecutionStep::Await(coroutine));
}
Phase::AwaitingEmbedding => {
let result = result.take().ok_or_else(|| {
PyRuntimeError::new_err(
"semantic execution expected an embedding result",
)
})?;
let seed = match result {
Ok(value) => PythonEmbedder::extract(value.into_bound(py))
.map_err(|_| Error::Unavailable),
Err(error) => {
if !error.is_instance_of::<PyException>(py) {
return Err(error);
}
Err(Error::Unavailable)
}
};
return self.backend_step(py, seed);
}
Phase::AwaitingBackend => {
let result = result.take().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution expected a backend result")
})?;
let value = match result {
Ok(value) => value,
Err(error) => return Err(error),
};
let more = matches!(
&self.operation,
SemanticOperation::StoreBatch(queue) if !queue.is_empty()
);
if more {
self.phase = Phase::Start;
continue;
}
return Ok(ExecutionStep::Return(value));
}
}
}
})
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
if let Some(embedder) = self.service.semantic_embedder() {
embedder.traverse(visit)?;
}
Ok(())
}
}
pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult<Bound<'_, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

249
litellm-rust/crates/python-bridge/src/cache/semantic_step.rs vendored Normal file
View file

@ -0,0 +1,249 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::SemanticCacheContext;
use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest};
use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context};
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use serde_json::Value;
use super::{cache_error, embedder::PythonEmbedder};
pub(super) enum Op {
Lookup,
Store(Value),
StoreBatch(Vec<Value>),
}
#[derive(Clone, Copy)]
enum State {
Start,
AwaitingEmbedding,
AwaitingStorage,
Done,
}
pub(super) struct SemanticEmbedExecution {
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
op: Op,
now: Option<Duration>,
prepared: Vec<Option<Vec<f32>>>,
index: usize,
state: State,
}
impl SemanticEmbedExecution {
pub(super) fn lookup(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Lookup,
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
}
}
pub(super) fn store(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
response: Value,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Store(response),
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
}
}
pub(super) fn store_batch(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
responses: Vec<Value>,
) -> Self {
Self {
backend,
embedder,
prepared: vec![None; requests.len()],
requests,
op: Op::StoreBatch(responses),
now: None,
index: 0,
state: State::Start,
}
}
fn start(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
if self.now.is_none() {
self.now = Some(super::request::now());
}
while self.index < self.requests.len() {
let request = &self.requests[self.index];
let enabled = match &self.op {
Op::Lookup => request.controls.reads(),
Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(),
};
if !enabled {
self.index += 1;
continue;
}
let Some(prompt) = prompt_from_context(&request.context) else {
self.index += 1;
continue;
};
let metadata = request.context.metadata.clone();
let awaitable = self
.embedder
.async_embed_awaitable(py, &prompt, &metadata)?;
self.state = State::AwaitingEmbedding;
return Ok(ExecutionStep::Await(awaitable.unbind()));
}
self.state = State::AwaitingStorage;
self.storage_step(py)
}
fn storage_step(&self, py: Python<'_>) -> PyResult<ExecutionStep> {
let requests = self.requests.clone();
let prepared = self.prepared.clone();
let backend = Arc::clone(&self.backend);
let now = self
.now
.ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?;
let awaitable = match &self.op {
Op::Lookup => {
let Some(request) = requests.into_iter().next() else {
return Err(PyRuntimeError::new_err(
"semantic lookup requires one request",
));
};
match prepared.into_iter().next().flatten() {
Some(values) => {
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
run_async(
py,
async move { cache.async_lookup(&request, now).await },
cache_error,
)?
}
None => {
let cache = Arc::new(ResponseCache::new(backend));
run_async(
py,
async move { cache.async_lookup(&request, now).await },
cache_error,
)?
}
}
}
Op::Store(response) => {
let Some(request) = requests.into_iter().next() else {
return Err(PyRuntimeError::new_err(
"semantic store requires one request",
));
};
let response = response.clone();
match prepared.into_iter().next().flatten() {
Some(values) => {
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
run_async(
py,
async move { cache.async_store(&request, response, now).await },
cache_error,
)?
}
None => {
let cache = Arc::new(ResponseCache::new(backend));
run_async(
py,
async move { cache.async_store(&request, response, now).await },
cache_error,
)?
}
}
}
Op::StoreBatch(responses) => {
let responses = responses.clone();
run_async(
py,
async move {
for ((request, response), prepared) in
requests.into_iter().zip(responses).zip(prepared)
{
let Some(values) = prepared else {
continue;
};
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = ResponseCache::new(Arc::new(backend));
cache.async_store(&request, response, now).await?;
}
Ok(())
},
cache_error,
)?
}
};
Ok(ExecutionStep::Await(awaitable.unbind()))
}
fn resume_py(
&mut self,
py: Python<'_>,
result: Option<PyResult<Py<PyAny>>>,
) -> PyResult<ExecutionStep> {
match (self.state, result) {
(State::Start, None) => self.start(py),
(State::AwaitingEmbedding, Some(Ok(value))) => {
let values = value.bind(py).extract::<Vec<f64>>()?;
self.prepared[self.index] =
Some(values.into_iter().map(|value| value as f32).collect());
self.index += 1;
self.start(py)
}
(State::AwaitingStorage, Some(Ok(value))) => {
self.state = State::Done;
Ok(ExecutionStep::Return(value))
}
(_, Some(Err(error))) => Err(error),
_ => Err(PyRuntimeError::new_err(
"invalid semantic cache execution state",
)),
}
}
}
impl ExecutionBody for SemanticEmbedExecution {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| self.resume_py(py, result))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.embedder.traverse(visit)
}
}
pub(super) fn drive_semantic<'py>(
py: Python<'py>,
body: SemanticEmbedExecution,
) -> PyResult<Bound<'py, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

25
litellm-rust/crates/secrets-hashicorp/Cargo.toml Normal file
View file

@ -0,0 +1,25 @@
[package]
name = "litellm-secrets-hashicorp"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-core-utils.workspace = true
litellm-secrets-types.workspace = true
moka.workspace = true
rustify.workspace = true
rustify_derive.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
vaultrs.workspace = true
veil.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile = "3"
tokio.workspace = true
wiremock = "0.6.5"

22
litellm-rust/crates/secrets-hashicorp/src/cert_login.rs Normal file
View file

@ -0,0 +1,22 @@
#[derive(Debug, rustify_derive::Endpoint)]
#[endpoint(path = "/auth/{self.mount}/login", method = "POST")]
pub struct CertLoginRequest {
#[endpoint(skip)]
pub mount: String,
#[endpoint(raw)]
body: Vec<u8>,
}
impl CertLoginRequest {
pub fn new(name: Option<&str>) -> Self {
let body: Vec<u8> = match name {
Some(name) => serde_json::to_vec(&serde_json::json!({ "name": name }))
.expect("json object serialization is infallible"),
None => b"{}".to_vec(),
};
Self {
mount: "cert".to_owned(),
body,
}
}
}

161
litellm-rust/crates/secrets-hashicorp/src/config.rs Normal file
View file

@ -0,0 +1,161 @@
use std::{path::PathBuf, time::Duration};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::SecretValue;
use crate::Error;
const DEFAULT_ADDRESS: &str = "http://127.0.0.1:8200";
const DEFAULT_MOUNT: &str = "secret";
const DEFAULT_APPROLE_MOUNT_PATH: &str = "approle";
const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400);
const HCP_VAULT_ADDR: &str = "HCP_VAULT_ADDR";
const HCP_VAULT_TOKEN: &str = "HCP_VAULT_TOKEN";
const HCP_VAULT_NAMESPACE: &str = "HCP_VAULT_NAMESPACE";
const HCP_VAULT_LOGIN_NAMESPACE: &str = "HCP_VAULT_LOGIN_NAMESPACE";
const HCP_VAULT_SECRET_NAMESPACE: &str = "HCP_VAULT_SECRET_NAMESPACE";
const HCP_VAULT_MOUNT_NAME: &str = "HCP_VAULT_MOUNT_NAME";
const HCP_VAULT_PATH_PREFIX: &str = "HCP_VAULT_PATH_PREFIX";
const HCP_VAULT_APPROLE_ROLE_ID: &str = "HCP_VAULT_APPROLE_ROLE_ID";
const HCP_VAULT_APPROLE_SECRET_ID: &str = "HCP_VAULT_APPROLE_SECRET_ID";
const HCP_VAULT_APPROLE_MOUNT_PATH: &str = "HCP_VAULT_APPROLE_MOUNT_PATH";
const HCP_VAULT_CLIENT_CERT: &str = "HCP_VAULT_CLIENT_CERT";
const HCP_VAULT_CLIENT_KEY: &str = "HCP_VAULT_CLIENT_KEY";
const HCP_VAULT_CERT_ROLE: &str = "HCP_VAULT_CERT_ROLE";
const HCP_VAULT_REFRESH_INTERVAL: &str = "HCP_VAULT_REFRESH_INTERVAL";
const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL";
#[derive(Clone, Debug)]
pub struct AppRoleAuth {
pub role_id: String,
pub secret_id: SecretValue,
pub mount_path: String,
}
#[derive(Clone, Debug)]
pub struct TlsCertAuth {
pub cert_path: PathBuf,
pub key_path: PathBuf,
pub role: Option<String>,
}
#[derive(Clone, Debug)]
pub struct HashicorpVaultConfig {
pub address: String,
pub token: Option<SecretValue>,
pub namespace: Option<String>,
pub login_namespace: Option<String>,
pub secret_namespace: Option<String>,
pub mount: String,
pub path_prefix: Option<String>,
pub approle: Option<AppRoleAuth>,
pub tls_cert: Option<TlsCertAuth>,
pub refresh_interval: Duration,
}
impl HashicorpVaultConfig {
pub fn from_environment(environment: &dyn Lookup) -> Result<Self, Error> {
let address: String = environment
.get(HCP_VAULT_ADDR)
.and_then(|value| nonempty(value.trim()))
.map(|value| value.trim_end_matches('/').to_owned())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| DEFAULT_ADDRESS.to_owned());
let token: Option<SecretValue> = environment
.get(HCP_VAULT_TOKEN)
.and_then(nonempty)
.map(SecretValue::new);
let namespace: Option<String> = path_component(environment.get(HCP_VAULT_NAMESPACE));
let login_namespace: Option<String> =
path_component(environment.get(HCP_VAULT_LOGIN_NAMESPACE));
let secret_namespace: Option<String> =
path_component(environment.get(HCP_VAULT_SECRET_NAMESPACE));
let mount: String = path_component(environment.get(HCP_VAULT_MOUNT_NAME))
.unwrap_or_else(|| DEFAULT_MOUNT.to_owned());
let path_prefix: Option<String> = path_component(environment.get(HCP_VAULT_PATH_PREFIX));
let approle: Option<AppRoleAuth> = match (
environment
.get(HCP_VAULT_APPROLE_ROLE_ID)
.and_then(nonempty),
environment
.get(HCP_VAULT_APPROLE_SECRET_ID)
.and_then(nonempty)
.map(SecretValue::new),
) {
(Some(role_id), Some(secret_id)) => Some(AppRoleAuth {
role_id,
secret_id,
mount_path: path_component(environment.get(HCP_VAULT_APPROLE_MOUNT_PATH))
.unwrap_or_else(|| DEFAULT_APPROLE_MOUNT_PATH.to_owned()),
}),
_ => None,
};
let tls_cert: Option<TlsCertAuth> = match (
environment.get(HCP_VAULT_CLIENT_CERT).and_then(nonempty),
environment.get(HCP_VAULT_CLIENT_KEY).and_then(nonempty),
) {
(Some(cert_path), Some(key_path)) => Some(TlsCertAuth {
cert_path: PathBuf::from(cert_path),
key_path: PathBuf::from(key_path),
role: environment.get(HCP_VAULT_CERT_ROLE).and_then(nonempty),
}),
_ => None,
};
let refresh_interval: Duration = refresh_interval(environment)?;
Ok(Self {
address,
token,
namespace,
login_namespace,
secret_namespace,
mount,
path_prefix,
approle,
tls_cert,
refresh_interval,
})
}
pub fn login_namespace(&self) -> Option<&str> {
self.login_namespace
.as_deref()
.or(self.namespace.as_deref())
}
pub fn secret_namespace(&self) -> Option<&str> {
self.secret_namespace
.as_deref()
.or(self.namespace.as_deref())
}
}
fn nonempty(value: impl AsRef<str>) -> Option<String> {
let value: &str = value.as_ref();
(!value.is_empty()).then(|| value.to_owned())
}
fn path_component(value: Option<String>) -> Option<String> {
value
.and_then(|value| nonempty(value.trim()))
.map(|value| value.trim_matches('/').to_owned())
.filter(|value| !value.is_empty())
}
fn refresh_interval(environment: &dyn Lookup) -> Result<Duration, Error> {
let value: Option<String> = environment
.get(HCP_VAULT_REFRESH_INTERVAL)
.and_then(nonempty)
.or_else(|| {
environment
.get(SECRET_MANAGER_REFRESH_INTERVAL)
.and_then(nonempty)
});
let Some(value) = value else {
return Ok(DEFAULT_REFRESH_INTERVAL);
};
let seconds: i64 = value.parse().map_err(|_| Error::RefreshInterval)?;
if seconds < 0 {
return Err(Error::RefreshInterval);
}
Ok(Duration::from_secs(seconds as u64))
}

34
litellm-rust/crates/secrets-hashicorp/src/error.rs Normal file
View file

@ -0,0 +1,34 @@
#[derive(thiserror::Error, veil::Redact)]
pub enum Error {
#[error("HashiCorp Vault requires an enterprise license")]
EnterpriseRequired,
#[error("invalid secret name")]
InvalidSecretName(#[from] litellm_secrets_types::Error),
#[error("HashiCorp Vault client failed")]
Client(
#[from]
#[redact]
vaultrs::error::ClientError,
),
#[error("HashiCorp Vault client settings are invalid: {message}")]
ClientSettings { message: String },
#[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")]
TlsIdentity {
path: std::path::PathBuf,
message: String,
},
#[error("HashiCorp Vault login returned HTTP {status}")]
LoginStatus { status: u16 },
#[error("HashiCorp Vault login response is malformed")]
MalformedLogin,
#[error("HashiCorp Vault authentication is not configured")]
NoAuthConfigured,
#[error("HashiCorp Vault returned HTTP {status}")]
Status { status: u16 },
#[error("HashiCorp Vault response payload is malformed")]
MalformedPayload,
#[error("HashiCorp Vault secret value is not a string")]
NonStringValue,
#[error("invalid HashiCorp Vault refresh interval")]
RefreshInterval,
}

10
litellm-rust/crates/secrets-hashicorp/src/lib.rs Normal file
View file

@ -0,0 +1,10 @@
#![forbid(unsafe_code)]
mod cert_login;
mod config;
mod error;
pub mod secret_manager;
pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth};
pub use error::Error;
pub use secret_manager::{HashicorpVault, SecretLocation};

359
litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs Normal file
View file

@ -0,0 +1,359 @@
use std::{
collections::HashMap,
fmt,
sync::Arc,
time::{Duration, Instant},
};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_types::{
BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name,
};
use moka::future::Cache;
use rustify::errors::ClientError as RustifyClientError;
use serde_json::Value;
use tokio::sync::Mutex;
use vaultrs::{
api,
auth::approle,
client::{Identity, VaultClient, VaultClientSettingsBuilder},
error::ClientError,
kv2,
};
use crate::{Error, HashicorpVaultConfig, TlsCertAuth, cert_login::CertLoginRequest};
const CACHE_CAPACITY: u64 = 200;
#[derive(Clone)]
struct CachedClient {
client: Arc<VaultClient>,
expires_at: Option<Instant>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SecretLocation {
pub namespace: Option<String>,
pub mount: String,
pub path: String,
}
#[derive(Clone)]
pub struct HashicorpVault {
config: HashicorpVaultConfig,
cache: Cache<String, SecretValue>,
auth_client: Arc<Mutex<Option<CachedClient>>>,
}
impl fmt::Debug for HashicorpVault {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("HashicorpVault")
.field("config", &self.config)
.finish_non_exhaustive()
}
}
impl HashicorpVault {
pub fn new(
environment: Arc<dyn Lookup + Send + Sync>,
enterprise_enabled: bool,
) -> Result<Self, Error> {
let config: HashicorpVaultConfig =
HashicorpVaultConfig::from_environment(environment.as_ref())?;
Self::from_config(config, enterprise_enabled)
}
pub fn from_config(
config: HashicorpVaultConfig,
enterprise_enabled: bool,
) -> Result<Self, Error> {
if !enterprise_enabled {
return Err(Error::EnterpriseRequired);
}
let cache: Cache<String, SecretValue> = Cache::builder()
.max_capacity(CACHE_CAPACITY)
.time_to_live(config.refresh_interval)
.build();
Ok(Self {
config,
cache,
auth_client: Arc::new(Mutex::new(None)),
})
}
pub fn secret_location(&self, secret_name: &str) -> Result<SecretLocation, Error> {
validate_secret_name(secret_name).map_err(Error::InvalidSecretName)?;
let path: String = [
self.config.path_prefix.clone(),
Some(secret_name.to_owned()),
]
.into_iter()
.flatten()
.collect::<Vec<String>>()
.join("/");
Ok(SecretLocation {
namespace: self.config.secret_namespace().map(str::to_owned),
mount: self.config.mount.clone(),
path,
})
}
pub fn config(&self) -> &HashicorpVaultConfig {
&self.config
}
pub async fn async_read_secret(&self, secret_name: &str) -> Result<Option<SecretValue>, Error> {
let location: SecretLocation = self.secret_location(secret_name)?;
let cache_key: String = cache_key(&location);
if let Some(value) = self.cache.get(&cache_key).await {
return Ok(Some(value));
}
let client: Arc<VaultClient> = self.vault_client().await?;
let data: HashMap<String, Value> =
match kv2::read(client.as_ref(), &location.mount, &location.path).await {
Ok(data) => data,
Err(error) if api_status(&error) == Some(404) => return Ok(None),
Err(error) => return Err(map_api_error(error, ErrorContext::Read)),
};
let Some(value) = data.get("key") else {
return Ok(None);
};
let value: &str = value.as_str().ok_or(Error::NonStringValue)?;
let value: SecretValue = SecretValue::new(value);
self.cache.insert(cache_key, value.clone()).await;
Ok(Some(value))
}
pub async fn async_write_secret(
&self,
secret_name: &str,
value: SecretValue,
description: Option<&str>,
) -> Result<Value, Error> {
let location: SecretLocation = self.secret_location(secret_name)?;
let cache_key: String = cache_key(&location);
let data: HashMap<String, Value> = match description {
Some(description) => [
("key".to_owned(), Value::String(value.expose().to_owned())),
(
"description".to_owned(),
Value::String(description.to_owned()),
),
]
.into_iter()
.collect(),
None => [("key".to_owned(), Value::String(value.expose().to_owned()))]
.into_iter()
.collect(),
};
let client: Arc<VaultClient> = self.vault_client().await?;
let metadata = kv2::set(client.as_ref(), &location.mount, &location.path, &data)
.await
.map_err(|error| map_api_error(error, ErrorContext::Secret))?;
self.cache.invalidate(&cache_key).await;
serde_json::to_value(metadata)
.map_err(|source| Error::Client(ClientError::JsonParseError { source }))
}
pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> {
let location: SecretLocation = self.secret_location(secret_name)?;
let cache_key: String = cache_key(&location);
let client: Arc<VaultClient> = self.vault_client().await?;
kv2::delete_latest(client.as_ref(), &location.mount, &location.path)
.await
.map_err(|error| map_api_error(error, ErrorContext::Secret))?;
self.cache.invalidate(&cache_key).await;
Ok(())
}
pub async fn async_rotate_secret(
&self,
current_name: &str,
new_name: &str,
value: &SecretValue,
) -> Result<Value, Error> {
async_rotate_secret(self, current_name, new_name, value).await
}
async fn vault_client(&self) -> Result<Arc<VaultClient>, Error> {
let mut cached = self.auth_client.lock().await;
if let Some(entry) = cached.as_ref()
&& entry
.expires_at
.is_none_or(|expires_at| expires_at > Instant::now())
{
return Ok(entry.client.clone());
}
let (client, expires_at): (VaultClient, Option<Instant>) =
match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) {
(Some(approle), _) => {
let login_client: VaultClient =
self.build_client(self.config.login_namespace(), "")?;
let auth = approle::login(
&login_client,
&approle.mount_path,
&approle.role_id,
approle.secret_id.expose(),
)
.await
.map_err(|error| map_api_error(error, ErrorContext::Login))?;
(
self.build_client(self.config.secret_namespace(), &auth.client_token)?,
token_expiry(auth.lease_duration),
)
}
(None, Some(tls)) => {
let login_client: VaultClient =
self.build_client(self.config.login_namespace(), "")?;
let endpoint: CertLoginRequest = CertLoginRequest::new(tls.role.as_deref());
let auth = api::auth(&login_client, endpoint)
.await
.map_err(|error| map_api_error(error, ErrorContext::Login))?;
(
self.build_client(self.config.secret_namespace(), &auth.client_token)?,
token_expiry(auth.lease_duration),
)
}
(None, None) => {
let token: SecretValue =
self.config.token.clone().ok_or(Error::NoAuthConfigured)?;
(
self.build_client(self.config.secret_namespace(), token.expose())?,
None,
)
}
};
let client: Arc<VaultClient> = Arc::new(client);
*cached = Some(CachedClient {
client: client.clone(),
expires_at,
});
Ok(client)
}
fn build_client(&self, namespace: Option<&str>, token: &str) -> Result<VaultClient, Error> {
let settings = VaultClientSettingsBuilder::default()
.address(&self.config.address)
.token(token.to_owned())
.namespace(namespace.map(str::to_owned))
.identity(identity_for(self.config.tls_cert.as_ref())?)
.ca_certs(Vec::new())
.verify(true)
.build()
.map_err(|message| Error::ClientSettings {
message: message.to_string(),
})?;
VaultClient::new(settings).map_err(Error::Client)
}
}
impl BaseSecretManager for HashicorpVault {
type Error = Error;
type WriteResponse = Value;
type DeleteResponse = ();
async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
HashicorpVault::async_read_secret(self, name).await
}
async fn async_write_secret(
&self,
name: &str,
value: &SecretValue,
description: Option<&str>,
) -> Result<Value, Error> {
HashicorpVault::async_write_secret(self, name, value.clone(), description).await
}
async fn async_delete_secret(
&self,
name: &str,
_recovery_window_in_days: i64,
) -> Result<(), Error> {
HashicorpVault::async_delete_secret(self, name).await
}
}
#[derive(Clone, Copy)]
enum ErrorContext {
Login,
Read,
Secret,
}
fn cache_key(location: &SecretLocation) -> String {
format!(
"{:?}/{}/{}",
location.namespace, location.mount, location.path
)
}
fn identity_for(tls: Option<&TlsCertAuth>) -> Result<Option<Identity>, Error> {
tls.map(|tls| {
let cert: Vec<u8> = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity {
path: tls.cert_path.clone(),
message: source.to_string(),
})?;
let key: Vec<u8> = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity {
path: tls.key_path.clone(),
message: source.to_string(),
})?;
Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| {
Error::TlsIdentity {
path: tls.cert_path.clone(),
message: source.to_string(),
}
})
})
.transpose()
}
fn map_api_error(error: ClientError, context: ErrorContext) -> Error {
match error {
ClientError::APIError { code, .. } => match context {
ErrorContext::Login => Error::LoginStatus { status: code },
ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code },
},
ClientError::JsonParseError { source } => match context {
ErrorContext::Login => Error::MalformedLogin,
ErrorContext::Read => Error::MalformedPayload,
ErrorContext::Secret => Error::Client(ClientError::JsonParseError { source }),
},
ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError => {
malformed_response(context)
}
ClientError::RestClientError { source } => match source {
RustifyClientError::ServerResponseError { code, .. } => match context {
ErrorContext::Login => Error::LoginStatus { status: code },
ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code },
},
RustifyClientError::ResponseParseError { .. } => malformed_response(context),
source => Error::Client(ClientError::RestClientError { source }),
},
error => Error::Client(error),
}
}
fn api_status(error: &ClientError) -> Option<u16> {
match error {
ClientError::APIError { code, .. } => Some(*code),
ClientError::RestClientError {
source: RustifyClientError::ServerResponseError { code, .. },
} => Some(*code),
_ => None,
}
}
fn malformed_response(context: ErrorContext) -> Error {
match context {
ErrorContext::Login => Error::MalformedLogin,
ErrorContext::Read => Error::MalformedPayload,
ErrorContext::Secret => Error::Client(ClientError::ResponseDataEmptyError),
}
}
fn token_expiry(lease_duration: u64) -> Option<Instant> {
(lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration))
}

602
litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs Normal file
View file

@ -0,0 +1,602 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_core_utils::settings::Lookup;
use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig};
use litellm_secrets_types::SecretValue;
use serde::Deserialize;
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_json, header, method, path},
};
fn config(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVaultConfig {
let mut environment_values: HashMap<String, String> = values
.iter()
.map(|(name, value)| ((*name).to_owned(), (*value).to_owned()))
.collect();
environment_values.insert("HCP_VAULT_ADDR".to_owned(), server.uri());
let environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| environment_values.get(name).cloned());
HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap()
}
fn manager(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVault {
HashicorpVault::from_config(config(server, values), true).unwrap()
}
fn auth_response(token: &str, lease_duration: u64) -> serde_json::Value {
json!({
"auth": {
"client_token": token,
"accessor": "",
"policies": [],
"token_policies": [],
"metadata": null,
"lease_duration": lease_duration,
"renewable": false,
"entity_id": "",
"token_type": "service",
"orphan": false
},
"lease_id": "",
"lease_duration": lease_duration,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})
}
fn read_response(data: serde_json::Value) -> serde_json::Value {
json!({
"data": {
"data": data,
"metadata": {
"created_time": "",
"deletion_time": "",
"custom_metadata": null,
"destroyed": false,
"version": 1
}
},
"lease_id": "",
"lease_duration": 0,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})
}
#[tokio::test]
async fn token_reads_use_vault_headers_and_cache_values() {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.and(header("X-Vault-Token", "token"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]);
assert_eq!(
manager
.async_read_secret("name")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
let requests = server.received_requests().await.unwrap();
assert!(
requests
.iter()
.all(|request| !request.headers.contains_key("X-Vault-Namespace"))
);
assert_eq!(
manager
.async_read_secret("name")
.await
.unwrap()
.unwrap()
.expose(),
"value"
);
}
#[tokio::test]
async fn namespace_mount_and_prefix_are_sanitized_in_the_url() {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/kv-prod/data/virtual-keys/name"))
.and(header("X-Vault-Namespace", "team-a"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(
&server,
&[
("HCP_VAULT_TOKEN", "token"),
("HCP_VAULT_SECRET_NAMESPACE", " /team-a/ "),
("HCP_VAULT_MOUNT_NAME", " /kv-prod/ "),
("HCP_VAULT_PATH_PREFIX", " /virtual-keys/ "),
],
);
let location = manager.secret_location("name").unwrap();
assert_eq!(location.namespace.as_deref(), Some("team-a"));
assert_eq!(location.mount, "kv-prod");
assert_eq!(location.path, "virtual-keys/name");
assert!(manager.async_read_secret("name").await.unwrap().is_some());
}
#[test]
fn trailing_address_slashes_are_removed() {
let environment: Arc<dyn Lookup + Send + Sync> = Arc::new(|name: &str| match name {
"HCP_VAULT_ADDR" => Some("http://vault.test:8200///".to_owned()),
"HCP_VAULT_TOKEN" => Some("token".to_owned()),
_ => None,
});
let config: HashicorpVaultConfig =
HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap();
let manager: HashicorpVault = HashicorpVault::from_config(config, true).unwrap();
assert_eq!(
manager.secret_location("name").unwrap(),
litellm_secrets_hashicorp::SecretLocation {
namespace: None,
mount: "secret".to_owned(),
path: "name".to_owned(),
}
);
}
#[rstest::rstest]
#[case("-1")]
#[case("not-a-number")]
fn invalid_refresh_intervals_are_rejected(#[case] value: &str) {
let environment: Arc<dyn Lookup + Send + Sync> = Arc::new(move |name: &str| match name {
"HCP_VAULT_REFRESH_INTERVAL" => Some(value.to_owned()),
_ => None,
});
assert!(matches!(
HashicorpVaultConfig::from_environment(environment.as_ref()),
Err(Error::RefreshInterval)
));
}
#[tokio::test]
async fn approle_login_uses_namespace_and_reuses_the_token() {
let server: MockServer = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/auth/custom-approle/login"))
.and(header("X-Vault-Namespace", "login-root"))
.and(body_json(json!({"role_id": "role", "secret_id": "secret"})))
.respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 3600)))
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.and(header("X-Vault-Token", "login-token"))
.and(header("X-Vault-Namespace", "secret-root"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(1)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name-2"))
.respond_with(ResponseTemplate::new(404).set_body_json(json!({"errors": ["missing"]})))
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(
&server,
&[
("HCP_VAULT_APPROLE_ROLE_ID", "role"),
("HCP_VAULT_APPROLE_SECRET_ID", "secret"),
("HCP_VAULT_APPROLE_MOUNT_PATH", "custom-approle"),
("HCP_VAULT_NAMESPACE", "secret-root"),
("HCP_VAULT_LOGIN_NAMESPACE", "login-root"),
],
);
assert!(manager.async_read_secret("name").await.unwrap().is_some());
assert!(manager.async_read_secret("name-2").await.unwrap().is_none());
}
#[tokio::test]
async fn approle_tokens_expire_after_the_vault_lease() {
let server: MockServer = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/auth/approle/login"))
.respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 1)))
.expect(2)
.mount(&server)
.await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(2)
.mount(&server)
.await;
let manager: HashicorpVault = manager(
&server,
&[
("HCP_VAULT_APPROLE_ROLE_ID", "role"),
("HCP_VAULT_APPROLE_SECRET_ID", "secret"),
("HCP_VAULT_REFRESH_INTERVAL", "0"),
],
);
assert!(manager.async_read_secret("first").await.unwrap().is_some());
tokio::time::sleep(Duration::from_secs(1) + Duration::from_millis(50)).await;
assert!(manager.async_read_secret("second").await.unwrap().is_some());
}
#[tokio::test]
async fn tls_login_posts_the_role_and_uses_the_client_identity() {
let server: MockServer = MockServer::start().await;
let directory: tempfile::TempDir = tempfile::tempdir().unwrap();
let cert_path = directory.path().join("client.crt");
let key_path = directory.path().join("client.key");
std::fs::write(&cert_path, TEST_CERTIFICATE).unwrap();
std::fs::write(&key_path, TEST_PRIVATE_KEY).unwrap();
Mock::given(method("POST"))
.and(path("/v1/auth/cert/login"))
.and(header("X-Vault-Namespace", "login-ns"))
.respond_with(ResponseTemplate::new(200).set_body_json(auth_response("cert-token", 0)))
.expect(2)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.and(header("X-Vault-Token", "cert-token"))
.and(header("X-Vault-Namespace", "secret-ns"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(2)
.mount(&server)
.await;
let role_values: HashMap<String, String> = HashMap::from([
("HCP_VAULT_ADDR".to_owned(), server.uri()),
(
"HCP_VAULT_CLIENT_CERT".to_owned(),
cert_path.to_str().unwrap().to_owned(),
),
(
"HCP_VAULT_CLIENT_KEY".to_owned(),
key_path.to_str().unwrap().to_owned(),
),
("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()),
(
"HCP_VAULT_LOGIN_NAMESPACE".to_owned(),
"login-ns".to_owned(),
),
(
"HCP_VAULT_SECRET_NAMESPACE".to_owned(),
"secret-ns".to_owned(),
),
]);
let role_environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| role_values.get(name).cloned());
let role_manager: HashicorpVault = HashicorpVault::new(role_environment, true).unwrap();
assert!(
role_manager
.async_read_secret("name")
.await
.unwrap()
.is_some()
);
let no_role_values: HashMap<String, String> = HashMap::from([
("HCP_VAULT_ADDR".to_owned(), server.uri()),
(
"HCP_VAULT_CLIENT_CERT".to_owned(),
cert_path.to_str().unwrap().to_owned(),
),
(
"HCP_VAULT_CLIENT_KEY".to_owned(),
key_path.to_str().unwrap().to_owned(),
),
(
"HCP_VAULT_LOGIN_NAMESPACE".to_owned(),
"login-ns".to_owned(),
),
(
"HCP_VAULT_SECRET_NAMESPACE".to_owned(),
"secret-ns".to_owned(),
),
]);
let no_role_environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| no_role_values.get(name).cloned());
let no_role_manager: HashicorpVault = HashicorpVault::new(no_role_environment, true).unwrap();
assert!(
no_role_manager
.async_read_secret("name")
.await
.unwrap()
.is_some()
);
let login_bodies: Vec<serde_json::Value> = server
.received_requests()
.await
.unwrap()
.iter()
.filter(|request| request.method.as_str() == "POST")
.map(|request| serde_json::from_slice(&request.body).unwrap())
.collect();
assert!(login_bodies.contains(&json!({"name": "vault-role"})));
assert!(login_bodies.contains(&json!({})));
}
#[rstest::rstest]
#[case::missing(404, json!({"errors": ["missing"]}), 0)]
#[case::malformed(200, json!({"data": "invalid"}), 1)]
#[case::missing_key(200, json!({}), 0)]
#[case::non_string(200, json!({"key": 1}), 2)]
#[tokio::test]
async fn read_responses_distinguish_absence_and_malformed_payloads(
#[case] status: u16,
#[case] body: serde_json::Value,
#[case] expected: u8,
) {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(ResponseTemplate::new(status).set_body_json(
if status == 200 && expected != 1 {
read_response(body)
} else {
body
},
))
.expect(1)
.mount(&server)
.await;
let result: Result<Option<SecretValue>, Error> =
manager(&server, &[("HCP_VAULT_TOKEN", "token")])
.async_read_secret("name")
.await;
match expected {
0 => assert!(result.unwrap().is_none()),
1 => assert!(matches!(result, Err(Error::MalformedPayload))),
2 => assert!(matches!(result, Err(Error::NonStringValue))),
_ => unreachable!(),
}
}
#[tokio::test]
async fn write_and_delete_invalidate_the_read_cache() {
let server: MockServer = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/name"))
.respond_with(
ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))),
)
.expect(2)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/v1/secret/data/name"))
.and(body_json(
json!({"data": {"key": "updated", "description": "description"}}),
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"data": {
"created_time": "",
"deletion_time": "",
"custom_metadata": null,
"destroyed": false,
"version": 2
},
"lease_id": "",
"lease_duration": 0,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("DELETE"))
.and(path("/v1/secret/data/name"))
.respond_with(ResponseTemplate::new(204))
.expect(1)
.mount(&server)
.await;
let manager: HashicorpVault = manager(&server, &[("HCP_VAULT_TOKEN", "token")]);
assert!(manager.async_read_secret("name").await.unwrap().is_some());
assert!(
manager
.async_write_secret("name", SecretValue::new("updated"), Some("description"))
.await
.is_ok()
);
assert!(manager.async_read_secret("name").await.unwrap().is_some());
manager.async_delete_secret("name").await.unwrap();
}
#[tokio::test]
async fn no_auth_and_invalid_names_fail_without_requests() {
let server: MockServer = MockServer::start().await;
let manager: HashicorpVault = manager(&server, &[]);
assert!(matches!(
manager.async_read_secret("name").await,
Err(Error::NoAuthConfigured)
));
assert!(matches!(
manager.async_read_secret("../name").await,
Err(Error::InvalidSecretName(_))
));
assert!(server.received_requests().await.unwrap().is_empty());
}
#[tokio::test]
async fn debug_output_redacts_authentication_values() {
let server: MockServer = MockServer::start().await;
let manager: HashicorpVault =
HashicorpVault::from_config(config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), true)
.unwrap();
let debug: String = format!("{manager:?}");
assert!(!debug.contains("token-value"));
assert!(!debug.contains("secret-id"));
}
#[derive(Deserialize)]
struct ParityCase {
env: HashMap<String, String>,
expected_secret_url: String,
expected_login_url: Option<String>,
expected_login_namespace: Option<String>,
expected_secret_namespace: Option<String>,
secret_name: String,
}
#[test]
fn configuration_matches_python_parity_fixture() {
let cases: Vec<ParityCase> = serde_json::from_str(include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../tests/test_litellm/secret_managers/hashicorp_vault_parity.json"
)))
.unwrap();
for case in cases {
let values: HashMap<String, String> = case.env.clone();
let environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(move |name: &str| values.get(name).cloned());
let config: HashicorpVaultConfig =
HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap();
let manager: HashicorpVault = HashicorpVault::from_config(config.clone(), true).unwrap();
let location = manager.secret_location(&case.secret_name).unwrap();
let namespace = location
.namespace
.as_deref()
.map(|namespace| format!("{namespace}/"))
.unwrap_or_default();
assert_eq!(
format!(
"{}/v1/{}{}/data/{}",
config.address, namespace, location.mount, location.path
),
case.expected_secret_url
);
let login_url = config.approle.as_ref().map_or_else(
|| {
config
.tls_cert
.as_ref()
.map(|_| format!("{}/v1/auth/cert/login", config.address))
},
|approle| {
Some(format!(
"{}/v1/auth/{}/login",
config.address, approle.mount_path
))
},
);
assert_eq!(login_url, case.expected_login_url);
assert_eq!(
manager.config().login_namespace(),
case.expected_login_namespace.as_deref()
);
assert_eq!(
manager.config().secret_namespace(),
case.expected_secret_namespace.as_deref()
);
}
}
#[tokio::test]
#[ignore]
async fn live_vault_round_trip() {
let environment: Arc<dyn Lookup + Send + Sync> =
Arc::new(litellm_core_utils::settings::ProcessEnvironment);
let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap();
let name: String = std::env::var("LITELLM_VAULT_LIVE_SECRET_NAME").unwrap();
let value: SecretValue = SecretValue::new("native-live-value");
let location = manager.secret_location(&name).unwrap();
println!(
"native provenance: {} vaultrs {} {:?} {} {}",
module_path!(),
manager.config().address,
location.namespace,
location.mount,
location.path
);
manager
.async_write_secret(&name, value.clone(), None)
.await
.unwrap();
assert_eq!(
manager.async_read_secret(&name).await.unwrap().unwrap(),
value
);
manager.async_delete_secret(&name).await.unwrap();
assert!(manager.async_read_secret(&name).await.unwrap().is_none());
}
const TEST_CERTIFICATE: &str = "-----BEGIN CERTIFICATE-----
MIIDDzCCAfegAwIBAgIUeMzLFLM/mRbPGbNAew5N2UTscocwDQYJKoZIhvcNAQEL
BQAwFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MB4XDTI2MDkyMTIwMjA1OVoXDTI2
MDkyMjIwMjA1OVowFzEVMBMGA1UEAwwMbGl0ZWxsbS10ZXN0MIIBIjANBgkqhkiG
9w0BAQEFAAOCAQ8AMIIBCgKCAQEAveYoSUJXybmkHmQsBfhBcv2Ob5Oy8ejZu+B3
vTnrPumW4ANi1XXKBSazRGB3fEtAgr+3KhKeHaSKEQeBwJkAEBfdmQv0tpXICwHs
1kFNtU0owy54HVW5/ia+LMszsFcPzVIoMnbUOuiKr9RaV7P+IEFzILPBVuV4DoYH
yocjD3+9QNqokWgNL8LK37JijmNEFVaKFz0X6SyL2VRDlfPWTEBK52Gp/pvDgA6G
eTSfyI+kCm9h5ECTYUAtmatk9WPVS8sWOqV1EXVanFyYBU+mDxoywAS1/6CHeIPh
bNmCOZjPoO9qWBJ7ZyGhOconBigXY8qnlXymev+44IPHrx4urwIDAQABo1MwUTAd
BgNVHQ4EFgQUvaZrZ6HKtbr3ekeZmgy4b5Pq95QwHwYDVR0jBBgwFoAUvaZrZ6HK
tbr3ekeZmgy4b5Pq95QwDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOC
AQEAEejrD8d1qDxW55XxQ4IC31rufoEvDV955jyvh2kALPaN/i5oWsBGI+UAQZna
aaoQXwzlmHrtDUBWl0LztVTUamIleUep2+PLLauqqt43vxppxMX8Jn2mnPO20YE/
hIzGx0jN/LBG8PDyLSvHdlgjP9ofA4Vg4rTQugdXRgOvlCE/epnH/MADcg9KYJtJ
C1RObCIkL3LcdUbjStJRCY/U/FeWcgyncEPz95OFDkbrlNDajb6o6CkYfouqvhTc
8XlgjjAVKIbAbRgbVu3elsquuFM97x2DzWDjkrMNmDt1FJ9ubK36gL6B3o0UMaoQ
00R7x/eqvH+EkWa/2ekW9lpleQ==
-----END CERTIFICATE-----
";
const TEST_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC95ihJQlfJuaQe
ZCwF+EFy/Y5vk7Lx6Nm74He9Oes+6ZbgA2LVdcoFJrNEYHd8S0CCv7cqEp4dpIoR
B4HAmQAQF92ZC/S2lcgLAezWQU21TSjDLngdVbn+Jr4syzOwVw/NUigydtQ66Iqv
1FpXs/4gQXMgs8FW5XgOhgfKhyMPf71A2qiRaA0vwsrfsmKOY0QVVooXPRfpLIvZ
VEOV89ZMQErnYan+m8OADoZ5NJ/Ij6QKb2HkQJNhQC2Zq2T1Y9VLyxY6pXURdVqc
XJgFT6YPGjLABLX/oId4g+Fs2YI5mM+g72pYEntnIaE5yicGKBdjyqeVfKZ6/7jg
g8evHi6vAgMBAAECggEAGdJjlP6b8Fa5bdaCM/ebcrbuuNZVJVbb0JPHxGfNSLs7
pE9hj5QaOdQW2Uviw3h6F61ZCzQH4xD+Iy2po5ZKb2XHYKnDB1bboj+LRGER337T
9aJqe9at2VTMVEv3Rdm40NsEk0QcPLxlK16NQFK90gYEUSSQPDAswJDSG2R/zHn+
vADI907mW/goEJHeLn8PWGlNlSiR6x+5JJtq+GXCzUzVvJYQSCLGxCSl2x2H+0g7
NhFI0zPpdzNmO/h+yhzaFb6Rp5U8+ZsnZ3qYjQ/03gw1myTDKJt1YaO9JvArnNYX
hcJQQ8Rt0bHhcrZA16bBOpqZlo5pKCicwI/netgN8QKBgQDcFz7AzdJ26sMSV32V
rwrMgIoggt8qDjO1ARwqW35A1TIge0FoW4M4KpsXQGGfT341uU1esXEcyZ/1L/5X
3ql2gX4DbOYLZLWYzZGR2hq33oi8HkhN98QrEwL9emSH8NqYX3Xxja3PrmCrSYJe
Zbnd9TIm2XkxyMoyXJu6M/QvnwKBgQDc4dzqTbxoGEGa5MuJoGmMwPnqgdG9UM5J
eExVnh7osxc2sOdsiPeRjjQTxs9v2kJwctC359OJoo9yGaaJeSghU4LEWJo1sqnA
fzSCLammYvtVAtniyNv5Mxk/6Uimi4NNDKaAKB+m4K2uSn3U9AmY7KPYMGaSbS9W
XSnobjxm8QKBgC8bPpAvvWs8ZhIn7bY659nLbUT2HeO3dHO6UBf0yzn/J6JyHxbB
93zvCZDZc8uQTRgcmCW7XtVlhjoJUqvl+Wlm39zF0xr/LCsPXKfWAb/2/lcdOCaP
8Emz4QD10EyUTYUtcWYJB/mafhBLRH8F0Nlj4J8WDu2L51MOJTqeYhZLAoGAWffN
icocAbJPlo22sdoa4+/+W5yBF8GAJMDRJtZ+9H1t6SLpQHYRkMIBSETkXUTjZvX9
Ocs9iIQkNW9pO/mTdO+VBfCo71JUfknR02xR+6m5gYjlws/ZeYlssXGN2/hbhNiw
QOcW7Vv6olFJK6Iy/oz0t6wPO3kpnN3Zogi0paECgYEAwo44M1DdYCtV0snhmYM9
5u0mPfYt5P2SVLXyUbr+vFTfrTL/WKnXIJgbsnj3Gvf+GIZv9tKcXhSNmEHQCYX4
X3w9iTPddCHuvZ1fpufi2TyArJh0OkoNtLXJHTKrHjf2N+61AQzFiv5WieJrdE+H
qr32PTUuVGPyO9LyTY4/RL0=
-----END PRIVATE KEY-----
";

2
litellm-rust/crates/secrets/Cargo.toml
View file

@ -9,6 +9,7 @@ repository.workspace = true
default = []
aws = ["dep:litellm-secrets-aws"]
google = ["dep:litellm-secrets-google"]
hashicorp = ["dep:litellm-secrets-hashicorp"]
azure = ["dep:litellm-secrets-azure"]
cyberark = ["dep:litellm-secrets-cyberark"]
@ -16,6 +17,7 @@ cyberark = ["dep:litellm-secrets-cyberark"]
litellm-secrets-types.workspace = true
litellm-secrets-aws = { workspace = true, optional = true }
litellm-secrets-google = { workspace = true, optional = true }
litellm-secrets-hashicorp = { workspace = true, optional = true }
litellm-secrets-azure = { workspace = true, optional = true }
litellm-secrets-cyberark = { workspace = true, optional = true }
litellm-core-utils.workspace = true

2
litellm-rust/crates/secrets/README.md
View file

@ -9,3 +9,5 @@ Backend failures propagate by default. To allow fallback during a backend failur
`get_secret` preserves value types. `get_secret_str` accepts a string default and rejects boolean or JSON values with `Error::TypeMismatch`. `get_secret_bool` accepts a boolean default and converts strings containing `true` or `false`, ignoring surrounding whitespace and ASCII case. Other strings and JSON values produce `Error::TypeMismatch`. Conversion failures never activate fallback or replace a found value with the default
Provider payloads remain strings unless explicitly selecting a field from an AWS primary JSON secret. Google caches only successfully decoded string payloads, so reads have identical values and types before and after caching. Confirmed absence and failed reads are not cached. AWS resource-not-found responses and Google HTTP 404 responses indicate absence. Other provider errors remain errors, and successful responses without the required payload are malformed responses rather than missing secrets
The HashiCorp Vault backend is enabled with the `hashicorp` feature and reads KV v2 values from `HCP_VAULT_*` environment variables. It supports static tokens, AppRole authentication, and TLS certificate authentication

3
litellm-rust/crates/secrets/src/error.rs
View file

@ -30,6 +30,9 @@ pub enum Error {
#[cfg(feature = "google")]
#[error(transparent)]
Google(#[from] litellm_secrets_google::Error),
#[cfg(feature = "hashicorp")]
#[error(transparent)]
Hashicorp(#[from] litellm_secrets_hashicorp::Error),
#[cfg(feature = "azure")]
#[error(transparent)]
Azure(#[from] litellm_secrets_azure::Error),

10
litellm-rust/crates/secrets/src/handler.rs
View file

@ -13,6 +13,8 @@ pub enum SecretManager {
GoogleKms(crate::google::GoogleKms),
#[cfg(feature = "google")]
GoogleSecretManager(crate::google::GoogleSecretManager),
#[cfg(feature = "hashicorp")]
HashicorpVault(crate::hashicorp::HashicorpVault),
#[cfg(feature = "azure")]
AzureKeyVault(crate::azure::AzureKeyVault),
#[cfg(feature = "cyberark")]
@ -31,6 +33,8 @@ impl SecretManager {
Self::GoogleKms(_) => KeyManagementSystem::GoogleKms,
#[cfg(feature = "google")]
Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager,
#[cfg(feature = "hashicorp")]
Self::HashicorpVault(_) => KeyManagementSystem::HashicorpVault,
#[cfg(feature = "azure")]
Self::AzureKeyVault(_) => KeyManagementSystem::AzureKeyVault,
#[cfg(feature = "cyberark")]
@ -86,6 +90,12 @@ pub async fn get_secret_from_manager(
.get_secret_from_google_secret_manager(secret_name)
.await
.map_err(Error::from),
#[cfg(feature = "hashicorp")]
SecretManager::HashicorpVault(client) => client
.async_read_secret(secret_name)
.await
.map(|value| value.map(Secret::String))
.map_err(Error::from),
#[cfg(feature = "azure")]
SecretManager::AzureKeyVault(client) => client
.get_secret_from_azure_key_vault(secret_name)

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

@ -23,3 +23,5 @@ pub use litellm_secrets_azure as azure;
pub use litellm_secrets_cyberark as cyberark;
#[cfg(feature = "google")]
pub use litellm_secrets_google as google;
#[cfg(feature = "hashicorp")]
pub use litellm_secrets_hashicorp as hashicorp;

142
litellm-rust/crates/secrets/tests/handler.rs
View file

@ -105,6 +105,148 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites
Err(Error::MissingCiphertext)
));
}
#[cfg(feature = "hashicorp")]
#[tokio::test]
async fn hashicorp_handler_resolves_found_missing_and_failed_values() {
use std::sync::Arc;
use litellm_core_utils::settings::Lookup;
use litellm_secrets::{
Error, FailurePolicy, KeyManagementSettings, SecretManager, SecretManagerState,
SecretResolver, hashicorp::HashicorpVault, hashicorp::HashicorpVaultConfig,
};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path},
};
let found_server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/v1/secret/data/KEY"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"data": {
"data": {"key": "remote"},
"metadata": {
"created_time": "",
"deletion_time": "",
"custom_metadata": null,
"destroyed": false,
"version": 1
}
},
"lease_id": "",
"lease_duration": 0,
"renewable": false,
"request_id": "",
"warnings": null,
"wrap_info": null
})))
.mount(&found_server)
.await;
let found_environment: Arc<dyn Lookup + Send + Sync> = Arc::new({
let address = found_server.uri();
move |name: &str| match name {
"HCP_VAULT_ADDR" => Some(address.clone()),
"HCP_VAULT_TOKEN" => Some("token".into()),
_ => None,
}
});
let found_config = HashicorpVaultConfig::from_environment(found_environment.as_ref()).unwrap();
let found_manager = HashicorpVault::from_config(found_config, true).unwrap();
let found_resolver = SecretResolver::new(
Arc::new(SecretManagerState::new(
SecretManager::HashicorpVault(found_manager),
KeyManagementSettings {
hosted_keys: Some(vec!["KEY".into()]),
..Default::default()
},
)),
Arc::new(|_: &str| None),
litellm_secrets::OidcResolver::default(),
);
assert_eq!(
found_resolver
.get_secret_str("KEY", None)
.await
.unwrap()
.unwrap()
.expose(),
"remote"
);
let missing_server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(404).set_body_json(serde_json::json!({"errors": ["missing"]})),
)
.mount(&missing_server)
.await;
let missing_environment: Arc<dyn Lookup + Send + Sync> = Arc::new({
let address = missing_server.uri();
move |name: &str| match name {
"HCP_VAULT_ADDR" => Some(address.clone()),
"HCP_VAULT_TOKEN" => Some("token".into()),
_ => None,
}
});
let missing_config =
HashicorpVaultConfig::from_environment(missing_environment.as_ref()).unwrap();
let missing_manager = HashicorpVault::from_config(missing_config, true).unwrap();
let missing_state = SecretManagerState::new(
SecretManager::HashicorpVault(missing_manager),
KeyManagementSettings {
hosted_keys: Some(vec!["KEY".into()]),
..Default::default()
},
);
let missing = litellm_secrets::get_secret_from_manager(
missing_state.backend().unwrap(),
"KEY",
missing_state.settings().unwrap(),
&|_: &str| None,
)
.await
.unwrap();
assert!(missing.is_none());
let failed_server = MockServer::start().await;
Mock::given(method("GET"))
.respond_with(
ResponseTemplate::new(500).set_body_json(serde_json::json!({"errors": ["failed"]})),
)
.mount(&failed_server)
.await;
let failed_environment: Arc<dyn Lookup + Send + Sync> = Arc::new({
let address = failed_server.uri();
move |name: &str| match name {
"HCP_VAULT_ADDR" => Some(address.clone()),
"HCP_VAULT_TOKEN" => Some("token".into()),
_ => None,
}
});
let failed_config =
HashicorpVaultConfig::from_environment(failed_environment.as_ref()).unwrap();
let failed_manager = HashicorpVault::from_config(failed_config, true).unwrap();
let failed_state = SecretManagerState::new(
SecretManager::HashicorpVault(failed_manager),
KeyManagementSettings {
hosted_keys: Some(vec!["KEY".into()]),
..Default::default()
},
);
let failed_resolver = SecretResolver::new(
Arc::new(failed_state),
Arc::new(|_: &str| None),
litellm_secrets::OidcResolver::default(),
)
.with_failure_policy(FailurePolicy::Propagate);
assert!(matches!(
failed_resolver.get_secret_str("KEY", None).await,
Err(Error::Hashicorp(
litellm_secrets::hashicorp::Error::Status { status: 500 }
))
));
}
#[cfg(feature = "azure")]
#[tokio::test]

12
litellm/constants.py
View file

@ -1749,6 +1749,12 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float(
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300"))
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30"))
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000"))
SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float(
os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5")
)
SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float(
os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5")
)
TOOL_SPEND_TOP_TOOLS: Final = 100
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
@ -1860,6 +1866,12 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"max_ui_session_budget",
"budget_rollover",
"mcp_tool_search",
"turn_off_message_logging",
"datadog_params",
"datadog_llm_observability_params",
"newrelic_params",
"pointfive_params",
"aws_sqs_callback_params",
]
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))

4
litellm/litellm_core_utils/logging_worker.py
View file

@ -484,9 +484,13 @@ class LoggingWorker:
so it correctly handles items that have been dequeued but whose
callback hasn't finished yet — ``queue.empty()`` would return True in
that window and cause us to skip the wait.
``start()`` runs first so a queue left behind by a previous event loop
is carried onto this one and drained here instead of joined forever.
"""
if self._queue is None:
return
self.start()
await self._queue.join()
async def clear_queue(self):

79
litellm/model_prices_and_context_window_backup.json
View file

@ -43037,21 +43037,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.83746e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.784544e-06,
"output_cost_per_token": 1.767492e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.4356e-08,
"cache_read_input_token_cost": 7.36455e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -70299,8 +70299,8 @@
"supports_web_search": true
},
"openrouter/meta-llama/llama-4-maverick": {
"input_cost_per_token": 2e-07,
"output_cost_per_token": 8e-07,
"input_cost_per_token": 1.875e-07,
"output_cost_per_token": 6.525e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 16384,
@ -73794,6 +73794,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_above_128k_tokens": 5e-07,
"litellm_provider": "openrouter",
@ -73815,6 +73816,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6-flash": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1e-07,
"litellm_provider": "openrouter",
@ -73855,6 +73857,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-2.0-code": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 5e-07,
"input_cost_per_token_above_128k_tokens": 1e-06,
"litellm_provider": "openrouter",
@ -76876,6 +76879,26 @@
"supports_vision": true,
"supports_web_search": true
},
"moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
@ -76975,5 +76998,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

11
litellm/proxy/auth/handle_jwt.py
View file

@ -14,7 +14,7 @@ import hashlib
import os
import re
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
@ -1602,6 +1602,7 @@ class JWTAuthManager:
team_object: LiteLLM_TeamTable | None,
route: str,
request_method: str | None = None,
team_allowed_routes: Collection[str] = (),
) -> bool:
normalized_request_method: Final = request_method.upper() if isinstance(request_method, str) else None
if not RouteChecks.is_auth_enforced_pass_through_route(
@ -1610,8 +1611,11 @@ class JWTAuthManager:
):
return True
if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes):
return True
# JWT team selection is team-scoped; key metadata is not available here,
# so passthrough access is granted only by the selected team's metadata.
# so beyond the JWT config grant above, only the selected team's metadata grants access.
return RouteChecks.check_passthrough_route_access(
route=route,
user_api_key_dict=UserAPIKeyAuth(team_metadata=(team_object.metadata or {}) if team_object else {}),
@ -1689,6 +1693,7 @@ class JWTAuthManager:
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes,
):
is_allowed = False
denied_auth_enforced_pass_through_route = True
@ -2584,6 +2589,7 @@ class JWTAuthManager:
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes,
):
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
@ -2653,6 +2659,7 @@ class JWTAuthManager:
team_object=team_object,
route=route,
request_method=request_method,
team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes,
):
JWTAuthManager._raise_team_passthrough_route_denial(route=route)
elif team_id is None:

35
litellm/proxy/auth/route_checks.py
View file

@ -278,7 +278,11 @@ class RouteChecks:
route=route,
method=RouteChecks._get_request_method(request=request),
):
RouteChecks._require_auth_pass_through_access(route=route, valid_token=valid_token)
RouteChecks._require_auth_pass_through_access(
route=route,
valid_token=valid_token,
jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token),
)
elif RouteChecks.is_llm_api_route(route=route):
pass
elif RouteChecks.is_info_route(route=route):
@ -689,16 +693,43 @@ class RouteChecks:
),
)
@staticmethod
def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool:
"""
Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do:
a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names
no path segment (``*``, ``/*``) is skipped.
"""
return any(
RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route)
for allowed_route in team_allowed_routes
if allowed_route.rstrip("*").strip("/")
)
@staticmethod
def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]:
"""``team_allowed_routes`` for team tokens built by JWT auth; JWT-mapped virtual keys stay key-scoped."""
if valid_token.jwt_claims is None or valid_token.token is not None or valid_token.team_id is None:
return ()
from litellm.proxy.proxy_server import jwt_handler
return jwt_handler.litellm_jwtauth.team_allowed_routes
@staticmethod
def _require_auth_pass_through_access(
route: str,
valid_token: UserAPIKeyAuth,
jwt_team_allowed_routes: Collection[str] = (),
) -> None:
"""
Require an explicit ``allowed_passthrough_routes`` match for auth=true pass-through.
Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the
key or team, or an explicit JWT ``team_allowed_routes`` entry.
"""
if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token):
return
if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes):
return
raise RouteChecks._auth_pass_through_denied_exception(route=route)
@staticmethod

15
litellm/proxy/config_resolvers/__init__.py
View file

@ -5,6 +5,17 @@ from litellm.proxy.config_resolvers._descriptors import (
FieldSource,
resolve_fields,
)
from litellm.proxy.config_resolvers.settings_store import SettingsStore, config_ownership_message
from litellm.proxy.config_resolvers.settings_store import (
SettingsStore,
config_ownership_message,
source_for,
)
__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "config_ownership_message", "resolve_fields")
__all__ = (
"FieldDescriptor",
"FieldSource",
"SettingsStore",
"config_ownership_message",
"resolve_fields",
"source_for",
)

7
litellm/proxy/config_resolvers/settings_store.py
View file

@ -176,3 +176,10 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def _resolution_for(self, key: str) -> Resolved:
yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT)
return resolve(yaml_value, self._db_value(key))
def source_for(settings: SettingsStore, key: str, default: object = None) -> FieldSource:
source: Final = settings.source(key)
if source == "unset":
return "default" if default is not None else "unset"
return source

170
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -4,7 +4,7 @@ Per-session auto-router benchmarks rollup.
At request time the spend writer builds one AutoRouterTurnTransaction per successful
auto-routed request (a request whose metadata carries a routing_decision) and queues it
on the prisma client. The spend-log flush job drains the queue into
LiteLLM_AutoRouterSession with one conditional upsert per turn: the statement classifies
key and user session rollups with one atomic statement per turn: each upsert classifies
the turn (same model, first visit, return to a model the session already used, out of
order) against the row's own columns, so nothing is read before the write and concurrent
pods compose. The benchmarks endpoint aggregates these rows and never touches
@ -35,10 +35,27 @@ if TYPE_CHECKING:
CACHE_TTL_5M_SECONDS: Final = 300
CACHE_TTL_1H_SECONDS: Final = 3600
AUTOROUTER_BENCHMARKS_SQL: Final = """
_SESSION_COLUMNS: Final = """
api_key, session_id, router_name, router_type, first_turn_at, last_turn_at,
last_model, models, turns, unordered_turns, covered_turns, cache_hits,
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns,
baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend,
savings_estimated_baseline_models
"""
AUTOROUTER_BENCHMARKS_SQL: Final = f"""
WITH windowed AS (
SELECT * FROM "LiteLLM_AutoRouterSession"
WHERE last_turn_at >= $1::timestamp
SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterSession"
WHERE $4::text IS NULL
AND last_turn_at >= $1::timestamp
AND first_turn_at < $2::timestamp
AND ($3::text IS NULL OR api_key = $3::text)
UNION ALL
SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterUserSession"
WHERE (($4::text IS NOT NULL AND user_id = $4::text) OR ($4::text IS NULL AND api_key = ''))
AND last_turn_at >= $1::timestamp
AND first_turn_at < $2::timestamp
AND ($3::text IS NULL OR api_key = $3::text)
),
@ -53,7 +70,7 @@ tier_maps AS (
)
SELECT
agg.*,
COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns
COALESCE(tier_maps.tier_turns, '{{}}'::jsonb) AS tier_turns
FROM (
SELECT
router_name,
@ -111,6 +128,7 @@ class AutoRouterTurnTransaction:
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
user_id: str = ""
class TurnCacheFacts(NamedTuple):
@ -214,10 +232,11 @@ def build_autorouter_turn_transaction(
if not isinstance(routing_decision, Mapping) or not routing_decision:
return None
router_name: Final = routing_decision.get("router_model_name") or payload.get("model_group")
api_key: Final = payload.get("api_key")
api_key: Final = payload.get("api_key") or ""
user_id: Final = payload.get("user") or ""
session_id: Final = payload.get("session_id")
model: Final = payload.get("model")
if not (isinstance(router_name, str) and router_name and api_key and session_id and model):
if not (isinstance(router_name, str) and router_name and (api_key or user_id) and session_id and model):
return None
turn_at: Final = _turn_time_utc(str(payload.get("startTime") or ""))
if turn_at is None:
@ -236,6 +255,7 @@ def build_autorouter_turn_transaction(
estimated_savings: Final = recorded_estimated_autorouter_savings(metadata)
return AutoRouterTurnTransaction(
api_key=api_key,
user_id=user_id,
session_id=bounded_session_id(session_id),
router_name=router_name,
router_type=str(routing_decision.get("router_type") or "unknown"),
@ -293,18 +313,18 @@ _RETURN_MISS: Final = (
_IDLE_SECONDS: Final = f"EXTRACT(EPOCH FROM {_TURN_AT}::timestamp) - (t.models -> {_MODEL} ->> 'at')::float8"
_CACHE_TOUCHED: Final = f"{_TOUCHED}::int = 1"
UPSERT_AUTOROUTER_SESSION_SQL: Final = f"""
INSERT INTO "LiteLLM_AutoRouterSession" AS t (
api_key, session_id, router_name, router_type, first_turn_at, last_turn_at,
last_model, models, turns, unordered_turns, covered_turns, cache_hits,
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns,
baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend,
savings_estimated_baseline_models
def _session_upsert_sql(*, user_scoped: bool) -> str:
table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession"
user_column: Final = "user_id, " if user_scoped else ""
user_value: Final = f"{_p('user_id')}::text, " if user_scoped else ""
required_identity: Final = _p("user_id" if user_scoped else "api_key")
return f"""
INSERT INTO "{table_name}" AS t (
{user_column}{_SESSION_COLUMNS}
)
VALUES (
{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
SELECT
{user_value}{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
{_MODEL}, jsonb_build_object({_MODEL}, jsonb_build_object('at', EXTRACT(EPOCH FROM {_TURN_AT}::timestamp), 'ttl', {_CACHE_TTL}::int)),
1, 0, {_COVERED}::int, {_CACHE_HIT}::int,
0, 0, 1, {_CACHE_HIT}::int,
@ -315,8 +335,8 @@ VALUES (
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA},
{_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8,
{_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA}
)
ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
WHERE {required_identity}::text <> ''
ON CONFLICT ({user_column}api_key, session_id, router_name) DO UPDATE SET
turns = t.turns + 1,
total_tokens = t.total_tokens + EXCLUDED.total_tokens,
spend = t.spend + EXCLUDED.spend,
@ -365,6 +385,17 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
"""
UPSERT_AUTOROUTER_SESSION_SQL: Final = f"""
WITH key_rollup AS (
{_session_upsert_sql(user_scoped=False)}
RETURNING 1
)
{_session_upsert_sql(user_scoped=True)}
"""
UPSERT_AUTOROUTER_USER_SESSION_SQL: Final = _session_upsert_sql(user_scoped=True)
def _as_sql_param(value: str | float | bool | datetime | None) -> str | float | None:
if isinstance(value, bool):
return int(value)
@ -377,18 +408,23 @@ def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float
return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS)
async def write_autorouter_turn(db: SupportsExecuteRaw, transaction: AutoRouterTurnTransaction) -> None:
await db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction))
async def write_autorouter_turn(
db: SupportsExecuteRaw,
transaction: AutoRouterTurnTransaction,
statement: str = UPSERT_AUTOROUTER_SESSION_SQL,
) -> None:
await db.execute_raw(statement, *_upsert_params(transaction))
async def _upsert_turn_with_retry(
prisma_client: PrismaClient,
transaction: AutoRouterTurnTransaction,
n_retry_times: int,
statement: str,
) -> None:
for attempt in range(n_retry_times + 1):
try:
await write_autorouter_turn(prisma_client.db, transaction)
await write_autorouter_turn(prisma_client.db, transaction, statement)
except DB_RETRY_SAFE_ERROR_TYPES:
if attempt >= n_retry_times:
raise
@ -397,6 +433,58 @@ async def _upsert_turn_with_retry(
return
def _session_partition(transaction: AutoRouterTurnTransaction) -> tuple[str, str, str, str]:
identity: Final = ("key", transaction.api_key) if transaction.api_key else ("user", transaction.user_id)
return (*identity, transaction.session_id, transaction.router_name)
async def _drain_session_partition(
prisma_client: PrismaClient,
transactions: tuple[AutoRouterTurnTransaction, ...],
n_retry_times: int,
statement: str,
) -> tuple[AutoRouterTurnTransaction, ...]:
for position, transaction in enumerate(transactions):
try:
await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times, statement)
except Exception as flush_err: # noqa: BLE001 # stop dependent turns without retrying an ambiguous write
verbose_proxy_logger.error(
"Spend tracking - auto-router session rollup flush failed for router %s; "
"%s of %s turn writes stopped in this partition: %s",
transaction.router_name,
len(transactions) - position,
len(transactions),
flush_err,
)
return transactions[position:]
return ()
async def _flush_session_partition(
prisma_client: PrismaClient,
transactions: tuple[AutoRouterTurnTransaction, ...],
n_retry_times: int,
) -> None:
failed_suffix: Final = await _drain_session_partition(
prisma_client, transactions, n_retry_times, UPSERT_AUTOROUTER_SESSION_SQL
)
if not failed_suffix or not failed_suffix[0].api_key:
return
failed_user: Final = failed_suffix[0].user_id
other_users: Final = sorted(
(
transaction
for transaction in failed_suffix[1:]
if transaction.user_id and transaction.user_id != failed_user
),
key=lambda transaction: transaction.user_id,
)
for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id):
await _drain_session_partition(
prisma_client, tuple(user_turns), n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL
)
async def flush_autorouter_turn_transactions(
prisma_client: PrismaClient,
transactions: Sequence[AutoRouterTurnTransaction],
@ -407,38 +495,20 @@ async def flush_autorouter_turn_transactions(
Statements run sequentially in per-session event order: a turn's classification
depends on the turns before it, and Postgres rejects one multi-row INSERT touching
the same key twice. Only ConnectError is retried, per statement, because it proves
that statement never reached the database. Any other failure drops the remaining
turns of THAT session only, with an error log, and the flush continues with the
next session: sessions are independent state machines, so one poisoned statement
must not discard unrelated sessions, and a repeated increment is worse than an
undercount. Callers must not add their own retry around this function.
that statement never reached the database. A failed write stops its key and user
histories for this batch. Other users sharing that key can still advance their
independent user histories, with the key projection disabled and the real key
identity preserved. The failed turn is never replayed. Callers must not add their
own retry around this function.
"""
if not transactions:
return
ordered: Final = sorted(
transactions,
key=lambda transaction: (
transaction.api_key,
transaction.session_id,
transaction.router_name,
transaction.turn_at,
),
key=lambda transaction: (*_session_partition(transaction), transaction.turn_at),
)
for session_key, session_group in groupby(
for _, session_group in groupby(
ordered,
key=lambda transaction: (transaction.api_key, transaction.session_id, transaction.router_name),
key=_session_partition,
):
session_turns = tuple(session_group)
for position, transaction in enumerate(session_turns):
try:
await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times)
except Exception as flush_err: # noqa: BLE001 # a statement failure drops only its session's remainder by design
verbose_proxy_logger.error(
"Spend tracking - auto-router session rollup flush failed for router %s; "
"%s of %s turn transactions dropped for one session: %s",
session_key[2],
len(session_turns) - position,
len(session_turns),
flush_err,
)
break
await _flush_session_partition(prisma_client, tuple(session_group), n_retry_times)

40
litellm/proxy/db/baseline_accounting.py
View file

@ -171,6 +171,7 @@ class _Change(BaseModel):
request_id: str
publication: BaselinePublication
api_key: str
user_id: str = ""
session_id: str
router_name: str
baseline_model: str
@ -256,42 +257,54 @@ SET publication = x.publication::text
FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
WHERE observations.request_id = x.request_id
"""
_UPDATE_SESSIONS: Final = """
def _session_correction_sql(*, user_scoped: bool) -> str:
table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession"
identity_columns: Final = ("user_id, " if user_scoped else "") + "api_key, session_id, router_name"
user_filter: Final = "WHERE user_id <> ''" if user_scoped else ""
user_match: Final = "session.user_id = totals.user_id AND " if user_scoped else ""
return f"""
WITH changes AS (
SELECT * FROM jsonb_to_recordset($1::jsonb) AS x(
api_key text, session_id text, router_name text, baseline_model text,
user_id text, api_key text, session_id text, router_name text, baseline_model text,
covered_delta int, actual_delta float8, savings_delta float8
)
{user_filter}
), totals AS (
SELECT api_key, session_id, router_name, SUM(covered_delta)::int AS covered_delta,
SELECT {identity_columns}, SUM(covered_delta)::int AS covered_delta,
SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta
FROM changes GROUP BY api_key, session_id, router_name
FROM changes GROUP BY {identity_columns}
), models AS (
SELECT api_key, session_id, router_name, jsonb_object_agg(baseline_model, delta) AS deltas
SELECT {identity_columns}, jsonb_object_agg(baseline_model, delta) AS deltas
FROM (
SELECT api_key, session_id, router_name, baseline_model, SUM(covered_delta)::int AS delta
FROM changes GROUP BY api_key, session_id, router_name, baseline_model
) grouped GROUP BY api_key, session_id, router_name
SELECT {identity_columns}, baseline_model, SUM(covered_delta)::int AS delta
FROM changes GROUP BY {identity_columns}, baseline_model
) grouped GROUP BY {identity_columns}
)
UPDATE "LiteLLM_AutoRouterSession" AS session
UPDATE "{table_name}" AS session
SET saved_spend = session.saved_spend + totals.savings_delta,
savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta,
savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta,
savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta,
savings_estimated_baseline_models = (
SELECT COALESCE(jsonb_object_agg(key, value), '{}'::jsonb) FROM (
SELECT COALESCE(jsonb_object_agg(key, value), '{{}}'::jsonb) FROM (
SELECT key, SUM(value::int)::int AS value FROM (
SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models)
UNION ALL SELECT * FROM jsonb_each_text(models.deltas)
) combined GROUP BY key HAVING SUM(value::int) > 0
) counts
)
FROM totals JOIN models USING (api_key, session_id, router_name)
WHERE session.api_key = totals.api_key AND session.session_id = totals.session_id
FROM totals JOIN models USING ({identity_columns})
WHERE {user_match}session.api_key = totals.api_key AND session.session_id = totals.session_id
AND session.router_name = totals.router_name
"""
_UPDATE_SESSIONS: Final = _session_correction_sql(user_scoped=False)
_UPDATE_USER_SESSIONS: Final = _session_correction_sql(user_scoped=True)
def _primary_transaction(client: PrismaClient) -> _TransactionManager:
primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db))
return primary.tx(timeout=_TRANSACTION_TIMEOUT)
@ -308,6 +321,7 @@ def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, n
request_id=record.observation.request_id,
publication=new,
api_key=record.api_key,
user_id=record.turn.user_id if record.turn is not None else "",
session_id=record.session_id,
router_name=record.router_name,
baseline_model=record.baseline_model,
@ -357,6 +371,8 @@ async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None:
serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":"))
await db.execute_raw(_UPDATE_LOGS, serialized)
await db.execute_raw(_UPDATE_SESSIONS, serialized)
if any(change.user_id for change in changes):
await db.execute_raw(_UPDATE_USER_SESSIONS, serialized)
for entity, table in DAILY_SPEND_TABLES.items():
if adjustments := tuple(
change.daily.adjustment(target, change.savings_delta, change.request_id)

61
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -1,5 +1,6 @@
import asyncio
import time
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final, Literal, TypeAlias
@ -40,6 +41,28 @@ class TableCleanupResult:
stop_reason: StopReason
class _RunProgress:
"""How far one cleanup run has got, reported if that run is cancelled"""
def __init__(self) -> None:
self.rows_deleted: int = 0
self.batches: int = 0
def record_batch(self, rows_deleted: int) -> None:
self.rows_deleted += rows_deleted
self.batches += 1
_run_progress: ContextVar[_RunProgress] = ContextVar("spend_log_cleanup_run_progress")
def _record_run_batch(rows_deleted: int) -> None:
"""Count a batch towards the run in progress, if a run is what issued it"""
progress: Final = _run_progress.get(None)
if progress is not None:
progress.record_batch(rows_deleted)
class _RemainingRow(BaseModel):
"""One row of the capped outstanding-rows probe, validated out of prisma's untyped result."""
@ -422,6 +445,7 @@ class SpendLogCleanup:
total_deleted += deleted_count
run_count += 1
_record_run_batch(deleted_count)
# Add a small sleep to prevent overwhelming the database
await asyncio.sleep(0.1)
@ -492,6 +516,18 @@ class SpendLogCleanup:
deadline=deadline,
)
async def _delete_old_autorouter_user_session_rows(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
return await self._delete_old_rows_batched(
prisma_client,
cutoff_date,
table_name="LiteLLM_AutoRouterUserSession",
key_columns=("user_id", "api_key", "session_id", "router_name"),
time_column="last_turn_at",
deadline=deadline,
)
async def _delete_old_health_check_rows(
self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
) -> TableCleanupResult:
@ -560,9 +596,17 @@ class SpendLogCleanup:
)
except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job
verbose_proxy_logger.warning("Auto-router baseline retention remains pending")
sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline)
sessions_result: Final = await self._delete_old_autorouter_session_rows(
prisma_client, session_cutoff, self._group_deadline(deadline, 2)
)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted)
return (sessions_result,)
user_sessions_result: Final = await self._delete_old_autorouter_user_session_rows(
prisma_client, session_cutoff, deadline
)
verbose_proxy_logger.info(
"Deleted %s expired auto-router user session rollup rows", user_sessions_result.rows_deleted
)
return (sessions_result, user_sessions_result)
async def _clean_health_checks(
self, prisma_client: PrismaClient, retention_seconds: int, deadline: float
@ -601,6 +645,9 @@ class SpendLogCleanup:
If no pod_lock_manager, runs cleanup without distributed locking.
"""
lock_acquired = False
run_started_at: Final = time.monotonic()
progress: Final = _RunProgress()
progress_token: Final = _run_progress.set(progress)
try:
verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now())
self._refresh_bounds()
@ -681,6 +728,15 @@ class SpendLogCleanup:
self._run_outcome(spend_log_results + session_results + health_check_results)
)
except asyncio.CancelledError:
verbose_proxy_logger.error(
"Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here",
time.monotonic() - run_started_at,
progress.rows_deleted,
progress.batches,
)
SpendLogCleanupMetrics.record_run("aborted")
raise
except Exception as e:
# .exception() captures the traceback; str(e) alone on a Prisma/DB
# timeout is often empty and gives operators no signal to diagnose.
@ -692,6 +748,7 @@ class SpendLogCleanup:
SpendLogCleanupMetrics.record_run("aborted")
return # Return after error handling
finally:
_run_progress.reset(progress_token)
# Only release the lock if it was actually acquired
if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache:
await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME)

11
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -789,14 +789,18 @@ async def get_auto_router_benchmarks(
] = None,
end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None,
api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None,
user_id: Annotated[
str | None, Query(min_length=1, description="Filter to one canonical internal user recorded on each turn")
] = None,
) -> AutoRouterBenchmarksResponse:
"""
Benchmarks for the auto-router dashboard: session shape, savings against the configured
baseline, and prompt-caching behaviour bucketed by what the router did.
Reads the LiteLLM_AutoRouterSession rollup, folded once per request at spend-write time,
so this endpoint never scans LiteLLM_SpendLogs. A session is in the window when it
overlaps it: its last turn is on or after start_date and its first turn is on or before
Reads session rollups folded once per request at spend-write time, so this endpoint
never scans LiteLLM_SpendLogs. A user filter selects only turns attributed to that
internal user when written; older key-only history remains outside user views. A session
is in the window when it overlaps it: its last turn is on or after start_date and its first turn is on or before
end_date. Overall hit rate is over telemetry-bearing turns; each bucket's hit rate is
over that bucket's turns.
@ -826,6 +830,7 @@ async def get_auto_router_benchmarks(
start_day.isoformat(),
(end_day + timedelta(days=1)).isoformat(),
api_key,
user_id,
)
rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ())
groups: Final = (

9
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -1078,6 +1078,9 @@ if MCP_AVAILABLE:
return {"servers": registry_servers}
## FastAPI Routes
def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]:
return ((server.server_name or server.alias or server.server_id).lower(), server.server_id)
def _get_user_mcp_management_mode() -> UserMCPManagementMode:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
@ -1228,10 +1231,12 @@ if MCP_AVAILABLE:
detail="You do not have permission to view MCP servers for this team.",
)
redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id)
redacted_mcp_servers = sorted(
await _get_team_scoped_mcp_server_list(sanitized_team_id), key=_mcp_server_display_order
)
else:
servers: Final = await _resolve_accessible_mcp_servers(user_api_key_dict)
redacted_mcp_servers = _redact_mcp_credentials_list(servers)
redacted_mcp_servers = sorted(_redact_mcp_credentials_list(servers), key=_mcp_server_display_order)
if connected_app_view is True and is_ui_session_credential(user_api_key_dict):
reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict)

32
litellm/proxy/management_endpoints/router_settings_endpoints.py
View file

@ -8,6 +8,8 @@ GET /router/fields - Get router settings field definitions without values (for U
"""
import inspect
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, get_args
from fastapi import APIRouter, Depends
@ -16,6 +18,7 @@ from pydantic import BaseModel, Field
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for
from litellm.router import Router
from litellm.types.management_endpoints import (
ROUTER_SETTINGS_FIELDS,
@ -30,6 +33,7 @@ class RouterSettingsResponse(BaseModel):
fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata")
current_values: dict[str, Any] = Field(description="Current values of router settings")
routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option")
source: dict[str, FieldSource] = Field(description="Source of each current router setting")
class RouterFieldsResponse(BaseModel):
@ -39,6 +43,18 @@ class RouterFieldsResponse(BaseModel):
routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option")
def _router_setting_source(
settings: SettingsStore,
key: str,
current_value: object,
field_default: object,
) -> FieldSource:
source: Final = source_for(settings, key, field_default)
if source != "unset":
return source
return "default" if current_value is not None else "unset"
def _get_routing_strategies_from_router_class() -> list[str]:
"""
Dynamically extract routing strategies from the Router class __init__ method.
@ -109,15 +125,29 @@ async def get_router_settings(
# Merge with config values (config takes precedence)
current_values.update(router_settings_from_config)
# Update field values with current values
for field in router_fields:
if field.field_name in current_values:
field.field_value = current_values[field.field_name]
field_defaults: Final[Mapping[str, object]] = MappingProxyType(
{field.field_name: field.field_default for field in router_fields}
)
source: Final[Mapping[str, FieldSource]] = MappingProxyType(
{
key: _router_setting_source(
proxy_config.router_settings,
key,
current_values[key],
field_defaults.get(key),
)
for key in current_values
}
)
return RouterSettingsResponse(
fields=router_fields,
current_values=current_values,
routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS,
source=source,
)
except Exception as e:
verbose_proxy_logger.error("Error fetching router settings: %s", e)

98
litellm/proxy/proxy_server.py
View file

@ -50,6 +50,7 @@ import anyio
import websockets
import websockets.exceptions
from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError
from pydantic.fields import FieldInfo, PydanticUndefined
from typing_extensions import NotRequired, ReadOnly, assert_never
from litellm._uuid import uuid
@ -456,7 +457,13 @@ from litellm.proxy.common_utils.user_api_key_cache import (
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.config_resolvers import SettingsStore, config_ownership_message, resolve_fields
from litellm.proxy.config_resolvers import (
FieldSource,
SettingsStore,
config_ownership_message,
resolve_fields,
source_for,
)
from litellm.proxy.config_resolvers.alerting import (
EMAIL_DESCRIPTORS,
MS_TEAMS_DESCRIPTORS,
@ -715,6 +722,11 @@ from litellm.proxy.route_llm_request import route_request
from litellm.proxy.route_priority import hot_routes_first
from litellm.proxy.search_endpoints.endpoints import router as search_router
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
from litellm.proxy.shutdown.scheduled_jobs import (
AwaitableAsyncIOExecutor,
pause_scheduled_jobs,
stop_in_flight_scheduler_jobs,
)
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
run_scheduled_daily_global_spend_reconcile,
@ -1488,6 +1500,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
if model_info_scheduler is not scheduler:
model_info_scheduler.shutdown(wait=False)
# Shutdown event - stop starting scheduled jobs; the ones already running keep the drain window
if scheduler is not None:
pause_scheduled_jobs(scheduler)
# Shutdown event - drain in-flight requests before tearing down dependencies
# so SIGTERM (rolling update, scale-down, liveness kill) doesn't drop them.
GracefulShutdownManager.start_shutdown()
@ -1527,6 +1543,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
await _drain_spend_event_producer_on_shutdown()
# Shutdown event - finish or cancel in-flight scheduled jobs before the shutdown flushes and the DB disconnect
if scheduler is not None and scheduler_executor is not None:
try:
await stop_in_flight_scheduler_jobs(scheduler, scheduler_executor)
except Exception as e:
verbose_proxy_logger.error("Error stopping in-flight scheduled jobs: %s", e)
await flush_spend_counters_on_shutdown()
await _flush_spend_logs_queue_on_shutdown()
@ -2529,6 +2552,7 @@ celery_app_conn: Final = None
celery_fn: Final = None # Redis Queue for handling requests
scheduler = None
scheduler_executor: AwaitableAsyncIOExecutor | None = None # rebind-ok: bound once the scheduler is built at startup
# Global variable for anthropic beta headers reload scheduling
last_anthropic_beta_headers_reload = None
@ -4951,6 +4975,12 @@ def _as_settings_mapping(value: object) -> Mapping[str, SettingsJsonValue]:
return _SETTINGS_MAPPING.validate_python(value)
def _get_field_default(field_info: FieldInfo) -> JsonValue:
if field_info.default is PydanticUndefined:
return None
return cast(JsonValue, field_info.default) # cast-ok: Pydantic field defaults are JSON values at runtime
def _bind_general_settings_store(settings: SettingsStore) -> None:
global general_settings
general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings
@ -10081,7 +10111,7 @@ class ProxyStartupEvent:
proxy_logging_obj: ProxyLogging,
) -> ProxyWorkerHeartbeat:
"""Initializes scheduled background jobs"""
global heuristic_v1_tuning_baselines, store_model_in_db, scheduler # rebind-ok: startup publishes the one read-only baseline snapshot
global heuristic_v1_tuning_baselines, store_model_in_db, scheduler, scheduler_executor # rebind-ok: startup publishes the one read-only baseline snapshot
# MEMORY LEAK FIX: Configure scheduler with optimized settings
# Memray analysis showed APScheduler's normalize() and _apply_jitter() causing
@ -10090,9 +10120,9 @@ class ProxyStartupEvent:
# 1. Remove/minimize jitter to avoid normalize() memory explosion
# 2. Use larger misfire_grace_time to prevent backlog calculations
# 3. Set replace_existing=True to avoid duplicate jobs
from apscheduler.executors.asyncio import AsyncIOExecutor
from apscheduler.jobstores.memory import MemoryJobStore
scheduler_executor = AwaitableAsyncIOExecutor() # rebind-ok: shutdown awaits the jobs this executor runs
scheduler = AsyncIOScheduler(
job_defaults={
"coalesce": APSCHEDULER_COALESCE,
@ -10105,7 +10135,7 @@ class ProxyStartupEvent:
jobstores={"default": MemoryJobStore()}, # explicitly use memory job store
# Use simple executor to minimize overhead
executors={
"default": AsyncIOExecutor(),
"default": scheduler_executor,
},
# Disable timezone awareness to reduce computation
timezone=None,
@ -16006,6 +16036,22 @@ async def model_settings():
#### ALERTING MANAGEMENT ENDPOINTS ####
def _nested_setting_source(
settings: SettingsStore,
db_values: Mapping[str, JsonValue],
parent_key: str,
field_name: str,
field_default: JsonValue,
) -> FieldSource:
unset_source: Final[FieldSource] = "default" if field_default is not None else "unset"
parent_value: Final = settings.config_value(parent_key)
if isinstance(parent_value, Mapping) and field_name in parent_value:
return "config"
if settings.owned_by_config(parent_key):
return unset_source
return "db" if field_name in db_values else unset_source
@router.get(
"/alerting/settings",
description="Return the configurable alerting param, description, and current value",
@ -16043,17 +16089,20 @@ async def alerting_settings(
where={"param_name": "general_settings"}
)
if db_general_settings is not None and db_general_settings.param_value is not None:
db_general_settings_dict: Final = dict(db_general_settings.param_value)
alerting_args_dict: dict = cast( # cast-ok: ConfigGeneralSettings validates alerting_args as a dict on write
dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {})
)
alerting_values: list | None = cast( # cast-ok: ConfigGeneralSettings validates alerting as a list on write
list[JsonValue] | None, db_general_settings_dict.get("alerting")
)
else:
alerting_args_dict = {}
alerting_values = None
db_general_settings_dict: Final[Mapping[str, JsonValue]] = MappingProxyType(
dict(db_general_settings.param_value) # mutable-ok: Prisma returns the JSON column as a plain dict
if db_general_settings is not None and db_general_settings.param_value is not None
else {}
)
alerting_args_value: Final = db_general_settings_dict.get("alerting_args")
alerting_args_dict: Final[Mapping[str, JsonValue]] = MappingProxyType(
alerting_args_value if isinstance(alerting_args_value, dict) else {}
)
alerting_values: Final = cast( # cast-ok: alerting is stored as a JSON list when present
list[JsonValue] | None, db_general_settings_dict.get("alerting")
)
settings: Final = proxy_config.settings
allowed_args: Final = MappingProxyType(
{
@ -16082,9 +16131,9 @@ async def alerting_settings(
is_slack_enabled = False
if general_settings.get("alerting") and isinstance(general_settings["alerting"], list):
if "slack" in general_settings["alerting"]:
is_slack_enabled = True
alerting: Final = settings.get("alerting")
if isinstance(alerting, list) and "slack" in alerting:
is_slack_enabled = True
_response_obj = ConfigList(
field_name="slack_alerting",
@ -16092,6 +16141,7 @@ async def alerting_settings(
field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.",
field_value=is_slack_enabled,
stored_in_db=True if alerting_values is not None else False,
source=source_for(settings, "alerting"),
field_default_value=None,
premium_field=False,
)
@ -16099,6 +16149,7 @@ async def alerting_settings(
for field_name, field_info in SlackAlertingArgs.model_fields.items():
if field_name in allowed_args:
field_default: JsonValue = _get_field_default(field_info)
_stored_in_db: bool | None = None
if field_name in alerting_args_dict:
_stored_in_db = True
@ -16109,9 +16160,16 @@ async def alerting_settings(
field_name=field_name,
field_type=allowed_args[field_name],
field_description=field_info.description or "",
field_value=_slack_alerting_args_dict.get(field_name, None),
field_value=_slack_alerting_args_dict.get(field_name, field_default),
stored_in_db=_stored_in_db,
field_default_value=field_info.default,
source=_nested_setting_source(
settings,
alerting_args_dict,
"alerting_args",
field_name,
field_default,
),
field_default_value=field_default,
premium_field=(True if field_name == "region_outage_alert_ttl" else False),
)
return_val.append(_response_obj)

41
litellm/proxy/schema.prisma
View file

@ -1623,6 +1623,47 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
model LiteLLM_AutoRouterUserSession {
user_id String
api_key String
session_id String
router_name String
router_type String
first_turn_at DateTime
last_turn_at DateTime
last_model String
models Json @default("{}")
turns Int @default(0)
unordered_turns Int @default(0)
covered_turns Int @default(0)
cache_hits Int @default(0)
same_model_turns Int @default(0)
same_model_hits Int @default(0)
first_visit_turns Int @default(0)
first_visit_hits Int @default(0)
return_turns Int @default(0)
return_hits Int @default(0)
return_expired_misses Int @default(0)
return_within_ttl_misses Int @default(0)
ttl_5m_turns Int @default(0)
ttl_1h_turns Int @default(0)
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
baseline_models Json @default("{}")
@@id([user_id, api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_user_session_last_turn")
@@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn")
}
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
// either direction. forward duplicates the requests the keys did not route through the
// router through it, answering whether they should adopt it; reverse duplicates the

79
litellm/proxy/shutdown/scheduled_jobs.py Normal file
View file

@ -0,0 +1,79 @@
# pyright: reportMissingTypeStubs=false # apscheduler ships no type information
import asyncio
from collections.abc import Collection
from typing import Final, Protocol
from apscheduler.executors.asyncio import AsyncIOExecutor
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS,
SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS,
)
class StoppableScheduler(Protocol):
"""The slice of ``AsyncIOScheduler`` shutdown uses, which ships no type information"""
@property
def running(self) -> bool: ...
def pause(self) -> None: ...
def shutdown(self, wait: bool = ...) -> None: ...
class AwaitableAsyncIOExecutor(AsyncIOExecutor): # pyright: ignore[reportUntypedBaseClass] # apscheduler ships no type information and is absent from the type-check env
"""``AsyncIOExecutor`` whose in-flight job tasks can be awaited after ``shutdown`` cancels them"""
_pending_futures: Collection["asyncio.Future[object]"]
def in_flight_jobs(self) -> tuple["asyncio.Future[object]", ...]:
"""The job tasks that are running right now, as a snapshot"""
return tuple(future for future in self._pending_futures if not future.done())
def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None:
"""Stop the scheduler from starting jobs that shutdown would only cancel; running jobs continue"""
if scheduler.running:
scheduler.pause()
async def stop_in_flight_scheduler_jobs(
scheduler: StoppableScheduler,
executor: AwaitableAsyncIOExecutor,
*,
finish_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS,
cancel_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS,
) -> None:
"""
Let in-flight jobs finish for up to finish_timeout_seconds, then stop the scheduler and wait, bounded by
cancel_timeout_seconds, for the jobs it cancels.
Must run before the database is disconnected: a write job that finishes needs its connection,
and a job's cancellation handler is what records the run's outcome.
"""
if not scheduler.running:
return
in_flight: Final = executor.in_flight_jobs()
if in_flight:
verbose_proxy_logger.info(
"Waiting up to %ss for %d in-flight scheduled job(s) to finish",
finish_timeout_seconds,
len(in_flight),
)
still_running: Final = (
(await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] if in_flight else frozenset()
)
scheduler.shutdown(wait=False)
if not still_running:
return
verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running))
_done, pending = await asyncio.wait(still_running, timeout=cancel_timeout_seconds)
if pending:
verbose_proxy_logger.warning(
"%d scheduled job(s) did not finish within %ss of cancellation; giving up on them",
len(pending),
cancel_timeout_seconds,
)

86
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -15,8 +15,8 @@ from typing import (
from urllib.parse import urlparse
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model
from pydantic.fields import FieldInfo
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model
from pydantic.fields import FieldInfo, PydanticUndefined
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
@ -25,6 +25,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys
from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for
from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError
from litellm.proxy.config_resolvers.sso import (
SSO_FIELD_ENV_VARS,
@ -35,7 +36,10 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_PERMISSIONS,
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.spend_tracking.ptu_feature_flag import (
PTU_COST_ATTRIBUTION_ENV_VAR,
is_ptu_cost_attribution_enabled,
)
from litellm.proxy.utils import invalidate_config_param
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -45,6 +49,7 @@ from litellm.repositories.table_repositories import (
UISettingsRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.secret_managers.main import get_secret
from litellm.types.mcp import MCPToolSearchSettings
from litellm.types.proxy.management_endpoints.ui_sso import (
DefaultTeamSSOParams,
@ -199,6 +204,11 @@ class SettingsResponse(BaseModel):
"""Schema information including descriptions and property types for UI display"""
class _SettingsWithSchema(BaseModel):
values: dict[str, object]
field_schema: dict[str, object]
class SSOSettingsResponse(SettingsResponse):
"""Response model for SSO settings"""
@ -330,6 +340,8 @@ class UISettings(BaseModel):
class UISettingsResponse(SettingsResponse):
"""Response model for UI settings"""
source: dict[str, FieldSource]
# Allowlist of UI settings that can be stored
ALLOWED_UI_SETTINGS_FIELDS: Final = {
@ -748,6 +760,25 @@ def _root_schema(settings_class: type[BaseModel]) -> _RootSchema:
)
def _model_field_default(settings_class: type[BaseModel], field_name: str) -> object:
field_info: Final = settings_class.model_fields.get(field_name)
if field_info is None or field_info.default is PydanticUndefined:
return None
return cast(object, field_info.default) # cast-ok: Pydantic field defaults are untyped
def _ui_setting_source(
key: str,
value: object,
settings: SettingsStore,
settings_class: type[BaseModel],
) -> FieldSource:
if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING:
configured_value: Final = get_secret(PTU_COST_ATTRIBUTION_ENV_VAR, None)
return "config" if configured_value is not None or value is True else "default"
return source_for(settings, key, _model_field_default(settings_class, key))
async def _get_settings_with_schema(
settings_key: str,
settings_class: type[BaseModel],
@ -1705,7 +1736,7 @@ async def get_ui_settings():
Get UI-specific configuration flags.
All authenticated users can fetch these settings for client-side behavior.
"""
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy.proxy_server import prisma_client, proxy_config
if prisma_client is None:
raise HTTPException(
@ -1730,20 +1761,43 @@ async def get_ui_settings():
await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL)
# Build config-like object for schema helper
config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}}
settings: Final = await _get_settings_with_schema(
settings_key="ui_settings",
settings_class=_get_effective_ui_settings_class(),
config=config,
effective_ui_settings: Final[Mapping[str, object]] = MappingProxyType(
{
**ui_settings,
**{key: proxy_config.settings[key] for key in ALLOWED_UI_SETTINGS_FIELDS if key in proxy_config.settings},
}
)
config: Final[Mapping[str, object]] = MappingProxyType(
{"litellm_settings": MappingProxyType({"ui_settings": effective_ui_settings})}
)
settings_class: Final = _get_effective_ui_settings_class()
resolved_settings: Final = _SettingsWithSchema.model_validate(
await _get_settings_with_schema(
settings_key="ui_settings",
settings_class=settings_class,
config=config,
)
)
values: Final[Mapping[str, object]] = MappingProxyType(
{
**resolved_settings.values,
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(),
}
)
source: Final[Mapping[str, FieldSource]] = MappingProxyType(
{
key: (
_ui_setting_source(key, values[key], proxy_config.settings, settings_class)
if key in proxy_config.settings or key not in ui_settings
else "db"
)
for key in values
}
)
return UISettingsResponse(
values={
**settings["values"],
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(),
},
field_schema=settings["field_schema"],
values=values,
field_schema=resolved_settings.field_schema,
source=source,
)

5
litellm/router_strategy/complexity_router/config.py
View file

@ -1321,7 +1321,7 @@ class ComplexityRouterConfig(BaseModel):
)
enable_context_window_escalation: bool = Field(
default=True,
default=False,
description=(
"Escalate a request off a tier whose models provably cannot hold its prompt, before "
"dispatch. The classifier scores complexity and never prompt size, so a long agentic "
@ -1331,7 +1331,8 @@ class ComplexityRouterConfig(BaseModel):
"moves to the lowest configured tier with a model whose declared window fits; when "
"only some of the tier's models fit, the pick is restricted to those and the tier "
"keeps the request. Models with no resolvable window are never escalated away from "
"and never escalated onto. Set false to dispatch on complexity alone, as before."
"and never escalated onto. Disabled by default: omit or set false to dispatch on "
"complexity alone; set true to enable context-window escalation."
),
)
context_window_escalation_buffer: float = Field(

127
litellm/rust_bridge/_native.pyi
View file

@ -93,6 +93,133 @@ class ResponsesWebSocketConnection:
def recv_text(self) -> Future[str | None]: ...
def close(self) -> Future[None]: ...
@final
class _CacheTestBinding:
@property
def kind(self) -> str: ...
def lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | Sequence[object] | None = None,
) -> object: ...
def store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> None: ...
def lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> object: ...
def async_lookup(
self,
request: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_store(
self,
request: object,
response: object,
*,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[None]: ...
def async_lookup_batch(
self,
requests: Sequence[object],
*,
callback_kwargs: Sequence[object] | None = None,
) -> Future[object]: ...
def async_store_batch(
self,
requests: Sequence[object],
responses: Sequence[object],
*,
callback_result: object = None,
callback_kwargs: Mapping[str, object] | None = None,
) -> Future[object]: ...
def async_flush(self) -> Future[None]: ...
def ping(self) -> Future[object]: ...
@final
class _CacheTestHandle:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@staticmethod
def memory(
*,
capacity: int = 200,
ttl_seconds: float = 600.0,
max_entry_bytes: int = 1048576,
) -> _CacheTestHandle: ...
@staticmethod
def redis(
url: str,
*,
ttl_seconds: float = 60.0,
namespace: str | None = None,
startup_nodes: Sequence[tuple[str, int]] | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def disk(directory: str) -> _CacheTestHandle: ...
@staticmethod
def qdrant_semantic(
url: str,
*,
collection_name: str,
similarity_threshold: float,
vector_size: int,
embedding_model: str = "text-embedding-3-small",
api_key: str | None = None,
embedding_api_key: str | None = None,
embedding_api_base: str | None = None,
embedding_timeout_seconds: float | None = None,
quantization: str = "binary",
) -> _CacheTestHandle: ...
@staticmethod
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
@staticmethod
def redis_semantic(backend: object) -> _CacheTestHandle: ...
@staticmethod
def valkey_semantic(
url: str,
similarity_threshold: float,
index_name: str,
embedder: object,
) -> _CacheTestHandle: ...
@staticmethod
def gcs(
bucket_name: str,
*,
gcs_path: str | None = None,
path_service_account: str | None = None,
endpoint: str | None = None,
token: str | None = None,
) -> _CacheTestHandle: ...
@staticmethod
def s3(
bucket: str,
*,
region: str,
endpoint_url: str | None = None,
key_prefix: str = "",
access_key_id: str | None = None,
secret_access_key: str | None = None,
session_token: str | None = None,
) -> _CacheTestHandle: ...
@property
def backend(self) -> str: ...
def _bind_facade(self, facade: object) -> None: ...
@final
class _CacheTestResolver:
def __new__(cls, namespace: object) -> _CacheTestResolver: ...
def resolve(self) -> _CacheTestBinding: ...
@final
class TokenCounter:
def __new__(cls, tokenizer_json: str) -> TokenCounter: ...

79
model_prices_and_context_window.json
View file

@ -43037,21 +43037,21 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"input_cost_per_token": 8.92272e-07,
"input_cost_per_token": 8.83746e-07,
"input_cost_per_token_cache_hit": 4.4e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.784544e-06,
"output_cost_per_token": 1.767492e-06,
"source": "https://openrouter.ai/api/v1/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 7.4356e-08,
"cache_read_input_token_cost": 7.36455e-08,
"supports_audio_input": false,
"supports_pdf_input": false,
"supports_vision": false,
@ -70299,8 +70299,8 @@
"supports_web_search": true
},
"openrouter/meta-llama/llama-4-maverick": {
"input_cost_per_token": 2e-07,
"output_cost_per_token": 8e-07,
"input_cost_per_token": 1.875e-07,
"output_cost_per_token": 6.525e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 16384,
@ -73794,6 +73794,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 2.5e-07,
"input_cost_per_token_above_128k_tokens": 5e-07,
"litellm_provider": "openrouter",
@ -73815,6 +73816,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-1.6-flash": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 7.5e-08,
"input_cost_per_token_above_128k_tokens": 1e-07,
"litellm_provider": "openrouter",
@ -73855,6 +73857,7 @@
"supports_web_search": false
},
"openrouter/bytedance-seed/seed-2.0-code": {
"deprecation_date": "2026-11-11",
"input_cost_per_token": 5e-07,
"input_cost_per_token_above_128k_tokens": 1e-06,
"litellm_provider": "openrouter",
@ -76876,6 +76879,26 @@
"supports_vision": true,
"supports_web_search": true
},
"moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
@ -76975,5 +76998,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

41
schema.prisma
View file

@ -1623,6 +1623,47 @@ model LiteLLM_AutoRouterSession {
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
}
model LiteLLM_AutoRouterUserSession {
user_id String
api_key String
session_id String
router_name String
router_type String
first_turn_at DateTime
last_turn_at DateTime
last_model String
models Json @default("{}")
turns Int @default(0)
unordered_turns Int @default(0)
covered_turns Int @default(0)
cache_hits Int @default(0)
same_model_turns Int @default(0)
same_model_hits Int @default(0)
first_visit_turns Int @default(0)
first_visit_hits Int @default(0)
return_turns Int @default(0)
return_hits Int @default(0)
return_expired_misses Int @default(0)
return_within_ttl_misses Int @default(0)
ttl_5m_turns Int @default(0)
ttl_1h_turns Int @default(0)
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
baseline_models Json @default("{}")
@@id([user_id, api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_user_session_last_turn")
@@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn")
}
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
// either direction. forward duplicates the requests the keys did not route through the
// router through it, answering whether they should adopt it; reverse duplicates the

3
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -89,6 +89,9 @@
- {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"}
- {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"}
- {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"}
- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"}
- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"}
- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"}
- {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"}
- {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"}
- {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"}

1
tests/e2e/coverage_registry/schema.py
View file

@ -58,6 +58,7 @@ LlmRoute = Literal[
"openai",
"together_ai",
"vertex",
"xiaomi_mimo",
]
LlmCapability = Literal[

214
tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py Normal file
View file

@ -0,0 +1,214 @@
"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions.
Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via
``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes
us is that the reasoning chain surfaces as ``reasoning_content``, tool calls
survive translation, and the cost header plus spend row follow the proxy's own
cost-map price for the row (read back from ``/model/info``, never pinned here).
Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call, unwrap
from lifecycle import ResourceManager
from models import (
ChatBody,
ChatMessage,
ChatResponse,
ChatTool,
ChatToolFunction,
CostMapEntry,
LiteLLMParamsBody,
OutMessage,
SpendLogRow,
)
from passthrough_client import PassthroughClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash")
ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number."
WEATHER_PROMPT = "What is the weather in Paris? Use the tool."
COUNTING_PROMPT = "Count from 1 to 50, one number per line."
WEATHER_TOOL = ChatTool(
function=ChatToolFunction(
name="get_weather",
description="Get the current weather for a location.",
parameters={
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
)
)
class _WeatherArgs(BaseModel):
location: str
class _StreamDelta(BaseModel):
content: str | None = None
reasoning_content: str | None = None
class _StreamChoice(BaseModel):
delta: _StreamDelta | None = None
class _StreamChunk(BaseModel):
choices: list[_StreamChoice] = []
def _approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
@pytest.fixture(scope="module")
def registry(client: PassthroughClient) -> dict[str, CostMapEntry]:
return client.proxy.model_cost_map()
def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]:
model = f"e2e-xiaomi-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY")
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model, resources.key()
def _message(response: ChatResponse) -> OutMessage:
assert response.choices, f"Xiaomi returned no choices: {response}"
message = response.choices[0].message
assert message is not None, f"Xiaomi choice has no message: {response}"
return message
def _deltas(result: StreamingResponse) -> list[_StreamDelta]:
require_successful_call(result)
assert result.is_streaming, f"response was not streamed: {result.headers}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}"
return [
choice.delta
for event in result.stream_events
for choice in _StreamChunk.model_validate_json(event).choices
if choice.delta is not None
]
@pytest.mark.parametrize("backend", BACKENDS)
class TestXiaomiMimoChatCompletions:
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged")
def test_cost_header_and_spend_row_match_the_registry_price(
self,
client: PassthroughClient,
resources: ResourceManager,
registry: dict[str, CostMapEntry],
backend: str,
) -> None:
price = registry.get(backend)
assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0"
assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}"
assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}"
model, key = _register(client, resources, backend)
result = client.proxy.transport.send(
"/chat/completions",
headers=client.proxy.transport.bearer(key),
json=ChatBody(
model=model,
messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")],
max_tokens=1024,
),
)
require_successful_call(result)
response = ChatResponse.model_validate_json(result.body)
message = _message(response)
assert message.content and "43" in message.content, f"answer lost: {message}"
assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}"
usage = response.usage
assert usage is not None and usage.prompt_tokens and usage.completion_tokens, (
f"response carries no usage, so the cost cannot be real: {result.body[:300]}"
)
header_cost = result.response_cost
assert header_cost is not None and header_cost > 0, (
f"x-litellm-response-cost header missing or non-positive: {result.headers}"
)
cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0
expected = (
(usage.prompt_tokens - cached) * price.input_cost_per_token
+ cached * (price.cache_read_input_token_cost or 0.0)
+ usage.completion_tokens * price.output_cost_per_token
)
assert _approx_equal(header_cost, expected), (
f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}"
)
def _priced(rows: list[SpendLogRow]) -> bool:
return any(row.spend is not None and row.spend > 0 for row in rows)
rows = client.proxy.poll_logs_for_key(key, predicate=_priced)
priced = [row for row in rows if row.spend is not None and row.spend > 0]
assert priced, f"no priced spend row landed for key {key}; got {rows}"
row = priced[0]
assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}"
assert row.spend is not None and _approx_equal(row.spend, header_cost), (
f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}"
)
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works")
def test_reasoning_and_answer_stream_as_deltas(
self, client: PassthroughClient, resources: ResourceManager, backend: str
) -> None:
model, key = _register(client, resources, backend)
deltas = _deltas(
client.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=COUNTING_PROMPT)],
max_tokens=2048,
stream=True,
),
)
)
reasoning = "".join(delta.reasoning_content or "" for delta in deltas)
content = "".join(delta.content or "" for delta in deltas)
assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}"
assert "50" in content, f"streamed answer lost: {content[:300]!r}"
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works")
def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None:
model, key = _register(client, resources, backend)
message = _message(
unwrap(
client.proxy.chat(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=WEATHER_PROMPT)],
tools=[WEATHER_TOOL],
max_tokens=1024,
),
)
)
)
assert message.tool_calls, f"{backend} dropped the tool call: {message}"
call = message.tool_calls[0]
assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}"
assert call.function.name == "get_weather", f"wrong tool called: {call}"
assert call.function.arguments, f"tool call carries no arguments: {call}"
args = _WeatherArgs.model_validate_json(call.function.arguments)
assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}"

98
tests/e2e/ui/helpers/roundTrip.ts
View file

@ -12,17 +12,111 @@ export async function captureRequestBody(
match: { method: string; urlIncludes: string },
action: () => Promise<void>,
): Promise<Record<string, any>> {
const pending = page.waitForRequest((req) => req.method() === match.method && req.url().includes(match.urlIncludes));
const pending = page.waitForRequest(
(req) =>
req.method() === match.method && req.url().includes(match.urlIncludes),
);
await action();
const request = await pending;
return JSON.parse(request.postData() ?? "{}") as Record<string, any>;
}
/** Reads an endpoint as the master key, so a failure is bad data and not an expired UI token. */
export async function readBack<T = any>(page: Page, endpoint: string): Promise<T> {
export async function readBack<T = any>(
page: Page,
endpoint: string,
): Promise<T> {
const res = await page.request.get(endpoint, {
headers: { Authorization: `Bearer ${masterKey()}` },
});
expect(res.ok(), `GET ${endpoint}`).toBe(true);
return (await res.json()) as T;
}
type OperationOutcome =
| { readonly status: "success" }
| { readonly status: "failure"; readonly error: unknown };
type RunFailure =
| { readonly status: "action_failure"; readonly error: unknown }
| { readonly status: "cleanup_failure"; readonly error: unknown }
| {
readonly status: "action_and_cleanup_failure";
readonly actionError: unknown;
readonly cleanupError: unknown;
};
function toRunFailure(
actionOutcome: OperationOutcome,
cleanupOutcome: OperationOutcome,
): RunFailure | null {
if (
actionOutcome.status === "failure" &&
cleanupOutcome.status === "failure"
) {
return {
status: "action_and_cleanup_failure",
actionError: actionOutcome.error,
cleanupError: cleanupOutcome.error,
};
}
if (actionOutcome.status === "failure") {
return { status: "action_failure", error: actionOutcome.error };
}
if (cleanupOutcome.status === "failure") {
return { status: "cleanup_failure", error: cleanupOutcome.error };
}
return null;
}
function raiseRunFailure(failure: RunFailure): never {
switch (failure.status) {
case "action_failure":
throw failure.error;
case "cleanup_failure":
throw failure.error;
case "action_and_cleanup_failure":
throw new AggregateError(
[failure.actionError, failure.cleanupError],
"Action and cleanup failed",
);
}
}
async function runAction(
action: () => void | Promise<void>,
): Promise<OperationOutcome> {
return Promise.resolve()
.then(action)
.then(
() => ({ status: "success" as const }),
(error: unknown) => ({ status: "failure" as const, error }),
);
}
async function runCleanup(
cleanup: () => boolean | Promise<boolean>,
): Promise<OperationOutcome> {
return Promise.resolve()
.then(cleanup)
.then(
(succeeded) =>
succeeded
? { status: "success" as const }
: {
status: "failure" as const,
error: new Error("Failed to clean up UI E2E resource"),
},
(error: unknown) => ({ status: "failure" as const, error }),
);
}
export async function runWithCleanup(
action: () => void | Promise<void>,
cleanup: () => boolean | Promise<boolean>,
): Promise<void> {
const actionOutcome = await runAction(action);
const cleanupOutcome = await runCleanup(cleanup);
const failure = toRunFailure(actionOutcome, cleanupOutcome);
if (failure !== null) raiseRunFailure(failure);
}

75
tests/e2e/ui/tests/prompts/addPrompt.spec.ts Normal file
View file

@ -0,0 +1,75 @@
import { test, expect } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page as DashboardPage } from "../../fixtures/pages";
import { navigateToPage } from "../../helpers/navigation";
import { runWithCleanup } from "../../helpers/roundTrip";
import { masterKey, uniqueSuffix } from "../../helpers/traffic";
test.use({ storageState: ADMIN_STORAGE_PATH });
test.describe("Prompt upload form", () => {
test("uploads a prompt file and reads the created prompt back", async ({
page,
}) => {
const promptId = `e2e-prompt-${uniqueSuffix()}`;
const promptContent = "Hello {{name}}";
await runWithCleanup(
async () => {
await navigateToPage(page, DashboardPage.Prompts);
await page.getByRole("button", { name: "Upload .prompt File" }).click();
await expect(
page.getByRole("dialog", { name: "Add New Prompt" }),
).toBeVisible();
await page.getByLabel("Prompt ID").fill(promptId);
await page.locator('input[type="file"]').setInputFiles({
name: "e2e.prompt",
mimeType: "text/plain",
buffer: Buffer.from(
`---\nmodel: fake-openai-gpt-4\n---\n${promptContent}\n`,
),
});
await expect(page.getByText("Selected: e2e.prompt")).toBeVisible();
await page.getByRole("button", { name: "Create Prompt" }).click();
await expect
.poll(async () => {
const response = await page.request.get(
`/prompts/${encodeURIComponent(promptId)}/info?environment=development`,
{
headers: { Authorization: `Bearer ${masterKey()}` },
},
);
return response.ok();
})
.toBe(true);
await expect
.poll(async () => {
const response = await page.request.get(
`/prompts/${encodeURIComponent(promptId)}/info?environment=development`,
{
headers: { Authorization: `Bearer ${masterKey()}` },
},
);
if (!response.ok()) return undefined;
const promptInfo = (await response.json()) as {
raw_prompt_template?: { content?: string };
};
return promptInfo.raw_prompt_template?.content;
})
.toBe(promptContent);
await expect(page.getByText(promptId, { exact: true })).toBeVisible();
},
async () => {
const response = await page.request.delete(
`/prompts/${encodeURIComponent(promptId)}?environment=development`,
{
headers: { Authorization: `Bearer ${masterKey()}` },
},
);
return response.ok();
},
);
});
});

84
tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts Normal file
View file

@ -0,0 +1,84 @@
import { test, expect } from "@playwright/test";
import { ADMIN_STORAGE_PATH } from "../../constants";
import { Page as DashboardPage } from "../../fixtures/pages";
import { navigateToPage } from "../../helpers/navigation";
import {
captureRequestBody,
readBack,
runWithCleanup,
} from "../../helpers/roundTrip";
import { masterKey, uniqueSuffix } from "../../helpers/traffic";
test.use({ storageState: ADMIN_STORAGE_PATH });
test.describe("Tag management", () => {
test("creates, edits, reopens, and reads back a tag", async ({ page }) => {
const tagName = `e2e-tag-${uniqueSuffix()}`;
const description = "synthetic tag description";
const updatedDescription = `${description} updated`;
await runWithCleanup(
async () => {
await navigateToPage(page, DashboardPage.TagManagement);
await page.getByRole("button", { name: "+ Create New Tag" }).click();
await expect(
page.getByRole("dialog", { name: "Create New Tag" }),
).toBeVisible();
await page.getByLabel("Tag Name").fill(tagName);
await page.getByLabel("Description").fill(description);
await page.getByRole("button", { name: "Create Tag" }).click();
await expect
.poll(async () => {
const response = await readBack<Array<{ name: string }>>(
page,
"/tag/list",
);
return response.some((tag) => tag.name === tagName);
})
.toBe(true);
await expect(page.getByText(tagName, { exact: true })).toBeVisible();
await page.getByText(tagName, { exact: true }).click();
await expect(page.getByText("Tag Name:")).toBeVisible();
await page.getByRole("button", { name: "Edit Tag" }).click();
await page.getByLabel("Description").fill(updatedDescription);
const updateBody = await captureRequestBody(
page,
{ method: "POST", urlIncludes: "/tag/update" },
() => page.getByRole("button", { name: "Save Changes" }).click(),
);
expect(updateBody).toMatchObject({
name: tagName,
description: updatedDescription,
});
await expect
.poll(async () => {
const infoResponse = await page.request.post("/tag/info", {
headers: { Authorization: `Bearer ${masterKey()}` },
data: { names: [tagName] },
});
expect(infoResponse.ok()).toBe(true);
const info = (await infoResponse.json()) as Record<
string,
{ description?: string }
>;
return info[tagName]?.description;
})
.toBe(updatedDescription);
},
async () => {
const response = await page.request.post("/tag/delete", {
headers: {
Authorization: `Bearer ${masterKey()}`,
"Content-Type": "application/json",
},
data: { name: tagName },
});
return response.ok();
},
);
});
});

171
tests/proxy_behavior/spend/test_autorouter_session_rollup.py
View file

@ -6,17 +6,24 @@ tests/test_litellm/proxy/db/test_autorouter_session_rollup.py.
"""
import asyncio
import time
import uuid
from datetime import datetime, timedelta, timezone
from typing import Final
from types import SimpleNamespace
from typing import Final, TypedDict, cast
import pytest
from prisma import Prisma
from prisma.errors import RawQueryError
from typing_extensions import ReadOnly
from litellm.proxy.db.autorouter_session_rollup import (
AUTOROUTER_BENCHMARKS_SQL,
UPSERT_AUTOROUTER_SESSION_SQL,
AutoRouterTurnTransaction,
flush_autorouter_turn_transactions,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
pytestmark = pytest.mark.asyncio(loop_scope="session")
@ -45,6 +52,7 @@ async def _turn(
tier: "str | None" = None,
baseline: "str | None" = None,
estimated: bool = True,
user_id: str = "",
) -> None:
touched: Final = 1 if (hit or ttl is not None or not covered) else 0
await db.execute_raw(
@ -68,6 +76,7 @@ async def _turn(
int(estimated),
spend if estimated else 0.0,
saved if estimated else 0.0,
user_id,
)
@ -217,7 +226,7 @@ async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers
assert row["savings_estimated_actual_spend"] == pytest.approx(0.01 * sum(writers))
assert row["savings_estimated_saved_spend"] == pytest.approx(0.02 * sum(writers))
groups: Final = await db.query_raw(
AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key
AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key, None
)
assert len(groups) == 1
assert groups[0]["classifier_cost"] == row["classifier_cost"]
@ -242,7 +251,7 @@ async def test_unknown_and_legacy_turns_preserve_actual_spend_without_entering_t
assert row["saved_spend"] == pytest.approx(-0.03)
assert row["savings_estimated_baseline_models"] == {"opus": 1}
groups: Final = await db.query_raw(
AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key
AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key, None
)
assert len(groups) == 1
for actual in (row, groups[0]):
@ -277,6 +286,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db):
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
None,
None,
)
matching = [row for row in rows if row["router_name"] == router]
assert len(matching) == 1
@ -304,6 +314,7 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
first_key,
None,
)
matching = [row for row in rows if row["router_name"] == router]
assert len(matching) == 1
@ -317,10 +328,160 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db):
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
f"k-{uuid.uuid4()}",
None,
)
assert [row for row in unknown_key_rows if row["router_name"] == router] == []
class _BenchmarkRow(TypedDict):
sessions: ReadOnly[int]
turns: ReadOnly[int]
same_model_turns: ReadOnly[int]
first_visit_turns: ReadOnly[int]
spend: ReadOnly[float]
saved_spend: ReadOnly[float]
tier_turns: ReadOnly[dict[str, int]]
cache_hits: ReadOnly[int]
savings_estimated_turns: ReadOnly[int]
savings_estimated_actual_spend: ReadOnly[float]
savings_estimated_saved_spend: ReadOnly[float]
async def _scoped_benchmarks(
db: Prisma, router: str, user_id: str | None = None, key: str | None = None
) -> tuple[_BenchmarkRow, ...]:
rows: Final = await db.query_raw(
AUTOROUTER_BENCHMARKS_SQL,
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
key,
user_id,
)
return tuple(cast(_BenchmarkRow, row) for row in rows if row["router_name"] == router)
async def test_users_keep_written_identity_across_shared_keys_and_keyless_sessions(db: Prisma) -> None:
router: Final = f"r-{uuid.uuid4()}"
alice: Final = f"u-{uuid.uuid4()}"
bob: Final = f"u-{uuid.uuid4()}"
first_key: Final = f"k-{uuid.uuid4()}"
second_key: Final = f"k-{uuid.uuid4()}"
await _legacy_turn(db, first_key, T0, router=router)
await _turn(db, first_key, "A", T0 + timedelta(seconds=10), router=router, user_id=alice, tier="simple")
await _turn(
db, first_key, "B", T0 + timedelta(seconds=20), router=router, user_id=bob, spend=0.03, saved=0.06, tier="complex"
)
await _turn(db, second_key, "C", T0, router=router, user_id=alice, spend=0.02, saved=0.04)
await _turn(db, "", "A", T0, router=router, user_id=alice, ttl=300)
await _turn(db, "", "A", T0 + timedelta(seconds=1), router=router, user_id=alice, hit=1)
await _turn(db, "", "B", T0, router=router, user_id=bob, spend=0.04, saved=0.08)
await _turn(db, second_key, "C", T0 - timedelta(days=40), router=router, user_id=alice, session_id="expired")
alice_rows: Final = await _scoped_benchmarks(db, router, user_id=alice)
bob_rows: Final = await _scoped_benchmarks(db, router, user_id=bob)
global_rows: Final = await _scoped_benchmarks(db, router)
key_rows: Final = await _scoped_benchmarks(db, router, key=first_key)
intersection: Final = await _scoped_benchmarks(db, router, user_id=alice, key=first_key)
assert len(alice_rows) == len(bob_rows) == len(global_rows) == len(key_rows) == len(intersection) == 1
assert (alice_rows[0]["sessions"], alice_rows[0]["turns"], alice_rows[0]["same_model_turns"]) == (3, 4, 1)
assert (bob_rows[0]["sessions"], bob_rows[0]["turns"], bob_rows[0]["first_visit_turns"]) == (2, 2, 2)
assert alice_rows[0]["spend"] == pytest.approx(0.05)
assert bob_rows[0]["spend"] == pytest.approx(0.07)
assert alice_rows[0]["tier_turns"] == {"simple": 1}
assert bob_rows[0]["tier_turns"] == {"complex": 1}
assert (alice_rows[0]["cache_hits"], bob_rows[0]["cache_hits"]) == (1, 0)
assert (global_rows[0]["sessions"], global_rows[0]["turns"]) == (4, 7)
assert (alice_rows[0]["savings_estimated_turns"], bob_rows[0]["savings_estimated_turns"]) == (4, 2)
assert global_rows[0]["savings_estimated_turns"] == 6
for scoped in (alice_rows[0], bob_rows[0]):
assert scoped["savings_estimated_actual_spend"] == pytest.approx(scoped["spend"])
assert scoped["savings_estimated_saved_spend"] == pytest.approx(scoped["saved_spend"])
assert global_rows[0]["spend"] == pytest.approx(alice_rows[0]["spend"] + bob_rows[0]["spend"] + 0.01)
assert global_rows[0]["saved_spend"] == pytest.approx(alice_rows[0]["saved_spend"] + bob_rows[0]["saved_spend"] + 0.02)
assert global_rows[0]["tier_turns"] == {"simple": 1, "complex": 1}
assert (key_rows[0]["sessions"], key_rows[0]["turns"]) == (1, 3)
assert key_rows[0]["spend"] == pytest.approx(0.05)
assert (intersection[0]["sessions"], intersection[0]["turns"]) == (1, 1)
assert intersection[0]["spend"] == pytest.approx(0.01)
assert await _scoped_benchmarks(db, router, user_id=bob, key=second_key) == ()
assert await _scoped_benchmarks(db, router, user_id=f"u-{uuid.uuid4()}") == ()
assert await _scoped_benchmarks(db, router, user_id="") == ()
async def test_a_failed_user_projection_rolls_back_the_keys_increment(db: Prisma) -> None:
key: Final = f"k-{uuid.uuid4()}"
user_id: Final = "".join(str(uuid.uuid4()) for _ in range(200))
await _turn(db, key, "A", T0)
before: Final = await _row(db, key)
with pytest.raises(RawQueryError, match=r"index row (requires|size)"):
await _turn(db, key, "B", T0 + timedelta(seconds=1), user_id=user_id)
assert await _row(db, key) == before
assert await db.query_raw('SELECT user_id FROM "LiteLLM_AutoRouterUserSession" WHERE user_id = $1', user_id) == []
first_user: Final = f"u-{uuid.uuid4()}"
second_user: Final = f"u-{uuid.uuid4()}"
turns: Final = tuple(
AutoRouterTurnTransaction(
api_key=key,
user_id=user,
session_id="s1",
router_name="auto-1",
router_type="complexity",
model=model,
turn_at=T0 + timedelta(seconds=second),
total_tokens=100,
spend=0.01,
saved_spend=0.02,
classifier_cost=0.0,
covered=True,
cache_hit=False,
cache_ttl_seconds=None,
cache_touched=False,
)
for user, model, second in (
(first_user, "A", 1),
(user_id, "B", 2),
(first_user, "B", 3),
(second_user, "C", 4),
(first_user, "B", 5),
(second_user, "C", 6),
(user_id, "A", 7),
)
)
await flush_autorouter_turn_transactions(SimpleNamespace(db=db), tuple(reversed(turns)), n_retry_times=0)
key_row: Final = await _row(db, key)
assert (key_row["turns"], key_row["last_model"], key_row["unordered_turns"]) == (2, "A", 0)
assert key_row["spend"] == pytest.approx(0.02)
user_rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterUserSession" WHERE api_key = $1', key)
by_user: Final = {row["user_id"]: row for row in user_rows}
assert set(by_user) == {first_user, second_user}
for user, count, model in ((first_user, 3, "B"), (second_user, 2, "C")):
row: Final = by_user[user]
assert (row["turns"], row["same_model_turns"], row["unordered_turns"], row["last_model"]) == (count, 1, 0, model)
assert row["spend"] == pytest.approx(count * 0.01)
assert row["saved_spend"] == pytest.approx(count * 0.02)
async def test_user_session_cleanup_keeps_another_users_recent_keyless_session(db: Prisma) -> None:
router: Final = f"r-{uuid.uuid4()}"
expired_user: Final = f"u-{uuid.uuid4()}"
recent_user: Final = f"u-{uuid.uuid4()}"
await _turn(db, "", "A", T0 - timedelta(days=1), router=router, user_id=expired_user)
await _turn(db, "", "A", T0 + timedelta(days=1), router=router, user_id=recent_user)
cleaner: Final = SpendLogCleanup(general_settings={})
await cleaner._delete_old_autorouter_user_session_rows(
SimpleNamespace(db=db), T0.replace(tzinfo=timezone.utc), time.monotonic() + 60
)
assert await db.query_raw(
'SELECT user_id, turns FROM "LiteLLM_AutoRouterUserSession" WHERE router_name = $1', router
) == [{"user_id": recent_user, "turns": 1}]
async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db):
key = f"k-{uuid.uuid4()}"
router = f"r-{uuid.uuid4()}"
@ -334,6 +495,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
None,
None,
)
matching = sorted(
(row for row in rows if row["router_name"] == router),
@ -418,6 +580,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db):
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
None,
None,
)
grouped = next(row for row in rows if row["router_name"] == router)
assert grouped["tier_turns"] == {"simple": 2, "complex": 1}
@ -446,6 +609,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
None,
None,
)
by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router}
assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}}
@ -461,6 +625,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db):
(T0 - timedelta(days=1)).isoformat(),
(T0 + timedelta(days=1)).isoformat(),
None,
None,
)
grouped = next(row for row in rows if row["router_name"] == router)
assert grouped["tier_turns"] == {}

67
tests/proxy_behavior/spend/test_baseline_accounting.py
View file

@ -56,7 +56,9 @@ def record() -> Callable[..., BaselineAccountingRecord]:
},
)
def create(label: str = "first", started: float = 10000.0, identical: bool = True) -> BaselineAccountingRecord:
def create(
label: str = "first", started: float = 10000.0, identical: bool = True, user_id: str = ""
) -> BaselineAccountingRecord:
return BaselineAccountingRecord(
scope="autorouter-baseline:v3:" + run * 2, api_key=run, session_id=run,
router_name="test-router", baseline_model="anthropic/claude-opus-5",
@ -76,6 +78,7 @@ def record() -> Callable[..., BaselineAccountingRecord]:
total_tokens=6230, spend=0.17, saved_spend=0.0, classifier_cost=0.0,
covered=True, cache_hit=False, cache_ttl_seconds=3600, cache_touched=True,
baseline_model="anthropic/claude-opus-5",
user_id=user_id,
),
daily=DailyBaselineAttribution(
date="2026-09-15", api_key=run, model="claude-opus-5", custom_llm_provider="anthropic",
@ -99,21 +102,36 @@ async def _session(db: Prisma, record: BaselineAccountingRecord):
return rows[0]
async def _user_sessions(db: Prisma, record: BaselineAccountingRecord) -> dict[str, dict[str, object]]:
rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterUserSession" WHERE api_key=$1', record.api_key)
return {str(row["user_id"]): row for row in rows}
async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None:
store: Final = _store(db)
late: Final = record("late", 10001.0)
early: Final = record("early", identical=False)
late: Final = record("late", 10001.0, user_id="late-user")
early: Final = record("early", identical=False, user_id="early-user")
await _log(db, late)
assert await store.append(late) == "recorded"
assert await store.project(late.scope) == "published"
before: Final = await _session(db, late)
assert before["savings_estimated_actual_spend"] == before["spend"] == 0.17
assert before["saved_spend"] == 0.0
before_users: Final = await _user_sessions(db, late)
assert set(before_users) == {"late-user"}
assert before_users["late-user"]["savings_estimated_turns"] == 1
assert before_users["late-user"]["savings_estimated_baseline_models"] == {late.baseline_model: 1}
await _log(db, early)
assert await store.append(early) == "recorded"
pending: Final = await _session(db, late)
assert pending["spend"] == 0.34 and pending["savings_estimated_turns"] == 0
assert pending["saved_spend"] == pending["savings_estimated_actual_spend"] == 0.0
pending_users: Final = await _user_sessions(db, late)
assert set(pending_users) == {"late-user", "early-user"}
for user in pending_users.values():
assert user["turns"] == 1 and user["spend"] == 0.17
assert user["savings_estimated_turns"] == user["savings_estimated_actual_spend"] == user["saved_spend"] == 0
assert user["savings_estimated_baseline_models"] == {}
waiting: Final = await db.query_raw('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', late.observation.request_id)
assert waiting[0]["metadata"]["autorouter_savings"] is None
assert waiting[0]["metadata"]["autorouter_savings_estimate"]["reason"] == "pending_projection"
@ -125,37 +143,69 @@ async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma,
assert logs[0]["spend"] == 0.17
assert logs[0]["metadata"]["autorouter_savings_estimate"]["provenance"] == "modeled"
assert after["saved_spend"] == pytest.approx(logs[0]["metadata"]["autorouter_savings"])
after_users: Final = await _user_sessions(db, late)
assert after_users["early-user"] == pending_users["early-user"]
for field in (
"saved_spend", "savings_estimated_turns", "savings_estimated_actual_spend",
"savings_estimated_saved_spend", "savings_estimated_baseline_models",
):
assert after_users["late-user"][field] == after[field]
assert after_users["late-user"]["turns"] == 1 and after_users["late-user"]["spend"] == 0.17
for table in ("DailyUserSpend", "DailyTeamSpend", "DailyOrganizationSpend", "DailyEndUserSpend", "DailyAgentSpend", "DailyTagSpend"):
rows: Final = await db.query_raw(f'SELECT spend,api_requests,autorouter_savings_spend FROM "LiteLLM_{table}" WHERE api_key=$1', late.api_key)
assert rows[0]["spend"] == rows[0]["api_requests"] == 0
assert rows[0]["autorouter_savings_spend"] == pytest.approx(after["saved_spend"])
async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None:
event: Final = record()
@pytest.mark.parametrize("attributed", [True, False])
async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent(
db: Prisma, record: Callable[..., BaselineAccountingRecord], attributed: bool
) -> None:
event: Final = record(user_id="first-user" if attributed else "")
other: Final = record("other", 10001.0, user_id="second-user" if attributed else "")
await _log(db, event)
assert await _store(db, after_commit=True).append(event) == "unavailable"
store: Final = _store(db)
assert set(await asyncio.gather(*(store.append(event) for _ in range(4)))) == {"recorded"}
await _log(db, other)
assert await store.append(other) == "recorded"
if not attributed:
await db.execute_raw(
'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET data=(data::jsonb #- \'{turn,user_id}\')::text WHERE scope=$1',
event.scope,
)
assert await store.project(event.scope) == "published"
assert await store.project(event.scope) == "unchanged"
session: Final = await _session(db, event)
assert session["turns"] == session["savings_estimated_turns"] == 1
assert session["spend"] == session["savings_estimated_actual_spend"] == 0.17
assert session["turns"] == session["savings_estimated_turns"] == 2
assert session["spend"] == session["savings_estimated_actual_spend"] == 0.34
users: Final = await _user_sessions(db, event)
assert set(users) == ({"first-user", "second-user"} if attributed else set())
for user in users.values():
assert user["turns"] == user["savings_estimated_turns"] == 1
assert user["spend"] == user["savings_estimated_actual_spend"] == 0.17
assert user["savings_estimated_baseline_models"] == {event.baseline_model: 1}
async def test_publication_rollback_keeps_dirty_revision_for_retry(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None:
event: Final = record()
event: Final = record(user_id="rollback-user")
await _log(db, event)
store: Final = _store(db)
assert await store.append(event) == "recorded"
assert await _store(db, before_commit=True).project(event.scope) == "unavailable"
session: Final = await _session(db, event)
assert session["spend"] == 0.17 and session["savings_estimated_turns"] == 0
before_users: Final = await _user_sessions(db, event)
assert before_users["rollback-user"]["spend"] == 0.17
assert before_users["rollback-user"]["savings_estimated_turns"] == 0
assert before_users["rollback-user"]["savings_estimated_baseline_models"] == {}
revisions: Final = await db.query_raw('SELECT revision,published_revision FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope=$1', event.scope)
assert revisions[0]["revision"] > revisions[0]["published_revision"]
assert await store.project(event.scope) == "published"
assert (await _session(db, event))["savings_estimated_turns"] == 1
after_users: Final = await _user_sessions(db, event)
assert after_users["rollback-user"]["turns"] == after_users["rollback-user"]["savings_estimated_turns"] == 1
assert after_users["rollback-user"]["spend"] == after_users["rollback-user"]["savings_estimated_actual_spend"] == 0.17
async def test_conflicting_duplicate_cannot_restore_an_observed_estimate(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None:
@ -196,6 +246,7 @@ async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_a
db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch,
) -> None:
import os
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation

47
tests/test_litellm/litellm_core_utils/test_logging_worker.py
View file

@ -180,6 +180,53 @@ class TestLoggingWorker:
assert sorted(fired) == ["first", "second"]
@pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"])
def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded):
"""
Regression: ``flush()`` from a new event loop used to ``join()`` the queue bound to the
previous loop, whose unfinished counter nothing on the new loop ever decrements. The first
such flush hung until pytest-timeout killed it and every later one raised
``RuntimeError: ... is bound to a different event loop`` from the queue's Event.
"""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()
async def enqueue_on_first_loop():
if stranded == "still_queued":
worker._ensure_queue()
worker.enqueue(callback())
return
worker.ensure_initialized_and_enqueue(callback())
asyncio.run(enqueue_on_first_loop())
assert worker._queue is not None
expected_shape = (1, 0) if stranded == "still_queued" else (0, 1)
assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape
assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed"
async def flush_twice_on_second_loop():
await asyncio.wait_for(worker.flush(), timeout=5)
await asyncio.wait_for(worker.flush(), timeout=5)
asyncio.run(flush_twice_on_second_loop())
assert callback.await_count == 1
def test_flush_starts_a_worker_when_the_queue_has_none(self):
"""``flush()`` must drain a queue that exists on the current loop without a running worker."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()
async def enqueue_then_flush():
worker._ensure_queue()
worker.enqueue(callback())
assert worker._worker_task is None, "precondition: nothing is draining the queue yet"
await asyncio.wait_for(worker.flush(), timeout=3)
asyncio.run(enqueue_then_flush())
assert callback.await_count == 1
def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self):
"""A callback raising CancelledError must not abort the atexit flush of later events."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)

184
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -291,6 +291,106 @@ async def test_find_team_with_model_access_uses_request_method_for_passthrough_a
assert "allowed_passthrough_routes" in exc_info.value.detail
_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = {
"test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": {
"endpoint_id": "test-uuid-1",
"path": "/model-host/v1/extractor",
"type": "subpath",
"auth": True,
},
}
@pytest.mark.asyncio
async def test_find_team_with_model_access_team_allowed_routes_wildcard_grants_auth_passthrough():
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes", "/model-host/*"])
team_without_passthrough_allowlist = LiteLLM_TeamTable(team_id="team-a", models=["all-proxy-models"], metadata={})
with (
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=team_without_passthrough_allowlist,
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
):
team_id, team_obj = await JWTAuthManager.find_team_with_model_access(
team_ids={"team-a"},
requested_model=None,
route="/model-host/v1/extractor/predict",
jwt_handler=jwt_handler,
prisma_client=None,
user_api_key_cache=MagicMock(),
parent_otel_span=None,
proxy_logging_obj=MagicMock(),
request_method="POST",
)
assert team_id == "team-a"
assert team_obj == team_without_passthrough_allowlist
@pytest.mark.asyncio
async def test_auth_builder_header_team_allows_auth_passthrough_for_team_allowed_routes_wildcard():
from litellm.proxy.utils import ProxyLogging
jwt_handler = JWTHandler()
user_api_key_cache = DualCache()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=LiteLLM_JWTAuth(
team_ids_jwt_field="groups",
user_id_jwt_field="sub",
team_allowed_routes=["openai_routes", "/model-host/*"],
),
)
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt,
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
return_value=LiteLLM_TeamTable(team_id="team-2", metadata={}),
),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(None, None, None, None, "user-1"),
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
):
mock_auth_jwt.return_value = {"sub": "user-1", "scope": "", "groups": ["team-1", "team-2"]}
result = await JWTAuthManager.auth_builder(
api_key="jwt-token",
jwt_handler=jwt_handler,
request_data={},
general_settings={},
route="/model-host/v1/extractor/predict",
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
request_headers={"x-litellm-team-id": "team-2"},
request_method="POST",
)
assert result["team_id"] == "team-2"
@pytest.mark.asyncio
async def test_auth_builder_proxy_admin_user_role():
"""Test that is_proxy_admin is True when user_object.user_role is PROXY_ADMIN"""
@ -6463,6 +6563,90 @@ async def test_auth_builder_db_fallback_enforces_passthrough_route_access():
assert "passthrough route" in exc_info.value.detail
async def _auth_builder_via_db_team_fallback(team_allowed_routes: list[str]):
user_id = "u_passthrough"
user_object = LiteLLM_UserTable(
user_id=user_id,
user_role=LitellmUserRoles.INTERNAL_USER,
teams=["team_no_passthrough"],
)
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_allowed_routes=team_allowed_routes)
async def fake_get_team(team_id, **kwargs):
return LiteLLM_TeamTable(team_id=team_id, metadata={})
with (
patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock, return_value={"sub": user_id, "scope": ""}),
patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock),
patch.object(jwt_handler, "get_rbac_role", return_value=None),
patch.object(jwt_handler, "get_scopes", return_value=[]),
patch.object(jwt_handler, "get_object_id", return_value=None),
patch.object(
JWTAuthManager,
"get_user_info",
new_callable=AsyncMock,
return_value=(user_id, "u@example.com", True),
),
patch.object(jwt_handler, "get_org_id", return_value=None),
patch.object(jwt_handler, "get_end_user_id", return_value=None),
patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None),
patch.object(
JWTAuthManager,
"get_objects",
new_callable=AsyncMock,
return_value=(user_object, None, None, None, user_id),
),
patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock),
patch.object(JWTAuthManager, "validate_object_id", return_value=True),
patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock),
patch(
"litellm.proxy.auth.handle_jwt.get_team_object",
new_callable=AsyncMock,
side_effect=fake_get_team,
),
patch(
"litellm.proxy.auth.handle_jwt.get_team_membership",
new_callable=AsyncMock,
return_value=None,
),
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
):
return await JWTAuthManager.auth_builder(
api_key="test_jwt_token",
jwt_handler=jwt_handler,
request_data={},
general_settings={"enforce_rbac": False},
route="/model-host/v1/extractor/predict",
prisma_client=None,
user_api_key_cache=None,
parent_otel_span=None,
proxy_logging_obj=None,
request_headers=None,
request_method="POST",
)
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_team_allowed_routes_wildcard_grants_auth_passthrough():
result = await _auth_builder_via_db_team_fallback(team_allowed_routes=["openai_routes", "/model-host/*"])
assert result["team_id"] == "team_no_passthrough"
@pytest.mark.asyncio
async def test_auth_builder_db_fallback_route_groups_alone_do_not_grant_auth_passthrough():
with pytest.raises(HTTPException) as exc_info:
await _auth_builder_via_db_team_fallback(team_allowed_routes=["openai_routes", "mapped_pass_through_routes"])
assert exc_info.value.status_code == 403, exc_info.value.detail
assert "allowed_passthrough_routes" in exc_info.value.detail
@pytest.mark.asyncio
async def test_sync_user_role_and_teams_singular_claim_reconciles_memberships():
"""When fallback_to_db_teams is on but the JWT carries a singular team claim

116
tests/test_litellm/proxy/auth/test_route_checks.py
View file

@ -1249,6 +1249,122 @@ def test_non_proxy_admin_allows_auth_pass_through_with_team_allowlist():
)
@pytest.mark.parametrize(
"route, team_allowed_routes, expected",
[
("/model-host/v1/extractor/predict", ["/model-host/*"], True),
("/model-host", ["/model-host/*"], False),
("/model-host/v1/extractor", ["/model-host/v1/extractor"], True),
("/model-host/v1/extractor/predict", ["/model-host/v1/extractor"], False),
("/other/v1/extractor", ["/model-host/*"], False),
("/model-host/v1/extractor", ["openai_routes", "llm_api_routes", "mapped_pass_through_routes"], False),
("/model-host/v1/extractor", ["*"], False),
("/model-host/v1/extractor", ["/*"], False),
("/model-host/v1/extractor", [], False),
],
)
def test_jwt_team_routes_grant_pass_through_only_for_explicit_paths(route, team_allowed_routes, expected):
assert (
RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes)
is expected
)
_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = {
"test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": {
"endpoint_id": "test-uuid-1",
"path": "/model-host/v1/extractor",
"type": "subpath",
"auth": True,
},
}
def _jwt_handler_with_team_allowed_routes(team_allowed_routes: list[str]):
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTHandler
jwt_handler: Final = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=team_allowed_routes)
return jwt_handler
def _check_model_host_route_as(valid_token: UserAPIKeyAuth, team_allowed_routes: list[str]) -> None:
with (
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
_AUTH_ENFORCED_MODEL_HOST_ROUTES,
),
patch("litellm.proxy.utils.get_server_root_path", return_value="/"),
patch(
"litellm.proxy.proxy_server.jwt_handler",
_jwt_handler_with_team_allowed_routes(team_allowed_routes),
),
):
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=None,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/model-host/v1/extractor/predict",
request=MagicMock(spec=Request),
valid_token=valid_token,
request_data={},
)
def test_non_proxy_admin_allows_auth_pass_through_for_jwt_team_allowed_routes_wildcard():
jwt_token: Final = UserAPIKeyAuth(
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id="team-a",
jwt_claims={"sub": "test_user"},
)
_check_model_host_route_as(jwt_token, team_allowed_routes=["openai_routes", "/model-host/*"])
def test_non_proxy_admin_denies_auth_pass_through_for_jwt_when_only_route_groups_configured():
jwt_token: Final = UserAPIKeyAuth(
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id="team-a",
jwt_claims={"sub": "test_user"},
)
with pytest.raises(HTTPException) as exc_info:
_check_model_host_route_as(jwt_token, team_allowed_routes=["openai_routes", "mapped_pass_through_routes"])
assert exc_info.value.status_code == 403, exc_info.value.detail
assert "allowed_passthrough_routes" in exc_info.value.detail
@pytest.mark.parametrize(
"api_key, team_id, jwt_claims",
[
("sk-test-key", "team-a", None),
("sk-test-key", "team-a", {"sub": "test_user"}),
(None, "team-a", None),
(None, None, {"sub": "test_user"}),
],
ids=["plain_virtual_key", "jwt_mapped_virtual_key", "keyless_non_jwt_caller", "jwt_without_team"],
)
def test_non_proxy_admin_jwt_team_allowed_routes_grant_pass_through_only_to_jwt_team_callers(
api_key, team_id, jwt_claims
):
caller: Final = UserAPIKeyAuth(
api_key=api_key,
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
team_id=team_id,
jwt_claims=jwt_claims,
)
with pytest.raises(HTTPException) as exc_info:
_check_model_host_route_as(caller, team_allowed_routes=["openai_routes", "/model-host/*"])
assert exc_info.value.status_code == 403, exc_info.value.detail
assert "allowed_passthrough_routes" in exc_info.value.detail
def test_virtual_key_without_llm_api_routes_cannot_access_pass_through():
"""
Test that virtual keys without llm_api_routes permission cannot access registered pass-through endpoints.

115
tests/test_litellm/proxy/db/test_autorouter_session_rollup.py
View file

@ -56,6 +56,31 @@ def _build(payload: dict | None = None, metadata: dict | None = None):
class TestBuildTransaction:
@pytest.mark.parametrize(
"api_key, user_id, included",
[
("hashed-key", "canonical-user", True),
("hashed-key", None, True),
("hashed-key", "", True),
("", "canonical-user", True),
("", None, False),
("", "", False),
],
)
def test_attribution_uses_the_canonical_user_even_without_a_key(
self, api_key: str, user_id: str | None, included: bool
) -> None:
transaction: Final = _build(
payload=_payload(api_key=api_key, user=user_id),
metadata=_metadata(user="client-user", user_api_key_user_id="metadata-user"),
)
if not included:
assert transaction is None
return
assert transaction is not None
assert transaction.api_key == api_key
assert transaction.user_id == (user_id or "")
def test_successful_auto_routed_turn_builds_every_field(self):
transaction = _build(
metadata=_metadata(
@ -205,23 +230,43 @@ class TestBuildTransaction:
class _FakeDB:
def __init__(self, failures: "list[Exception] | None" = None, poison_session: str | None = None):
def __init__(
self,
failures: "list[Exception] | None" = None,
poison_session: str | None = None,
poison_user: str | None = None,
commit_then_error_users: frozenset[str] = frozenset(),
):
self.calls: list[tuple] = []
self.attempts: list[tuple[str, tuple[object, ...]]] = []
self._failures = list(failures or [])
self._poison_session = poison_session
self._poison_user = poison_user
self._commit_then_error_users = commit_then_error_users
async def execute_raw(self, sql: str, *params: object) -> int:
self.attempts.append((sql, params))
if self._poison_session is not None and params[1] == self._poison_session:
raise RuntimeError("index row size exceeds btree maximum")
if self._poison_user is not None and params[19] == self._poison_user:
raise RuntimeError("index row size exceeds btree maximum")
if self._failures:
raise self._failures.pop(0)
self.calls.append((sql, params))
if params[19] in self._commit_then_error_users:
raise RuntimeError("commit succeeded but acknowledgement was lost")
return 1
class _FakeClient:
def __init__(self, failures: "list[Exception] | None" = None, poison_session: str | None = None):
self.db = _FakeDB(failures, poison_session)
def __init__(
self,
failures: "list[Exception] | None" = None,
poison_session: str | None = None,
poison_user: str | None = None,
commit_then_error_users: frozenset[str] = frozenset(),
):
self.db = _FakeDB(failures, poison_session, poison_user, commit_then_error_users)
def _transaction(
@ -229,9 +274,11 @@ def _transaction(
at: datetime = datetime(2026, 8, 1, 12, 0, 0),
tier: str | None = "medium",
baseline_model: str | None = "anthropic/claude-opus-5",
api_key: str = "k1",
user_id: str = "",
) -> AutoRouterTurnTransaction:
return AutoRouterTurnTransaction(
api_key="k1",
api_key=api_key,
session_id=session_id,
router_name="live-auto",
router_type="complexity",
@ -247,6 +294,7 @@ def _transaction(
cache_touched=False,
tier=tier,
baseline_model=baseline_model,
user_id=user_id,
)
@ -261,7 +309,7 @@ class TestFlush:
def test_params_marshal_in_statement_order(self):
client = _FakeClient()
asyncio.run(flush_autorouter_turn_transactions(client, [_transaction()]))
asyncio.run(flush_autorouter_turn_transactions(client, [_transaction(user_id="canonical-user")]))
sql, params = client.db.calls[0]
assert sql == UPSERT_AUTOROUTER_SESSION_SQL
assert params == (
@ -284,8 +332,65 @@ class TestFlush:
0,
0.0,
0.0,
"canonical-user",
)
def test_a_keys_turns_stay_chronological_when_its_canonical_user_changes(self) -> None:
client: Final = _FakeClient()
earlier: Final = _transaction(user_id="z-user", at=datetime(2026, 8, 1, 12, 0, 0))
later: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 10))
asyncio.run(flush_autorouter_turn_transactions(client, [later, earlier]))
assert [(params[5], params[19]) for _, params in client.db.calls] == [
("2026-08-01T12:00:00", "z-user"),
("2026-08-01T12:00:10", "a-user"),
]
def test_one_keyless_users_failed_session_does_not_drop_another_users_turn(self) -> None:
client: Final = _FakeClient(poison_user="a-user")
failed: Final = _transaction(api_key="", user_id="a-user")
other: Final = _transaction(api_key="", user_id="b-user", at=datetime(2026, 8, 1, 12, 0, 10))
asyncio.run(flush_autorouter_turn_transactions(client, [other, failed]))
assert [(params[0], params[1], params[19]) for _, params in client.db.calls] == [("", "s1", "b-user")]
def test_uncertain_commits_quarantine_only_the_key_and_each_failed_user(self) -> None:
client: Final = _FakeClient(commit_then_error_users=frozenset({"a-failed", "c-failed"}))
turns: Final = tuple(
_transaction(user_id=user, at=datetime(2026, 8, 1, 12, 0, second), api_key=key)
for user, second, key in (
("b-healthy", 0, "k1"),
("a-failed", 1, "k1"),
("b-healthy", 2, "k1"),
("c-failed", 3, "k1"),
("b-healthy", 4, "k1"),
("d-healthy", 5, "k1"),
("c-failed", 6, "k1"),
("d-healthy", 7, "k1"),
("a-failed", 8, "k1"),
("", 9, "k1"),
("z-other", 10, "k2"),
)
)
asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns))))
assert client.db.attempts == client.db.calls
assert [
(params[0], params[19], params[5])
for sql, params in client.db.calls
if sql == UPSERT_AUTOROUTER_SESSION_SQL
] == [
("k1", "b-healthy", "2026-08-01T12:00:00"),
("k1", "a-failed", "2026-08-01T12:00:01"),
("k2", "z-other", "2026-08-01T12:00:10"),
]
assert [params[19] for _, params in client.db.attempts].count("a-failed") == 1
assert [params[19] for _, params in client.db.attempts].count("c-failed") == 1
for user, seconds in (("b-healthy", (2, 4)), ("c-failed", (3,)), ("d-healthy", (5, 7))):
assert [
(params[0], params[5])
for sql, params in client.db.calls
if sql != UPSERT_AUTOROUTER_SESSION_SQL and params[19] == user
] == [("k1", f"2026-08-01T12:00:{second:02d}") for second in seconds]
def test_a_connect_error_retries_the_same_statement(self):
client = _FakeClient(failures=[httpx.ConnectError("boom")])
asyncio.run(flush_autorouter_turn_transactions(client, [_transaction()]))

42
tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py
View file

@ -5,6 +5,7 @@ Unit tests for auto router management endpoints
from collections.abc import Mapping, Sequence
from functools import partial
from pathlib import Path
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
@ -781,17 +782,43 @@ class TestAutoRouterBenchmarks:
assert _summed_agg_row([complexity, quality]).tier_turns == {}
@pytest.mark.asyncio
async def test_non_admin_roles_cannot_read_benchmarks(self):
@pytest.mark.parametrize("user_id", [None, "own-user", "other-user"])
async def test_non_admin_roles_cannot_read_benchmarks(self, user_id: str | None):
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
with pytest.raises(HTTPException) as err:
await get_auto_router_benchmarks(
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-x"),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-x", user_id="own-user"
),
start_date="2026-08-01",
end_date="2026-08-02",
user_id=user_id,
)
assert err.value.status_code == 403
@pytest.mark.asyncio
async def test_an_empty_user_filter_is_rejected_before_querying_deployment_data(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
import httpx
from fastapi import FastAPI
from litellm.proxy import proxy_server
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
query: Final = AsyncMock(return_value=[])
monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=SimpleNamespace(query_raw=query)))
app: Final = FastAPI()
app.get("/auto_router/benchmarks")(get_auto_router_benchmarks)
app.dependency_overrides[user_api_key_auth] = lambda: ADMIN
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response: Final = await client.get("/auto_router/benchmarks", params={"user_id": ""})
assert response.status_code == 422
query.assert_not_awaited()
@pytest.mark.asyncio
async def test_a_reversed_window_is_rejected(self, monkeypatch: pytest.MonkeyPatch):
from litellm.proxy import proxy_server
@ -807,7 +834,11 @@ class TestAutoRouterBenchmarks:
assert err.value.status_code == 400
@pytest.mark.asyncio
async def test_endpoint_returns_groups_and_totals_from_the_rollup(self, monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
@pytest.mark.parametrize("user_id", [None, "selected-user"])
async def test_endpoint_returns_groups_and_totals_from_the_rollup(
self, monkeypatch: pytest.MonkeyPatch, role: LitellmUserRoles, user_id: str | None
):
from litellm.proxy import proxy_server
from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks
@ -822,12 +853,13 @@ class TestAutoRouterBenchmarks:
monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})())
response = await get_auto_router_benchmarks(
user_api_key_dict=ADMIN,
user_api_key_dict=UserAPIKeyAuth(user_role=role, api_key="sk-admin", user_id="viewer"),
start_date="2026-07-01",
end_date="2026-08-01",
api_key="key-hash",
user_id=user_id,
)
assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash")
assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash", user_id)
assert response.routers_in_scope == 1
assert response.groups[0].router_name == "live-auto"
assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0

56
tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
View file

@ -6,7 +6,7 @@ import logging
from contextlib import ExitStack
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import List, Optional, cast
from typing import Final, List, Optional, cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -1614,6 +1614,60 @@ class TestTeamScopedMCPServerAccess:
result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id")
assert len(result) == 1
class TestFetchAllMCPServersOrdering:
def test_display_order_is_case_insensitive_name_then_id(self) -> None:
servers: Final = (
LiteLLM_MCPServerTable(server_id="s-2", server_name="GitHub", alias="aaa", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="s-1", alias="github", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="s-0", server_name="Slack", alias="zzz", transport=MCPTransport.http),
LiteLLM_MCPServerTable(server_id="confluence", server_name="", alias="", transport=MCPTransport.http),
)
ordered: Final = sorted(servers, key=mgmt_endpoints._mcp_server_display_order)
assert [s.server_id for s in ordered] == ["confluence", "s-1", "s-2", "s-0"]
@pytest.mark.parametrize("team_id", [None, "team-1"])
@pytest.mark.parametrize("reverse", [False, True])
@pytest.mark.asyncio
async def test_list_is_sorted_by_display_name_regardless_of_resolution_order(
self, team_id: str | None, reverse: bool
) -> None:
mock_user_auth: Final = generate_mock_user_api_key_auth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin_user",
)
servers: Final = (
generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"),
generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"),
generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"),
)
resolved: Final = list(reversed(servers) if reverse else servers)
mock_manager: Final = MagicMock()
mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved)
with (
patch( # test-quality-ok: the route reads a module-global manager with no injection seam
"litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager",
mock_manager,
),
patch( # test-quality-ok: admin view is derived from module-global proxy settings
"litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view",
return_value=True,
),
patch( # test-quality-ok: auth contexts need a live prisma client
"litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts",
AsyncMock(return_value=[mock_user_auth]),
),
patch( # test-quality-ok: isolate the route's ordering from team database resolution
"litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list",
AsyncMock(return_value=resolved),
),
):
result: Final = await mgmt_endpoints.fetch_all_mcp_servers(
user_api_key_dict=mock_user_auth, team_id=team_id
)
assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"]
@pytest.mark.asyncio
async def test_restricted_virtual_key_cannot_use_team_id_filter(self):
"""Restricted virtual keys must not bypass access limits via team_id."""

340
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -1284,11 +1284,11 @@ class TestUpdateModel:
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
patch(
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
side_effect=lambda value: value,
),
patch(
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(
return_value=ReconcileOutcome(still_desired=None, live_after=None)
@ -4615,7 +4615,7 @@ class TestPatchModelBlockedAuthGate:
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
patch(
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(
return_value=ReconcileOutcome(still_desired=None, live_after=None)
@ -7284,3 +7284,337 @@ class TestTeamMemberAutoRouterWrites:
assert json.loads(written["model_info"])["member_auto_router"] is True
assert appended.await_args.kwargs["data"].models == ["new-personal-router"]
assert appended.await_args.kwargs["data"].team_id == "member-team"
class TestModelManagementActorEdges:
@pytest.mark.asyncio
async def test_add_model_rejects_non_team_internal_user(self):
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER)
prisma: Final = MagicMock()
deployment: Final = Deployment(
model_name="internal-model",
litellm_params=LiteLLM_Params(model="openai/test-model"),
model_info=ModelInfo(id="internal-model-id"),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
):
with pytest.raises(ProxyException) as exc_info:
await add_new_model(model_params=deployment, user_api_key_dict=actor)
assert str(exc_info.value.code) == "403"
assert "permission" in str(exc_info.value).lower()
prisma.db.litellm_proxymodeltable.create.assert_not_called()
@pytest.mark.asyncio
async def test_add_model_rejects_proxy_admin_viewer(self):
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
actor: Final = UserAPIKeyAuth(
user_id="view-only-user", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
)
prisma: Final = MagicMock()
deployment: Final = Deployment(
model_name="view-only-model",
litellm_params=LiteLLM_Params(model="openai/test-model"),
model_info=ModelInfo(id="view-only-model-id"),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
):
with pytest.raises(ProxyException) as exc_info:
await add_new_model(model_params=deployment, user_api_key_dict=actor)
assert str(exc_info.value.code) == "403"
assert "view-only" in str(exc_info.value).lower()
prisma.db.litellm_proxymodeltable.create.assert_not_called()
@pytest.mark.asyncio
async def test_add_model_requires_database_storage(self):
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
prisma: Final = MagicMock()
deployment: Final = Deployment(
model_name="database-disabled-model",
litellm_params=LiteLLM_Params(model="openai/test-model"),
model_info=ModelInfo(id="database-disabled-model-id"),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", False), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam
):
with pytest.raises(ProxyException) as exc_info:
await add_new_model(model_params=deployment, user_api_key_dict=actor)
assert str(exc_info.value.code) == "500"
assert "STORE_MODEL_IN_DB" in str(exc_info.value)
prisma.db.litellm_proxymodeltable.create.assert_not_called()
@pytest.mark.asyncio
async def test_legacy_model_update_persists_changed_field(self):
from litellm.proxy.management_endpoints.model_management_endpoints import update_model
model_id: Final = "legacy-update-model-id"
existing_row: Final = MagicMock()
existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30}
existing_row.model_dump.return_value = {
"model_name": "legacy-update-model",
"litellm_params": existing_row.litellm_params,
"model_info": {"id": model_id},
}
existing_row.model_dump_json.return_value = "{}"
updated_row: Final = MagicMock()
updated_row.model_dump_json.return_value = "{}"
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
router: Final = MagicMock()
router.get_model_ids.return_value = [model_id]
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
side_effect=lambda value: value,
),
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
await update_model(
model_params=updateDeployment(
litellm_params=updateLiteLLMParams(timeout=42),
model_info=ModelInfo(id=model_id),
),
user_api_key_dict=actor,
)
written: Final = json.loads(
prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]
)
assert written["timeout"] == 42
assert written["model"] == "openai/test-model"
@pytest.mark.asyncio
async def test_legacy_model_update_explicit_null_preserves_existing_field(self):
from litellm.proxy.management_endpoints.model_management_endpoints import update_model
model_id: Final = "legacy-null-model-id"
existing_row: Final = MagicMock()
existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30}
existing_row.model_dump.return_value = {
"model_name": "legacy-null-model",
"litellm_params": existing_row.litellm_params,
"model_info": {"id": model_id},
}
existing_row.model_dump_json.return_value = "{}"
updated_row: Final = MagicMock()
updated_row.model_dump_json.return_value = "{}"
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
router: Final = MagicMock()
router.get_model_ids.return_value = [model_id]
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
side_effect=lambda value: value,
),
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
):
await update_model(
model_params=updateDeployment(
litellm_params=updateLiteLLMParams(timeout=None),
model_info=ModelInfo(id=model_id),
),
user_api_key_dict=actor,
)
written: Final = json.loads(
prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]
)
assert written["timeout"] == 30
@pytest.mark.asyncio
async def test_patch_model_rejects_config_file_model(self):
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import patch_model
model_id: Final = "config-model-id"
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None)
prisma.db.litellm_proxymodeltable.update = AsyncMock()
router: Final = MagicMock()
router.get_deployment.return_value = Deployment(
model_name="config-model",
litellm_params=LiteLLM_Params(model="openai/test-model"),
model_info=ModelInfo(id=model_id),
)
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam
):
with pytest.raises(ProxyException) as exc_info:
await patch_model(
model_id=model_id,
patch_data=updateDeployment(
litellm_params=updateLiteLLMParams(timeout=42),
model_info=ModelInfo(id=model_id),
),
user_api_key_dict=actor,
)
assert str(exc_info.value.code) == "400"
assert "Cannot edit config-based model" in str(exc_info.value)
prisma.db.litellm_proxymodeltable.update.assert_not_awaited()
@contextlib.contextmanager
def _client_for(self, actor: UserAPIKeyAuth) -> Iterator[TestClient]:
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy.proxy_server import app
app.dependency_overrides[proxy_server.user_api_key_auth] = lambda: actor
try:
yield TestClient(app)
finally:
app.dependency_overrides.pop(proxy_server.user_api_key_auth, None)
def test_post_model_new_binds_to_actor_guard(self):
actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER)
prisma: Final = MagicMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
self._client_for(actor) as client,
):
response: Final = client.post(
"/model/new",
json={
"model_name": "internal-model",
"litellm_params": {"model": "openai/test-model"},
"model_info": {"id": "internal-model-id"},
},
)
assert response.status_code == 403
assert "permission" in response.text.lower()
prisma.db.litellm_proxymodeltable.create.assert_not_called()
def test_post_legacy_model_update_binds_to_persistence(self):
model_id: Final = "legacy-route-model-id"
existing_row: Final = LiteLLM_ProxyModelTable(
model_id=model_id,
model_name="legacy-route-model",
litellm_params={"model": "openai/test-model", "timeout": 30},
model_info={"id": model_id},
created_by="admin",
updated_by="admin",
)
updated_row: Final = LiteLLM_ProxyModelTable(
model_id=model_id,
model_name="legacy-route-model",
litellm_params={"model": "openai/test-model", "timeout": 42},
model_info={"id": model_id},
created_by="admin",
updated_by="admin",
)
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row)
prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row)
router: Final = MagicMock()
router.get_model_ids.return_value = [model_id]
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
side_effect=lambda value: value,
),
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
),
patch( # test-quality-ok: [TQ008] audit logging is outside the persistence contract
"litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log",
new=AsyncMock(return_value=None),
),
self._client_for(actor) as client,
):
response: Final = client.post(
"/model/update",
json={
"litellm_params": {"timeout": 42},
"model_info": {"id": model_id},
},
)
assert response.status_code == 200, response.text
written: Final = json.loads(
prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"]
)
assert written["timeout"] == 42
def test_patch_config_model_binds_to_patch_route(self):
model_id: Final = "config-route-model-id"
prisma: Final = MagicMock()
prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None)
prisma.db.litellm_proxymodeltable.update = AsyncMock()
router: Final = MagicMock()
router.get_deployment.return_value = Deployment(
model_name="config-route-model",
litellm_params=LiteLLM_Params(model="openai/test-model"),
model_info=ModelInfo(id=model_id),
)
actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
self._client_for(actor) as client,
):
response: Final = client.patch(
f"/model/{model_id}/update",
json={
"litellm_params": {"timeout": 42},
"model_info": {"id": model_id},
},
)
assert response.status_code == 400
assert "Cannot edit config-based model" in response.text
prisma.db.litellm_proxymodeltable.update.assert_not_awaited()

48
tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py
View file

@ -4,6 +4,8 @@ Tests for router settings management endpoints.
Tests the GET endpoints for router settings and router fields.
"""
from collections.abc import Mapping
from typing import Any, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -15,12 +17,23 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.router_settings_endpoints import (
get_router_settings,
)
from litellm.proxy.config_resolvers import SettingsStore
from litellm.proxy.proxy_server import app
from litellm.router import Router
client = TestClient(app)
class _StubProxyConfig:
def __init__(self, router_settings: SettingsStore, config_router_settings: Mapping[str, Any]) -> None:
self.router_settings: Final = router_settings
self._config_router_settings: Final = dict(config_router_settings)
async def get_config(self, config_file_path: str | None = None) -> dict[str, Any]:
del config_file_path
return {"router_settings": dict(self._config_router_settings)}
class TestRouterSettingsEndpoints:
"""Test suite for router settings endpoints"""
@ -75,6 +88,31 @@ class TestRouterSettingsEndpoints:
assert isinstance(routing_strategy_field["options"], list)
assert len(routing_strategy_field["options"]) > 0
@pytest.mark.asyncio
async def test_get_router_settings_reports_sources(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
store = SettingsStore("router_settings")
store.load_yaml({"routing_strategy": "simple-shuffle"})
store.apply_db_row("router_settings", {"num_retries": 3})
monkeypatch.setattr(
proxy_server,
"proxy_config",
_StubProxyConfig(
store,
{"routing_strategy": "simple-shuffle", "num_retries": 3},
),
)
monkeypatch.setattr(proxy_server, "llm_router", None)
admin_user = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-x"
)
response = await get_router_settings(user_api_key_dict=admin_user)
assert response.source["routing_strategy"] == "config"
assert response.source["num_retries"] == "db"
@pytest.mark.asyncio
async def test_get_router_settings_includes_routing_groups_from_live_router(
self, monkeypatch
@ -102,12 +140,10 @@ class TestRouterSettingsEndpoints:
)
monkeypatch.setattr(proxy_server, "llm_router", llm_router)
async def fake_get_config(self, config_file_path=None):
return {}
monkeypatch.setattr(
proxy_server.ProxyConfig, "get_config", fake_get_config, raising=True
proxy_server,
"proxy_config",
_StubProxyConfig(SettingsStore("router_settings"), {}),
)
admin_user = UserAPIKeyAuth(
@ -116,6 +152,8 @@ class TestRouterSettingsEndpoints:
response = await get_router_settings(user_api_key_dict=admin_user)
assert response.current_values.get("routing_groups") == groups
assert response.current_values["timeout"] is not None
assert response.source["timeout"] == "default"
rg_field = next(f for f in response.fields if f.field_name == "routing_groups")
assert rg_field.field_value == groups

176
tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py
View file

@ -11,13 +11,18 @@ Pins (PR2):
from __future__ import annotations
from collections.abc import Callable, Mapping
from contextlib import AbstractContextManager
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi.testclient import TestClient
import litellm
from litellm.proxy import proxy_server
from litellm.proxy._types import LitellmUserRoles
from litellm.proxy.config_resolvers.settings_rules import JsonValue
from litellm.proxy.config_resolvers.settings_store import SettingsStore
from .conftest import normalize # type: ignore[import-not-found]
@ -179,6 +184,177 @@ def test_model_settings_method_not_allowed(client, auth_as):
# ---------------------------------------------------------------------------
def _alerting_client(
monkeypatch: pytest.MonkeyPatch,
*,
yaml_values: Mapping[str, JsonValue],
db_row: Mapping[str, JsonValue],
live_args: Mapping[str, JsonValue],
) -> "SettingsStore":
pc = MagicMock()
row = MagicMock()
row.param_value = db_row
pc.db.litellm_config.find_first = AsyncMock(return_value=row)
monkeypatch.setattr(proxy_server, "prisma_client", pc)
logging_obj = MagicMock()
args_model = MagicMock()
args_model.model_dump = MagicMock(return_value=live_args)
logging_obj.slack_alerting_instance.alerting_args = args_model
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj)
store = SettingsStore("general_settings")
store.load_yaml(yaml_values)
store.apply_db_row("general_settings", db_row)
monkeypatch.setattr(proxy_server.proxy_config, "settings", store)
monkeypatch.setattr(proxy_server, "general_settings", store)
return store
def test_alerting_settings_reports_sources(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
_alerting_client(
monkeypatch,
yaml_values={
"alerting": ["slack"],
"alerting_args": {"daily_report_frequency": 3, "report_check_interval": 300},
},
db_row={"alerting_args": {"daily_report_frequency": 7, "outage_alert_ttl": 4242}},
live_args={"daily_report_frequency": 3},
)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/alerting/settings")
assert response.status_code == 200
by_name = {entry["field_name"]: entry for entry in response.json()}
assert by_name["slack_alerting"]["source"] == "config"
assert by_name["daily_report_frequency"]["source"] == "config"
assert by_name["report_check_interval"]["source"] == "config"
assert by_name["outage_alert_ttl"]["source"] == "default"
assert by_name["budget_alert_ttl"]["source"] == "default"
def test_alerting_settings_reports_db_source_when_the_file_omits_alerting_args(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
store = _alerting_client(
monkeypatch,
yaml_values={"alerting": ["slack"]},
db_row={
"alerting_args": {
"outage_alert_ttl": 4242,
"region_outage_alert_ttl": [],
"report_check_interval": None,
}
},
live_args={"outage_alert_ttl": 4242},
)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/alerting/settings")
assert response.status_code == 200
by_name = {entry["field_name"]: entry for entry in response.json()}
assert store.source("alerting_args") == "db"
assert by_name["outage_alert_ttl"]["source"] == "db"
assert by_name["region_outage_alert_ttl"]["source"] == "db"
assert by_name["report_check_interval"]["source"] == "db"
assert by_name["budget_alert_ttl"]["source"] == "default"
def test_alerting_settings_reports_config_source_when_db_disagrees(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.proxy.config_resolvers import SettingsStore
db_alerting_args = {"daily_report_frequency": 7}
pc = MagicMock()
row = MagicMock()
row.param_value = {"alerting_args": db_alerting_args}
pc.db.litellm_config.find_first = AsyncMock(return_value=row)
monkeypatch.setattr(proxy_server, "prisma_client", pc)
logging_obj = MagicMock()
args_model = MagicMock()
args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 3})
logging_obj.slack_alerting_instance.alerting_args = args_model
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj)
store = SettingsStore("general_settings")
store.load_yaml({"alerting_args": {"daily_report_frequency": 3}})
store.apply_db_row("general_settings", {"alerting_args": db_alerting_args})
monkeypatch.setattr(proxy_server.proxy_config, "settings", store)
monkeypatch.setattr(proxy_server, "general_settings", store)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/alerting/settings")
assert response.status_code == 200
by_name = {entry["field_name"]: entry for entry in response.json()}
assert store.source("alerting_args") == "config"
assert by_name["daily_report_frequency"]["field_value"] == 3
assert by_name["daily_report_frequency"]["source"] == "config"
@pytest.mark.parametrize("db_alerting_args", [None, []])
def test_alerting_settings_handles_empty_db_args(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
monkeypatch: pytest.MonkeyPatch,
db_alerting_args: JsonValue,
) -> None:
from litellm.proxy.config_resolvers import SettingsStore
pc = MagicMock()
row = MagicMock()
row.param_value = {"alerting_args": db_alerting_args}
pc.db.litellm_config.find_first = AsyncMock(return_value=row)
monkeypatch.setattr(proxy_server, "prisma_client", pc)
logging_obj = MagicMock()
args_model = MagicMock()
args_model.model_dump = MagicMock(return_value={})
logging_obj.slack_alerting_instance.alerting_args = args_model
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj)
store = SettingsStore("general_settings")
store.load_yaml({"alerting_args": {"report_check_interval": 300}})
monkeypatch.setattr(proxy_server.proxy_config, "settings", store)
monkeypatch.setattr(proxy_server, "general_settings", store)
with auth_as(LitellmUserRoles.PROXY_ADMIN):
response = client.get("/alerting/settings")
assert response.status_code == 200
by_name = {entry["field_name"]: entry for entry in response.json()}
assert by_name["report_check_interval"]["source"] == "config"
assert by_name["budget_alert_ttl"]["source"] == "default"
@pytest.mark.parametrize(
("field_default", "expected"),
[(43200, "default"), (None, "unset")],
)
def test_nested_setting_source_without_a_config_or_db_value(field_default: JsonValue, expected: str) -> None:
store = SettingsStore("general_settings")
store.load_yaml({})
assert (
proxy_server._nested_setting_source(store, {}, "alerting_args", "budget_alert_ttl", field_default) == expected
)
def test_alerting_settings_no_db_error(client, auth_as, no_prisma):
"""Pins ``GET /alerting/settings`` (error: db not connected)."""
with auth_as(LitellmUserRoles.PROXY_ADMIN):

158
tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py Normal file
View file

@ -0,0 +1,158 @@
import asyncio
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from datetime import datetime, timedelta
import pytest
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.shutdown.scheduled_jobs import (
AwaitableAsyncIOExecutor,
pause_scheduled_jobs,
stop_in_flight_scheduler_jobs,
)
class _Job:
"""A scheduled job that blocks until cancelled, or for ``work_seconds``, and records what it observed"""
def __init__(self, swallow_cancellation: bool = False, work_seconds: float | None = None) -> None:
self.started = asyncio.Event()
self.events: list[str] = []
self.swallow_cancellation = swallow_cancellation
self.work_seconds = work_seconds
async def run(self) -> None:
self.started.set()
try:
if self.work_seconds is None:
await asyncio.Event().wait()
else:
await asyncio.sleep(self.work_seconds)
self.events.append("committed")
except asyncio.CancelledError:
self.events.append("cancelled")
if self.swallow_cancellation:
await asyncio.Event().wait()
raise
finally:
self.events.append("finished")
@asynccontextmanager
async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOScheduler, AwaitableAsyncIOExecutor]]:
"""A started scheduler with every job in flight, stopped on the way out whatever the test did"""
executor = AwaitableAsyncIOExecutor()
scheduler = AsyncIOScheduler(executors={"default": executor})
for index, job in enumerate(jobs):
scheduler.add_job(job.run, id=f"job-{index}", next_run_time=datetime.now())
scheduler.start()
try:
for job in jobs:
await asyncio.wait_for(job.started.wait(), timeout=5)
yield scheduler, executor
finally:
if scheduler.running:
scheduler.shutdown(wait=False)
stragglers = executor.in_flight_jobs()
for straggler in stragglers:
straggler.cancel()
await asyncio.gather(*stragglers, return_exceptions=True)
@pytest.mark.asyncio
async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns():
"""The job's own CancelledError handler records how a run ended, so shutdown must wait for it"""
job = _Job()
async with _running_scheduler(job) as (scheduler, executor):
await stop_in_flight_scheduler_jobs(scheduler, executor)
assert job.events == ["cancelled", "finished"]
assert scheduler.running is False
assert executor.in_flight_jobs() == ()
@pytest.mark.asyncio
async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled():
"""A spend write cancelled mid-commit drops the rows it popped, so short jobs get to finish first"""
write = _Job(work_seconds=0.2)
stuck = _Job()
async with _running_scheduler(write, stuck) as (scheduler, executor):
await stop_in_flight_scheduler_jobs(scheduler, executor, finish_timeout_seconds=2.0)
assert write.events == ["committed", "finished"]
assert stuck.events == ["cancelled", "finished"]
assert scheduler.running is False
@pytest.mark.asyncio
async def test_every_in_flight_job_is_cancelled_not_only_the_first():
first, second = _Job(), _Job()
async with _running_scheduler(first, second) as (scheduler, executor):
await stop_in_flight_scheduler_jobs(scheduler, executor)
assert first.events == ["cancelled", "finished"]
assert second.events == ["cancelled", "finished"]
@pytest.mark.asyncio
async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(caplog):
"""A job that swallows CancelledError must not hold the pod past its termination grace period"""
job = _Job(swallow_cancellation=True)
async with _running_scheduler(job) as (scheduler, executor):
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await stop_in_flight_scheduler_jobs(scheduler, executor, cancel_timeout_seconds=0.05)
assert job.events == ["cancelled"]
assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text
@pytest.mark.asyncio
async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler():
async with _running_scheduler() as (scheduler, executor):
await stop_in_flight_scheduler_jobs(scheduler, executor)
await asyncio.sleep(0)
assert scheduler.running is False
@pytest.mark.asyncio
async def test_a_scheduler_that_never_started_is_left_alone():
"""The proxy runs without a scheduler when it has no database"""
executor = AwaitableAsyncIOExecutor()
scheduler = AsyncIOScheduler(executors={"default": executor})
await stop_in_flight_scheduler_jobs(scheduler, executor)
assert scheduler.running is False
@pytest.mark.asyncio
async def test_pausing_stops_new_jobs_from_starting_but_leaves_running_ones_alone():
"""A job due during the shutdown drain would only be cancelled, so it must not start at all"""
running = _Job()
async with _running_scheduler(running) as (scheduler, executor):
late = _Job()
scheduler.add_job(late.run, id="late", next_run_time=datetime.now() + timedelta(seconds=0.1))
pause_scheduled_jobs(scheduler)
await asyncio.sleep(0.3)
assert late.started.is_set() is False
assert running.events == []
assert scheduler.running is True
await stop_in_flight_scheduler_jobs(scheduler, executor)
assert running.events == ["cancelled", "finished"]
assert late.started.is_set() is False
@pytest.mark.asyncio
async def test_pausing_a_scheduler_that_never_started_is_a_no_op():
scheduler = AsyncIOScheduler(executors={"default": AwaitableAsyncIOExecutor()})
pause_scheduled_jobs(scheduler)
assert scheduler.running is False

60
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -11766,6 +11766,66 @@ def test_prompt_caching_settings_propagate_on_config_reload(monkeypatch, field_n
assert getattr(litellm, field_name) == db_value
@pytest.mark.asyncio
async def test_db_stored_datadog_redaction_settings_apply_before_logger_init(monkeypatch: pytest.MonkeyPatch):
"""A DB-only litellm_settings row that pairs success_callback: ["datadog"] with
datadog_params.turn_off_message_logging: true must build the DataDogLogger redacted, the
same as the identical block in YAML. Regression for the redaction keys being absent from
the safe-override allowlist while the callback half of the row was honoured."""
import litellm.proxy.proxy_server as ps
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.litellm_core_utils import litellm_logging
monkeypatch.setenv("DD_API_KEY", "test-key")
monkeypatch.setenv("DD_SITE", "us5.datadoghq.com")
monkeypatch.setattr(litellm, "datadog_params", None)
monkeypatch.setattr(litellm, "turn_off_message_logging", False)
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(litellm_logging, "_in_memory_loggers", [])
db_row = {
"success_callback": ["datadog"],
"datadog_params": {"turn_off_message_logging": True},
"turn_off_message_logging": True,
}
pc = ps.ProxyConfig()
pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", db_row))
pc._add_callbacks_from_db_config({"litellm_settings": db_row})
datadog_loggers = [cb for cb in litellm.success_callback if isinstance(cb, DataDogLogger)]
assert len(datadog_loggers) == 1
assert datadog_loggers[0].turn_off_message_logging is True
assert litellm.turn_off_message_logging is True
@pytest.mark.parametrize(
"field_name",
[
"datadog_params",
"datadog_llm_observability_params",
"newrelic_params",
"pointfive_params",
"aws_sqs_callback_params",
],
)
def test_db_stored_callback_params_propagate_to_litellm_module(monkeypatch: pytest.MonkeyPatch, field_name: str):
"""Every callback init params block stored in the DB litellm_settings row must land on the
litellm module before the matching logger is built, so the DB row behaves like YAML."""
import litellm.proxy.proxy_server as ps
monkeypatch.setattr(litellm, field_name, None)
db_value = {"turn_off_message_logging": True}
pc = ps.ProxyConfig()
pc._apply_litellm_settings_db_values(pc._prepared_db_settings_values("litellm_settings", {field_name: db_value}))
assert getattr(litellm, field_name) == db_value
def test_get_config_list_marks_untouched_prompt_caching_flag_as_not_set(monkeypatch):
"""The flag defaults to False rather than None, so a plain 'is not None' check would
report the default as 'In Config' and imply an admin had set it."""

136
tests/test_litellm/proxy/test_spend_log_cleanup.py
View file

@ -7,6 +7,7 @@ import math
import time
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -793,18 +794,20 @@ async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup():
tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
assert any('"LiteLLM_SpendLogs"' in sql for sql in tables)
assert not any('"LiteLLM_AutoRouterSession"' in sql for sql in tables)
assert not any('"LiteLLM_AutoRouterUserSession"' in sql for sql in tables)
assert not any('"LiteLLM_HealthCheckTable"' in sql for sql in tables)
@pytest.mark.asyncio
async def test_session_retention_alone_cleans_only_the_session_rollup():
client = _mock_prisma_for_retention([0])
async def test_session_retention_alone_cleans_both_session_rollups():
client = _mock_prisma_for_retention([0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_autorouter_session_retention_period": "365d"})
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(client)
tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
assert len(tables) == 1
assert len(tables) == 2
assert '"LiteLLM_AutoRouterSession"' in tables[0]
assert '"LiteLLM_AutoRouterUserSession"' in tables[1]
@pytest.mark.asyncio
@ -825,7 +828,7 @@ async def test_health_check_retention_alone_cleans_only_the_health_check_table()
@pytest.mark.asyncio
async def test_each_retention_key_cuts_off_at_its_own_horizon():
client = _mock_prisma_for_retention([0, 0, 0, 0])
client = _mock_prisma_for_retention([0, 0, 0, 0, 0])
cleaner = SpendLogCleanup(
general_settings={
"maximum_spend_logs_retention_period": "7d",
@ -839,6 +842,8 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
(
"LiteLLM_AutoRouterSession"
if '"LiteLLM_AutoRouterSession"' in call[0][0]
else "LiteLLM_AutoRouterUserSession"
if '"LiteLLM_AutoRouterUserSession"' in call[0][0]
else "LiteLLM_HealthCheckTable"
if '"LiteLLM_HealthCheckTable"' in call[0][0]
else "logs"
@ -848,6 +853,7 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
now = datetime.now(timezone.utc)
assert (now - cutoffs["logs"]).days == 7
assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365
assert cutoffs["LiteLLM_AutoRouterUserSession"] == cutoffs["LiteLLM_AutoRouterSession"]
assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30
@ -1417,3 +1423,125 @@ def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(st
"""
results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons)
assert SpendLogCleanup._run_outcome(results) == expected
_OTHER_OUTCOMES: Final = ("completed", "budget_exhausted", "batch_cap_reached", "skipped_locked", "skipped_disabled")
def _runs_recorded(outcome: str) -> float:
"""The real ``litellm_spend_log_cleanup_runs_total`` sample for one outcome, 0 when unset"""
from prometheus_client import REGISTRY
return REGISTRY.get_sample_value("litellm_spend_log_cleanup_runs_total", {"outcome": outcome}) or 0.0
@pytest.mark.asyncio
async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch):
"""A run cut short by shutdown must leave its outcome and how far it got behind"""
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
mock_logger = MagicMock()
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
aborted_runs_before = _runs_recorded("aborted")
other_runs_before = {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES}
third_batch_reached = asyncio.Event()
async def _execute_raw(sql, *args):
if third_batch_reached.is_set():
raise AssertionError("no batch may be issued after the cancelled one")
if _execute_raw.calls < 2:
_execute_raw.calls += 1
return 150
third_batch_reached.set()
await asyncio.Event().wait()
_execute_raw.calls = 0
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = _execute_raw
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = MagicMock()
cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
cleaner.pod_lock_manager.release_lock = AsyncMock()
run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client))
await asyncio.wait_for(third_batch_reached.wait(), timeout=5)
run.cancel()
with pytest.raises(asyncio.CancelledError):
await run
assert _runs_recorded("aborted") == aborted_runs_before + 1
assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before
cleaner.pod_lock_manager.release_lock.assert_awaited_once()
mock_logger.exception.assert_not_called()
(error_call,) = mock_logger.error.call_args_list
rendered = error_call[0][0] % error_call[0][1:]
assert rendered.startswith("Spend log cleanup cancelled after ")
assert "s (rows_deleted=300, batches=2)" in rendered
@pytest.mark.asyncio
async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch):
"""The scheduler holds one cleaner for the life of the process, so progress must not carry over"""
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
mock_logger = MagicMock()
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()])
with pytest.raises(asyncio.CancelledError):
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
(error_call,) = mock_logger.error.call_args_list
rendered = error_call[0][0] % error_call[0][1:]
assert "(rows_deleted=150, batches=1)" in rendered
@pytest.mark.asyncio
async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch):
"""With APSCHEDULER_MAX_INSTANCES above one, two runs share the cleaner but not their progress"""
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
mock_logger = MagicMock()
monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
first_batch_done = asyncio.Event()
second_run_done = asyncio.Event()
async def _slow_execute_raw(sql, *args):
first_batch_done.set()
await second_run_done.wait()
return 100
slow_client = MagicMock()
_wire_tx(slow_client.db)
slow_client.db.execute_raw = _slow_execute_raw
fast_client = MagicMock()
_wire_tx(fast_client.db)
fast_client.db.execute_raw = AsyncMock(side_effect=[150, 150, 0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = None
slow_run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(slow_client))
await asyncio.wait_for(first_batch_done.wait(), timeout=5)
await cleaner.cleanup_old_spend_logs(fast_client)
second_run_done.set()
await asyncio.sleep(0)
slow_run.cancel()
with pytest.raises(asyncio.CancelledError):
await slow_run
(error_call,) = mock_logger.error.call_args_list
rendered = error_call[0][0] % error_call[0][1:]
assert "(rows_deleted=100, batches=1)" in rendered

84
tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py
View file

@ -1342,6 +1342,48 @@ class TestProxySettingEndpoints:
where={"id": "ui_settings"}
)
def test_get_ui_settings_reports_sources(self, monkeypatch: pytest.MonkeyPatch) -> None:
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy import proxy_server
from litellm.proxy.config_resolvers import SettingsStore
mock_prisma = MagicMock()
mock_db_record = MagicMock()
mock_db_record.ui_settings = {
"disable_model_add_for_internal_users": True,
"require_auth_for_public_ai_hub": True,
}
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(
return_value=mock_db_record
)
monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma)
store = SettingsStore("general_settings")
store.load_yaml(
{
"disable_model_add_for_internal_users": False,
"forward_client_headers_to_llm_api": True,
}
)
store.apply_db_row(
"ui_settings",
{"disable_model_add_for_internal_users": True},
)
monkeypatch.setattr(proxy_server.proxy_config, "settings", store)
monkeypatch.setattr(proxy_server, "general_settings", store)
response = client.get("/get/ui_settings")
assert response.status_code == 200
data = response.json()
assert data["values"]["disable_model_add_for_internal_users"] is False
assert data["values"]["forward_client_headers_to_llm_api"] is True
assert data["values"]["require_auth_for_public_ai_hub"] is True
assert data["source"]["disable_model_add_for_internal_users"] == "config"
assert data["source"]["forward_client_headers_to_llm_api"] == "config"
assert data["source"]["require_auth_for_public_ai_hub"] == "db"
def test_get_ui_settings_schema_description_preserved_with_extensions(
self, mock_auth, monkeypatch
):
@ -3477,6 +3519,7 @@ class TestPtuCostAttributionUISetting:
assert response.status_code == 200
assert response.json()["values"]["enable_ptu_cost_attribution"] is False
assert response.json()["source"]["enable_ptu_cost_attribution"] == "default"
def test_reported_true_once_the_env_var_is_set(self, mock_auth, monkeypatch):
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
@ -3488,6 +3531,47 @@ class TestPtuCostAttributionUISetting:
assert response.status_code == 200
assert response.json()["values"]["enable_ptu_cost_attribution"] is True
assert response.json()["source"]["enable_ptu_cost_attribution"] == "config"
def test_reported_config_when_secret_manager_enables_the_flag(
self, mock_auth: None, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
monkeypatch.setattr(
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.is_ptu_cost_attribution_enabled",
lambda: True,
)
self._mock_prisma(monkeypatch)
response = client.get("/get/ui_settings")
assert response.status_code == 200
assert response.json()["values"]["enable_ptu_cost_attribution"] is True
assert response.json()["source"]["enable_ptu_cost_attribution"] == "config"
def test_reported_config_when_secret_manager_disables_the_flag(
self, mock_auth: None, monkeypatch: pytest.MonkeyPatch
) -> None:
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
monkeypatch.setattr(
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.is_ptu_cost_attribution_enabled",
lambda: False,
)
monkeypatch.setattr(
"litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_secret",
lambda *_args: False,
)
self._mock_prisma(monkeypatch)
response = client.get("/get/ui_settings")
assert response.status_code == 200
assert response.json()["values"]["enable_ptu_cost_attribution"] is False
assert response.json()["source"]["enable_ptu_cost_attribution"] == "config"
def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch):
"""A row written before the allowlist existed must not be able to turn the feature on."""

70
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -13,6 +13,7 @@ import time
from collections.abc import AsyncIterator, Mapping, Sequence
from copy import deepcopy
from functools import partial
from types import MappingProxyType
from typing import Dict, Final, List, Literal
from unittest.mock import AsyncMock, MagicMock, patch
@ -90,6 +91,7 @@ from litellm.types.router import (
TaggedPreRoutingStrategy,
)
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.management_endpoints.auto_router_endpoints import RequestComplexityRouterConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@ -6681,6 +6683,7 @@ class TestTierModelAffinity:
litellm_router_instance=_windowed_router(_SMALL, _BIG),
complexity_router_config={
"tiers": {"SIMPLE": ["small-model", "big-model"]},
"enable_context_window_escalation": True,
"adaptive": adaptive,
"deployment_affinity": True,
"session_affinity": False,
@ -13878,8 +13881,12 @@ _CJK_TURNS = [
]
def _tier_config(**overrides) -> Dict:
return {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, **overrides}
def _tier_config(**overrides: object) -> dict[str, object]:
return {
"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"},
"enable_context_window_escalation": True,
**overrides,
}
class TestContextWindowEscalation:
@ -13938,7 +13945,7 @@ class TestContextWindowEscalation:
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG),
complexity_router_config={"tiers": {"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}},
complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}),
)
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
@ -13975,7 +13982,7 @@ class TestContextWindowEscalation:
},
]
),
complexity_router_config={"tiers": {"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}},
complexity_router_config=_tier_config(tiers={"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}),
)
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
@ -14026,7 +14033,7 @@ class TestContextWindowEscalation:
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=_windowed_router(*deployments),
complexity_router_config={"tiers": tiers},
complexity_router_config=_tier_config(tiers=tiers),
)
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
@ -14035,19 +14042,37 @@ class TestContextWindowEscalation:
assert result.model == expected_model
@pytest.mark.asyncio
async def test_the_disabled_gate_dispatches_on_complexity_alone(self):
"""The escape hatch: enable_context_window_escalation false restores today's behavior."""
router = ComplexityRouter(
@pytest.mark.parametrize("enabled", (None, False, True), ids=("omitted", "disabled", "enabled"))
@pytest.mark.parametrize("serialized", (False, True), ids=("config", "http-json"))
async def test_context_window_escalation_requires_opt_in(self, enabled: bool | None, serialized: bool) -> None:
setting: Final = (
MappingProxyType({"enable_context_window_escalation": enabled})
if enabled is not None
else MappingProxyType({})
)
raw_config: Final = RequestComplexityRouterConfig.model_validate(
MappingProxyType(
{"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting}
)
)
config: Final = (
RequestComplexityRouterConfig.model_validate_json(raw_config.model_dump_json())
if serialized
else raw_config
)
router: Final = ComplexityRouter(
model_name="test-router",
litellm_router_instance=_windowed_router(_SMALL, _BIG),
complexity_router_config=_tier_config(enable_context_window_escalation=False),
complexity_router_config=config.model_dump(exclude_unset=not serialized, exclude_none=True),
)
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
result: Final = await router.async_pre_routing_hook(
model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS
)
assert result is not None
assert result.model == "small-model"
assert "context_escalated" not in result.routing_decision
assert result.model == ("big-model" if enabled else "small-model")
assert result.routing_decision.get("context_escalated", False) is (enabled is True)
@pytest.mark.asyncio
async def test_out_of_band_system_and_tools_count_against_the_window(self):
@ -14156,7 +14181,7 @@ class TestContextWindowEscalation:
},
]
),
complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["small-model", "mid-model"]}},
complexity_router_config=_tier_config(adaptive=True, tiers={"SIMPLE": ["small-model", "mid-model"]}),
)
result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS)
@ -14186,7 +14211,7 @@ class TestContextWindowEscalation:
},
]
),
complexity_router_config={"tiers": {"SIMPLE": "cop-pool", "COMPLEX": "big-model"}},
complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}),
)
real_get_llm_provider = litellm.get_llm_provider
copilot_resolutions: List = []
@ -14219,7 +14244,7 @@ class TestContextWindowEscalation:
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}},
"complexity_router_config": _tier_config(),
},
},
{
@ -15139,7 +15164,12 @@ class TestHealthFallbackDispatch:
) -> None:
from litellm.types.router import RouterRateLimitError
router: Final = self._router(config={"tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"}})
router: Final = self._router(
config={
"tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"},
"enable_context_window_escalation": True,
}
)
router.add_deployment(
Deployment(
model_name="large",
@ -15216,7 +15246,13 @@ class TestHealthFallbackDispatch:
@pytest.mark.asyncio
@pytest.mark.parametrize("default_fits", [True, False])
async def test_modality_default_must_also_fit_context(self, default_fits: bool) -> None:
router: Final = self._router(config={"modality_routing": True, "tiers": {"SIMPLE": "primary"}})
router: Final = self._router(
config={
"modality_routing": True,
"tiers": {"SIMPLE": "primary"},
"enable_context_window_escalation": True,
}
)
for deployment in router.model_list:
deployment["model_info"]["supports_vision"] = deployment["model_name"] == "fallback"
deployment["model_info"]["max_input_tokens"] = 10000 if default_fits else 10

54
tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py Normal file
View file

@ -0,0 +1,54 @@
from datetime import datetime, timedelta
from typing import Final
from litellm import Router
from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict
MODEL_GROUP: Final = "lowest-tpm-router"
HIGH_USAGE_DEPLOYMENT_ID: Final = "highest-usage"
LOW_USAGE_DEPLOYMENT_ID: Final = "lowest-usage"
def _deployment(deployment_id: str) -> DeploymentTypedDict:
params: LiteLLMParamsTypedDict = {
"model": "gpt-4o",
"api_key": "key",
"mock_response": f"from {deployment_id}",
}
return {
"model_name": MODEL_GROUP,
"litellm_params": params,
"model_info": {"id": deployment_id},
}
def test_usage_based_routing_v1_selects_the_lowest_recorded_tpm() -> None:
router: Final = Router(
model_list=[
_deployment(HIGH_USAGE_DEPLOYMENT_ID),
_deployment(LOW_USAGE_DEPLOYMENT_ID),
],
routing_strategy="usage-based-routing",
num_retries=0,
)
usage_by_deployment: Final = {
HIGH_USAGE_DEPLOYMENT_ID: 100,
LOW_USAGE_DEPLOYMENT_ID: 1,
}
now: Final = datetime.now()
cache_keys: Final = tuple(
f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}"
for offset in range(60)
)
for cache_key in cache_keys:
router.cache.set_cache(
key=cache_key, value=usage_by_deployment, ttl=float("inf")
)
deployment: Final = router.get_available_deployment(
model=MODEL_GROUP,
messages=[{"role": "user", "content": "test"}],
)
assert deployment["model_info"]["id"] == LOW_USAGE_DEPLOYMENT_ID

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