scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 145

Merge remote-tracking branch 'origin/main' into litellm_mcp_ui_prompts_resources
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	tests/e2e/ui/tests/mcp/mcpTools.spec.ts
This commit is contained in:
joshua 2026-09-22 03:41:57 +00:00
commit 109dc8cc1c
390 changed files with 37293 additions and 2594 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

57
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json
View file

@ -6267,6 +6267,63 @@
],
"title": "Spend update queue sizes (litellm_<queue>_size)",
"type": "timeseries"
},
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"description": "Requests that carried usage but were logged at $0 on a model whose pricing entry has a non-zero rate, by requested model and reason",
"fieldConfig": {
"defaults": {
"color": {
"mode": "palette-classic"
},
"custom": {
"drawStyle": "line",
"fillOpacity": 10,
"lineWidth": 1,
"showPoints": "never",
"spanNulls": false
},
"unit": "short"
},
"overrides": []
},
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 430
},
"id": 110,
"options": {
"legend": {
"calcs": [],
"displayMode": "list",
"placement": "bottom",
"showLegend": true
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${DS_PROMETHEUS}"
},
"editorMode": "code",
"expr": "sum(rate(litellm_zero_cost_requests_total[$__rate_interval])) by (requested_model, reason)",
"legendFormat": "{{requested_model}} / {{reason}}",
"range": true,
"refId": "A"
}
],
"title": "litellm_zero_cost_requests rate",
"type": "timeseries"
}
],
"preload": false,

11
docker/README.md
View file

@ -2,6 +2,17 @@
This guide provides instructions for building and running the LiteLLM application using Docker and Docker Compose.
> **Just want to run LiteLLM?** This guide builds from source. To run the published
> image instead, use `docker-compose.quickstart.yml` in this directory — the
> two-service stack (gateway + Postgres) that the
> [Docker quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start) documents:
>
> ```bash
> curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
> printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
> docker compose -f docker-compose.quickstart.yml up -d
> ```
## Prerequisites
- Docker

41
docker/docker-compose.quickstart.yml Normal file
View file

@ -0,0 +1,41 @@
# LiteLLM quickstart stack: the gateway plus a Postgres database that stores
# models, virtual keys, and spend logs. Used by
# https://docs.litellm.ai/docs/proxy/docker_quick_start
#
# curl -sSLO https://github.com/BerriAI/litellm/raw/main/docker/docker-compose.quickstart.yml
# printf 'LITELLM_MASTER_KEY=sk-%s\nLITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" "$(openssl rand -hex 32)" > .env
# docker compose -f docker-compose.quickstart.yml up -d
#
# Compose reads .env from this directory. Keep it: regenerating LITELLM_SALT_KEY
# makes credentials already stored in the database unreadable. For anything
# beyond local evaluation, pin the image to a specific release tag.
services:
litellm:
image: docker.litellm.ai/berriai/litellm:main-stable
ports:
- "4000:4000"
environment:
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env - see the header of this file}
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?set it in .env - see the header of this file}
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
STORE_MODEL_IN_DB: "True"
depends_on:
db:
condition: service_healthy
db:
image: postgres:16
environment:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_DB: litellm
healthcheck:
test: ["CMD-SHELL", "pg_isready -U litellm"]
interval: 5s
timeout: 5s
retries: 10
volumes:
- postgres_data:/var/lib/postgresql/data
volumes:
postgres_data:

2
gateway/routes/allowlist.py
View file

@ -96,10 +96,12 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/assemblyai/",
"/eu.assemblyai/",
"/deepgram/",
"/fal_ai/",
"/langfuse/",
"/vllm/",
"/mistral/",
"/typesafe/",
"/openrouter/",
"/nvidia_nim/",
"/groq/",
"/voyage/",

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];

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

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260921000000_add_password_reset_columns/migration.sql Normal file
View file

@ -0,0 +1,3 @@
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "password_reset_required" BOOLEAN;
ALTER TABLE "LiteLLM_UserTable" ADD COLUMN IF NOT EXISTS "last_breach_check_at" TIMESTAMP(3);

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

@ -73,6 +73,7 @@ model LiteLLM_AgentsTable {
static_headers Json? @default("{}")
extra_headers String[] @default([])
agent_access_groups String[] @default([])
access_group_ids String[] @default([])
object_permission_id String?
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
spend Float @default(0.0)
@ -246,6 +247,8 @@ model LiteLLM_UserTable {
organization_id String?
object_permission_id String?
password String?
password_reset_required Boolean?
last_breach_check_at DateTime?
teams String[] @default([])
user_role String?
max_budget Float?
@ -1621,6 +1624,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

542
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"
@ -569,6 +650,49 @@ dependencies = [
"tracing",
]
[[package]]
name = "axum"
version = "0.8.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
dependencies = [
"axum-core",
"bytes",
"futures-util",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"itoa",
"matchit",
"memchr",
"mime",
"percent-encoding",
"pin-project-lite",
"serde_core",
"sync_wrapper",
"tower",
"tower-layer",
"tower-service",
]
[[package]]
name = "axum-core"
version = "0.5.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
dependencies = [
"bytes",
"futures-core",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"mime",
"pin-project-lite",
"sync_wrapper",
"tower-layer",
"tower-service",
]
[[package]]
name = "azure_core"
version = "1.1.0"
@ -980,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"
@ -1377,6 +1511,18 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "fallible-iterator"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
[[package]]
name = "fallible-streaming-iterator"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
[[package]]
name = "fancy-regex"
version = "0.17.0"
@ -1428,6 +1574,12 @@ version = "1.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
[[package]]
name = "foldhash"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
[[package]]
name = "form_urlencoded"
version = "1.2.2"
@ -1892,11 +2044,34 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
dependencies = [
"foldhash",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
dependencies = [
"allocator-api2",
"equivalent",
"foldhash",
]
[[package]]
name = "hashlink"
version = "0.12.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a596f1b20ed2cc5ecac41a164aaebc7258057060f06c0cf7a2ba3991ee7990fb"
dependencies = [
"hashbrown 0.17.1",
]
[[package]]
name = "heck"
@ -2277,6 +2452,12 @@ version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "iter-read"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071ed4cc1afd86650602c7b11aa2e1ce30762a1c27193201cb5cee9c6ebb1294"
[[package]]
name = "itertools"
version = "0.13.0"
@ -2435,6 +2616,17 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libsqlite3-sys"
version = "0.38.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@ -2540,6 +2732,36 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-cache-disk"
version = "0.1.0"
dependencies = [
"litellm-cache",
"py_literal",
"rand 0.8.7",
"rstest",
"rusqlite",
"serde-pickle",
"serde_json",
"tempfile",
"tokio",
]
[[package]]
name = "litellm-cache-gcs"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-cache",
"percent-encoding",
"reqwest 0.12.28",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-cache-memory"
version = "0.1.0"
@ -2550,6 +2772,26 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-qdrant-semantic"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-response",
"qdrant-client",
"reqwest 0.12.28",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"tokio-stream",
"tonic",
"tonic-prost",
"uuid",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
@ -2562,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"
@ -2578,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"
@ -2740,12 +3028,19 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-azure-blob",
"litellm-cache-disk",
"litellm-cache-gcs",
"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",
@ -2756,12 +3051,17 @@ dependencies = [
"litellm-types",
"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",
]
[[package]]
@ -2778,6 +3078,7 @@ dependencies = [
"litellm-secrets-azure",
"litellm-secrets-cyberark",
"litellm-secrets-google",
"litellm-secrets-hashicorp",
"litellm-secrets-types",
"moka",
"reqwest 0.12.28",
@ -2875,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"
@ -2966,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"
@ -2988,6 +3318,22 @@ version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c"
[[package]]
name = "matchit"
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"
@ -3601,6 +3947,27 @@ dependencies = [
"serde",
]
[[package]]
name = "qdrant-client"
version = "1.19.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dddc19df129bad7346ebd027288621ab1ac7e52678371f906b9a8622d7aaf87e"
dependencies = [
"anyhow",
"derive_builder",
"futures",
"parking_lot",
"prost",
"prost-types",
"semver",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"tonic",
"tonic-prost",
]
[[package]]
name = "quick-error"
version = "1.2.3"
@ -4033,6 +4400,16 @@ dependencies = [
"windows-sys 0.52.0",
]
[[package]]
name = "rsqlite-vfs"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c"
dependencies = [
"hashbrown 0.16.1",
"thiserror 2.0.19",
]
[[package]]
name = "rstest"
version = "0.26.1"
@ -4073,6 +4450,21 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "rusqlite"
version = "0.40.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3"
dependencies = [
"bitflags 2.13.1",
"fallible-iterator",
"fallible-streaming-iterator",
"hashlink",
"libsqlite3-sys",
"smallvec",
"sqlite-wasm-rs",
]
[[package]]
name = "rustc-hash"
version = "2.1.3"
@ -4088,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"
@ -4330,6 +4756,19 @@ dependencies = [
"serde_derive",
]
[[package]]
name = "serde-pickle"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b641fdc8bcf2781ee78b30c599700d64ad4f412976143e4c5d0b9df906bb4843"
dependencies = [
"byteorder",
"iter-read",
"num-bigint 0.4.8",
"num-traits",
"serde",
]
[[package]]
name = "serde_core"
version = "1.0.229"
@ -4429,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"
@ -4545,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"
@ -4557,6 +5013,18 @@ dependencies = [
"unicode-segmentation",
]
[[package]]
name = "sqlite-wasm-rs"
version = "0.5.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75"
dependencies = [
"cc",
"js-sys",
"rsqlite-vfs",
"wasm-bindgen",
]
[[package]]
name = "sse-stream"
version = "0.2.6"
@ -4615,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"
@ -4646,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"
@ -4968,8 +5459,12 @@ version = "0.14.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
dependencies = [
"async-trait",
"axum",
"base64 0.22.1",
"bytes",
"flate2",
"h2 0.4.15",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
@ -4979,6 +5474,7 @@ dependencies = [
"percent-encoding",
"pin-project",
"rustls-native-certs",
"socket2 0.6.5",
"sync_wrapper",
"tokio",
"tokio-rustls 0.26.4",
@ -5060,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",
@ -5142,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",
]
@ -5244,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"
@ -5303,6 +5806,31 @@ 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"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
[[package]]
name = "veil"
version = "0.3.0"
@ -5746,7 +6274,7 @@ dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure",
"synstructure 0.13.2",
]
[[package]]
@ -5787,7 +6315,7 @@ dependencies = [
"proc-macro2",
"quote",
"syn 2.0.119",
"synstructure",
"synstructure 0.13.2",
]
[[package]]

12
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,7 +33,12 @@ 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" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }
@ -50,9 +56,14 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
qdrant-client = { version = "1.19.0", default-features = false }
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"] }
@ -69,6 +80,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
percent-encoding = "2.3"
webpki-roots = "1"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"

8
litellm-rust/crates/auth-gcp/src/lib.rs
View file

@ -128,6 +128,14 @@ impl VertexAuth {
}
}
pub async fn access_token(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<String, Error> {
self.load_provider(config, env_lookup).await?.token().await
}
pub async fn validate_environment(
&self,
headers: Vec<(String, String)>,

19
litellm-rust/crates/cache-disk/Cargo.toml Normal file
View file

@ -0,0 +1,19 @@
[package]
name = "litellm-cache-disk"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
py_literal = "0.4.0"
rand.workspace = true
rusqlite = { version = "0.40", features = ["bundled"] }
serde-pickle = "1.2"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile = "3.27.0"

10
litellm-rust/crates/cache-disk/src/adapter.rs Normal file
View file

@ -0,0 +1,10 @@
use litellm_cache::Error;
use crate::StoredValue;
pub trait ValueAdapter: Send + Sync + 'static {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error>;
fn write(&self, payload: Vec<u8>) -> StoredValue;
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error>;
fn counter_value(&self, value: f64) -> StoredValue;
}

301
litellm-rust/crates/cache-disk/src/cache.rs Normal file
View file

@ -0,0 +1,301 @@
use std::{
path::Path,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
};
use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
pub struct DiskCache<S, D = DiskcacheSqliteStore, A = PythonDiskCacheAdapter> {
store: Arc<D>,
adapter: Arc<A>,
codec: S,
}
impl<S: CacheCodec> DiskCache<S> {
pub fn open(directory: impl AsRef<Path>, codec: S) -> Result<Self, Error> {
Ok(Self {
store: Arc::new(DiskcacheSqliteStore::open(directory)?),
adapter: Arc::new(PythonDiskCacheAdapter),
codec,
})
}
}
impl<S: CacheCodec, D: DiskStore> DiskCache<S, D, PythonDiskCacheAdapter> {
pub fn with_store(store: D, codec: S) -> Self {
Self {
store: Arc::new(store),
adapter: Arc::new(PythonDiskCacheAdapter),
codec,
}
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
pub fn with_adapter(store: D, adapter: A, codec: S) -> Self {
Self {
store: Arc::new(store),
adapter: Arc::new(adapter),
codec,
}
}
pub fn directory(&self) -> &Path {
self.store.directory()
}
fn decode_stored(&self, value: StoredValue) -> Result<Option<S::Value>, Error> {
let Some(bytes) = self.adapter.read(value)? else {
return Ok(None);
};
self.codec.decode(&bytes).map(Some)
}
async fn run_blocking<T, F>(store: Arc<D>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&D) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || operation(&store))
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D, A> {
type Value = S::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> {
let value = self.adapter.write(self.codec.encode(&value)?);
let expire_time = context.ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
self.store.set(key, value, expire_time, unix_now())
}
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.store
.get(key, unix_now())?
.map(|value| self.decode_stored(value))
.transpose()
.map(|value| value.flatten())
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let value = self.adapter.write(self.codec.encode(&value)?);
let ttl = context.ttl;
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
let expire_time = ttl.map(|ttl| unix_now() + ttl.as_secs_f64());
store.set(&key, value, expire_time, unix_now())
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = key.to_string();
let value = Self::run_blocking(Arc::clone(&self.store), move |store| {
store.get(&key, unix_now())
})
.await?;
value
.map(|value| self.decode_stored(value))
.transpose()
.map(|value| value.flatten())
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: ExactCacheContext,
) -> Result<(), Error> {
let entries = entries
.into_iter()
.map(|(key, value)| {
self.codec
.encode(&value)
.map(|value| (key, self.adapter.write(value)))
})
.collect::<Result<Vec<_>, _>>()?;
let expire_after = context.ttl;
Self::run_blocking(Arc::clone(&self.store), move |store| {
for (key, value) in entries {
let expire_time = expire_after.map(|ttl| unix_now() + ttl.as_secs_f64());
store.set(&key, value, expire_time, unix_now())?;
}
Ok(())
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let result = Self::run_blocking(Arc::clone(&self.store), |store| {
store.probe().map(|_| CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Disk cache connection test successful".into(),
error: None,
})
})
.await;
Ok(match result {
Ok(result) => result,
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Disk cache connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
fn batch_get_cache(
&self,
keys: &[String],
context: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
keys.iter()
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let values = Self::run_blocking(Arc::clone(&self.store), move |store| {
keys.into_iter()
.map(|key| store.get(&key, unix_now()).map(|value| (key, value)))
.collect::<Result<Vec<_>, _>>()
})
.await?;
values
.into_iter()
.map(|(_, value)| match value {
None => Ok(BatchEntry::Miss),
Some(value) => match self.decode_stored(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
},
})
.collect()
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DeleteCache for DiskCache<S, D, A> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.store.pop(key, unix_now()).map(|_| ())
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = key.to_string();
Self::run_blocking(Arc::clone(&self.store), move |store| {
store.pop(&key, unix_now()).map(|_| ())
})
.await
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D, A> {
fn flush_cache(&self) -> Result<(), Error> {
self.store.clear()
}
async fn async_flush_cache(&self) -> Result<(), Error> {
Self::run_blocking(Arc::clone(&self.store), |store| store.clear()).await
}
}
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
for DiskCache<S, D, A>
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
increment(
self.adapter.as_ref(),
self.store.as_ref(),
key,
amount,
context.ttl,
)
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = key.to_string();
let adapter = Arc::clone(&self.adapter);
Self::run_blocking(Arc::clone(&self.store), move |store| {
increment(adapter.as_ref(), store, &key, amount, context.ttl)
})
.await
}
}
fn increment<A: ValueAdapter, D: DiskStore>(
adapter: &A,
store: &D,
key: &str,
amount: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let mut result = None;
let mut apply = |current: Option<StoredValue>| {
let initial = adapter.counter_seed(current)?;
let value = initial + amount;
let stored = adapter.counter_value(value);
result = Some(value);
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
};
store.update(key, unix_now(), &mut apply)?;
result.ok_or(Error::InvalidEntry)
}
fn unix_now() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs_f64()
}

11
litellm-rust/crates/cache-disk/src/lib.rs Normal file
View file

@ -0,0 +1,11 @@
mod adapter;
mod cache;
mod python;
mod sqlite;
mod store;
pub use adapter::ValueAdapter;
pub use cache::DiskCache;
pub use python::PythonDiskCacheAdapter;
pub use sqlite::DiskcacheSqliteStore;
pub use store::{DiskStore, StoredValue};

77
litellm-rust/crates/cache-disk/src/python/mod.rs Normal file
View file

@ -0,0 +1,77 @@
mod value;
use litellm_cache::Error;
use py_literal::Value;
use crate::{StoredValue, ValueAdapter};
#[derive(Clone, Copy, Debug, Default)]
pub struct PythonDiskCacheAdapter;
impl PythonDiskCacheAdapter {
fn python_get_cache(value: StoredValue) -> Result<Option<Value>, Error> {
let value = match value {
StoredValue::Bytes(value) => Value::Bytes(value),
StoredValue::Text(value) => Value::String(value),
StoredValue::Integer(value) => Value::Integer(value.into()),
StoredValue::Float(value) => Value::Float(value),
StoredValue::Pickle(value) => value::from_pickle(&value)?,
};
if !value::is_truthy(&value) {
return Ok(None);
}
match value {
Value::String(text) => Ok(Some(
value::from_json_text(&text).unwrap_or(Value::String(text)),
)),
Value::Bytes(bytes) => match std::str::from_utf8(&bytes) {
Ok(text) => Ok(Some(
value::from_json_text(text).unwrap_or(Value::Bytes(bytes)),
)),
Err(_) => Ok(Some(Value::Bytes(bytes))),
},
value => Ok(Some(value)),
}
}
}
impl ValueAdapter for PythonDiskCacheAdapter {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, Error> {
match value {
StoredValue::Text(value) => Ok((!value.is_empty()).then(|| value.into_bytes())),
StoredValue::Bytes(value) => Ok((!value.is_empty()).then_some(value)),
value => {
let Some(value) = Self::python_get_cache(value)? else {
return Ok(None);
};
value::to_json(&value).map(Some)
}
}
}
fn write(&self, payload: Vec<u8>) -> StoredValue {
StoredValue::Bytes(payload)
}
fn counter_seed(&self, value: Option<StoredValue>) -> Result<f64, Error> {
let Some(value) = value else {
return Ok(0.0);
};
let Some(value) = Self::python_get_cache(value)? else {
return Ok(0.0);
};
Ok(if value::is_int(&value) {
value::to_f64(&value).unwrap_or(0.0)
} else {
0.0
})
}
fn counter_value(&self, value: f64) -> StoredValue {
if value.fract() == 0.0 && value >= i64::MIN as f64 && value <= i64::MAX as f64 {
StoredValue::Integer(value as i64)
} else {
StoredValue::Float(value)
}
}
}

173
litellm-rust/crates/cache-disk/src/python/value.rs Normal file
View file

@ -0,0 +1,173 @@
use litellm_cache::Error;
use py_literal::Value;
use serde_json::{Map, Number};
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
let value = serde_pickle::value_from_slice(bytes, Default::default())
.map_err(|_| Error::InvalidEntry)?;
from_pickle_value(value)
}
fn from_pickle_value(value: serde_pickle::Value) -> Result<Value, Error> {
match value {
serde_pickle::Value::None => Ok(Value::None),
serde_pickle::Value::Bool(value) => Ok(Value::Boolean(value)),
serde_pickle::Value::I64(value) => integer(value.to_string()),
serde_pickle::Value::Int(value) => integer(value.to_string()),
serde_pickle::Value::F64(value) => Ok(Value::Float(value)),
serde_pickle::Value::String(value) => Ok(Value::String(value)),
serde_pickle::Value::Bytes(value) => Ok(Value::Bytes(value)),
serde_pickle::Value::List(values) => values
.into_iter()
.map(from_pickle_value)
.collect::<Result<Vec<_>, _>>()
.map(Value::List),
serde_pickle::Value::Tuple(values) => values
.into_iter()
.map(from_pickle_value)
.collect::<Result<Vec<_>, _>>()
.map(Value::Tuple),
serde_pickle::Value::Set(values) => values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()
.map(Value::Set),
serde_pickle::Value::FrozenSet(values) => values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()
.map(Value::Set),
serde_pickle::Value::Dict(values) => values
.into_iter()
.map(|(key, value)| Ok((from_pickle_hashable(key)?, from_pickle_value(value)?)))
.collect::<Result<Vec<_>, Error>>()
.map(Value::Dict),
}
}
fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Error> {
Ok(match value {
serde_pickle::HashableValue::None => Value::None,
serde_pickle::HashableValue::Bool(value) => Value::Boolean(value),
serde_pickle::HashableValue::I64(value) => integer(value.to_string())?,
serde_pickle::HashableValue::Int(value) => integer(value.to_string())?,
serde_pickle::HashableValue::F64(value) => Value::Float(value),
serde_pickle::HashableValue::Bytes(value) => Value::Bytes(value),
serde_pickle::HashableValue::String(value) => Value::String(value),
serde_pickle::HashableValue::Tuple(values) => Value::Tuple(
values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()?,
),
serde_pickle::HashableValue::FrozenSet(values) => Value::Set(
values
.into_iter()
.map(from_pickle_hashable)
.collect::<Result<Vec<_>, _>>()?,
),
})
}
fn integer(value: String) -> Result<Value, Error> {
value.parse().map_err(|_| Error::InvalidEntry)
}
pub(crate) fn from_json(value: serde_json::Value) -> Value {
match value {
serde_json::Value::Null => Value::None,
serde_json::Value::Bool(value) => Value::Boolean(value),
serde_json::Value::Number(value) => {
if value.is_i64() || value.is_u64() {
integer(value.to_string())
.unwrap_or(Value::Float(value.as_f64().unwrap_or(f64::NAN)))
} else {
Value::Float(value.as_f64().unwrap_or(f64::NAN))
}
}
serde_json::Value::String(value) => Value::String(value),
serde_json::Value::Array(values) => {
Value::List(values.into_iter().map(from_json).collect())
}
serde_json::Value::Object(values) => Value::Dict(
values
.into_iter()
.map(|(key, value)| (Value::String(key), from_json(value)))
.collect(),
),
}
}
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
serde_json::from_str(value)
.map(from_json)
.map_err(|_| Error::InvalidEntry)
}
pub(crate) fn is_truthy(value: &Value) -> bool {
match value {
Value::None => false,
Value::Boolean(value) => *value,
Value::Integer(value) => value.to_string() != "0",
Value::Float(value) => *value != 0.0,
Value::Complex(value) => value.re != 0.0 || value.im != 0.0,
Value::String(value) => !value.is_empty(),
Value::Bytes(value) => !value.is_empty(),
Value::Tuple(value) | Value::List(value) | Value::Set(value) => !value.is_empty(),
Value::Dict(value) => !value.is_empty(),
}
}
pub(crate) fn is_int(value: &Value) -> bool {
matches!(value, Value::Integer(_) | Value::Boolean(_))
}
pub(crate) fn to_f64(value: &Value) -> Option<f64> {
match value {
Value::Integer(value) => value.to_string().parse().ok(),
Value::Boolean(value) => Some(if *value { 1.0 } else { 0.0 }),
_ => None,
}
}
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
}
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
Ok(match value {
Value::None => serde_json::Value::Null,
Value::Boolean(value) => serde_json::Value::Bool(*value),
Value::Integer(value) => serde_json::Value::Number(
value
.to_string()
.parse::<Number>()
.map_err(|_| Error::InvalidEntry)?,
),
Value::Float(value) => {
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
}
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
Value::String(value) => serde_json::Value::String(value.clone()),
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
serde_json::Value::Array(
values
.iter()
.map(to_json_value)
.collect::<Result<Vec<_>, _>>()?,
)
}
Value::Dict(values) => {
let values = values
.iter()
.map(|(key, value)| {
let Value::String(key) = key else {
return Err(Error::InvalidEntry);
};
Ok((key.clone(), to_json_value(value)?))
})
.collect::<Result<Map<String, serde_json::Value>, _>>()?;
serde_json::Value::Object(values)
}
})
}

817
litellm-rust/crates/cache-disk/src/sqlite.rs Normal file
View file

@ -0,0 +1,817 @@
use std::{
collections::HashMap,
fs::{self, OpenOptions},
io::Write,
path::{Path, PathBuf},
sync::Mutex,
};
use litellm_cache::Error;
use rand::RngCore;
use rusqlite::{Connection, OptionalExtension, params, types::Value};
use crate::{DiskStore, StoredValue};
const MODE_RAW: i64 = 1;
const MODE_BINARY: i64 = 2;
const MODE_TEXT: i64 = 3;
const MODE_PICKLE: i64 = 4;
const DEFAULT_DISK_MIN_FILE_SIZE: i64 = 2_i64.pow(15);
const DEFAULT_SIZE_LIMIT: i64 = 2_i64.pow(30);
const DEFAULT_CULL_LIMIT: i64 = 10;
pub struct DiskcacheSqliteStore {
directory: PathBuf,
connection: Mutex<Connection>,
min_file_size: usize,
eviction_policy: String,
size_limit: i64,
cull_limit: i64,
statistics: bool,
}
struct StoredColumns {
size: i64,
mode: i64,
filename: Option<String>,
value: Option<Value>,
}
struct Row {
rowid: i64,
mode: i64,
filename: Option<String>,
value: Value,
}
impl DiskcacheSqliteStore {
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
let directory = directory.as_ref().to_path_buf();
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
let database = directory.join("cache.db");
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
connection
.busy_timeout(std::time::Duration::from_secs(60))
.map_err(|_| Error::Unavailable)?;
let mut settings = read_settings(&connection)?;
for (key, value) in default_settings() {
settings.entry(key).or_insert(value);
}
for (key, value) in settings
.iter()
.filter(|(key, _)| key.starts_with("sqlite_"))
{
apply_pragma(&connection, key, value)?;
}
connection
.execute_batch(
"CREATE TABLE IF NOT EXISTS Settings (
key TEXT NOT NULL UNIQUE,
value
)",
)
.map_err(|_| Error::Unavailable)?;
for (key, value) in &settings {
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
connection
.execute(
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
}
}
for (key, value) in [
("count", Value::Integer(0)),
("size", Value::Integer(0)),
("hits", Value::Integer(0)),
("misses", Value::Integer(0)),
] {
connection
.execute(
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
}
connection
.execute_batch(
"CREATE TABLE IF NOT EXISTS Cache (
rowid INTEGER PRIMARY KEY,
key BLOB,
raw INTEGER,
store_time REAL,
expire_time REAL,
access_time REAL,
access_count INTEGER DEFAULT 0,
tag BLOB,
size INTEGER DEFAULT 0,
mode INTEGER DEFAULT 0,
filename TEXT,
value BLOB
);
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
)
.map_err(|_| Error::Unavailable)?;
let eviction_policy = setting_string(&settings, "eviction_policy")
.unwrap_or_else(|| "least-recently-stored".to_string());
match eviction_policy.as_str() {
"none" => {}
"least-recently-stored" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
)
.map_err(|_| Error::Unavailable)?;
}
"least-recently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
)
.map_err(|_| Error::Unavailable)?;
}
"least-frequently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
)
.map_err(|_| Error::Unavailable)?;
}
_ => return Err(Error::Unavailable),
}
connection
.execute_batch(
"CREATE TRIGGER IF NOT EXISTS Settings_count_insert
AFTER INSERT ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value + 1
WHERE key = \"count\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_count_delete
AFTER DELETE ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value - 1
WHERE key = \"count\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_insert
AFTER INSERT ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value + NEW.size
WHERE key = \"size\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_update
AFTER UPDATE ON Cache FOR EACH ROW BEGIN
UPDATE Settings
SET value = value + NEW.size - OLD.size
WHERE key = \"size\"; END;
CREATE TRIGGER IF NOT EXISTS Settings_size_delete
AFTER DELETE ON Cache FOR EACH ROW BEGIN
UPDATE Settings SET value = value - OLD.size
WHERE key = \"size\"; END;",
)
.map_err(|_| Error::Unavailable)?;
let min_file_size = setting_i64(&settings, "disk_min_file_size")
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
.try_into()
.map_err(|_| Error::Unavailable)?;
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
Ok(Self {
directory,
connection: Mutex::new(connection),
min_file_size,
eviction_policy,
size_limit,
cull_limit,
statistics,
})
}
fn set_locked(
&self,
connection: &Connection,
key: &str,
columns: StoredColumns,
expire_time: Option<f64>,
now: f64,
) -> Result<Vec<String>, Error> {
let mut cleanup = Vec::new();
if let Some(old_filename) = connection
.query_row(
"SELECT filename FROM Cache WHERE key = ? AND raw = 1",
params![key],
|row| row.get::<_, Option<String>>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?
.flatten()
{
cleanup.push(old_filename);
}
let (size, mode, filename, value) =
(columns.size, columns.mode, columns.filename, columns.value);
let rowid = connection
.query_row(
"SELECT rowid FROM Cache WHERE key = ? AND raw = 1",
params![key],
|row| row.get::<_, i64>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?;
if let Some(rowid) = rowid {
connection
.execute(
"UPDATE Cache SET store_time = ?, expire_time = ?, access_time = ?,
access_count = 0, tag = NULL, size = ?, mode = ?, filename = ?, value = ?
WHERE rowid = ?",
params![now, expire_time, now, size, mode, filename, value, rowid],
)
.map_err(|_| Error::Unavailable)?;
} else {
connection
.execute(
"INSERT INTO Cache(
key, raw, store_time, expire_time, access_time, access_count,
tag, size, mode, filename, value
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
params![key, now, expire_time, now, size, mode, filename, value],
)
.map_err(|_| Error::Unavailable)?;
}
cleanup.extend(self.cull(connection, now)?);
Ok(cleanup)
}
fn cull(&self, connection: &Connection, now: f64) -> Result<Vec<String>, Error> {
if self.cull_limit <= 0 {
return Ok(Vec::new());
}
let mut cleanup = Vec::new();
let expired = connection
.prepare(
"SELECT rowid, filename FROM Cache
WHERE expire_time IS NOT NULL AND expire_time < ?
ORDER BY expire_time LIMIT ?",
)
.map_err(|_| Error::Unavailable)?
.query_map(params![now, self.cull_limit], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
for (_, filename) in &expired {
if let Some(filename) = filename {
cleanup.push(filename.clone());
}
}
for (rowid, _) in &expired {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
}
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
if remaining <= 0 || self.volume(connection)? < self.size_limit {
return Ok(cleanup);
}
let order = match self.eviction_policy.as_str() {
"none" => return Ok(cleanup),
"least-recently-stored" => "store_time",
"least-recently-used" => "access_time",
"least-frequently-used" => "access_count",
_ => return Err(Error::Unavailable),
};
let rows = connection
.prepare(&format!(
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
))
.map_err(|_| Error::Unavailable)?
.query_map(params![remaining], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
for (_, filename) in &rows {
if let Some(filename) = filename {
cleanup.push(filename.clone());
}
}
for (rowid, _) in rows {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
}
Ok(cleanup)
}
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
let page_count: i64 = connection
.query_row("PRAGMA page_count", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
let page_size: i64 = connection
.query_row("PRAGMA page_size", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
let size: i64 = connection
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
row.get(0)
})
.map_err(|_| Error::Unavailable)?;
Ok(page_count.saturating_mul(page_size).saturating_add(size))
}
}
impl DiskStore for DiskcacheSqliteStore {
fn directory(&self) -> &Path {
&self.directory
}
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
if !self.statistics && !has_get_update(&self.eviction_policy) {
return row
.map(|row| fetch_row(&self.directory, row))
.transpose()
.map(|value| value.flatten());
}
transactional(&connection, |connection| {
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
let Some(row) = row else {
if self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
}
return Ok(None);
};
let rowid = row.rowid;
let value = fetch_row(&self.directory, row);
let hit = value.as_ref().is_ok_and(Option::is_some);
if hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
[],
)
.map_err(|_| Error::Unavailable)?;
} else if !hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
}
if has_get_update(&self.eviction_policy) && hit {
let update = match self.eviction_policy.as_str() {
"least-recently-used" => "UPDATE Cache SET access_time = ? WHERE rowid = ?",
"least-frequently-used" => {
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
}
_ => return Err(Error::Unavailable),
};
if self.eviction_policy == "least-recently-used" {
connection
.execute(update, params![now, rowid])
.map_err(|_| Error::Unavailable)?;
} else {
connection
.execute(update, params![rowid])
.map_err(|_| Error::Unavailable)?;
}
}
value
})
}
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error> {
let columns = store_value(&self.directory, self.min_file_size, value)?;
let new_filename = columns.filename.clone();
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let result = transactional(&connection, |connection| {
self.set_locked(connection, key, columns, expire_time, now)
});
match result {
Ok(cleanup) => {
cleanup_files(&self.directory, cleanup);
Ok(())
}
Err(error) => {
if let Some(filename) = new_filename {
remove_file(&self.directory, &filename);
}
Err(error)
}
}
}
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let selected = transactional(&connection, |connection| {
let row = connection
.query_row(
"SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1
AND (expire_time IS NULL OR expire_time > ?)",
params![key, now],
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?;
let Some(row) = row else {
return Ok(None);
};
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
.map_err(|_| Error::Unavailable)?;
Ok(Some(row))
})?;
let Some(row) = selected else {
return Ok(None);
};
let filename = row.filename.clone();
let result = fetch_row(&self.directory, row)?;
if let Some(filename) = filename {
remove_file(&self.directory, &filename);
}
Ok(result)
}
fn clear(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let mut last_rowid = 0_i64;
loop {
let batch = transactional(&connection, |connection| {
let rows = connection
.prepare(
"SELECT rowid, filename FROM Cache
WHERE rowid > ? ORDER BY rowid LIMIT 100",
)
.map_err(|_| Error::Unavailable)?
.query_map(params![last_rowid], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
if rows.is_empty() {
return Ok(rows);
}
let ids = rows
.iter()
.map(|(rowid, _)| rowid.to_string())
.collect::<Vec<_>>()
.join(",");
connection
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
.map_err(|_| Error::Unavailable)?;
Ok(rows)
})?;
if batch.is_empty() {
return Ok(());
}
last_rowid = batch.last().map(|(rowid, _)| *rowid).unwrap_or(last_rowid);
cleanup_files(
&self.directory,
batch
.into_iter()
.filter_map(|(_, filename)| filename)
.collect(),
);
}
}
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let mut created_filename = None;
let result = transactional(&connection, |connection| {
let current = connection
.query_row(
"SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1
AND (expire_time IS NULL OR expire_time > ?)",
params![key, now],
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?
.map(|row| fetch_row(&self.directory, row))
.transpose()?
.flatten();
let (value, expire_time) = apply(current)?;
let columns = store_value(&self.directory, self.min_file_size, value)?;
created_filename = columns.filename.clone();
let cleanup = self.set_locked(connection, key, columns, expire_time, now)?;
Ok(cleanup)
});
match result {
Ok(cleanup) => {
cleanup_files(&self.directory, cleanup);
Ok(())
}
Err(error) => {
if let Some(filename) = created_filename {
remove_file(&self.directory, &filename);
}
Err(error)
}
}
}
fn probe(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
connection
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0),
)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
}
fn default_settings() -> HashMap<String, Value> {
HashMap::from([
("statistics".to_string(), Value::Integer(0)),
("tag_index".to_string(), Value::Integer(0)),
(
"eviction_policy".to_string(),
Value::Text("least-recently-stored".to_string()),
),
("size_limit".to_string(), Value::Integer(DEFAULT_SIZE_LIMIT)),
("cull_limit".to_string(), Value::Integer(DEFAULT_CULL_LIMIT)),
("sqlite_auto_vacuum".to_string(), Value::Integer(1)),
("sqlite_cache_size".to_string(), Value::Integer(8192)),
(
"sqlite_journal_mode".to_string(),
Value::Text("wal".to_string()),
),
(
"sqlite_mmap_size".to_string(),
Value::Integer(2_i64.pow(26)),
),
("sqlite_synchronous".to_string(), Value::Integer(1)),
(
"disk_min_file_size".to_string(),
Value::Integer(DEFAULT_DISK_MIN_FILE_SIZE),
),
("disk_pickle_protocol".to_string(), Value::Integer(5)),
])
}
fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Error> {
let mut statement = match connection.prepare("SELECT key, value FROM Settings") {
Ok(statement) => statement,
Err(_) => return Ok(HashMap::new()),
};
statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
.map_err(|_| Error::Unavailable)?
.collect::<Result<HashMap<_, _>, _>>()
.map_err(|_| Error::Unavailable)
}
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
match value {
Value::Integer(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
Value::Text(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
_ => Err(Error::Unavailable),
}
}
fn setting_i64(settings: &HashMap<String, Value>, key: &str) -> Option<i64> {
match settings.get(key) {
Some(Value::Integer(value)) => Some(*value),
_ => None,
}
}
fn setting_string(settings: &HashMap<String, Value>, key: &str) -> Option<String> {
match settings.get(key) {
Some(Value::Text(value)) => Some(value.clone()),
_ => None,
}
}
fn has_get_update(policy: &str) -> bool {
matches!(policy, "least-recently-used" | "least-frequently-used")
}
fn row_from_query(row: &rusqlite::Row<'_>) -> rusqlite::Result<Row> {
Ok(Row {
rowid: row.get(0)?,
mode: row.get(2)?,
filename: row.get(3)?,
value: row.get(4)?,
})
}
fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
match row.mode {
MODE_RAW => match row.value {
Value::Blob(value) => Ok(Some(StoredValue::Bytes(value))),
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
Value::Null => Err(Error::InvalidEntry),
},
MODE_BINARY | MODE_PICKLE => {
let bytes = match row.value {
Value::Blob(value) => value,
Value::Null => {
let Some(value) = read_file(directory, row.filename.as_deref())? else {
return Ok(None);
};
value
}
_ => return Err(Error::InvalidEntry),
};
Ok(Some(if row.mode == MODE_BINARY {
StoredValue::Bytes(bytes)
} else {
StoredValue::Pickle(bytes)
}))
}
MODE_TEXT => {
let bytes = match row.value {
Value::Null => {
let Some(value) = read_file(directory, row.filename.as_deref())? else {
return Ok(None);
};
value
}
Value::Blob(value) => value,
Value::Text(value) => value.into_bytes(),
_ => return Err(Error::InvalidEntry),
};
Ok(Some(StoredValue::Text(
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
)))
}
_ => Err(Error::InvalidEntry),
}
}
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
let Some(filename) = filename else {
return Err(Error::InvalidEntry);
};
match fs::read(directory.join(filename)) {
Ok(value) => Ok(Some(value)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(_) => Err(Error::Unavailable),
}
}
fn store_value(
directory: &Path,
min_file_size: usize,
value: StoredValue,
) -> Result<StoredColumns, Error> {
match value {
StoredValue::Integer(value) => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Integer(value)),
}),
StoredValue::Float(value) => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Real(value)),
}),
StoredValue::Text(value) if value.chars().count() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Text(value)),
}),
StoredValue::Text(value) => {
let bytes = value.into_bytes();
let filename = write_file(directory, &bytes)?;
Ok(StoredColumns {
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_TEXT,
filename: Some(filename),
value: None,
})
}
StoredValue::Bytes(value) if value.len() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_RAW,
filename: None,
value: Some(Value::Blob(value)),
}),
StoredValue::Bytes(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_BINARY,
filename: Some(filename),
value: None,
})
}
StoredValue::Pickle(value) if value.len() < min_file_size => Ok(StoredColumns {
size: 0,
mode: MODE_PICKLE,
filename: None,
value: Some(Value::Blob(value)),
}),
StoredValue::Pickle(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
mode: MODE_PICKLE,
filename: Some(filename),
value: None,
})
}
}
}
fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
let mut random = [0_u8; 16];
rand::rngs::OsRng.fill_bytes(&mut random);
let hex = random
.iter()
.map(|byte| format!("{byte:02x}"))
.collect::<String>();
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
let path = directory.join(&filename);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
}
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(path)
.map_err(|_| Error::Unavailable)?;
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
Ok(filename)
}
fn cleanup_files(directory: &Path, filenames: Vec<String>) {
for filename in filenames {
remove_file(directory, &filename);
}
}
fn remove_file(directory: &Path, filename: &str) {
let path = directory.join(filename);
let _ = fs::remove_file(&path);
}
fn transactional<T>(
connection: &Connection,
operation: impl FnOnce(&Connection) -> Result<T, Error>,
) -> Result<T, Error> {
connection
.execute_batch("BEGIN IMMEDIATE")
.map_err(|_| Error::Unavailable)?;
match operation(connection) {
Ok(value) => {
connection
.execute_batch("COMMIT")
.map_err(|_| Error::Unavailable)?;
Ok(value)
}
Err(error) => {
let _ = connection.execute_batch("ROLLBACK");
Err(error)
}
}
}

33
litellm-rust/crates/cache-disk/src/store.rs Normal file
View file

@ -0,0 +1,33 @@
use std::path::Path;
use litellm_cache::Error;
#[derive(Clone, Debug, PartialEq)]
pub enum StoredValue {
Bytes(Vec<u8>),
Text(String),
Integer(i64),
Float(f64),
Pickle(Vec<u8>),
}
pub trait DiskStore: Send + Sync + 'static {
fn directory(&self) -> &Path;
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn set(
&self,
key: &str,
value: StoredValue,
expire_time: Option<f64>,
now: f64,
) -> Result<(), Error>;
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error>;
fn clear(&self) -> Result<(), Error>;
fn update(
&self,
key: &str,
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error>;
fn probe(&self) -> Result<(), Error>;
}

431
litellm-rust/crates/cache-disk/tests/cache.rs Normal file
View file

@ -0,0 +1,431 @@
use std::{
fs,
path::{Path, PathBuf},
sync::Arc,
thread,
time::Duration,
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
FlushCache, JsonCodec,
};
use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
use rstest::{fixture, rstest};
use rusqlite::Connection;
use serde_json::{Value, json};
use tempfile::TempDir;
struct Sandbox {
directory: TempDir,
}
#[fixture]
fn sandbox() -> Sandbox {
Sandbox {
directory: tempfile::tempdir().unwrap(),
}
}
impl Sandbox {
fn store(&self) -> DiskcacheSqliteStore {
DiskcacheSqliteStore::open(self.directory.path()).unwrap()
}
fn cache<V>(&self) -> DiskCache<JsonCodec<V>>
where
JsonCodec<V>: CacheCodec,
{
DiskCache::open(self.directory.path(), JsonCodec::new()).unwrap()
}
fn db(&self) -> Connection {
Connection::open(self.directory.path().join("cache.db")).unwrap()
}
fn value_files(&self) -> Vec<PathBuf> {
fn visit(directory: &Path, files: &mut Vec<PathBuf>) {
for entry in fs::read_dir(directory).unwrap() {
let path = entry.unwrap().path();
if path.is_dir() {
visit(&path, files);
} else if path.extension().is_some_and(|extension| extension == "val") {
files.push(path);
}
}
}
let mut files = Vec::new();
visit(self.directory.path(), &mut files);
files
}
}
#[rstest]
fn relative_store_directory_is_absolutized(sandbox: Sandbox) {
let relative = PathBuf::from(format!(
".litellm-cache-disk-{}",
sandbox
.directory
.path()
.file_name()
.unwrap()
.to_string_lossy()
));
let store = DiskcacheSqliteStore::open(&relative).unwrap();
assert!(store.directory().is_absolute());
assert!(store.directory().ends_with(&relative));
let directory = store.directory().to_path_buf();
drop(store);
fs::remove_dir_all(directory).unwrap();
}
#[derive(Clone, Copy, Debug, Default)]
struct TextAdapter;
impl ValueAdapter for TextAdapter {
fn read(&self, value: StoredValue) -> Result<Option<Vec<u8>>, litellm_cache::Error> {
match value {
StoredValue::Text(value) => Ok(Some(value.into_bytes())),
_ => Ok(None),
}
}
fn write(&self, payload: Vec<u8>) -> StoredValue {
StoredValue::Text(String::from_utf8(payload).unwrap())
}
fn counter_seed(&self, _: Option<StoredValue>) -> Result<f64, litellm_cache::Error> {
Ok(0.0)
}
fn counter_value(&self, value: f64) -> StoredValue {
if value.fract() == 0.0 {
StoredValue::Integer(value as i64)
} else {
StoredValue::Float(value)
}
}
}
#[rstest]
fn roundtrip_persists_and_reopens(sandbox: Sandbox) {
let context = ExactCacheContext::default();
let opened = sandbox.cache::<Value>();
opened
.set_cache("key", json!({"answer": 42}), &context)
.unwrap();
assert_eq!(
opened.get_cache("key", &context).unwrap(),
Some(json!({"answer": 42}))
);
drop(opened);
let reopened = sandbox.cache::<Value>();
assert_eq!(
reopened.get_cache("key", &context).unwrap(),
Some(json!({"answer": 42}))
);
}
#[rstest]
fn ttl_and_expired_culling_match_cache_contract(sandbox: Sandbox) {
let store = sandbox.store();
store
.set(
"expired",
StoredValue::Bytes(b"old".to_vec()),
Some(10.0),
0.0,
)
.unwrap();
assert_eq!(store.get("expired", 10.0).unwrap(), None);
store
.set("new", StoredValue::Bytes(b"new".to_vec()), None, 11.0)
.unwrap();
assert_eq!(
sandbox
.db()
.query_row("SELECT COUNT(*) FROM Cache", [], |row| row.get::<_, i64>(0))
.unwrap(),
1
);
assert_eq!(
sandbox
.db()
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0)
)
.unwrap(),
1
);
}
#[rstest]
fn batch_preserves_order_and_classifies_misses_and_invalid_values(sandbox: Sandbox) {
let store = sandbox.store();
store
.set(
"hit",
StoredValue::Bytes(br#"{"ok":true}"#.to_vec()),
None,
0.0,
)
.unwrap();
store
.set(
"invalid",
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
None,
0.0,
)
.unwrap();
let entries = sandbox
.cache::<Value>()
.batch_get_cache(
&["hit".into(), "missing".into(), "invalid".into()],
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"ok": true})),
BatchEntry::Miss,
BatchEntry::Invalid
]
);
}
#[rstest]
#[case(StoredValue::Bytes(Vec::new()))]
#[case(StoredValue::Text(String::new()))]
#[case(StoredValue::Integer(0))]
#[case(StoredValue::Float(0.0))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]))]
#[case(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]))]
fn falsy_values_are_misses(sandbox: Sandbox, #[case] value: StoredValue) {
sandbox.store().set("key", value, None, 0.0).unwrap();
assert_eq!(
sandbox
.cache::<Value>()
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
None
);
}
#[rstest]
#[case(Some(StoredValue::Integer(2)), 1.5, 3.5, "real")]
#[case(Some(StoredValue::Integer(2)), 1.0, 3.0, "integer")]
#[case(Some(StoredValue::Float(3.5)), 1.0, 1.0, "integer")]
#[case(Some(StoredValue::Text("not a number".into())), 2.0, 2.0, "integer")]
#[case(Some(StoredValue::Text("5".into())), 2.0, 7.0, "integer")]
#[case(Some(StoredValue::Text("3.5".into())), 2.0, 2.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0, 2.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 1.0, 3.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 1.0, 1.0, "integer")]
#[case(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 4.0, 4.0, "integer")]
fn counters_follow_python_initialization(
sandbox: Sandbox,
#[case] initial: Option<StoredValue>,
#[case] amount: f64,
#[case] expected: f64,
#[case] sqlite_type: &str,
) {
if let Some(initial) = initial {
sandbox.store().set("counter", initial, None, 0.0).unwrap();
}
let cache = sandbox.cache::<f64>();
assert_eq!(
cache
.increment_cache("counter", amount, ExactCacheContext::default())
.unwrap(),
expected
);
assert_eq!(
sandbox
.db()
.query_row(
"SELECT typeof(value) FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, String>(0)
)
.unwrap(),
sqlite_type
);
}
#[rstest]
fn counters_are_atomic_across_concurrent_callers(sandbox: Sandbox) {
let cache = Arc::new(sandbox.cache::<f64>());
let workers = (0..8)
.map(|_| {
let cache = Arc::clone(&cache);
thread::spawn(move || {
for _ in 0..25 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
})
})
.collect::<Vec<_>>();
for worker in workers {
worker.join().unwrap();
}
assert_eq!(
cache
.increment_cache("counter", 0.0, ExactCacheContext::default())
.unwrap(),
200.0
);
}
#[rstest]
fn fractional_then_integer_increment_follows_python_behavior(sandbox: Sandbox) {
let cache = sandbox.cache::<f64>();
assert_eq!(
cache
.increment_cache("counter", 3.5, ExactCacheContext::default())
.unwrap(),
3.5
);
assert_eq!(
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap(),
1.0
);
}
#[rstest]
fn increment_ttl_replacement_clears_expiry_without_ttl(sandbox: Sandbox) {
let cache = sandbox.cache::<f64>();
cache
.increment_cache(
"counter",
1.0,
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
},
)
.unwrap();
assert!(
sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0)
)
.unwrap()
);
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
assert!(
!sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0)
)
.unwrap()
);
}
#[rstest]
fn custom_adapter_controls_storage_and_reads(sandbox: Sandbox) {
let cache = DiskCache::with_adapter(sandbox.store(), TextAdapter, JsonCodec::<Value>::new());
cache
.set_cache("key", json!({"answer": 42}), &ExactCacheContext::default())
.unwrap();
assert!(matches!(
sandbox.store().get("key", 0.0).unwrap(),
Some(StoredValue::Text(_))
));
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"answer": 42}))
);
}
#[rstest]
fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox) {
let large = vec![b'x'; 32 * 1024];
sandbox
.store()
.set("large", StoredValue::Bytes(large.clone()), None, 0.0)
.unwrap();
assert_eq!(sandbox.value_files().len(), 1);
sandbox
.store()
.set(
"large",
StoredValue::Bytes(vec![b'y'; 32 * 1024]),
None,
0.0,
)
.unwrap();
assert_eq!(sandbox.value_files().len(), 1);
sandbox.store().pop("large", 0.0).unwrap();
assert!(sandbox.value_files().is_empty());
sandbox
.store()
.set("a", StoredValue::Bytes(large.clone()), None, 0.0)
.unwrap();
sandbox
.store()
.set("b", StoredValue::Bytes(large), None, 0.0)
.unwrap();
sandbox.store().clear().unwrap();
assert!(sandbox.value_files().is_empty());
}
#[rstest]
#[tokio::test]
async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
};
cache
.async_set_cache("a", json!(1), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(
vec![("b".into(), json!(2)), ("c".into(), json!(3))],
context.clone(),
)
.await
.unwrap();
assert_eq!(
cache.async_get_cache("a", &context).await.unwrap(),
Some(json!(1))
);
assert_eq!(
cache
.async_batch_get_cache(vec!["c".into(), "missing".into()], context.clone())
.await
.unwrap(),
vec![BatchEntry::Hit(json!(3)), BatchEntry::Miss]
);
cache.async_delete_cache("a").await.unwrap();
cache.async_flush_cache().await.unwrap();
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}

113
litellm-rust/crates/cache-disk/tests/python_compat.rs Normal file
View file

@ -0,0 +1,113 @@
use litellm_cache::Error;
use litellm_cache_disk::{PythonDiskCacheAdapter, StoredValue, ValueAdapter};
use rstest::rstest;
enum ReadExpectation {
Bytes(&'static [u8]),
Miss,
Invalid,
}
#[rstest]
#[case::pickled_dictionary_with_string_keys(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e]),
ReadExpectation::Bytes(br#"{"a":1}"#)
)]
#[case::pickled_list_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x65, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_tuple_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4b, 0x01, 0x4b, 0x02, 0x86, 0x94, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_set_of_integers(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8f, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x90, 0x2e]),
ReadExpectation::Bytes(br#"[1,2]"#)
)]
#[case::pickled_response_envelope(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x28, 0x8c, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x94, 0x47, 0x3f, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x94, 0x8c, 0x08, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x20, 0x31, 0x7d, 0x94, 0x75, 0x2e]),
ReadExpectation::Bytes(br#"{"response":"{\"a\": 1}","timestamp":1.5}"#)
)]
#[case::non_json_text(
StoredValue::Text("not json".into()),
ReadExpectation::Bytes(b"not json")
)]
#[case::json_text(
StoredValue::Text("{\"a\": 1}".into()),
ReadExpectation::Bytes(br#"{"a": 1}"#)
)]
#[case::non_utf8_bytes(
StoredValue::Bytes(vec![0xff, 0xfe]),
ReadExpectation::Bytes(&[0xff, 0xfe])
)]
#[case::integer_seven(StoredValue::Integer(7), ReadExpectation::Bytes(b"7"))]
#[case::float_one_point_five(StoredValue::Float(1.5), ReadExpectation::Bytes(b"1.5"))]
#[case::pickled_true(
StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e]),
ReadExpectation::Bytes(b"true")
)]
#[case::pickled_negative_integer(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0xfd, 0xff, 0xff, 0xff, 0x2e]),
ReadExpectation::Bytes(b"-3")
)]
#[case::pickled_bytes(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x43, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x94, 0x2e]),
ReadExpectation::Invalid
)]
#[case::pickled_dictionary_with_integer_key(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x4b, 0x01, 0x8c, 0x01, 0x61, 0x94, 0x73, 0x2e]),
ReadExpectation::Invalid
)]
#[case::pickled_complex(
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x2e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x69, 0x6e, 0x73, 0x94, 0x8c, 0x07, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x94, 0x93, 0x94, 0x47, 0x3f, 0xf0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x86, 0x94, 0x52, 0x94, 0x2e]),
ReadExpectation::Invalid
)]
#[case::truncated_pickle(
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
ReadExpectation::Invalid
)]
#[case::empty_bytes(StoredValue::Bytes(Vec::new()), ReadExpectation::Miss)]
#[case::empty_text(StoredValue::Text(String::new()), ReadExpectation::Miss)]
#[case::zero_integer(StoredValue::Integer(0), ReadExpectation::Miss)]
#[case::zero_float(StoredValue::Float(0.0), ReadExpectation::Miss)]
#[case::pickled_none(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_false(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_zero(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_zero_float(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_string(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_list(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_dictionary(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), ReadExpectation::Miss)]
#[case::pickled_empty_tuple(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]), ReadExpectation::Miss)]
fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation) {
let result = PythonDiskCacheAdapter.read(row);
match expected {
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
}
}
#[rstest]
#[case::integer_two(Some(StoredValue::Integer(2)), 2.0)]
#[case::float_three_point_five(Some(StoredValue::Float(3.5)), 0.0)]
#[case::text_not_a_number(Some(StoredValue::Text("not a number".into())), 0.0)]
#[case::text_five(Some(StoredValue::Text("5".into())), 5.0)]
#[case::text_three_point_five(Some(StoredValue::Text("3.5".into())), 0.0)]
#[case::pickled_true(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0)]
#[case::pickled_two(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 2.0)]
#[case::pickled_dictionary(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 0.0)]
#[case::missing(None, 0.0)]
#[case::pickled_none(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 0.0)]
fn python_counter_seed_cases(#[case] row: Option<StoredValue>, #[case] expected: f64) {
assert_eq!(PythonDiskCacheAdapter.counter_seed(row).unwrap(), expected);
}
#[rstest]
#[case::integer_three(3.0, StoredValue::Integer(3))]
#[case::fractional_three_point_five(3.5, StoredValue::Float(3.5))]
#[case::negative_zero(-0.0, StoredValue::Integer(0))]
#[case::large_float(1e300, StoredValue::Float(1e300))]
fn python_counter_value_cases(#[case] value: f64, #[case] expected: StoredValue) {
assert_eq!(PythonDiskCacheAdapter.counter_value(value), expected);
}

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

@ -0,0 +1,20 @@
[package]
name = "litellm-cache-gcs"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
futures-util.workspace = true
litellm-auth-gcp.workspace = true
litellm-auth-types.workspace = true
litellm-cache.workspace = true
percent-encoding.workspace = true
reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]
serde_json.workspace = true
tokio.workspace = true
wiremock = "0.6.5"

260
litellm-rust/crates/cache-gcs/src/cache.rs Normal file
View file

@ -0,0 +1,260 @@
use std::{future::Future, sync::Arc, time::Duration};
use futures_util::future::try_join_all;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext,
FlushCache,
};
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
use reqwest::Client;
use crate::{GcpTokenSource, TokenSource};
pub const DEFAULT_ENDPOINT: &str = "https://storage.googleapis.com";
const OBJECT_NAME_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC
.remove(b'-')
.remove(b'_')
.remove(b'.')
.remove(b'~');
pub fn key_prefix(gcs_path: Option<&str>) -> String {
match gcs_path {
Some(path) if !path.is_empty() => format!("{}/", path.trim_end_matches('/')),
_ => String::new(),
}
}
#[derive(Clone, Debug)]
pub struct GcsConfig {
pub bucket_name: String,
pub gcs_path: Option<String>,
pub path_service_account: Option<String>,
pub endpoint: String,
}
impl GcsConfig {
pub fn new(bucket_name: impl Into<String>) -> Self {
Self {
bucket_name: bucket_name.into(),
gcs_path: None,
path_service_account: None,
endpoint: DEFAULT_ENDPOINT.to_string(),
}
}
}
pub struct GcsCache<S: CacheCodec> {
config: GcsConfig,
key_prefix: String,
client: Client,
token: Arc<dyn TokenSource>,
codec: S,
}
impl<S: CacheCodec> GcsCache<S> {
pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
Self::with_token_source(config, codec, token)
}
pub fn with_token_source(
config: GcsConfig,
codec: S,
token: Arc<dyn TokenSource>,
) -> Result<Self, Error> {
let client = Client::builder().build().map_err(|_| Error::Unavailable)?;
let key_prefix = key_prefix(config.gcs_path.as_deref());
Ok(Self {
config,
key_prefix,
client,
token,
codec,
})
}
pub fn bucket_name(&self) -> &str {
&self.config.bucket_name
}
pub fn key_prefix(&self) -> &str {
&self.key_prefix
}
pub fn path_service_account(&self) -> Option<&str> {
self.config.path_service_account.as_deref()
}
pub fn object_name(&self, key: &str) -> String {
format!("{}{}", self.key_prefix, key)
}
fn encoded_object_name(&self, key: &str) -> String {
percent_encode(self.object_name(key).as_bytes(), OBJECT_NAME_ENCODE_SET).to_string()
}
fn endpoint(&self, path: &str) -> String {
format!("{}{}", self.config.endpoint.trim_end_matches('/'), path)
}
async fn async_set(&self, key: &str, value: S::Value) -> Result<(), Error> {
let token = self.token.bearer_token().await?;
let payload = self.codec.encode(&value)?;
let url = self.endpoint(&format!(
"/upload/storage/v1/b/{}/o?uploadType=media&name={}",
self.config.bucket_name,
self.encoded_object_name(key)
));
let response = self
.client
.post(url)
.bearer_auth(token)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.body(payload)
.send()
.await
.map_err(|_| Error::Unavailable)?;
if !response.status().is_success() {
return Err(Error::Unavailable);
}
Ok(())
}
async fn async_get(&self, key: &str) -> Result<Option<S::Value>, Error> {
let token = self.token.bearer_token().await?;
let url = self.endpoint(&format!(
"/storage/v1/b/{}/o/{}?alt=media",
self.config.bucket_name,
self.encoded_object_name(key)
));
let response = self
.client
.get(url)
.bearer_auth(token)
.header(reqwest::header::CONTENT_TYPE, "application/json")
.send()
.await
.map_err(|_| Error::Unavailable)?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if !response.status().is_success() {
return Err(Error::Unavailable);
}
let body = response.bytes().await.map_err(|_| Error::Unavailable)?;
self.codec
.decode(&body)
.map(Some)
.map_err(|_| Error::InvalidEntry)
}
fn run_sync<T, F>(future: F) -> Result<T, Error>
where
F: Future<Output = Result<T, Error>> + Send,
T: Send,
{
let run = || {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|_| Error::Unavailable)
.and_then(|runtime| runtime.block_on(future))
};
if let Ok(handle) = tokio::runtime::Handle::try_current() {
if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
return tokio::task::block_in_place(run);
}
return std::thread::scope(|scope| {
scope
.spawn(run)
.join()
.map_err(|_| Error::Unavailable)
.and_then(|result| result)
});
}
run()
}
}
impl<S: CacheCodec> BaseCache for GcsCache<S> {
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
None
}
fn set_cache(&self, key: &str, value: Self::Value, _: &Self::Context) -> Result<(), Error> {
Self::run_sync(self.async_set(key, value))
}
fn get_cache(&self, key: &str, _: &Self::Context) -> Result<Option<Self::Value>, Error> {
Self::run_sync(self.async_get(key))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
_: Self::Context,
) -> Result<(), Error> {
self.async_set(key, value).await
}
async fn async_get_cache(
&self,
key: &str,
_: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get(key).await
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: Self::Context,
) -> Result<(), Error> {
try_join_all(entries.into_iter().map(|(key, value)| {
let context = context.clone();
async move { self.async_set_cache(&key, value, context).await }
}))
.await
.map(|_| ())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<S: CacheCodec> BatchCache for GcsCache<S> {
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: Self::Context,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
try_join_all(keys.into_iter().map(|key| {
let context = context.clone();
async move {
match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}
}))
.await
}
}
impl<S: CacheCodec> FlushCache for GcsCache<S> {
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
}

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

@ -0,0 +1,5 @@
mod cache;
mod token;
pub use cache::{DEFAULT_ENDPOINT, GcsCache, GcsConfig, key_prefix};
pub use token::{GcpTokenSource, StaticTokenSource, TokenSource};

44
litellm-rust/crates/cache-gcs/src/token.rs Normal file
View file

@ -0,0 +1,44 @@
use std::{future::Future, pin::Pin};
use litellm_auth_gcp::{VertexAuth, VertexConfig};
use litellm_auth_types::{InputSource, SecretValue, Sourced};
use litellm_cache::Error;
pub trait TokenSource: Send + Sync + 'static {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>>;
}
pub struct GcpTokenSource {
auth: VertexAuth,
config: VertexConfig,
}
impl GcpTokenSource {
pub fn new(path_service_account: Option<String>) -> Self {
let credentials = path_service_account
.map(|path| Sourced::new(SecretValue::new(path), InputSource::Deployment));
Self {
auth: VertexAuth::default(),
config: VertexConfig::new(credentials, None, None),
}
}
}
impl TokenSource for GcpTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async move {
self.auth
.access_token(&self.config, &|name| std::env::var(name).ok())
.await
.map_err(|_| Error::Unavailable)
})
}
}
pub struct StaticTokenSource(pub String);
impl TokenSource for StaticTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async move { Ok(self.0.clone()) })
}
}

324
litellm-rust/crates/cache-gcs/tests/cache.rs Normal file
View file

@ -0,0 +1,324 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache,
JsonCodec,
};
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix};
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_bytes, header, method, path, query_param},
};
fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
GcsConfig {
bucket_name: "bucket".into(),
gcs_path: gcs_path.map(str::to_string),
path_service_account: None,
endpoint: server.uri(),
}
}
fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache<JsonCodec<serde_json::Value>> {
GcsCache::with_token_source(
config(server, gcs_path),
JsonCodec::new(),
Arc::new(StaticTokenSource("tok".into())),
)
.unwrap()
}
#[tokio::test]
async fn set_writes_encoded_object_and_headers() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.and(header("authorization", "Bearer tok"))
.and(header("content-type", "application/json"))
.and(body_bytes(br#"{"value":"entry"}"#))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
cache(&server, Some("cache/"))
.set_cache(
"team:a b/c",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(
requests[0].url.query(),
Some("uploadType=media&name=cache%2Fteam%3Aa%20b%2Fc")
);
}
#[tokio::test]
async fn get_maps_statuses_and_decode_failures() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.and(query_param("alt", "media"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/missing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/server-error"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/invalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
let cache = cache(&server, None);
assert_eq!(
cache
.get_cache("hit", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
assert_eq!(
cache
.get_cache("missing", &ExactCacheContext::default())
.unwrap(),
None
);
assert_eq!(
cache
.get_cache("server-error", &ExactCacheContext::default())
.unwrap_err(),
Error::Unavailable
);
assert_eq!(
cache
.get_cache("invalid", &ExactCacheContext::default())
.unwrap_err(),
Error::InvalidEntry
);
}
#[test]
fn key_prefix_normalizes_paths() {
assert_eq!(key_prefix(None), "");
assert_eq!(key_prefix(Some("a/b/")), "a/b/");
assert_eq!(key_prefix(Some("a/b")), "a/b/");
assert_eq!(key_prefix(Some("")), "");
}
#[tokio::test]
async fn object_names_use_python_quote_encoding() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.respond_with(ResponseTemplate::new(200))
.expect(2)
.mount(&server)
.await;
let cache = cache(&server, Some("p/"));
cache
.set_cache(
"a~b-c_d.e/f g%h",
json!({"value": "punctuation"}),
&ExactCacheContext::default(),
)
.unwrap();
cache
.set_cache(
"ключ",
json!({"value": "utf8"}),
&ExactCacheContext::default(),
)
.unwrap();
let requests = server.received_requests().await.unwrap();
let queries: Vec<_> = requests
.iter()
.filter_map(|request| request.url.query())
.collect();
assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h"));
assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87"));
}
#[tokio::test]
async fn ignores_ttl_and_writes_pipeline_concurrently() {
let server = MockServer::start().await;
for key in ["one", "two", "three"] {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.and(query_param("name", key))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
}
let cache = cache(&server, None);
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))),
None
);
cache
.async_set_cache_pipeline(
vec![
("one".into(), json!({"key": "one"})),
("two".into(), json!({"key": "two"})),
("three".into(), json!({"key": "three"})),
],
ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))),
)
.await
.unwrap();
}
#[tokio::test]
async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/missing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/invalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
assert_eq!(
cache(&server, None)
.async_batch_get_cache(
vec!["hit".into(), "missing".into(), "invalid".into()],
ExactCacheContext::default(),
)
.await
.unwrap(),
vec![
BatchEntry::Hit(json!({"value": "entry"})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
);
}
#[tokio::test]
async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() {
let server = MockServer::start().await;
let cache = cache(&server, None);
assert_eq!(cache.flush_cache(), Ok(()));
assert_eq!(cache.disconnect().await, Ok(()));
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
}
#[test]
fn sync_operations_work_without_an_active_runtime() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.unwrap();
let server = runtime.block_on(MockServer::start());
runtime.block_on(
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server),
);
runtime.block_on(
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server),
);
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
}
#[tokio::test(flavor = "multi_thread")]
async fn sync_operations_work_inside_a_multi_thread_runtime() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
}
struct FailingTokenSource;
impl TokenSource for FailingTokenSource {
fn bearer_token(
&self,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String, Error>> + Send + '_>>
{
Box::pin(async { Err(Error::Unavailable) })
}
}
#[tokio::test]
async fn token_source_failure_skips_http() {
let server = MockServer::start().await;
let cache = GcsCache::with_token_source(
config(&server, None),
JsonCodec::<serde_json::Value>::new(),
Arc::new(FailingTokenSource),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap_err(),
Error::Unavailable
);
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}

24
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml Normal file
View file

@ -0,0 +1,24 @@
[package]
name = "litellm-cache-qdrant-semantic"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
futures-util.workspace = true
litellm-cache.workspace = true
qdrant-client = { workspace = true, features = ["serde"] }
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
uuid.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
rstest.workspace = true
tonic = "0.14"
tonic-prost = "0.14"
tokio-stream = "0.1"

75
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs Normal file
View file

@ -0,0 +1,75 @@
use std::time::Duration;
use litellm_cache::Error;
use reqwest::Client;
use serde_json::Value;
use crate::Embedder;
pub struct OpenAiEmbedder {
client: Client,
api_base: String,
api_key: String,
model: String,
timeout: Option<Duration>,
}
pub struct OpenAiEmbedderConfig {
pub api_base: String,
pub api_key: String,
pub model: String,
pub timeout: Option<Duration>,
}
impl OpenAiEmbedder {
pub fn new(client: Client, config: OpenAiEmbedderConfig) -> Self {
Self {
client,
api_base: config.api_base.trim_end_matches('/').to_owned(),
api_key: config.api_key,
model: config.model,
timeout: config.timeout,
}
}
}
impl Embedder for OpenAiEmbedder {
fn model(&self) -> &str {
&self.model
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
let request = self
.client
.post(format!("{}/embeddings", self.api_base))
.bearer_auth(&self.api_key)
.json(&serde_json::json!({
"model": self.model,
"input": input,
"encoding_format": "float",
}));
let response = if let Some(timeout) = self.timeout {
request.timeout(timeout)
} else {
request
}
.send()
.await
.map_err(|_| Error::Unavailable)?
.error_for_status()
.map_err(|_| Error::Unavailable)?;
let body: Value = response.json().await.map_err(|_| Error::Unavailable)?;
body.get("data")
.and_then(Value::as_array)
.and_then(|data| data.first())
.and_then(|item| item.get("embedding"))
.and_then(Value::as_array)
.and_then(|embedding| {
embedding
.iter()
.map(|value| value.as_f64().map(|value| value as f32))
.collect::<Option<Vec<_>>>()
})
.ok_or(Error::Unavailable)
}
}

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

@ -0,0 +1,7 @@
mod embedder;
mod prompt;
mod semantic;
pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig};
pub use prompt::prompt_from_messages;
pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization};

59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs Normal file
View file

@ -0,0 +1,59 @@
use serde_json::Value;
fn search_results_text(search_results: Option<&Value>) -> String {
let Some(Value::Array(results)) = search_results else {
return String::new();
};
results
.iter()
.filter_map(Value::as_object)
.flat_map(|result| {
let source = result
.get("source")
.and_then(Value::as_str)
.map(str::to_owned);
let title = result
.get("title")
.and_then(Value::as_str)
.map(str::to_owned);
let content = result
.get("content")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
let citations = result
.get("citations")
.filter(|value| !value.is_null())
.map(|value| serde_json::to_string(value).unwrap_or_default());
source
.into_iter()
.chain(title)
.chain(content)
.chain(citations)
})
.collect()
}
pub fn prompt_from_messages(messages: &[Value]) -> String {
messages
.iter()
.filter_map(Value::as_object)
.map(|message| {
let content = match message.get("content") {
Some(Value::String(content)) => content.clone(),
Some(Value::Array(parts)) => parts
.iter()
.filter_map(Value::as_object)
.filter_map(|part| part.get("text").and_then(Value::as_str))
.collect(),
_ => String::new(),
};
format!(
"{content}{}",
search_results_text(message.get("search_results"))
)
})
.collect()
}

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

@ -0,0 +1,262 @@
use std::future::Future;
use futures_util::future::try_join_all;
use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
use qdrant_client::{
Payload, Qdrant,
qdrant::{
BinaryQuantizationBuilder, CompressionRatio, Condition, CreateCollectionBuilder,
CreateFieldIndexCollectionBuilder, Distance, FieldType, Filter, PointStruct,
ProductQuantizationBuilder, QuantizationSearchParamsBuilder, ScalarQuantizationBuilder,
SearchParamsBuilder, SearchPointsBuilder, UpsertPointsBuilder, VectorParamsBuilder,
},
};
use serde_json::{Map, Value, json};
use uuid::Uuid;
use crate::prompt_from_messages;
pub trait Embedder: Send + Sync + 'static {
fn model(&self) -> &str;
fn embed(&self, input: &str) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
#[derive(Clone, Debug, PartialEq)]
pub enum Quantization {
Binary,
Scalar,
Product,
}
pub struct QdrantSemanticConfig {
pub collection_name: String,
pub similarity_threshold: f64,
pub vector_size: u64,
pub quantization: Quantization,
}
pub struct QdrantSemanticCache<E: Embedder, C: CacheCodec> {
client: Qdrant,
embedder: E,
codec: C,
config: QdrantSemanticConfig,
runtime: tokio::runtime::Handle,
}
impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
pub async fn connect(
client: Qdrant,
embedder: E,
codec: C,
config: QdrantSemanticConfig,
runtime: tokio::runtime::Handle,
) -> Result<Self, Error> {
let exists = client
.collection_exists(config.collection_name.clone())
.await
.map_err(|_| Error::Unavailable)?;
if !exists {
client
.create_collection(
CreateCollectionBuilder::new(config.collection_name.clone())
.vectors_config(VectorParamsBuilder::new(
config.vector_size,
Distance::Cosine,
))
.quantization_config(quantization(&config.quantization)),
)
.await
.map_err(|_| Error::Unavailable)?;
}
let _ = client
.create_field_index(CreateFieldIndexCollectionBuilder::new(
config.collection_name.clone(),
"litellm_cache_key".to_owned(),
FieldType::Keyword,
))
.await;
Ok(Self {
client,
embedder,
codec,
config,
runtime,
})
}
pub fn collection_name(&self) -> &str {
&self.config.collection_name
}
pub fn similarity_threshold(&self) -> f64 {
self.config.similarity_threshold
}
pub fn vector_size(&self) -> u64 {
self.config.vector_size
}
pub fn embedder(&self) -> &E {
&self.embedder
}
fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
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(messages))
}
async fn set(
&self,
key: &str,
value: C::Value,
context: &SemanticCacheContext,
) -> Result<(), Error> {
let prompt = Self::prompt(context)?;
let vector = self.embedder.embed(&prompt).await?;
let response =
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
let payload = Payload::try_from(json!({
"litellm_cache_key": key,
"text": prompt,
"response": response,
}))
.map_err(|_| Error::InvalidEntry)?;
self.client
.upsert_points(
UpsertPointsBuilder::new(
self.collection_name(),
vec![PointStruct::new(
Uuid::new_v4().to_string(),
vector,
payload,
)],
)
.wait(true),
)
.await
.map_err(|_| Error::Unavailable)?;
Ok(())
}
async fn get(
&self,
key: &str,
context: &SemanticCacheContext,
) -> Result<Option<C::Value>, Error> {
let prompt = Self::prompt(context)?;
let vector = self.embedder.embed(&prompt).await?;
let result = self
.client
.search_points(
SearchPointsBuilder::new(self.collection_name(), vector, 1)
.with_payload(true)
.filter(Filter::must([Condition::matches(
"litellm_cache_key",
key.to_owned(),
)]))
.params(
SearchParamsBuilder::default().quantization(
QuantizationSearchParamsBuilder::default()
.ignore(false)
.rescore(true)
.oversampling(3.0),
),
),
)
.await
.map_err(|_| Error::Unavailable)?;
let Some(point) = result.result.into_iter().next() else {
return Ok(None);
};
let payload: Map<String, Value> = Payload::from(point.payload).into();
if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) {
return Ok(None);
}
if f64::from(point.score) < self.config.similarity_threshold {
return Ok(None);
}
let response = payload
.get("response")
.and_then(Value::as_str)
.ok_or(Error::InvalidEntry)?;
self.codec.decode(response.as_bytes()).map(Some)
}
}
fn quantization(value: &Quantization) -> qdrant_client::qdrant::quantization_config::Quantization {
match value {
Quantization::Binary => BinaryQuantizationBuilder::new(false).into(),
Quantization::Scalar => ScalarQuantizationBuilder::default()
.quantile(0.99)
.always_ram(false)
.into(),
Quantization::Product => ProductQuantizationBuilder::new(CompressionRatio::X16.into())
.always_ram(false)
.into(),
}
}
impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
type Value = C::Value;
type Context = SemanticCacheContext;
fn get_ttl(&self, _: &Self::Context) -> Option<std::time::Duration> {
None
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
self.runtime.block_on(self.set(key, value, context))
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.runtime.block_on(self.get(key, context))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
self.set(key, value, &context).await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key, context).await
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: Self::Context,
) -> Result<(), Error> {
try_join_all(entries.into_iter().map(|(key, value)| {
let context = context.clone();
async move { self.async_set_cache(&key, value, context).await }
}))
.await
.map(|_| ())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}

166
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs Normal file
View file

@ -0,0 +1,166 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::Error;
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig};
use serde_json::Value;
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
struct TestHttpServer {
address: std::net::SocketAddr,
request: Arc<Mutex<Option<Vec<u8>>>>,
task: tokio::task::JoinHandle<()>,
}
impl TestHttpServer {
async fn response(status: &str, body: &str) -> Self {
Self::response_after(status, body, Duration::ZERO).await
}
async fn response_after(status: &str, body: &str, delay: Duration) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let request = Arc::new(Mutex::new(None));
let captured = request.clone();
let status = status.to_owned();
let body = body.to_owned();
let task = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let request_bytes = read_request(&mut stream).await;
*captured.lock().unwrap() = Some(request_bytes);
tokio::time::sleep(delay).await;
let response = format!(
"HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
body.len()
);
stream.write_all(response.as_bytes()).await.unwrap();
});
Self {
address,
request,
task,
}
}
fn base_url(&self) -> String {
format!("http://{}", self.address)
}
}
impl Drop for TestHttpServer {
fn drop(&mut self) {
self.task.abort();
}
}
async fn read_request(stream: &mut tokio::net::TcpStream) -> Vec<u8> {
let mut bytes = Vec::new();
let header_end = loop {
let mut chunk = [0_u8; 1024];
let count = stream.read(&mut chunk).await.unwrap();
assert_ne!(count, 0);
bytes.extend_from_slice(&chunk[..count]);
if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
break end + 4;
}
};
let headers = String::from_utf8_lossy(&bytes[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
line.split_once(':')
.filter(|(name, _)| name.eq_ignore_ascii_case("content-length"))
.map(|(_, value)| value.trim())
})
.unwrap()
.parse::<usize>()
.unwrap();
while bytes.len() < header_end + content_length {
let mut chunk = [0_u8; 1024];
let count = stream.read(&mut chunk).await.unwrap();
assert_ne!(count, 0);
bytes.extend_from_slice(&chunk[..count]);
}
bytes
}
fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
OpenAiEmbedderConfig {
api_base: base,
api_key: "test-key".to_owned(),
model: "test-model".to_owned(),
timeout,
}
}
#[tokio::test]
async fn posts_embeddings_request_and_parses_vector() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(
format!("{}/", server.base_url()),
Some(Duration::from_secs(1)),
),
);
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
let request = server.request.lock().unwrap().clone().unwrap();
let request_text = String::from_utf8(request).unwrap();
assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n"));
assert!(request_text.contains("\r\nauthorization: Bearer test-key\r\n"));
let body = request_text.split("\r\n\r\n").nth(1).unwrap();
let body: Value = serde_json::from_str(body).unwrap();
assert_eq!(body["model"], "test-model");
assert_eq!(body["input"], "hello");
assert_eq!(body["encoding_format"], "float");
}
#[tokio::test]
async fn status_and_timeout_errors_are_unavailable() {
let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::response_after(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
Duration::from_millis(500),
)
.await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(server.base_url(), Some(Duration::from_millis(200))),
);
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::response_after(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
Duration::from_millis(100),
)
.await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(server.base_url(), Some(Duration::from_secs(1))),
);
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
}
#[tokio::test]
async fn uses_the_injected_client() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
let client = reqwest::Client::builder()
.user_agent("litellm-embedder-test")
.build()
.unwrap();
let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
let request = server.request.lock().unwrap().clone().unwrap();
let request_text = String::from_utf8(request).unwrap();
assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n"));
}

38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs Normal file
View file

@ -0,0 +1,38 @@
use litellm_cache_qdrant_semantic::prompt_from_messages;
use serde_json::json;
#[test]
fn prompt_matches_python_message_content_rules() {
let messages = vec![
json!({"role": "user", "content": "hello"}),
json!({
"role": "user",
"content": [
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": {"url": "ignored"}},
{"type": "text", "text": "!"},
],
}),
];
assert_eq!(prompt_from_messages(&messages), "helloworld!");
}
#[test]
fn prompt_includes_search_result_text_and_compact_citations() {
let messages = vec![json!({
"role": "tool",
"content": null,
"search_results": [{
"source": "source",
"title": "title",
"content": [{"text": "body"}],
"citations": {"page": 1, "section": "intro"},
}],
})];
assert_eq!(
prompt_from_messages(&messages),
r#"sourcetitlebody{"page":1,"section":"intro"}"#
);
}

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

@ -0,0 +1,422 @@
#[path = "support/mod.rs"]
mod support;
use std::{collections::HashMap, sync::Arc, time::Duration};
use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
use litellm_cache_qdrant_semantic::{
Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
};
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
};
use qdrant_client::Payload;
use qdrant_client::{
Qdrant,
qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams},
};
use serde_json::{Value as JsonValue, json};
use support::{FakeQdrant, FakeState, StoredPoint};
#[derive(Clone)]
struct FixedEmbedder {
vectors: Arc<HashMap<String, Vec<f32>>>,
}
impl FixedEmbedder {
fn new(vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>) -> Self {
Self {
vectors: Arc::new(
vectors
.into_iter()
.map(|(prompt, vector)| (prompt.to_owned(), vector))
.collect(),
),
}
}
}
impl Embedder for FixedEmbedder {
fn model(&self) -> &str {
"fixed"
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
}
}
fn config(quantization: Quantization) -> QdrantSemanticConfig {
QdrantSemanticConfig {
collection_name: "semantic".to_owned(),
similarity_threshold: 0.9,
vector_size: 2,
quantization,
}
}
fn context(prompt: &str) -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": prompt}])),
..Default::default()
}
}
fn value(response: JsonValue) -> CacheEntry {
CacheEntry {
timestamp: Some(1.0),
response,
}
}
async fn connect(
server: &FakeQdrant,
vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>,
) -> QdrantSemanticCache<FixedEmbedder, ResponseCacheCodec> {
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new(vectors),
ResponseCacheCodec,
config(Quantization::Binary),
tokio::runtime::Handle::current(),
)
.await
.unwrap()
}
#[tokio::test(flavor = "multi_thread")]
#[expect(
deprecated,
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
)]
async fn connect_sets_collection_quantization_and_index() {
for (quantization, expected) in [
(Quantization::Binary, 0),
(Quantization::Scalar, 1),
(Quantization::Product, 2),
] {
let server = FakeQdrant::start(FakeState::default()).await;
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new([]),
ResponseCacheCodec,
config(quantization),
tokio::runtime::Handle::current(),
)
.await
.unwrap();
let state = server.state.lock().unwrap();
let request = &state.created_collections[0];
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) =
request
.vectors_config
.as_ref()
.and_then(|config| config.config.clone())
else {
panic!("missing vector params");
};
assert_eq!(size, 2);
assert_eq!(distance, Distance::Cosine as i32);
let quantization_config = request
.quantization_config
.as_ref()
.unwrap()
.quantization
.unwrap();
match (expected, quantization_config) {
(0, qdrant::quantization_config::Quantization::Binary(binary)) => {
assert_eq!(binary.always_ram, Some(false));
}
(1, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
assert_eq!(scalar.quantile, Some(0.99));
assert_eq!(scalar.always_ram, Some(false));
}
(2, qdrant::quantization_config::Quantization::Product(product)) => {
assert_eq!(product.compression, CompressionRatio::X16 as i32);
assert_eq!(product.always_ram, Some(false));
}
_ => panic!("unexpected quantization"),
}
assert!(state.index_creations >= 1);
assert_eq!(state.field_indexes[0].collection_name, "semantic");
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
assert_eq!(
state.field_indexes[0].field_type,
Some(qdrant::FieldType::Keyword as i32)
);
server.stop();
}
}
#[tokio::test(flavor = "multi_thread")]
async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
let server = FakeQdrant::start(FakeState {
collections: ["semantic".to_owned()].into_iter().collect(),
fail_field_index: true,
..Default::default()
})
.await;
let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
let state = server.state.lock().unwrap();
assert!(state.created_collections.is_empty());
assert!(state.index_creations >= 1);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn async_and_sync_set_get_store_exact_payload() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
let ctx = context("hello");
let entry = value(json!({"answer": 42}));
cache
.async_set_cache("key", entry.clone(), ctx.clone())
.await
.unwrap();
assert_eq!(
cache.async_get_cache("key", &ctx).await.unwrap().as_ref(),
Some(&entry)
);
{
let state = server.state.lock().unwrap();
let payload = &state.points[0].payload;
let mut payload_keys = payload.keys().cloned().collect::<Vec<_>>();
payload_keys.sort();
assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]);
assert_eq!(payload["litellm_cache_key"], Value::from("key"));
assert_eq!(
payload["response"],
Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap())
);
}
let sync_entry = entry.clone();
let sync_cache = cache.clone();
let sync_ctx = ctx.clone();
tokio::task::spawn_blocking(move || {
sync_cache
.set_cache("sync", sync_entry.clone(), &sync_ctx)
.unwrap();
assert_eq!(
sync_cache.get_cache("sync", &sync_ctx).unwrap(),
Some(sync_entry)
);
})
.await
.unwrap();
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn misses_and_payload_validation_are_safe() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(
&server,
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
)
.await;
let entry = value(json!({"answer": 1}));
cache
.async_set_cache("key", entry, context("hello"))
.await
.unwrap();
assert_eq!(
cache
.async_get_cache("other", &context("hello"))
.await
.unwrap(),
None
);
assert_eq!(
cache
.async_get_cache("key", &context("near"))
.await
.unwrap(),
None
);
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({
"litellm_cache_key": 99,
"response": "{}",
}))
.unwrap()
.into(),
});
assert_eq!(
cache
.async_get_cache("99", &context("hello"))
.await
.unwrap(),
None
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await;
let empty = SemanticCacheContext::default();
assert_eq!(
cache
.async_set_cache("key", value(json!({})), empty.clone())
.await,
Err(Error::MissingPrompt)
);
assert_eq!(
cache.async_get_cache("key", &empty).await,
Err(Error::MissingPrompt)
);
assert_eq!(
cache.async_get_cache("key", &context("unknown")).await,
Err(Error::Unavailable)
);
cache
.async_set_cache(
"ttl",
value(json!({"ttl": true})),
context("one").with_ttl(Some(Duration::from_secs(1))),
)
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(1_100)).await;
assert!(
cache
.async_get_cache(
"ttl",
&context("one").with_ttl(Some(Duration::from_secs(1))),
)
.await
.unwrap()
.is_some()
);
cache
.async_set_cache_pipeline(
vec![
("one".to_owned(), value(json!({"n": 1}))),
("two".to_owned(), value(json!({"n": 2}))),
],
context("one"),
)
.await
.unwrap();
assert!(
cache
.async_get_cache("one", &context("one"))
.await
.unwrap()
.is_some()
);
assert!(
cache
.async_get_cache("two", &context("one"))
.await
.unwrap()
.is_some()
);
assert_eq!(
server.state.lock().unwrap().upsert_waits,
vec![Some(true), Some(true), Some(true)]
);
assert_eq!(cache.get_ttl(&context("one")), None);
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn response_payloads_decode_and_invalid_entries_fail() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
for (key, response) in [
("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")),
("garbage", json!("not json")),
("missing", json!("unused")),
] {
let mut payload = serde_json::Map::new();
payload.insert("litellm_cache_key".to_owned(), json!(key));
if key != "missing" {
payload.insert("response".to_owned(), response);
}
server.insert_point(StoredPoint {
id: Some(PointId::from(key.len() as u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(JsonValue::Object(payload))
.unwrap()
.into(),
});
}
assert_eq!(
cache
.async_get_cache("python", &context("hello"))
.await
.unwrap(),
Some(value(json!({"a": 1})))
);
assert_eq!(
cache.async_get_cache("garbage", &context("hello")).await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_get_cache("missing", &context("hello")).await,
Err(Error::InvalidEntry)
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn response_cache_facade_turns_invalid_entry_into_miss() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
let request = ResponseCacheRequest::<SemanticCacheContext>::new(CacheKeyInput {
preset: Some("key".to_owned()),
..Default::default()
})
.with_context(context("hello"));
let response = json!({"answer": 42});
let facade = ResponseCache::new(cache.clone());
facade
.async_store(&request, response.clone(), Duration::from_secs(1))
.await
.unwrap();
assert_eq!(
facade
.async_lookup(&request, Duration::from_secs(1))
.await
.unwrap(),
Some(response)
);
{
let mut state = server.state.lock().unwrap();
state.points[0]
.payload
.insert("response".to_owned(), Value::from("not json"));
}
assert_eq!(
facade
.async_lookup(&request, Duration::from_secs(1))
.await
.unwrap(),
None
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn stopped_qdrant_server_maps_to_unavailable() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
server.stop();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(
cache.async_get_cache("key", &context("hello")).await,
Err(Error::Unavailable)
);
}

342
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs Normal file
View file

@ -0,0 +1,342 @@
use std::{
collections::{HashMap, HashSet},
net::SocketAddr,
sync::{Arc, Mutex},
};
use qdrant_client::qdrant::collections_server::CollectionsServer;
use qdrant_client::qdrant::{
self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse,
CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId,
PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors,
collections_server::Collections,
points_server::{Points, PointsServer},
};
use tokio::sync::oneshot;
use tokio_stream::wrappers::TcpListenerStream;
use tonic::{Request, Response, Status, transport::Server};
#[derive(Clone, Debug)]
pub struct StoredPoint {
pub id: Option<PointId>,
pub vector: Vec<f32>,
pub payload: HashMap<String, Value>,
}
#[derive(Default)]
pub struct FakeState {
pub collections: HashSet<String>,
pub created_collections: Vec<CreateCollection>,
pub field_indexes: Vec<CreateFieldIndexCollection>,
pub points: Vec<StoredPoint>,
pub upsert_waits: Vec<Option<bool>>,
pub index_creations: usize,
pub fail_field_index: bool,
}
#[derive(Clone)]
pub struct FakeQdrant {
pub state: Arc<Mutex<FakeState>>,
pub address: SocketAddr,
shutdown: Arc<Mutex<Option<oneshot::Sender<()>>>>,
}
impl FakeQdrant {
pub async fn start(state: FakeState) -> Self {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let state = Arc::new(Mutex::new(state));
let service = FakeService {
state: state.clone(),
};
let (shutdown_tx, shutdown_rx) = oneshot::channel();
tokio::spawn(async move {
Server::builder()
.add_service(CollectionsServer::new(service.clone()))
.add_service(PointsServer::new(service))
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
let _ = shutdown_rx.await;
})
.await
.unwrap();
});
Self {
state,
address,
shutdown: Arc::new(Mutex::new(Some(shutdown_tx))),
}
}
pub fn url(&self) -> String {
format!("http://{}", self.address)
}
pub fn stop(&self) {
self.shutdown
.lock()
.unwrap()
.take()
.unwrap()
.send(())
.unwrap();
}
pub fn insert_point(&self, point: StoredPoint) {
self.state.lock().unwrap().points.push(point);
}
}
#[derive(Clone)]
struct FakeService {
state: Arc<Mutex<FakeState>>,
}
macro_rules! unimplemented_collections {
($($name:ident, $request:ty, $response:ty);* $(;)?) => {
$(
fn $name<'life0, 'async_trait>(
&'life0 self,
_: Request<$request>,
) -> std::pin::Pin<
Box<
dyn std::future::Future<
Output = Result<Response<$response>, Status>,
> + Send
+ 'async_trait,
>,
>
where
'life0: 'async_trait,
Self: 'async_trait,
{
Box::pin(async { Err(Status::unimplemented(stringify!($name))) })
}
)*
};
}
macro_rules! unimplemented_points {
($($name:ident, $request:ty, $response:ty);* $(;)?) => {
$(
fn $name<'life0, 'async_trait>(
&'life0 self,
_: Request<$request>,
) -> std::pin::Pin<
Box<
dyn std::future::Future<
Output = Result<Response<$response>, Status>,
> + Send
+ 'async_trait,
>,
>
where
'life0: 'async_trait,
Self: 'async_trait,
{
Box::pin(async { Err(Status::unimplemented(stringify!($name))) })
}
)*
};
}
#[tonic::async_trait]
impl Collections for FakeService {
async fn create(
&self,
request: Request<CreateCollection>,
) -> Result<Response<CollectionOperationResponse>, Status> {
let request = request.into_inner();
let mut state = self.state.lock().unwrap();
state.collections.insert(request.collection_name.clone());
state.created_collections.push(request);
Ok(Response::new(CollectionOperationResponse {
result: true,
..Default::default()
}))
}
async fn collection_exists(
&self,
request: Request<CollectionExistsRequest>,
) -> Result<Response<CollectionExistsResponse>, Status> {
let exists = self
.state
.lock()
.unwrap()
.collections
.contains(&request.into_inner().collection_name);
Ok(Response::new(CollectionExistsResponse {
result: Some(CollectionExists { exists }),
..Default::default()
}))
}
unimplemented_collections!(
get, qdrant::GetCollectionInfoRequest, qdrant::GetCollectionInfoResponse;
list, qdrant::ListCollectionsRequest, qdrant::ListCollectionsResponse;
update, qdrant::UpdateCollection, qdrant::CollectionOperationResponse;
delete, qdrant::DeleteCollection, qdrant::CollectionOperationResponse;
update_aliases, qdrant::ChangeAliases, qdrant::CollectionOperationResponse;
list_collection_aliases, qdrant::ListCollectionAliasesRequest, qdrant::ListAliasesResponse;
list_aliases, qdrant::ListAliasesRequest, qdrant::ListAliasesResponse;
collection_cluster_info, qdrant::CollectionClusterInfoRequest, qdrant::CollectionClusterInfoResponse;
update_collection_cluster_setup, qdrant::UpdateCollectionClusterSetupRequest, qdrant::UpdateCollectionClusterSetupResponse;
create_shard_key, qdrant::CreateShardKeyRequest, qdrant::CreateShardKeyResponse;
delete_shard_key, qdrant::DeleteShardKeyRequest, qdrant::DeleteShardKeyResponse;
list_shard_keys, qdrant::ListShardKeysRequest, qdrant::ListShardKeysResponse;
);
}
#[tonic::async_trait]
impl Points for FakeService {
async fn create_field_index(
&self,
request: Request<CreateFieldIndexCollection>,
) -> Result<Response<PointsOperationResponse>, Status> {
let mut state = self.state.lock().unwrap();
state.index_creations += 1;
state.field_indexes.push(request.into_inner());
if state.fail_field_index {
return Err(Status::internal("field index failure"));
}
Ok(Response::new(PointsOperationResponse::default()))
}
async fn upsert(
&self,
request: Request<qdrant::UpsertPoints>,
) -> Result<Response<PointsOperationResponse>, Status> {
let request = request.into_inner();
let mut state = self.state.lock().unwrap();
state.upsert_waits.push(request.wait);
for point in request.points {
let stored = StoredPoint {
id: point.id.clone(),
vector: dense_vector(point.vectors)?,
payload: point.payload,
};
if let Some(existing) = state
.points
.iter_mut()
.find(|existing| existing.id == stored.id)
{
*existing = stored;
} else {
state.points.push(stored);
}
}
Ok(Response::new(PointsOperationResponse::default()))
}
async fn search(
&self,
request: Request<SearchPoints>,
) -> Result<Response<SearchResponse>, Status> {
let request = request.into_inner();
let key_filter = keyword_filter(request.filter.as_ref());
let state = self.state.lock().unwrap();
let mut results = state
.points
.iter()
.filter(|point| {
key_filter.as_ref().is_none_or(|(field, expected)| {
point
.payload
.get(field)
.and_then(|value| {
let value: serde_json::Value = value.clone().into();
value
.as_str()
.map(str::to_owned)
.or_else(|| value.as_i64().map(|value| value.to_string()))
})
.is_some_and(|value| value == *expected)
})
})
.map(|point| ScoredPoint {
id: point.id.clone(),
payload: point.payload.clone(),
score: cosine(&request.vector, &point.vector),
..Default::default()
})
.collect::<Vec<_>>();
results.sort_by(|left, right| right.score.total_cmp(&left.score));
results.truncate(request.limit as usize);
Ok(Response::new(SearchResponse {
result: results,
..Default::default()
}))
}
unimplemented_points!(
delete, qdrant::DeletePoints, qdrant::PointsOperationResponse;
get, qdrant::GetPoints, qdrant::GetResponse;
update_vectors, qdrant::UpdatePointVectors, qdrant::PointsOperationResponse;
delete_vectors, qdrant::DeletePointVectors, qdrant::PointsOperationResponse;
set_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse;
overwrite_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse;
delete_payload, qdrant::DeletePayloadPoints, qdrant::PointsOperationResponse;
clear_payload, qdrant::ClearPayloadPoints, qdrant::PointsOperationResponse;
delete_field_index, qdrant::DeleteFieldIndexCollection, qdrant::PointsOperationResponse;
create_vector_name, qdrant::CreateVectorNameRequest, qdrant::PointsOperationResponse;
delete_vector_name, qdrant::DeleteVectorNameRequest, qdrant::PointsOperationResponse;
search_batch, qdrant::SearchBatchPoints, qdrant::SearchBatchResponse;
search_groups, qdrant::SearchPointGroups, qdrant::SearchGroupsResponse;
scroll, qdrant::ScrollPoints, qdrant::ScrollResponse;
recommend, qdrant::RecommendPoints, qdrant::RecommendResponse;
recommend_batch, qdrant::RecommendBatchPoints, qdrant::RecommendBatchResponse;
recommend_groups, qdrant::RecommendPointGroups, qdrant::RecommendGroupsResponse;
discover, qdrant::DiscoverPoints, qdrant::DiscoverResponse;
discover_batch, qdrant::DiscoverBatchPoints, qdrant::DiscoverBatchResponse;
count, qdrant::CountPoints, qdrant::CountResponse;
update_batch, qdrant::UpdateBatchPoints, qdrant::UpdateBatchResponse;
query, qdrant::QueryPoints, qdrant::QueryResponse;
query_batch, qdrant::QueryBatchPoints, qdrant::QueryBatchResponse;
query_groups, qdrant::QueryPointGroups, qdrant::QueryGroupsResponse;
facet, qdrant::FacetCounts, qdrant::FacetResponse;
search_matrix_pairs, qdrant::SearchMatrixPoints, qdrant::SearchMatrixPairsResponse;
search_matrix_offsets, qdrant::SearchMatrixPoints, qdrant::SearchMatrixOffsetsResponse;
);
}
fn dense_vector(vectors: Option<Vectors>) -> Result<Vec<f32>, Status> {
let Some(Vectors {
vectors_options:
Some(qdrant::vectors::VectorsOptions::Vector(Vector {
vector: Some(qdrant::vector::Vector::Dense(qdrant::DenseVector { data })),
..
})),
}) = vectors
else {
return Err(Status::invalid_argument("expected dense vector"));
};
Ok(data)
}
fn keyword_filter(filter: Option<&Filter>) -> Option<(String, String)> {
filter?
.must
.iter()
.find_map(|condition| match condition.condition_one_of.as_ref()? {
qdrant::condition::ConditionOneOf::Field(field) => {
let qdrant::r#match::MatchValue::Keyword(value) =
field.r#match.as_ref()?.match_value.as_ref()?
else {
return None;
};
Some((field.key.clone(), value.clone()))
}
_ => None,
})
}
fn cosine(left: &[f32], right: &[f32]) -> f32 {
let dot = left
.iter()
.zip(right)
.map(|(left, right)| left * right)
.sum::<f32>();
let left_norm = left.iter().map(|value| value * value).sum::<f32>().sqrt();
let right_norm = right.iter().map(|value| value * value).sum::<f32>().sqrt();
dot / (left_norm * right_norm)
}

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

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

@ -1,21 +1,21 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, 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 {
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: ExactCacheContext,
pub context: C,
pub max_age: Option<Duration>,
}
impl ResponseCacheRequest {
impl<C: CacheContext + Default> ResponseCacheRequest<C> {
pub fn new(key: CacheKeyInput) -> Self {
Self {
key,
@ -26,17 +26,35 @@ impl ResponseCacheRequest {
default_on: true,
..Default::default()
},
context: ExactCacheContext::default(),
context: C::default(),
max_age: None,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
impl<C: CacheContext> ResponseCacheRequest<C> {
pub fn with_context<D: CacheContext>(self, context: D) -> ResponseCacheRequest<D> {
ResponseCacheRequest {
key: self.key,
controls: self.controls,
context,
max_age: self.max_age,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
where
B::Context: Default + PartialEq,
{
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
impl<B> ResponseCache<B>
where
B: BaseCache<Value = CacheEntry>,
B::Context: Default + PartialEq,
{
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -45,8 +63,12 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
&self.backend
}
pub fn backend_arc(&self) -> &Arc<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&ExactCacheContext::default())
self.backend.get_ttl(&B::Context::default())
}
pub async fn async_flush(&self) -> Result<(), Error>
@ -62,7 +84,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -81,7 +103,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -101,7 +123,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -126,7 +148,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -153,7 +175,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -172,7 +194,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -193,7 +215,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value)>,
now: Duration,
) -> Result<(), Error> {
self.async_store_entries(
@ -209,7 +231,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
/// the freshness of its original response.
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
) -> Result<(), Error> {
let writable = entries
.into_iter()
@ -249,8 +271,8 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
}
fn partial_hits(
requests: &[ResponseCacheRequest],
readable: Vec<(usize, &ResponseCacheRequest)>,
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {

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

@ -1,12 +1,15 @@
use std::{
sync::{
Arc,
Arc, Mutex,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_cache::{BaseCache, CacheCodec, Error};
use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
@ -30,6 +33,82 @@ fn request() -> ResponseCacheRequest {
})
}
struct SemanticBackend {
entries: Mutex<Vec<(String, CacheEntry)>>,
contexts: Mutex<Vec<SemanticCacheContext>>,
}
impl BaseCache for SemanticBackend {
type Value = CacheEntry;
type Context = SemanticCacheContext;
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
None
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
self.contexts.lock().unwrap().push(context.clone());
self.entries.lock().unwrap().push((key.to_owned(), value));
Ok(())
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.contexts.lock().unwrap().push(context.clone());
Ok(self
.entries
.lock()
.unwrap()
.iter()
.find(|(entry_key, _)| entry_key == key)
.map(|(_, entry)| entry.clone()))
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "ok".into(),
error: None,
})
}
}
#[test]
fn semantic_context_reaches_backend_for_store_and_lookup() {
let backend = Arc::new(SemanticBackend {
entries: Mutex::new(Vec::new()),
contexts: Mutex::new(Vec::new()),
});
let cache = ResponseCache::new(backend.clone());
let context = SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "hello"}])),
..Default::default()
};
let request = request().with_context(context.clone());
let response = json!({"answer": 42});
cache
.store(&request, response.clone(), Duration::from_secs(100))
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::from_secs(100)).unwrap(),
Some(response)
);
assert_eq!(
backend.contexts.lock().unwrap().as_slice(),
&[context.clone(), context]
);
}
#[tokio::test]
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
let clock = Arc::new(AtomicU64::new(100));

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/error.rs vendored
View file

@ -6,4 +6,8 @@ pub enum Error {
InvalidEntry,
#[error("flushing Redis requires an explicit namespace")]
UnscopedFlush,
#[error("operation is not supported by this cache")]
UnsupportedOperation,
#[error("semantic cache requires request messages")]
MissingPrompt,
}

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

@ -8,7 +8,7 @@ mod error;
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};

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

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

@ -24,9 +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
@ -38,8 +46,11 @@ litellm-host-python.workspace = true
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
tokio = { workspace = true, features = ["sync"] }
url.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
[dev-dependencies]
serde.workspace = true
@ -47,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)

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

@ -1,11 +1,14 @@
use std::time::Duration;
use std::{path::PathBuf, time::Duration};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::CacheType;
use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use pyo3::{
exceptions::{PyTypeError, PyValueError},
exceptions::{PyAttributeError, PyTypeError, PyValueError},
prelude::*,
types::{PyAny, PyDict, PyList, PyString},
types::{PyAny, PyBool, PyDict, PyList, PyString},
};
use super::{native::NativeResponseCache, request::duration};
@ -26,6 +29,10 @@ pub(super) struct MemoryCacheConfig {
pub(super) max_entry_bytes: usize,
}
pub(super) struct DiskCacheConfig {
pub(super) directory: PathBuf,
}
#[derive(Debug, PartialEq)]
pub(super) enum RedisProtocol {
Resp2,
@ -75,11 +82,31 @@ pub(super) struct RedisCacheConfig {
pub(super) connection: RedisConnectionConfig,
}
#[derive(Debug, PartialEq)]
pub(super) struct GcsCacheConfig {
pub(super) bucket_name: String,
pub(super) key_prefix: String,
pub(super) path_service_account: Option<String>,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
#[allow(
dead_code,
reason = "embedding settings are projected so drift falls back to Python"
)]
pub(super) struct RedisSemanticCacheConfig {
pub(super) redis_url: String,
pub(super) index_name: String,
pub(super) similarity_threshold: f64,
pub(super) embedding_model: String,
pub(super) embedding_max_input_tokens: Option<u64>,
pub(super) embedding_timeout: Option<f64>,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
@ -91,10 +118,45 @@ struct RedisClientProjection<'py> {
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct ValkeySemanticCacheConfig {
pub(super) similarity_threshold: f64,
pub(super) index_name: String,
pub(super) embedding_model: String,
pub(super) connection: RedisConnectionConfig,
}
pub(super) struct QdrantSemanticCacheConfig {
pub(super) grpc_url: String,
pub(super) api_key: Option<String>,
pub(super) collection_name: String,
pub(super) similarity_threshold: f64,
pub(super) vector_size: u64,
pub(super) embedding: OpenAiEmbedderConfig,
pub(super) quantization: Quantization,
}
impl QdrantSemanticCacheConfig {
pub(super) fn to_qdrant_config(&self) -> QdrantSemanticConfig {
QdrantSemanticConfig {
collection_name: self.collection_name.clone(),
similarity_threshold: self.similarity_threshold,
vector_size: self.vector_size,
quantization: self.quantization.clone(),
}
}
}
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
RedisSemantic(Box<RedisSemanticCacheConfig>),
QdrantSemantic(Box<QdrantSemanticCacheConfig>),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
@ -109,6 +171,13 @@ pub(super) enum UnsupportedCacheConfig {
RedisCredentials,
RedisConnection,
RedisOption,
S3Client,
S3Credentials,
S3Option,
GcsBucket,
DiskStore,
QdrantEndpoint,
SemanticEmbedding,
}
impl UnsupportedCacheConfig {
@ -119,6 +188,15 @@ impl UnsupportedCacheConfig {
Self::RedisCredentials => "native Redis credentials require Python",
Self::RedisConnection => "native Redis connection type is not implemented",
Self::RedisOption => "native Redis configuration requires Python",
Self::S3Client => "native S3 client type is not implemented",
Self::S3Credentials => "native S3 credentials require Python",
Self::S3Option => "native S3 configuration requires Python",
Self::GcsBucket => "native GCS cache requires a configured bucket name",
Self::DiskStore => "native disk cache requires the built-in diskcache store",
Self::QdrantEndpoint => {
"native Qdrant requires the default REST port so the gRPC port can be derived"
}
Self::SemanticEmbedding => "native semantic embedding requires Python",
}
}
}
@ -161,21 +239,54 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::S3) => match project_s3(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::S3(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Gcs) => match project_gcs(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Gcs(backend),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Disk) => match project_disk(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Disk(backend),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::QdrantSemantic) => match project_qdrant_semantic(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::QdrantSemantic(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::AzureBlob(backend),
}))
}),
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::Gcs,
)
| None => Ok(CacheConfigProjection::Unsupported(
Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::RedisSemantic(Box::new(backend)),
}))
}),
None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
}
@ -185,9 +296,17 @@ impl NativeCacheConfig {
let default_ttl = match &self.backend {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::AzureBlob(_) => None,
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO),
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_)
| CacheBackendConfig::RedisSemantic(_)
| CacheBackendConfig::QdrantSemantic(_) => None,
};
if service.default_ttl() != default_ttl {
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
&& service.default_ttl() != default_ttl
{
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
@ -212,6 +331,115 @@ impl NativeCacheConfig {
CacheBackendConfig::Redis(config) => (service.namespace()
!= config.namespace.as_deref())
.then_some("facade and native backend namespaces must match"),
CacheBackendConfig::S3(_) if service.kind() != "s3" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => {
Some("facade and native backend buckets must match")
}
CacheBackendConfig::S3(config)
if service.key_prefix() != Some(config.key_prefix.as_str()) =>
{
Some("facade and native backend key prefixes must match")
}
CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => {
Some("facade and native backend regions must match")
}
CacheBackendConfig::S3(config)
if service.endpoint()
!= config
.endpoint
.as_ref()
.map(|endpoint| endpoint.url.as_str()) =>
{
Some("facade and native backend endpoints must match")
}
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Gcs(config)
if service
.gcs_backend()
.is_none_or(|backend| backend.bucket_name() != config.bucket_name) =>
{
Some("facade and native backend buckets must match")
}
CacheBackendConfig::Gcs(config)
if service
.gcs_backend()
.is_none_or(|backend| backend.key_prefix() != config.key_prefix) =>
{
Some("facade and native backend key prefixes must match")
}
CacheBackendConfig::Gcs(config)
if service.gcs_backend().is_none_or(|backend| {
backend.path_service_account() != config.path_service_account.as_deref()
}) =>
{
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::ValkeySemantic(config) => {
if service.kind() != "valkey-semantic" {
return Some("facade and native backend types must match");
}
let Some((threshold, index_name)) = service.semantic_config() else {
return Some("facade and native backend types must match");
};
(threshold != config.similarity_threshold || index_name != config.index_name)
.then_some("facade and native semantic settings must match")
}
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::Disk(config) => {
let Some(directory) = service.directory() else {
return Some("facade and native backend types must match");
};
let native = std::fs::canonicalize(directory).ok();
let facade = std::fs::canonicalize(&config.directory).ok();
(native != facade).then_some("facade and native backend directories must match")
}
CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::RedisSemantic(config)
if service.index_name() != Some(config.index_name.as_str()) =>
{
Some("facade and native backend index names must match")
}
CacheBackendConfig::RedisSemantic(config)
if service.similarity_threshold()
!= Some(f64::from(config.similarity_threshold as f32)) =>
{
Some("facade and native backend similarity thresholds must match")
}
CacheBackendConfig::RedisSemantic(_) => None,
CacheBackendConfig::QdrantSemantic(config) if service.kind() != "qdrant_semantic" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::QdrantSemantic(config)
if service.collection_name() != Some(config.collection_name.as_str()) =>
{
Some("facade and native backend collections must match")
}
CacheBackendConfig::QdrantSemantic(config)
if service.similarity_threshold() != Some(config.similarity_threshold) =>
{
Some("facade and native backend similarity thresholds must match")
}
CacheBackendConfig::QdrantSemantic(config)
if service.vector_size() != Some(config.vector_size) =>
{
Some("facade and native backend vector sizes must match")
}
CacheBackendConfig::QdrantSemantic(config)
if service.embedding_model() != Some(config.embedding.model.as_str()) =>
{
Some("facade and native backend embedding models must match")
}
CacheBackendConfig::QdrantSemantic(_) => None,
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
@ -225,6 +453,105 @@ impl NativeCacheConfig {
}
}
#[inline(never)]
fn project_qdrant_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<QdrantSemanticCacheConfig, UnsupportedCacheConfig>> {
let rest_url = backend.getattr("qdrant_api_base")?.extract::<String>()?;
let parsed = match url::Url::parse(&rest_url) {
Ok(value) => value,
Err(_) => return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)),
};
if !matches!(parsed.scheme(), "http" | "https")
|| (!parsed.path().is_empty() && parsed.path() != "/")
|| parsed.query().is_some()
|| parsed.host_str().is_none()
|| parsed.port() != Some(6333)
{
return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint));
}
let mut grpc_url = parsed;
if grpc_url.set_port(Some(6334)).is_err() {
return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint));
}
grpc_url.set_path("");
grpc_url.set_query(None);
if optional_attribute(backend, "embedding_max_input_tokens")?
.is_some_and(|value| !value.is_none())
{
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
}
let configured_model = backend.getattr("embedding_model")?.extract::<String>()?;
let embedding_model = configured_model
.strip_prefix("openai/")
.unwrap_or(&configured_model)
.to_owned();
if !embedding_model.starts_with("text-embedding-") {
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
}
let proxy_server = py_sys_module(backend.py())?;
if let Some(proxy_server) = proxy_server {
let router = proxy_server.getattr("llm_router")?;
let model_list = proxy_server.getattr("llm_model_list")?;
let embedding_router = backend.py().import("litellm.caching._embedding_router")?;
if !embedding_router
.getattr("resolve_embedding_router")?
.call1((configured_model.as_str(), router, model_list))?
.is_none()
{
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
}
}
let litellm = backend.py().import("litellm")?;
for name in ["api_key", "openai_key", "api_base"] {
if !litellm.getattr(name)?.is_none() {
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
}
}
let Ok(embedding_api_key) = std::env::var("OPENAI_API_KEY") else {
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
};
if embedding_api_key.is_empty() {
return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding));
}
let embedding_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 timeout = optional_attribute(backend, "embedding_timeout")?
.map(|value| value.extract::<Option<f64>>())
.transpose()?
.flatten()
.map(duration)
.transpose()?;
Ok(Ok(QdrantSemanticCacheConfig {
grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(),
api_key: optional_string(backend.getattr("qdrant_api_key")?)?,
collection_name: backend.getattr("collection_name")?.extract()?,
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
vector_size: backend.getattr("vector_size")?.extract::<u64>()?,
embedding: OpenAiEmbedderConfig {
api_base: embedding_api_base,
api_key: embedding_api_key,
model: embedding_model,
timeout,
},
quantization: Quantization::Binary,
}))
}
fn py_sys_module(py: Python<'_>) -> PyResult<Option<Bound<'_, PyAny>>> {
match py
.import("sys")?
.getattr("modules")?
.get_item("litellm.proxy.proxy_server")
{
Ok(module) => Ok(Some(module)),
Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(py) => Ok(None),
Err(error) => Err(error),
}
}
#[inline(never)]
fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConfig> {
let client = backend.getattr("container_client")?;
@ -240,6 +567,27 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConf
})
}
#[inline(never)]
pub(super) fn project_redis_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<RedisSemanticCacheConfig> {
Ok(RedisSemanticCacheConfig {
redis_url: backend.getattr("_redis_url")?.extract::<String>()?,
index_name: backend
.getattr("_index_name")?
.extract::<Option<String>>()?
.unwrap_or_else(|| "litellm_semantic_cache_index".into()),
similarity_threshold: backend.getattr("similarity_threshold")?.extract::<f64>()?,
embedding_model: backend.getattr("embedding_model")?.extract::<String>()?,
embedding_max_input_tokens: backend
.getattr("embedding_max_input_tokens")?
.extract::<Option<u64>>()?,
embedding_timeout: backend
.getattr("embedding_timeout")?
.extract::<Option<f64>>()?,
})
}
#[inline(never)]
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
@ -252,6 +600,38 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
})
}
#[inline(never)]
fn project_gcs(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<GcsCacheConfig, UnsupportedCacheConfig>> {
let bucket_name = match backend.getattr("bucket_name")?.extract::<Option<String>>() {
Ok(Some(bucket_name)) if !bucket_name.is_empty() => bucket_name,
_ => return Ok(Err(UnsupportedCacheConfig::GcsBucket)),
};
Ok(Ok(GcsCacheConfig {
bucket_name,
key_prefix: backend.getattr("key_prefix")?.extract::<String>()?,
path_service_account: backend
.getattr("path_service_account")?
.extract::<Option<String>>()?,
}))
}
#[inline(never)]
fn project_disk(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<DiskCacheConfig, UnsupportedCacheConfig>> {
let store = backend.getattr("disk_cache")?;
if !instance_class_is(&store, "diskcache.core", "Cache")?
|| !instance_class_is(&store.getattr("_disk")?, "diskcache.core", "Disk")?
{
return Ok(Err(UnsupportedCacheConfig::DiskStore));
}
Ok(Ok(DiskCacheConfig {
directory: PathBuf::from(store.getattr("directory")?.extract::<String>()?),
}))
}
#[inline(never)]
fn project_redis(
backend: &Bound<'_, PyAny>,
@ -346,6 +726,77 @@ fn project_redis(
}))
}
#[inline(never)]
fn project_s3(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<S3CacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("s3_client")?;
if !instance_class_is(&client, "botocore.client", "S3")? {
return Ok(Err(UnsupportedCacheConfig::S3Client));
}
let meta = client.getattr("meta")?;
let Some(region) = optional_string(meta.getattr("region_name")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let Some(endpoint_url) = optional_string(meta.getattr("endpoint_url")?)? else {
return Ok(Err(UnsupportedCacheConfig::S3Option));
};
let client_config = meta.getattr("config")?;
for name in ["s3", "proxies", "client_cert"] {
if optional_attribute(&client_config, name)?.is_some_and(|value| !value.is_none()) {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
}
let signature = match optional_attribute(&client_config, "signature_version")? {
Some(value) => value.extract::<Option<String>>()?,
None => None,
};
if signature.as_deref() != Some("s3v4") {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let insecure = endpoint_url.starts_with("http://");
let verify = optional_attribute_chain(&client, &["_endpoint", "http_session", "_verify"])?;
let verified = verify
.and_then(|value| value.cast::<PyBool>().ok().map(|value| value.is_true()))
.unwrap_or(false);
if !verified && !insecure {
return Ok(Err(UnsupportedCacheConfig::S3Option));
}
let credentials = optional_attribute_chain(&client, &["_request_signer", "_credentials"])?
.ok_or(UnsupportedCacheConfig::S3Credentials);
let credentials = match credentials {
Ok(credentials) if !credentials.is_none() => credentials,
_ => return Ok(Err(UnsupportedCacheConfig::S3Credentials)),
};
let auth = if credentials.getattr("method")?.extract::<String>()?.as_str() == "explicit" {
AwsAuthConfig {
access_key_id: credentials
.getattr("access_key")?
.extract::<Option<String>>()?,
secret_access_key: credentials
.getattr("secret_key")?
.extract::<Option<String>>()?,
session_token: credentials.getattr("token")?.extract::<Option<String>>()?,
region_name: Some(region.clone()),
..Default::default()
}
} else {
AwsAuthConfig {
region_name: Some(region.clone()),
..Default::default()
}
};
let default_endpoint = endpoint_url == format!("https://s3.{region}.amazonaws.com")
|| (region == "us-east-1" && endpoint_url == "https://s3.amazonaws.com");
Ok(Ok(S3CacheConfig {
bucket: backend.getattr("bucket_name")?.extract::<String>()?,
key_prefix: backend.getattr("key_prefix")?.extract::<String>()?,
region,
endpoint: (!default_endpoint).then_some(S3Endpoint { url: endpoint_url }),
auth,
}))
}
#[inline(never)]
fn project_standalone_client<'py>(
client: &Bound<'py, PyAny>,
@ -380,7 +831,7 @@ fn project_standalone_client<'py>(
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
source: &Bound<'_, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
@ -468,6 +919,71 @@ fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[inline(never)]
fn project_valkey_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<ValkeySemanticCacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("sync_client")?;
let pool = client.getattr("connection_pool")?;
let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
for key in ["credential_provider", "redis_connect_func"] {
if has_value(&resolved, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
if is_tls {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let connection = RedisConnectionConfig {
host: required_string(&resolved, "host")?,
port: u16::try_from(required_i64(&resolved, "port")?)
.map_err(|_| PyValueError::new_err("invalid Redis port"))?,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol: RedisProtocol::Resp2,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
read_timeout: None,
connect_timeout: None,
socket_keepalive: None,
health_check_interval: Duration::ZERO,
client_name: None,
tls: None,
};
if connection.host.is_empty() {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
Ok(Ok(ValkeySemanticCacheConfig {
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
index_name: backend.getattr("index_name")?.extract()?,
embedding_model: backend.getattr("embedding_model")?.extract()?,
connection,
}))
}
#[inline(never)]
fn project_connection_pool<'py>(
pool: &Bound<'py, PyAny>,
) -> PyResult<Result<(Bound<'py, PyDict>, bool), UnsupportedCacheConfig>> {
if !instance_class_is(pool, "redis.connection", "ConnectionPool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
let connection_class = resolved
.get_item("connection_class")?
.unwrap_or(pool.getattr("connection_class")?);
let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? {
false
} else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
true
} else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
Ok(Ok((resolved, is_tls)))
}
#[inline(never)]
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {
@ -549,6 +1065,31 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult<O
}
}
#[inline(never)]
fn optional_attribute<'py>(
value: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
match value.getattr(name) {
Ok(value) => Ok(Some(value)),
Err(error) if error.is_instance_of::<PyAttributeError>(value.py()) => Ok(None),
Err(error) => Err(error),
}
}
#[inline(never)]
fn optional_attribute_chain<'py>(
value: &Bound<'py, PyAny>,
names: &[&str],
) -> PyResult<Option<Bound<'py, PyAny>>> {
names
.iter()
.try_fold(Some(value.clone()), |current, name| match current {
Some(current) => optional_attribute(&current, name),
None => Ok(None),
})
}
#[inline(never)]
fn optional_string(value: Bound<'_, PyAny>) -> PyResult<Option<String>> {
Ok(value
@ -637,12 +1178,13 @@ mod tests {
use pyo3::{prelude::*, types::PyDict};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_redis_semantic::RedisSemanticConfig;
use super::{
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig,
RedisProtocol,
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
use crate::cache::{embedder::PythonEmbedder, native::NativeResponseCache};
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
@ -715,6 +1257,49 @@ mod tests {
});
}
#[test]
fn redis_semantic_service_mismatch_accepts_backend_precision_threshold() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(_redis_url='redis://127.0.0.1/', _index_name='semantic_idx', similarity_threshold=0.8, embedding_model='text-embedding-3-small', embedding_max_input_tokens=None, embedding_timeout=None)\n\
facade = SimpleNamespace(type='redis-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let backend = facade.getattr("cache").unwrap();
let embedder = PythonEmbedder::new(backend.clone().unbind());
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Redis semantic cache should be supported");
};
let CacheBackendConfig::RedisSemantic(config) = config.backend else {
panic!("expected Redis semantic configuration");
};
let service = NativeResponseCache::redis_semantic(
&config.redis_url,
embedder,
RedisSemanticConfig {
index_name: config.index_name.clone(),
similarity_threshold: config.similarity_threshold as f32,
},
)
.unwrap();
let matching_config = NativeCacheConfig {
policy: CachePolicy {
mode: "default-on".into(),
ttl: None,
namespace: None,
supported_call_types: None,
redis_flush_size: None,
semantic_cache_scope: "key".into(),
},
backend: CacheBackendConfig::RedisSemantic(config),
};
assert_eq!(matching_config.service_mismatch(&service), None);
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
@ -758,6 +1343,87 @@ mod tests {
});
}
#[test]
fn projects_valkey_semantic_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.max_connections = 12\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Valkey semantic cache should be supported");
};
let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else {
panic!("expected Valkey semantic configuration");
};
assert_eq!(valkey.similarity_threshold, 0.85);
assert_eq!(valkey.index_name, "semantic_idx");
assert_eq!(valkey.embedding_model, "text-embedding-3-small");
assert_eq!(valkey.connection.host, "cache.internal");
assert_eq!(valkey.connection.port, 6390);
assert_eq!(valkey.connection.database, 2);
assert_eq!(valkey.connection.pool_size, 12);
assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2);
assert!(valkey.connection.tls.is_none());
});
}
#[test]
fn valkey_semantic_tls_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("TLS Valkey semantic cache should stay on Python");
};
assert_eq!(
reason.message(),
"native Redis connection type is not implemented"
);
});
}
#[test]
fn valkey_semantic_dynamic_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic Valkey authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn dynamic_redis_auth_stays_on_python() {
Python::initialize();
@ -822,6 +1488,71 @@ mod tests {
});
}
#[test]
fn projects_gcs_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(bucket_name='bucket', key_prefix='cache/', path_service_account='credentials.json')\n\
facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("GCS cache should be supported");
};
let CacheBackendConfig::Gcs(gcs) = config.backend else {
panic!("expected GCS configuration");
};
assert_eq!(
gcs,
GcsCacheConfig {
bucket_name: "bucket".into(),
key_prefix: "cache/".into(),
path_service_account: Some("credentials.json".into()),
}
);
let matching = NativeResponseCache::gcs(
litellm_cache_gcs::GcsConfig {
bucket_name: "bucket".into(),
gcs_path: Some("cache/".into()),
path_service_account: Some("credentials.json".into()),
endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(),
},
Some("token".into()),
)
.unwrap();
let matching_config = NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Gcs(gcs),
};
assert_eq!(matching_config.service_mismatch(&matching), None);
});
}
#[test]
fn rejects_gcs_without_a_bucket_name() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\
facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("GCS cache without a bucket should be unsupported");
};
assert!(matches!(&reason, UnsupportedCacheConfig::GcsBucket));
assert_eq!(
reason.message(),
"native GCS cache requires a configured bucket name"
);
});
}
#[test]
fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() {
Python::initialize();

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

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

@ -31,9 +31,15 @@ struct RedisPoolGuard {
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: Option<usize>,
client_name: &'static str,
attributes: RedisPoolAttributes,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -41,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,
@ -65,9 +77,12 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
max_connections: None,
};
const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL;
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
disk_store: Option<DiskStoreGuard>,
connection: ConnectionGuard,
}
@ -150,7 +165,9 @@ impl ObjectGuard {
return Ok(false);
}
for (name, value) in &expected.attributes {
if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
if (instance.contains(name)? && !self.config_names.contains(&name.as_str()))
|| !attributes.get_item(name)?.is(value.bind(py))
{
return Ok(false);
}
}
@ -171,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(),
@ -180,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)
@ -218,6 +238,26 @@ impl RedisPoolGuard {
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let store = backend.getattr("disk_cache")?;
Ok(Self {
reference: store.clone().unbind(),
directory: store.getattr("directory")?.extract()?,
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let store = backend.getattr("disk_cache")?;
Ok(self.reference.bind(py).is(&store)
&& self.directory == store.getattr("directory")?.extract::<String>()?)
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -246,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,
})
}
@ -261,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),
}
}
@ -269,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),
}
}
}
@ -290,16 +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_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")?;
@ -319,6 +410,15 @@ impl FacadeGuard {
if let Some(message) = config.service_mismatch(service) {
return Err(PyTypeError::new_err(message));
}
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,
@ -343,8 +443,30 @@ impl FacadeGuard {
"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()?,
connection: ConnectionGuard::capture(kind, cluster, &backend)?,
})
}
@ -357,12 +479,20 @@ impl FacadeGuard {
if !self.backend.matches(py, &backend)? {
return Ok(false);
}
if let Some(guard) = &self.disk_store
&& !guard.matches(py, &backend)?
{
return Ok(false);
}
self.connection.matches(py, &backend)
}
pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.outer.traverse(&visit)?;
self.backend.traverse(&visit)?;
if let Some(guard) = &self.disk_store {
guard.traverse(&visit)?;
}
self.connection.traverse(&visit)
}
}

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

@ -1,8 +1,27 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
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::*};
use litellm_http::ClientVariant;
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyRuntimeError, PyTypeError},
prelude::*,
types::PyDict,
};
use url::Url;
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
use super::{
cache_error,
config::{QdrantSemanticCacheConfig, project_redis_semantic},
embedder::PythonEmbedder,
facade::FacadeGuard,
native::NativeResponseCache,
request::duration,
};
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
@ -64,6 +83,199 @@ 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(
py: Python<'_>,
bucket_name: String,
gcs_path: Option<String>,
path_service_account: Option<String>,
endpoint: Option<String>,
token: Option<String>,
) -> PyResult<Self> {
let config = GcsConfig {
bucket_name,
gcs_path,
path_service_account,
endpoint: endpoint.unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()),
};
let service = release_gil(py, move || NativeResponseCache::gcs(config, token))
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (directory))]
fn disk(py: Python<'_>, directory: String) -> PyResult<Self> {
let service =
release_gil(py, move || NativeResponseCache::disk(&directory)).map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (url, *, collection_name, similarity_threshold, vector_size, embedding_model="text-embedding-3-small", api_key=None, embedding_api_key=None, embedding_api_base=None, embedding_timeout_seconds=None, quantization="binary"))]
#[expect(
clippy::too_many_arguments,
reason = "the test handle exposes the complete Qdrant constructor"
)]
fn qdrant_semantic(
py: Python<'_>,
url: String,
collection_name: String,
similarity_threshold: f64,
vector_size: u64,
embedding_model: &str,
api_key: Option<String>,
embedding_api_key: Option<String>,
embedding_api_base: Option<String>,
embedding_timeout_seconds: Option<f64>,
quantization: &str,
) -> PyResult<Self> {
let parsed = Url::parse(&url).map_err(|_| {
pyo3::exceptions::PyValueError::new_err(
"native Qdrant requires the default REST port so the gRPC port can be derived",
)
})?;
if !matches!(parsed.scheme(), "http" | "https")
|| (!parsed.path().is_empty() && parsed.path() != "/")
|| parsed.query().is_some()
|| parsed.host_str().is_none()
|| 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",
));
}
let mut grpc_url = parsed;
grpc_url.set_port(Some(6334)).map_err(|_| {
pyo3::exceptions::PyValueError::new_err(
"native Qdrant requires the default REST port so the gRPC port can be derived",
)
})?;
grpc_url.set_path("");
grpc_url.set_query(None);
let embedding_api_key = embedding_api_key
.or_else(|| {
std::env::var("OPENAI_API_KEY")
.ok()
.filter(|value| !value.is_empty())
})
.ok_or_else(|| {
pyo3::exceptions::PyValueError::new_err(
"native semantic embedding requires an OpenAI API key",
)
})?;
let embedding_api_base = embedding_api_base.unwrap_or_else(|| {
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 {
"binary" => Quantization::Binary,
"scalar" => Quantization::Scalar,
"product" => Quantization::Product,
_ => {
return Err(pyo3::exceptions::PyValueError::new_err(
"unsupported Qdrant quantization",
));
}
};
let config = QdrantSemanticCacheConfig {
grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(),
api_key,
collection_name,
similarity_threshold,
vector_size,
embedding: OpenAiEmbedderConfig {
api_base: embedding_api_base,
api_key: embedding_api_key,
model: embedding_model.to_owned(),
timeout: embedding_timeout_seconds.map(duration).transpose()?,
},
quantization,
};
let http_config = crate::http::call_config(py, &PyDict::new(py), true)?;
let client = crate::http::pool()
.client(&http_config, ClientVariant::Provider)
.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)
.await
.map_err(cache_error)
})?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[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> {
@ -79,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()
@ -87,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 {
@ -104,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)?;
}

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

@ -1,16 +1,19 @@
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::{
exceptions::{PyRuntimeError, PyValueError},
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
prelude::*,
};
@ -21,6 +24,7 @@ pub(crate) use self::{
fn cache_error(error: Error) -> PyErr {
match error {
Error::InvalidEntry => PyValueError::new_err(error.to_string()),
Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()),
_ => PyRuntimeError::new_err(error.to_string()),
}
}

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

File diff suppressed because it is too large Load diff

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

@ -1,9 +1,11 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
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::Value;
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
@ -12,24 +14,52 @@ struct RequestInput {
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
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> {
#[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> {
let mut request = ResponseCacheRequest::new(input.key);
if let Some(controls) = input.controls {
request.controls = controls;
}
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
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>> {
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<NativeRequest>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)

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]

4
litellm/__init__.py
View file

@ -2146,6 +2146,10 @@ if TYPE_CHECKING:
from .llms.edenai.videos.transformation import (
EdenAIVideoConfig as EdenAIVideoConfig,
)
from .llms.fal_ai.chat.transformation import (
FalAIChatConfig as FalAIChatConfig,
FalAIError as FalAIError,
)
from .llms.ovhcloud.chat.transformation import (
OVHCloudChatConfig as OVHCloudChatConfig,
)

4
litellm/_lazy_imports_registry.py
View file

@ -335,6 +335,8 @@ LLM_CONFIG_NAMES: Final = (
"EdenAITextToSpeechConfig",
"EdenAIImageGenerationConfig",
"EdenAIVideoConfig",
"FalAIChatConfig",
"FalAIError",
"OVHCloudChatConfig",
"OVHCloudEmbeddingConfig",
"CometAPIEmbeddingConfig",
@ -1251,6 +1253,8 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"EdenAITextToSpeechConfig": (".llms.edenai.text_to_speech.transformation", "EdenAITextToSpeechConfig"),
"EdenAIImageGenerationConfig": (".llms.edenai.image_generation.transformation", "EdenAIImageGenerationConfig"),
"EdenAIVideoConfig": (".llms.edenai.videos.transformation", "EdenAIVideoConfig"),
"FalAIChatConfig": (".llms.fal_ai.chat.transformation", "FalAIChatConfig"),
"FalAIError": (".llms.fal_ai.chat.transformation", "FalAIError"),
"OVHCloudChatConfig": (".llms.ovhcloud.chat.transformation", "OVHCloudChatConfig"),
"OVHCloudEmbeddingConfig": (
".llms.ovhcloud.embedding.transformation",

1
litellm/anthropic_interface/exceptions/exceptions.py
View file

@ -25,6 +25,7 @@ class AnthropicErrorDetail(TypedDict):
type: AnthropicErrorType
message: str
provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]]
litellm_call_id: NotRequired[ReadOnly[str]]
class AnthropicErrorResponse(TypedDict, total=False):

5
litellm/caching/qdrant_semantic_cache.py
View file

@ -12,6 +12,7 @@ import ast
import asyncio
import json
import os
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
import litellm
@ -36,6 +37,8 @@ from ._embedding_router import (
)
from .base_cache import BaseCache
_WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"})
if TYPE_CHECKING:
from litellm.router import Router
@ -313,6 +316,7 @@ class QdrantSemanticCache(BaseCache):
self.sync_client.put(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points",
headers=self.headers,
params=_WAIT_FOR_INDEXING,
json=data,
)
@ -422,6 +426,7 @@ class QdrantSemanticCache(BaseCache):
await self.async_client.put(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points",
headers=self.headers,
params=_WAIT_FOR_INDEXING,
json=data,
)

15
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))
@ -2121,3 +2133,6 @@ BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY: Final = "batch_enqueued_token_limit"
# Shared read-only empty mapping, for defaulting optional Mapping parameters without
# constructing a fresh mutable dict at each call site.
EMPTY_MAPPING: Final = MappingProxyType({})
# API endpoint for breached password k-anonymity search
HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range"

43
litellm/cost_calculator.py
View file

@ -948,7 +948,7 @@ def _extract_service_tier(source: object) -> str | None:
return None
def _get_usage_object(
def get_usage_object(
completion_response: object,
) -> Usage | None:
usage_obj: Final = cast(
@ -1336,7 +1336,7 @@ def completion_cost(
cache_creation_input_tokens: int | None = None
cache_read_input_tokens: int | None = None
audio_transcription_file_duration: float = 0.0
provider_usage_object: Final = _get_usage_object(completion_response=completion_response)
provider_usage_object: Final = get_usage_object(completion_response=completion_response)
cost_per_token_usage_object: Final[Usage | None] = (
_without_provider_stated_cost(provider_usage_object) if custom_pricing else provider_usage_object
)
@ -2033,6 +2033,45 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn
return None
def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None:
raw_entry: Final = litellm.model_cost.get(key)
return raw_entry if isinstance(raw_entry, Mapping) else None
def pricing_entry_for_cost_calc(
model: str | None,
completion_response: object | None,
custom_llm_provider: str | None,
custom_pricing: bool | None,
base_model: str | None,
router_model_id: str | None,
region_name: str | None,
litellm_logging_obj: LitellmLoggingObject | None,
) -> tuple[str, Mapping[str, object]] | None:
deployment_entry: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id)
deployment_key: Final = router_model_id or model
if deployment_entry is not None and deployment_key is not None:
registered_entry: Final = _raw_cost_map_entry(router_model_id) if router_model_id is not None else None
return deployment_key, registered_entry or deployment_entry
selected_model: Final = _select_model_name_for_cost_calc(
model=model,
completion_response=completion_response,
base_model=base_model,
custom_pricing=custom_pricing,
custom_llm_provider=custom_llm_provider,
router_model_id=router_model_id,
region_name=region_name,
)
candidates: Final = (selected_model, _get_response_model(completion_response), model)
resolved: Final = next(
(info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info),
None,
)
if resolved is None:
return None
return resolved["key"], _raw_cost_map_entry(resolved["key"]) or resolved
def ocr_cost(
model: str,
custom_llm_provider: str | None,

47
litellm/integrations/prometheus.py
View file

@ -10,6 +10,7 @@ import sys
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import replace
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast
from pydantic import BaseModel
@ -66,6 +67,7 @@ from litellm.types.proxy.carried_budget_state import (
from litellm.types.utils import (
StandardLoggingGuardrailInformation,
StandardLoggingPayload,
StandardLoggingZeroCostDiagnostic,
)
if TYPE_CHECKING:
@ -713,6 +715,15 @@ class PrometheusLogger(CustomLogger):
labelnames=self.get_labels_for_metric("litellm_requests_metric"),
)
self.litellm_zero_cost_requests_total = self._counter_factory(
name="litellm_zero_cost_requests_total",
documentation=(
"Requests that carried usage but were logged at $0 on a model whose pricing entry "
"has a non-zero rate, by reason (missing_pricing_key, pricing_not_applied, cost_calculation_error)"
),
labelnames=self.get_labels_for_metric("litellm_zero_cost_requests_total"),
)
# Cache metrics
self.litellm_cache_hits_metric = self._counter_factory(
name="litellm_cache_hits_metric",
@ -1410,6 +1421,11 @@ class PrometheusLogger(CustomLogger):
enum_values=enum_values,
label_context=label_context,
)
self._increment_zero_cost_requests_metric(
zero_cost_diagnostic=standard_logging_payload.get("zero_cost_diagnostic"),
enum_values=enum_values,
label_context=label_context,
)
# input, output, total token metrics
self._increment_token_metrics(
@ -1983,6 +1999,30 @@ class PrometheusLogger(CustomLogger):
amount=float(response_cost),
)
def _increment_zero_cost_requests_metric(
self,
zero_cost_diagnostic: StandardLoggingZeroCostDiagnostic | None,
enum_values: UserAPIKeyLabelValues,
label_context: PrometheusLabelFactoryContext,
) -> None:
if zero_cost_diagnostic is None:
return
supported_labels: Final = self.get_labels_for_metric("litellm_zero_cost_requests_total")
reason_label: Final = (
MappingProxyType({ZERO_COST_REASON_LABEL: zero_cost_diagnostic["reason"]})
if ZERO_COST_REASON_LABEL in supported_labels
else MappingProxyType({})
)
labels: Final = MappingProxyType(
{
**prometheus_label_factory(
supported_enum_labels=supported_labels, enum_values=enum_values, label_context=label_context
),
**reason_label,
}
)
self.litellm_zero_cost_requests_total.labels(**labels).inc()
@staticmethod
def _get_remaining_from_v3_rate_limit_headers(
standard_logging_payload: StandardLoggingPayload | None,
@ -2333,6 +2373,8 @@ class PrometheusLogger(CustomLogger):
team_alias=user_api_team_alias,
user=user_id,
model_id=standard_logging_payload.get("model_id", ""),
requested_model=standard_logging_payload.get("model_group"),
api_provider=standard_logging_payload.get("custom_llm_provider"),
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
@ -2345,6 +2387,11 @@ class PrometheusLogger(CustomLogger):
"litellm_llm_api_failed_requests_metric",
enum_values,
)
self._increment_zero_cost_requests_metric(
zero_cost_diagnostic=standard_logging_payload.get("zero_cost_diagnostic"),
enum_values=enum_values,
label_context=PrometheusLabelFactoryContext(enum_values),
)
self.set_llm_deployment_failure_metrics(kwargs)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,

32
litellm/litellm_core_utils/agentic_followup_kwargs.py Normal file
View file

@ -0,0 +1,32 @@
from collections.abc import Collection, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import Final
def build_agentic_followup_kwargs(
*,
request_kwargs: Mapping[str, object],
patch_kwargs: Mapping[str, object],
request_params: Collection[str],
depth: int,
max_loops: int,
fingerprints: Sequence[str],
fingerprint: str,
) -> Mapping[str, object]:
"""Kwargs for an agentic follow-up call: the request's kwargs overlaid by the plan's, never repeating a key already sent as a request param"""
seen: Final = [*fingerprints, fingerprint] # mutable-ok: the chat loop's settings reader only accepts a list
return MappingProxyType(
{
key: value
for key, value in chain(
((k, v) for k, v in request_kwargs.items() if k not in request_params),
((k, v) for k, v in patch_kwargs.items() if k not in request_params),
(
("_agentic_loop_depth", depth + 1),
("max_agentic_loops", max_loops),
("_agentic_loop_fingerprints", seen),
),
)
}
)

46
litellm/litellm_core_utils/chat_completion_agentic_loop.py
View file

@ -2,10 +2,13 @@
import json
from collections.abc import Mapping
from itertools import chain
from types import MappingProxyType
from typing import Final, cast
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
from litellm.litellm_core_utils.agentic_loop_settings import (
DEFAULT_MAX_AGENTIC_LOOPS,
validated_max_agentic_loops,
@ -117,13 +120,25 @@ def _wrap_response_as_fake_stream(
)
def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
metadata = kwargs_for_followup.get("litellm_metadata")
metadata = dict(metadata) if isinstance(metadata, dict) else {}
for key, value in kwargs_for_followup.items():
if key.startswith("_agentic_loop") or key == "max_agentic_loops" or is_interception_internal_key(key):
metadata[key] = value
kwargs_for_followup["litellm_metadata"] = metadata
def _with_agentic_loop_metadata(kwargs_for_followup: Mapping[str, object]) -> Mapping[str, object]:
metadata: Final = kwargs_for_followup.get("litellm_metadata")
return MappingProxyType(
{
**kwargs_for_followup,
"litellm_metadata": dict( # mutable-ok: the follow-up call's logging and proxy hooks write into litellm_metadata in place
chain(
metadata.items() if isinstance(metadata, dict) else (),
(
(key, value)
for key, value in kwargs_for_followup.items()
if key.startswith("_agentic_loop")
or key == "max_agentic_loops"
or is_interception_internal_key(key)
),
)
),
}
)
def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]:
@ -165,14 +180,17 @@ async def _execute_chat_completion_agentic_plan(
if "tool_choice" not in patch.optional_params:
optional_params_for_followup.pop("tool_choice", None)
kwargs_for_followup: Final = _filter_followup_kwargs(kwargs)
kwargs_for_followup.update(
{k: v for k, v in _filter_followup_kwargs(patch.kwargs).items() if k not in optional_params_for_followup}
kwargs_for_followup: Final = _with_agentic_loop_metadata(
build_agentic_followup_kwargs(
request_kwargs=_filter_followup_kwargs(kwargs),
patch_kwargs=_filter_followup_kwargs(patch.kwargs),
request_params=frozenset((*optional_params_for_followup, "model", "messages")),
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
)
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
_add_agentic_loop_metadata(kwargs_for_followup)
try:
response_followup = await litellm.acompletion(

40
litellm/litellm_core_utils/core_helpers.py
View file

@ -359,6 +359,46 @@ def get_litellm_metadata_from_kwargs(kwargs: dict):
return {}
def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object:
if isinstance(user_api_key_auth, Mapping):
return user_api_key_auth.get("budget_reservation")
return getattr(user_api_key_auth, "budget_reservation", None)
def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None:
stamped: Final = metadata.get("user_api_key_budget_reservation")
if isinstance(stamped, dict):
return stamped
on_auth_object: Final = _budget_reservation_on_auth_object(metadata.get("user_api_key_auth"))
return on_auth_object if isinstance(on_auth_object, dict) else None
def _stamp_budget_reservation_callback_bound(litellm_params: Mapping[str, object], callback_bound: bool) -> None:
for metadata_variable_name in ("metadata", "litellm_metadata"):
metadata = litellm_params.get(metadata_variable_name)
if not isinstance(metadata, Mapping):
continue
budget_reservation = budget_reservation_from_metadata(metadata)
if budget_reservation is not None:
budget_reservation["callback_bound"] = callback_bound
def bind_budget_reservation_to_callbacks(litellm_params: Mapping[str, object]) -> None:
"""Mark the request's budget reservation as owned by the success callbacks of this call.
The proxy releases any reservation still unbound when the request ends; one bound here
is left for the cost callback, which may finish after the response has been sent. Bind
only where a success handler is guaranteed to run: a logging object merely existing is
not that, since the proxy builds one for every route before calling anything.
"""
_stamp_budget_reservation_callback_bound(litellm_params, True)
def unbind_budget_reservation_from_callbacks(litellm_params: Mapping[str, object]) -> None:
"""Hand a failed call's reservation back to the request-end release: failure handlers never settle it."""
_stamp_budget_reservation_callback_bound(litellm_params, False)
def reconstruct_model_name(
model_name: str,
custom_llm_provider: str | None,

2
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -46,6 +46,8 @@ OPTIONAL_KWARGS_KEYS: Final = (
"bucket_name",
"s3_endpoint_url",
"s3_region_name",
"s3_access_key_id",
"s3_secret_access_key",
"vertex_credentials",
"vertex_project",
"vertex_location",

3
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -859,6 +859,9 @@ def _get_openai_compatible_provider_info(
elif custom_llm_provider == "edenai":
api_base = litellm.EdenAIChatConfig.get_api_base(api_base) # rebind-ok: chain resolves in place
dynamic_api_key = litellm.EdenAIChatConfig.get_api_key(api_key) # rebind-ok: chain resolves in place
elif custom_llm_provider == "fal_ai":
api_base = litellm.FalAIChatConfig.get_api_base(api_base) # rebind-ok: chain resolves in place
dynamic_api_key = litellm.FalAIChatConfig.get_api_key(api_key) # rebind-ok: chain resolves in place
elif custom_llm_provider == "aiml":
(
api_base,

21
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -17,14 +17,16 @@ import random
import sys
import threading
import time
from collections.abc import Awaitable, Callable
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass, replace
from datetime import datetime, timezone
from importlib.resources import files
from pathlib import Path
from types import MappingProxyType
from typing import Final, Protocol
import httpx
from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm import verbose_logger
@ -37,6 +39,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
)
FALLBACK_GENERALIZATIONS_KEY: Final = "fallback_generalizations"
_CATALOG_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]])
_CLI_ENTRYPOINT_NAMES: Final = frozenset({"lite", "litellm-proxy"})
@ -88,6 +91,18 @@ class GetModelCostMap:
"""Load the local backup model cost map bundled with the package."""
return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map
_loaded_catalog: Mapping[str, Mapping[str, object]] = MappingProxyType({})
@classmethod
def loaded_model_cost_map(cls) -> Mapping[str, Mapping[str, object]]:
"""The catalog as last loaded (bundled or remote), untouched by ``register_model`` or router registrations."""
return cls._loaded_catalog
@classmethod
def _snapshot_loaded_catalog(cls, model_cost: Mapping[str, object]) -> None:
raw: Final = _CATALOG_ADAPTER.validate_python(model_cost)
cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()})
@classmethod
def _get_backup_model_count(cls) -> int:
"""Return the number of models in the local backup (cached int)."""
@ -533,7 +548,9 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMapReloaded:
_cost_map_source_info.source_revision = loaded.revision
_cost_map_source_info.etag = loaded.etag
return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map))
finalized: Final = _finalize_model_cost_map(loaded.model_cost_map)
GetModelCostMap._snapshot_loaded_catalog(finalized) # pyright: ignore[reportPrivateUsage] # same module
return replace(loaded, model_cost_map=finalized)
def adopt_model_cost_map(

163
litellm/litellm_core_utils/litellm_logging.py
View file

@ -50,6 +50,8 @@ from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
_select_model_name_for_cost_calc,
get_usage_object,
pricing_entry_for_cost_calc,
)
from litellm.exceptions import (
BudgetExceededError,
@ -89,6 +91,10 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
InteractionsUsageObjectTransformation,
)
from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import (
diagnose_zero_cost,
zero_cost_warning,
)
from litellm.litellm_core_utils.logging_utils import (
truncate_base64_in_messages,
truncate_base64_in_messages_async,
@ -157,6 +163,7 @@ from litellm.types.utils import (
StandardLoggingPayloadStatusFields,
StandardLoggingPromptManagementMetadata,
StandardLoggingVectorStoreRequest,
StandardLoggingZeroCostDiagnostic,
TextCompletionResponse,
TranscriptionResponse,
Usage,
@ -614,6 +621,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape
## TIME TO FIRST TOKEN LOGGING ##
self.completion_start_time: datetime.datetime | None = None
self.zero_cost_warned: bool = False
self._llm_caching_handler: LLMCachingHandler | None = None
# INITIAL LITELLM_PARAMS
@ -1757,17 +1765,26 @@ class Logging(LiteLLMLoggingBaseClass):
if transformed_result is not None:
result = transformed_result
result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({})
result_additional_headers: Final = (
result_hidden_params.get("additional_headers")
if isinstance(result_hidden_params, dict)
else getattr(result_hidden_params, "additional_headers", None)
priced_result: Final = (
result.response
if isinstance(result, (ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent))
else result
)
if isinstance(result, (BaseModel, HttpxBinaryResponseContent)) and hasattr(result, "_hidden_params"):
result_hidden_params: Final = getattr(priced_result, "_hidden_params", None) or MappingProxyType({})
if isinstance(priced_result, (BaseModel, HttpxBinaryResponseContent)) and hasattr(
priced_result, "_hidden_params"
):
hidden_params: Final = result_hidden_params
if (
"response_cost" in hidden_params and hidden_params["response_cost"] is not None
): # use cost if already calculated
self._record_zero_cost_diagnostic(
priced_result,
hidden_params["response_cost"],
litellm_model_name=litellm_model_name,
router_model_id=router_model_id or hidden_params.get("model_id"),
)
return hidden_params["response_cost"]
elif router_model_id is None and "model_id" in hidden_params: # use model_id if not already set
router_model_id = hidden_params["model_id"]
@ -1779,18 +1796,7 @@ class Logging(LiteLLMLoggingBaseClass):
router_model_id = self.get_router_model_id()
## RESPONSE COST ##
spilled_over: Final = is_spilled_over_ptu_request(
model_info=_deployment_model_info(self.litellm_params if hasattr(self, "litellm_params") else None),
response_headers=self.model_call_details.get("response_headers"),
additional_headers=result_additional_headers,
)
custom_pricing: Final = (
False
if spilled_over
else use_custom_pricing_for_model(
litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None)
)
)
custom_pricing: Final = self._custom_pricing_for(priced_result)
prompt = self._prompt_for_cost_calculation()
@ -1799,7 +1805,7 @@ class Logging(LiteLLMLoggingBaseClass):
try:
response_cost_calculator_kwargs: Final = {
"response_object": result,
"response_object": priced_result,
"model": litellm_model_name or self.model,
"cache_hit": cache_hit,
"custom_llm_provider": self.model_call_details.get("custom_llm_provider", None),
@ -1842,9 +1848,18 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug("response_cost: %s", response_cost)
additional_response_cost: Final[object] = self.model_call_details.get("additional_response_cost")
if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0:
return (response_cost or 0.0) + additional_response_cost
return response_cost
total_response_cost: Final = (
(response_cost or 0.0) + additional_response_cost
if isinstance(additional_response_cost, (int, float)) and additional_response_cost > 0
else response_cost
)
self._record_zero_cost_diagnostic(
priced_result,
total_response_cost,
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
)
return total_response_cost
except Exception as e: # error calculating cost
debug_info = StandardLoggingModelCostFailureDebugInformation(
error_str=str(e),
@ -1858,9 +1873,108 @@ class Logging(LiteLLMLoggingBaseClass):
)
verbose_logger.debug("response_cost_failure_debug_information: %s", debug_info)
self.model_call_details["response_cost_failure_debug_information"] = debug_info
self._record_zero_cost_diagnostic(
priced_result,
None,
calculation_failed=True,
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
)
return None
def _record_zero_cost_diagnostic(
self,
result: object,
response_cost: float | None,
*,
calculation_failed: bool = False,
litellm_model_name: str | None = None,
router_model_id: str | None = None,
) -> None:
if response_cost is None and not calculation_failed:
return
if self.model_call_details.get("cache_hit") is True:
self.model_call_details["zero_cost_diagnostic"] = None
return
try:
finding: Final = self._zero_cost_finding(
result,
response_cost,
calculation_failed=calculation_failed,
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
)
except Exception as e: # noqa: BLE001 # the pricing helpers raise plain Exception and a diagnostic must never break cost tracking
verbose_logger.debug("zero_cost_diagnostic skipped: %s", e)
return
self.model_call_details["zero_cost_diagnostic"] = finding[0] if finding is not None else None
if finding is None or self.zero_cost_warned:
return
self.zero_cost_warned = True
verbose_logger.warning(finding[1])
def _zero_cost_finding(
self,
result: object,
response_cost: float | None,
*,
calculation_failed: bool,
litellm_model_name: str | None,
router_model_id: str | None,
) -> tuple[StandardLoggingZeroCostDiagnostic, str] | None:
metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(self.litellm_params)
if response_cost or is_unbilled_non_inference_call(self.call_type, metadata, result):
return None
usage: Final = get_usage_object(completion_response=result)
if usage is None:
return None
model: Final = litellm_model_name or self.model
custom_llm_provider: Final = self.model_call_details.get("custom_llm_provider")
pricing: Final = pricing_entry_for_cost_calc(
model=model,
completion_response=result,
custom_llm_provider=custom_llm_provider,
custom_pricing=self._custom_pricing_for(result),
base_model=_get_base_model_from_metadata(model_call_details=self.model_call_details),
router_model_id=router_model_id or self.get_router_model_id(),
region_name=_resolve_mantle_region_for_cost(
custom_llm_provider=custom_llm_provider,
litellm_params=self.model_call_details.get("litellm_params"),
),
litellm_logging_obj=self,
)
if pricing is None:
return None
diagnostic: Final = diagnose_zero_cost(
usage=usage, pricing_model=pricing[0], pricing_entry=pricing[1], calculation_failed=calculation_failed
)
if diagnostic is None:
return None
model_group: Final = metadata.get("model_group")
return diagnostic, zero_cost_warning(
diagnostic,
model_group=model_group if isinstance(model_group, str) else None,
model=model,
custom_llm_provider=custom_llm_provider,
usage=usage,
)
def _custom_pricing_for(self, result: object) -> bool:
litellm_params: Final = getattr(self, "litellm_params", None)
result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({})
additional_headers: Final = (
result_hidden_params.get("additional_headers")
if isinstance(result_hidden_params, dict)
else getattr(result_hidden_params, "additional_headers", None)
)
spilled_over: Final = is_spilled_over_ptu_request(
model_info=_deployment_model_info(litellm_params),
response_headers=self.model_call_details.get("response_headers"),
additional_headers=additional_headers,
)
return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params)
def _prompt_for_cost_calculation(self) -> str:
"""
The raw input string is only priced directly for text-to-speech, which bills per character.
@ -2205,6 +2319,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = 0.0
elif "response_cost" in hidden_params:
self.model_call_details["response_cost"] = hidden_params["response_cost"]
self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"])
elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0:
# Preserve response_cost if already calculated (e.g., by pass-through
# handlers like Gemini/Vertex which call completion_cost directly).
@ -5378,7 +5493,7 @@ def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str,
"""Access groups the auth layer stamped onto this request, from whichever metadata field carries them.
Detached internal sub-calls only inherit the identity keys, so the auth object is the
fallback there, exactly as _get_budget_reservation_from_metadata does for reservations.
fallback there, exactly as budget_reservation_from_metadata does for reservations.
"""
for metadata_variable_name in ("metadata", "litellm_metadata"):
metadata = litellm_params.get(metadata_variable_name)
@ -6499,6 +6614,7 @@ def get_standard_logging_object_payload(
error_str=error_str,
error_information=error_information,
response_cost_failure_debug_info=kwargs.get("response_cost_failure_debug_information"),
zero_cost_diagnostic=kwargs.get("zero_cost_diagnostic"),
guardrail_information=metadata.get("standard_logging_guardrail_information", None),
standard_built_in_tools_params=standard_built_in_tools_params,
)
@ -6677,6 +6793,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
response_cost=response_cost,
autorouter_savings=None,
response_cost_failure_debug_info=None,
zero_cost_diagnostic=None,
status="success",
total_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT + DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT),
prompt_tokens=int(DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT),

146
litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py Normal file
View file

@ -0,0 +1,146 @@
from collections.abc import Mapping
from functools import reduce
from typing import Final
from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm.types.utils import StandardLoggingZeroCostDiagnostic, Usage
ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total"
_TEXT_INPUT_RATE: Final = "input_cost_per_token"
_AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token"
_TEXT_OUTPUT_RATE: Final = "output_cost_per_token"
_AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token"
_RATE_KEY_MARKERS: Final = ("cost", "pricing")
_NESTED_PRICING: Final = TypeAdapter(Mapping[str, object] | tuple[object, ...])
_MAX_PRICING_DEPTH: Final = 4
def _audio_tokens(details: object) -> int:
audio_tokens: Final = getattr(details, "audio_tokens", None)
return audio_tokens if isinstance(audio_tokens, int) and audio_tokens > 0 else 0
def _tokens(value: object) -> int:
return value if isinstance(value, int) and value > 0 else 0
def used_pricing_keys(usage: Usage) -> tuple[str, ...]:
prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details)
completion_audio: Final = _audio_tokens(usage.completion_tokens_details)
prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio
completion_text: Final = _tokens(usage.completion_tokens) - completion_audio
components: Final = (
(_TEXT_INPUT_RATE, prompt_text),
(_AUDIO_INPUT_RATE, prompt_audio),
(_TEXT_OUTPUT_RATE, completion_text),
(_AUDIO_OUTPUT_RATE, completion_audio),
)
return tuple(key for key, count in components if count > 0)
def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None:
try:
return _NESTED_PRICING.validate_python(value)
except ValidationError:
return None
def _is_rate_key(key: str) -> bool:
return any(marker in key for marker in _RATE_KEY_MARKERS)
def _rate_values(value: object) -> tuple[object, ...]:
nested: Final = _nested_pricing(value)
if isinstance(nested, Mapping):
return tuple(child for key, child in nested.items() if _is_rate_key(key))
if nested is None:
return (value,)
return nested
def _expand_rate_values(values: tuple[object, ...], _depth: int) -> tuple[object, ...]:
return tuple(nested for value in values for nested in _rate_values(value))
def _is_positive_number(value: object) -> bool:
return not isinstance(value, bool) and isinstance(value, (int, float)) and value > 0
def _declares_a_rate(pricing_entry: Mapping[str, object]) -> bool:
leaves: Final = reduce(_expand_rate_values, range(_MAX_PRICING_DEPTH), (pricing_entry,))
return any(_is_positive_number(leaf) for leaf in leaves)
def _is_explicit_zero(value: object) -> bool:
return not isinstance(value, bool) and isinstance(value, (int, float)) and value == 0
def diagnose_zero_cost(
usage: Usage,
pricing_model: str,
pricing_entry: Mapping[str, object],
calculation_failed: bool,
) -> StandardLoggingZeroCostDiagnostic | None:
used_keys: Final = used_pricing_keys(usage)
if not used_keys:
return None
missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None)
if not missing_keys and all(_is_explicit_zero(pricing_entry[key]) for key in used_keys):
return None
if not _declares_a_rate(pricing_entry):
return None
if calculation_failed:
return StandardLoggingZeroCostDiagnostic(
reason="cost_calculation_error", pricing_model=pricing_model, missing_pricing_keys=()
)
if missing_keys:
return StandardLoggingZeroCostDiagnostic(
reason="missing_pricing_key", pricing_model=pricing_model, missing_pricing_keys=missing_keys
)
return StandardLoggingZeroCostDiagnostic(
reason="pricing_not_applied", pricing_model=pricing_model, missing_pricing_keys=()
)
def _cause(diagnostic: StandardLoggingZeroCostDiagnostic) -> str:
reason: Final = diagnostic["reason"]
match reason:
case "missing_pricing_key":
return (
f"pricing entry '{diagnostic['pricing_model']}' has no {', '.join(diagnostic['missing_pricing_keys'])}. "
"Set the missing rate in the deployment's model_info or in the model cost map, "
"or set every rate to 0 to mark the model free"
)
case "pricing_not_applied":
return (
f"pricing entry '{diagnostic['pricing_model']}' declares non-zero rates for this usage, "
"but the cost calculator returned $0"
)
case "cost_calculation_error":
return (
f"cost calculation raised for pricing entry '{diagnostic['pricing_model']}', "
"see response_cost_failure_debug_information"
)
case _:
return assert_never(reason)
def zero_cost_warning(
diagnostic: StandardLoggingZeroCostDiagnostic,
*,
model_group: str | None,
model: str,
custom_llm_provider: str | None,
usage: Usage,
) -> str:
request: Final = (
f"model_group={model_group or model} model={model} provider={custom_llm_provider or 'unknown'} "
f"prompt_tokens={_tokens(usage.prompt_tokens)} completion_tokens={_tokens(usage.completion_tokens)}"
)
return (
f"Billable request priced at $0 and logged as such ({request}): {_cause(diagnostic)}. "
f'Counted in {ZERO_COST_COUNTER_NAME}{{reason="{diagnostic["reason"]}"}}'
)

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