scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

Merge branch 'main' into fix/router-cooldown-extension

Browse source
This commit is contained in:
Ambuj Upadhyay 2026-10-27 07:36:07 +05:30
commit 6541003141
373 changed files with 25587 additions and 11224 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
.agents/skills/database-access/SKILL.md Normal file
View file

@ -0,0 +1,52 @@
---
name: database-access
description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope
---
# Database access
## Adding new database access
- Put Python database access behind an operation in `litellm/repositories`
- Give callers typed inputs and outputs, not Prisma clients, models, or query dictionaries
- Keep query execution and backend selection inside the repository
- Prisma is fine for the implementation, adding a query does not require a Rust port
- Native Rust callers use `litellm-rust/crates/db` directly
## Asked to move existing access behind a repository
- Identify the calls being moved and their transaction boundary
- Move the existing Prisma implementation into the relevant repository
- Replace those direct database calls with repository operations
- Preserve return values, errors, and transaction behavior, with focused regression tests
- Keep this an extraction, not a backend migration or unrelated cleanup
## An equivalent Rust or SQLx query already exists
- Leave it and the current backend selection unchanged unless integration is explicitly requested
- Its existence alone does not justify switching a Prisma caller to Rust
- Repository extraction can keep using Prisma even when the Rust query is available
## Asked to port a Prisma query to Rust or SQLx
- Identify the existing inputs, outputs, errors, and transaction boundary
- Check `litellm-rust/crates/db` for an equivalent query and reuse it if its behavior matches
- Otherwise follow `litellm-rust/crates/db/AGENTS.md`
- Put static SQL in `queries/` and compile-check it with SQLx
- Keep typed inputs and reusable row projections in `src/`
- Refresh SQLx query metadata after SQL changes
- Keep Python conversion in the Python bridge, not in `litellm-rust/crates/db`
- When caller integration is requested, expose the Rust operation through the bridge and call it from the repository
- A query-only port can land without switching callers
- Verify behavior against real PostgreSQL before switching callers
- Check relevant null handling, conflicts, ordering, errors, and transaction semantics
- Remove the old implementation only when no callers need it
## Migration context and boundaries
- Long-term direction is to centralize Python database access in repositories, then replace their Prisma implementations incrementally with Rust APIs
- Mixed Prisma and Rust implementations are expected during the transition, unrelated existing calls can stay as they are
- Both backends may have separate pools through PgBouncer, size them together and verify pooling-mode compatibility
- PgBouncer does not share transactions between Prisma and SQLx, keep each atomic unit of work on one backend and transaction
- Never retry a failed Rust write through Prisma when it may already have committed
- Prisma owns schema migrations today, eventual SQLx migration ownership and Prisma removal are separate planned steps, not part of a query port

319
.github/scripts/auto_update_price_and_context_window_file.py vendored
View file

@ -1,319 +0,0 @@
import asyncio
import aiohttp
import json
import math
from typing import Any
# Asynchronously fetch data from a given URL
async def fetch_data(url):
try:
# Create an asynchronous session
async with aiohttp.ClientSession() as session:
# Send a GET request to the URL
async with session.get(url) as resp:
# Raise an error if the response status is not OK
resp.raise_for_status()
# Parse the response JSON
resp_json = await resp.json()
print("Fetch the data from URL.")
# Return the 'data' field from the JSON response
return resp_json['data']
except Exception as e:
# Print an error message if fetching data fails
print("Error fetching data from URL:", e)
return None
FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models"
FRIENDLI_PROVIDER = "friendliai"
INHERITABLE_BASE_KEYS = (
"supports_pdf_input",
"supports_assistant_prefill",
"supports_adaptive_thinking",
"supports_output_config",
)
REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max")
def _find_base_model_entry(base_model: str, local_data: dict) -> str | None:
if not base_model:
return None
bm_tail = base_model.split("/")[-1].lower()
if base_model in local_data:
return base_model
for key in local_data:
if key.startswith("sample_spec") or key == "fallback_generalizations":
continue
if key.split("/")[-1].lower() == bm_tail:
return key
return None
def _reasoning_effort_levels(reasoning_options: list) -> list:
offered = {
val
for opt in reasoning_options or []
if opt.get("type") == "effort"
for val in opt.get("values", [])
}
return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered]
def _valid_token_price(value: object) -> bool:
try:
price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except
except (TypeError, ValueError):
return False
return math.isfinite(price) and price >= 0
def _has_valid_token_prices(pricing: dict | None) -> bool:
prices = pricing or {}
return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output"))
def _pricing(pricing: dict) -> dict:
out: dict[str, Any] = {}
if not pricing:
return out
if "input" in pricing:
out["input_cost_per_token"] = float(pricing["input"])
if "output" in pricing:
out["output_cost_per_token"] = float(pricing["output"])
if "input_cache_read" in pricing and pricing["input_cache_read"] is not None:
out["cache_read_input_token_cost"] = float(pricing["input_cache_read"])
return out
def _modality_flags(input_mods: list) -> dict:
mods = input_mods or []
has_image = "image" in mods
return {
"supports_vision": has_image,
"supports_image_input": has_image,
"supports_video_input": "video" in mods,
}
def transform_friendli_data(data: list, local_data: dict) -> dict:
transformed: dict[str, dict] = {}
if not data:
return transformed
for model in data:
# An unpriced row must never wholesale-replace an already priced local entry:
# missing prices cost-calculate as zero, silently zeroing tracked spend
if not _has_valid_token_prices(model.get("pricing")):
continue
model_id = model["id"]
base_model = model.get("base_model") or ""
entry: dict[str, Any] = {
"litellm_provider": FRIENDLI_PROVIDER,
}
base_key = _find_base_model_entry(base_model, local_data)
if base_key:
base_entry = local_data[base_key]
for k in INHERITABLE_BASE_KEYS:
if k in base_entry:
entry[k] = base_entry[k]
ctx = model.get("context_length")
if ctx is not None:
entry["max_input_tokens"] = int(ctx)
max_out = model.get("max_completion_tokens")
if max_out is not None:
entry["max_output_tokens"] = int(max_out)
entry["max_tokens"] = int(max_out)
pricing = _pricing(model.get("pricing", {}))
entry.update(pricing)
entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing
reasoning = model.get("reasoning") is True
entry["supports_reasoning"] = reasoning
if reasoning:
entry["reasoning_effort_levels"] = _reasoning_effort_levels(
model.get("reasoning_options", [])
)
func = model.get("functionality", {})
entry["supports_function_calling"] = func.get("tool_call") is True
entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True
is_struct = func.get("structured_output") is True
entry["supports_response_schema"] = is_struct
entry["supports_native_structured_output"] = is_struct
entry["supports_system_messages"] = func.get("system_messages") is True
entry["supports_tool_choice"] = func.get("tool_choice") is True
entry.update(_modality_flags(model.get("input_modalities", [])))
entry["mode"] = model.get("mode", "chat")
desc = model.get("description")
if desc:
entry["comment"] = desc
dep = model.get("deprecation_date")
if dep:
entry["deprecation_date"] = dep.split("T")[0]
entry["source"] = FRIENDLI_API_URL
transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry
return transformed
# Synchronize local data with remote data
def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()):
# Update existing keys in local_data with values from remote_data
# (replace_keys entries are swapped wholesale so a field the remote catalog
# dropped, e.g. cache pricing, cannot survive as a stale value)
for key in (set(local_data) & set(remote_data)):
if key in replace_keys:
local_data[key] = remote_data[key]
else:
local_data[key].update(remote_data[key])
# Add new keys from remote_data to local_data
for key in (set(remote_data) - set(local_data)):
local_data[key] = remote_data[key]
# Write data to the json file
def write_to_file(file_path, data):
try:
# Open the file in write mode
with open(file_path, "w") as file:
# Dump the data as JSON into the file
json.dump(data, file, indent=4)
print("Values updated successfully.")
except Exception as e:
# Print an error message if writing to file fails
print("Error updating JSON file:", e)
# Update the existing models and add the missing models for OpenRouter
def transform_openrouter_data(data):
transformed = {}
if not data:
return transformed
for row in data:
# Add the fields 'max_tokens' and 'input_cost_per_token'
obj = {
"max_tokens": row["context_length"],
"input_cost_per_token": float(row["pricing"]["prompt"]),
}
# Add 'max_output_tokens' as a field if it is not None
if "top_provider" in row and "max_completion_tokens" in row["top_provider"] and row["top_provider"]["max_completion_tokens"] is not None:
obj['max_output_tokens'] = int(row["top_provider"]["max_completion_tokens"])
# Add the field 'output_cost_per_token'
obj.update({
"output_cost_per_token": float(row["pricing"]["completion"]),
})
# Add field 'input_cost_per_image' if it exists and is non-zero
if "pricing" in row and "image" in row["pricing"] and float(row["pricing"]["image"]) != 0.0:
obj['input_cost_per_image'] = float(row["pricing"]["image"])
# Add the fields 'litellm_provider' and 'mode'
obj.update({
"litellm_provider": "openrouter",
"mode": "chat"
})
# Add the 'supports_vision' field if the modality is 'multimodal'
if row.get('architecture', {}).get('modality') == 'multimodal':
obj['supports_vision'] = True
# Use a composite key to store the transformed object
transformed[f'openrouter/{row["id"]}'] = obj
return transformed
# Update the existing models and add the missing models for Vercel AI Gateway
def transform_vercel_ai_gateway_data(data):
transformed = {}
if not data:
return transformed
for row in data:
# Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync
if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any(
row.get("pricing", {}).get(k) is None for k in ("input", "output")
):
continue
obj = {
"max_tokens": row["context_window"],
"input_cost_per_token": float(row["pricing"]["input"]),
"output_cost_per_token": float(row["pricing"]["output"]),
'max_output_tokens': row['max_tokens'],
'max_input_tokens': row["context_window"],
}
# Handle cache pricing if available
if "pricing" in row:
if "input_cache_read" in row["pricing"] and row["pricing"]["input_cache_read"] is not None:
obj['cache_read_input_token_cost'] = float(f"{float(row['pricing']['input_cache_read']):e}")
if "input_cache_write" in row["pricing"] and row["pricing"]["input_cache_write"] is not None:
obj['cache_creation_input_token_cost'] = float(f"{float(row['pricing']['input_cache_write']):e}")
mode = "embedding" if "embedding" in row["id"].lower() else "chat"
obj.update({"litellm_provider": "vercel_ai_gateway", "mode": mode})
transformed[f'vercel_ai_gateway/{row["id"]}'] = obj
return transformed
# Load local data from a specified file
def load_local_data(file_path):
try:
# Open the file in read mode
with open(file_path, "r") as file:
# Load and return the JSON data
return json.load(file)
except FileNotFoundError:
# Print an error message if the file is not found
print("File not found:", file_path)
return None
except json.JSONDecodeError as e:
# Print an error message if JSON decoding fails
print("Error decoding JSON:", e)
return None
def main():
local_file_path = "model_prices_and_context_window.json" # Path to the local data file
openrouter_url = "https://openrouter.ai/api/v1/models" # URL to fetch OpenRouter data
vercel_ai_gateway_url = "https://ai-gateway.vercel.sh/v1/models" # URL to fetch Vercel AI Gateway data
# Load local data from file
local_data = load_local_data(local_file_path)
# Fetch OpenRouter data
openrouter_data = asyncio.run(fetch_data(openrouter_url))
# Transform the fetched OpenRouter data
openrouter_data = transform_openrouter_data(openrouter_data)
# Fetch Vercel AI Gateway data
vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url))
# Transform the fetched Vercel AI Gateway data
vercel_data = transform_vercel_ai_gateway_data(vercel_data)
friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL))
friendli_data = transform_friendli_data(friendli_data, local_data)
# Combine both datasets
all_remote_data = {**openrouter_data, **vercel_data, **friendli_data}
# If both local and openrouter data are available, synchronize and save
if local_data and all_remote_data:
sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data))
write_to_file(local_file_path, local_data)
else:
print("Failed to fetch model data from either local file or URL.")
# Entry point of the script
if __name__ == "__main__":
main()

36
.github/workflows/auto_update_price_and_context_window.yml vendored
View file

@ -1,36 +0,0 @@
name: Updates model_prices_and_context_window.json and Create Pull Request
on:
schedule:
- cron: "0 0 * * 0" # Run every Sundays at midnight
#- cron: "0 0 * * *" # Run daily at midnight
permissions:
contents: write
pull-requests: write
jobs:
auto_update_price_and_context_window:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
- name: Create Pull Request
run: |
git add model_prices_and_context_window.json
git commit -m "Update model_prices_and_context_window.json file: $(date +'%Y-%m-%d')"
gh pr create --title "Update model_prices_and_context_window.json file" \
--body "Automated update for model_prices_and_context_window.json" \
--head auto-update-price-and-context-window-$(date +'%Y-%m-%d') \
--base main
env:
GH_TOKEN: ${{ secrets.GH_TOKEN }}

65
.github/workflows/sync-together-ai-models.yml vendored
View file

@ -1,65 +0,0 @@
name: Sync Together AI model registry
on:
schedule:
- cron: "30 6 * * *"
workflow_dispatch:
permissions:
contents: write
pull-requests: write
jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
env:
BASE_BRANCH: ${{ github.event.repository.default_branch }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ env.BASE_BRANCH }}
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Look for an already-open sync PR
id: existing
run: |
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
if [ -n "$open_pr" ]; then
echo "An open sync PR already exists on branch $open_pr; skipping this run."
fi
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
- name: Run the sync
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
if git diff --quiet; then
echo "Registry already in sync; no PR needed."
exit 0
fi
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base "$BASE_BRANCH"
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

4
.github/workflows/test-mcp-dependency-resolution.yml vendored
View file

@ -45,7 +45,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@ -132,7 +132,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:

8
.github/workflows/test-merge-smoke.yml vendored
View file

@ -33,7 +33,7 @@ jobs:
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14", "3.15"]
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
steps:
@ -76,6 +76,12 @@ jobs:
- name: Import and CLI checks
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py cli
- name: Routing and UUID compatibility
run: >-
$IN_NETNS .venv/bin/python -m pytest
tests/unit/router_strategy/test_complexity_router.py::TestCustomDimensions
tests/unit/test_uuid_helper.py -q
- name: Proxy startup check
run: $IN_NETNS .venv/bin/python .github/scripts/run_merge_smoke.py proxy-startup --diagnostics-dir "${RUNNER_TEMP}/smoke-diagnostics"

106
.github/workflows/test-rust-db.yml vendored Normal file
View file

@ -0,0 +1,106 @@
name: LiteLLM Rust DB
on:
push:
paths:
- "litellm-rust/crates/db/**"
- "litellm-rust/Cargo.toml"
- "litellm-rust/Cargo.lock"
- "litellm-rust/clippy.toml"
- "litellm-proxy-extras/litellm_proxy_extras/migrations/**"
- "**/schema.prisma"
- ".cargo/**"
- "rust-toolchain.toml"
- "pyproject.toml"
- "uv.lock"
- ".github/actions/setup-uv-with-retries/**"
- ".github/workflows/test-rust-db.yml"
pull_request:
branches:
- main
- "litellm_**"
paths:
- "litellm-rust/crates/db/**"
- "litellm-rust/Cargo.toml"
- "litellm-rust/Cargo.lock"
- "litellm-rust/clippy.toml"
- "litellm-proxy-extras/litellm_proxy_extras/migrations/**"
- "**/schema.prisma"
- ".cargo/**"
- "rust-toolchain.toml"
- "pyproject.toml"
- "uv.lock"
- ".github/actions/setup-uv-with-retries/**"
- ".github/workflows/test-rust-db.yml"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
env:
CARGO_TERM_COLOR: always
jobs:
db-queries:
runs-on: ubuntu-latest
timeout-minutes: 30
services:
postgres:
image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36
env:
POSTGRES_DB: litellm
POSTGRES_PASSWORD: postgres
POSTGRES_USER: postgres
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U postgres -d litellm"
--health-interval 10s
--health-timeout 5s
--health-retries 5
env:
DATABASE_URL: postgresql://postgres:postgres@localhost:5432/litellm
SQLX_OFFLINE: "false"
defaults:
run:
working-directory: litellm-rust
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- run: rustup toolchain install --no-self-update
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
workspaces: litellm-rust
cache-on-failure: true
- name: Install Python dependencies
working-directory: .
run: uv sync --frozen --extra extra-proxy
- name: Deploy Prisma migrations
working-directory: .
run: uv run --no-sync prisma migrate deploy --schema litellm-proxy-extras/litellm_proxy_extras/schema.prisma
- name: Install SQLx CLI
run: cargo install sqlx-cli --version 0.9.0 --locked --no-default-features -F postgres
- name: Check SQLx query metadata
working-directory: litellm-rust/crates/db
run: cargo sqlx prepare --check
- name: Test database queries
run: cargo test -p litellm-db --locked -- --include-ignored

6
ARCHITECTURE.md
View file

@ -332,12 +332,12 @@ Each translation is isolated in its own file, making it easy to test and modify
|--------------|----------|------------------|
| `/v1/chat/completions` | Anthropic | `llms/anthropic/chat/transformation.py` |
| `/v1/chat/completions` | Bedrock Converse | `llms/bedrock/chat/converse_transformation.py` |
| `/v1/chat/completions` | Bedrock Invoke | `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py` |
| `/v1/chat/completions` | Bedrock Invoke | `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py` |
| `/v1/chat/completions` | Gemini | `llms/gemini/chat/transformation.py` |
| `/v1/chat/completions` | Vertex AI | `llms/vertex_ai/gemini/transformation.py` |
| `/v1/chat/completions` | OpenAI | `llms/openai/chat/gpt_transformation.py` |
| `/v1/messages` (passthrough) | Anthropic | `llms/anthropic/pass_through/messages/transformation.py` |
| `/v1/messages` (passthrough) | Bedrock | `llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py` |
| `/v1/messages` (passthrough) | Bedrock | `llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py` |
| `/v1/messages` (passthrough) | Vertex AI | `llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py` |
| Passthrough endpoints | All | `proxy/pass_through_endpoints/llm_provider_handlers/` |
@ -346,7 +346,7 @@ Each translation is isolated in its own file, making it easy to test and modify
If `/v1/messages` → Bedrock Converse prompt caching isn't working but Bedrock Invoke works:
1. **Bedrock Converse translation**: `llms/bedrock/chat/converse_transformation.py`
2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py`
2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py`
3. Compare how each handles `cache_control` in `transform_request()`
### How translations work

4
Dockerfile
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

11
backend/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@ -124,7 +124,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/backend /app/backend
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \

76
cookbook/ai_coding_tool_guides/chatgpt_subscription_passthrough.md Normal file
View file

@ -0,0 +1,76 @@
# Use your ChatGPT subscription through LiteLLM
Sign into ChatGPT in Codex or Pi and send model requests through your existing LiteLLM proxy. The client keeps your login and refreshes it. Each request uses that client's ChatGPT account, while a separate LiteLLM virtual key controls gateway access and usage attribution.
An administrator adds a `chatgpt/` model using the usual model configuration. No ChatGPT login needs to be added to the proxy. This uses the same separate-key convention as provider passthrough, including Claude Code subscription forwarding.
## Add a model to the proxy
Reuse your existing proxy. If you need one, follow the [proxy quickstart](https://docs.litellm.ai/docs/proxy/docker_quick_start). Merge this entry into its model list and restart the proxy:
```yaml
model_list:
- model_name: gpt-6-luna
litellm_params:
model: chatgpt/gpt-6-luna
model_info:
mode: responses
```
Use a model available to the caller's subscription. The virtual key must allow the exposed model name. Leave saved ChatGPT logins and API credentials unset for a client-only deployment. Explicit cross-provider fallbacks can incur separate API charges; do not configure them for subscription-only traffic.
## Configure Codex
Use your existing LiteLLM URL and a virtual key assigned to your gateway user. Export that key as `LITELLM_API_KEY` in the environment that launches Codex. Keep the value out of committed configuration.
Add this provider to `~/.codex/config.toml`, preserving your other settings:
```toml
model = "gpt-6-luna"
model_provider = "litellm_subscription"
[model_providers.litellm_subscription]
name = "LiteLLM subscription"
base_url = "https://litellm.example.com/v1"
wire_api = "responses"
supports_websockets = false
requires_openai_auth = true
env_http_headers = { "x-litellm-api-key" = "LITELLM_API_KEY" }
```
Replace `https://litellm.example.com` with your proxy URL. A local proxy normally uses `http://localhost:4000/v1`. Launch Codex, sign into ChatGPT if needed, and send a prompt with a model available to your subscription. The selected model must match the name exposed by the gateway.
For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway).
## Native compatibility and Pi
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
### Configure Pi
Sign into the built-in `openai-codex` provider in Pi. Merge this override into `~/.pi/agent/models.json` and export your LiteLLM virtual key as `LITELLM_API_KEY` before launching Pi:
```json
{
"providers": {
"openai-codex": {
"baseUrl": "https://litellm.example.com/backend-api/codex",
"headers": {
"x-litellm-api-key": "LITELLM_API_KEY"
}
}
}
}
```
Keep Pi's subscription login as the provider credential. The header value names an environment variable, not a key to paste into the file.
## Verify account usage
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.
This flow does not connect accounts through the dashboard or pool subscription allowances. Use a Codex build that supports ChatGPT authentication with a custom model provider. The integration forwards HTTP and SSE requests; it does not add native WebSocket or other Codex backend routes.

4
docker/Dockerfile.database
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

4
docker/Dockerfile.non_root
View file

@ -1,8 +1,8 @@
# syntax=docker/dockerfile:1.7
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43

11
gateway/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
ARG PGBOUNCER_VERSION=1.26.0
@ -140,7 +140,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer

79
litellm-rust/Cargo.lock generated
View file

@ -876,6 +876,15 @@ dependencies = [
"time",
]
[[package]]
name = "backon"
version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
]
[[package]]
name = "base16ct"
version = "0.2.0"
@ -1248,7 +1257,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
dependencies = [
"bytes",
"futures-core",
"memchr",
"pin-project-lite",
"tokio",
"tokio-util",
]
[[package]]
@ -3642,6 +3655,14 @@ dependencies = [
"strum",
]
[[package]]
name = "litellm-circuit-breaker"
version = "0.1.0"
dependencies = [
"proptest",
"rstest",
]
[[package]]
name = "litellm-config"
version = "0.1.0"
@ -3655,6 +3676,8 @@ dependencies = [
"strum",
"tempfile",
"thiserror 2.0.19",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -3695,8 +3718,15 @@ dependencies = [
name = "litellm-db"
version = "0.1.0"
dependencies = [
"chrono",
"macro_rules_attribute",
"rstest",
"serde",
"serde_json",
"sqlx",
"strum",
"thiserror 2.0.19",
"tokio",
]
[[package]]
@ -3736,6 +3766,7 @@ dependencies = [
"futures-util",
"http-body-util",
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-config",
"litellm-gateway-auth",
"litellm-gateway-inference",
@ -3744,6 +3775,8 @@ dependencies = [
"litellm-http",
"litellm-inference",
"litellm-llms",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secrets",
"litellm-tracing",
"rstest",
@ -4275,6 +4308,28 @@ dependencies = [
"thiserror 2.0.19",
]
[[package]]
name = "litellm-redis-client"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-circuit-breaker",
"litellm-redis-testing",
"redis",
"rstest",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-redis-testing"
version = "0.1.0"
dependencies = [
"testcontainers-modules",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-router"
version = "0.1.0"
@ -5750,22 +5805,34 @@ dependencies = [
[[package]]
name = "redis"
version = "1.7.0"
version = "1.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
checksum = "7acd809e80afa754c63f586955699d8990efe1c4b6613183ee5f027a11fbc8b5"
dependencies = [
"arc-swap",
"arcstr",
"async-lock",
"backon",
"bytes",
"cfg-if",
"combine",
"crc16",
"futures-channel",
"futures-util",
"itoa",
"log",
"num-bigint 0.5.1",
"percent-encoding",
"pin-project-lite",
"rand 0.10.2",
"rustls 0.23.42",
"rustls-native-certs",
"ryu",
"sha1_smol",
"socket2 0.6.5",
"tokio",
"tokio-rustls 0.26.4",
"tokio-util",
"url",
"xxhash-rust",
]
@ -8202,18 +8269,18 @@ dependencies = [
[[package]]
name = "webpki-root-certs"
version = "1.0.10"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb"
checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "webpki-roots"
version = "1.0.10"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9"
checksum = "f5f280bbb2b287bc8e7de7e3852945a0ff50d1bbcb5ed42b6c7a568c546df171"
dependencies = [
"rustls-pki-types",
]

7
litellm-rust/Cargo.toml
View file

@ -69,6 +69,9 @@ litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface"
litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
litellm-host-python = { path = "crates/host-python" }
litellm-python-compat = { path = "crates/python-compat" }
litellm-circuit-breaker = { path = "crates/circuit-breaker" }
litellm-redis-client = { path = "crates/redis-client" }
litellm-redis-testing = { path = "crates/redis-testing" }
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
tracing = "0.1"
@ -88,6 +91,9 @@ rand = "0.8"
macro_rules_attribute = "0.2.3"
schemars = "1"
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
redis = "1.8.0"
redis-test = "1.0.4"
testcontainers-modules = "0.15.0"
qdrant-client = { version = "1.19.0", default-features = false }
uuid = { version = "1", features = ["v4"] }
rstest = "0.26.1"
@ -104,6 +110,7 @@ serde_with = { version = "=3.21.0", default-features = false, features = ["std",
sha2 = "0.10"
syn = { version = "2", default-features = false }
sqlx = { version = "0.9.0", default-features = false, features = ["json", "macros", "postgres", "runtime-tokio", "chrono", "tls-rustls-ring-native-roots"] }
chrono = { version = "0.4.45", default-features = false, features = ["std"] }
proc-macro2 = "1"
quote = "1"
subtle = "2"

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

@ -8,12 +8,12 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
redis = { workspace = true, features = ["tls-rustls"] }
sha2.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

4
litellm-rust/crates/cache-redis/Cargo.toml
View file

@ -7,12 +7,12 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] }
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true

4
litellm-rust/crates/cache-response/Cargo.toml
View file

@ -17,8 +17,8 @@ litellm-cache-gcs.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
redis = "1.7.0"
redis-test = "1.0.4"
redis.workspace = true
redis-test.workspace = true
rstest.workspace = true
tokio.workspace = true
wiremock.workspace = true

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

@ -8,13 +8,13 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
redis = { workspace = true, features = ["tls-rustls"] }
sha2.workspace = true
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

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

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

12
litellm-rust/crates/circuit-breaker/Cargo.toml Normal file
View file

@ -0,0 +1,12 @@
[package]
name = "litellm-circuit-breaker"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
[dev-dependencies]
proptest.workspace = true
rstest.workspace = true

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

@ -0,0 +1,444 @@
use std::{
sync::Arc,
time::{Duration, Instant},
};
use crate::{Machine, Observer, Shared};
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Config {
pub failure_threshold: u32,
pub recovery_timeout: Duration,
pub timeout_min_duration: Duration,
}
impl Default for Config {
fn default() -> Self {
Self {
failure_threshold: 5,
recovery_timeout: Duration::from_secs(60),
timeout_min_duration: Duration::from_secs(5),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Failure {
Hard,
Timeout,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum State {
Closed {
hard: u32,
timeouts: u32,
streak_started: Option<Instant>,
},
Open {
since: Instant,
},
HalfOpen,
}
impl State {
const CLOSED: Self = Self::Closed {
hard: 0,
timeouts: 0,
streak_started: None,
};
}
#[derive(Clone, Copy, Debug)]
pub enum Event {
Admit(Instant),
Succeeded,
Failed(Failure, Instant),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Output {
Admitted,
Probe,
Rejected,
Recorded,
}
impl Machine for Config {
type State = State;
type Event = Event;
type Output = Output;
fn step(&self, state: &State, event: Event) -> (State, Output) {
match (state, event) {
(State::Closed { .. }, Event::Admit(_)) => (state.clone(), Output::Admitted),
(State::Open { since }, Event::Admit(now))
if now.duration_since(*since) > self.recovery_timeout =>
{
(State::HalfOpen, Output::Probe)
}
(State::Open { .. } | State::HalfOpen, Event::Admit(_)) => {
(state.clone(), Output::Rejected)
}
(State::Closed { .. } | State::HalfOpen, Event::Succeeded) => {
(State::CLOSED, Output::Recorded)
}
(State::Open { .. }, Event::Succeeded) => (state.clone(), Output::Recorded),
(State::Closed { hard, .. }, Event::Failed(Failure::Hard, now)) => {
let hard = hard + 1;
let next = match hard >= self.failure_threshold {
true => State::Open { since: now },
false => State::Closed {
hard,
timeouts: 0,
streak_started: None,
},
};
(next, Output::Recorded)
}
(
State::Closed {
hard,
timeouts,
streak_started,
},
Event::Failed(Failure::Timeout, now),
) => {
let timeouts = timeouts + 1;
let started = streak_started.unwrap_or(now);
let opens = timeouts >= self.failure_threshold
&& now.duration_since(started) >= self.timeout_min_duration;
let next = match opens {
true => State::Open { since: now },
false => State::Closed {
hard: *hard,
timeouts,
streak_started: Some(started),
},
};
(next, Output::Recorded)
}
(State::Open { .. } | State::HalfOpen, Event::Failed(_, now)) => {
(State::Open { since: now }, Output::Recorded)
}
}
}
}
#[derive(Clone)]
pub struct Breaker(Arc<Shared<Config>>);
#[derive(Debug, PartialEq, Eq)]
pub struct Rejected;
pub struct Admission {
breaker: Breaker,
generation: u64,
probe: bool,
settled: bool,
}
impl Breaker {
pub fn new(config: Config, observer: Observer<State>) -> Self {
Self(Arc::new(Shared::new(config, State::CLOSED, observer)))
}
pub fn admit(&self, now: Instant) -> Result<Admission, Rejected> {
let (output, generation) = self.0.fire(Event::Admit(now));
let probe = match output {
Output::Admitted => false,
Output::Probe => true,
Output::Rejected | Output::Recorded => return Err(Rejected),
};
Ok(Admission {
breaker: self.clone(),
generation,
probe,
settled: false,
})
}
pub fn failed(&self, failure: Failure, now: Instant) {
self.0.fire(Event::Failed(failure, now));
}
pub fn state(&self) -> State {
self.0.state()
}
}
impl Admission {
pub fn probe(&self) -> bool {
self.probe
}
pub fn succeeded(mut self) {
self.settled = true;
self.breaker.0.fire_if(self.generation, Event::Succeeded);
}
pub fn failed(mut self, failure: Failure, now: Instant) {
self.settled = true;
self.breaker.failed(failure, now);
}
}
impl Drop for Admission {
fn drop(&mut self) {
if self.probe && !self.settled {
self.breaker.0.fire_if(
self.generation,
Event::Failed(Failure::Hard, Instant::now()),
);
}
}
}
#[cfg(test)]
mod tests {
use std::{
collections::HashMap,
sync::{Arc, Mutex},
time::{Duration, Instant},
};
use proptest::prelude::*;
use rstest::rstest;
use super::{Admission, Breaker, Config, Failure, State};
use crate::Transition;
#[derive(Clone, Copy, Debug)]
enum Step {
Fail(f64, Failure),
Admit(f64, bool),
Succeed(usize),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Phase {
Closed,
Open,
HalfOpen,
}
use Phase::{Closed as C, HalfOpen as H, Open as O};
#[rstest]
fn dropping_an_old_probe_does_not_undo_a_newer_recovery() {
let breaker = Breaker::new(
Config {
failure_threshold: 1,
recovery_timeout: Duration::from_secs(1),
..Config::default()
},
Box::new(|_| {}),
);
let now = Instant::now();
let pending = breaker.admit(now).unwrap();
breaker.failed(Failure::Hard, now);
let old_probe = breaker.admit(now + Duration::from_secs(2)).unwrap();
assert!(old_probe.probe());
pending.failed(Failure::Hard, now + Duration::from_secs(2));
let new_probe = breaker.admit(now + Duration::from_secs(4)).unwrap();
assert!(new_probe.probe());
new_probe.succeeded();
assert_eq!(breaker.state(), State::CLOSED);
drop(old_probe);
assert_eq!(breaker.state(), State::CLOSED);
assert!(!breaker.admit(now + Duration::from_secs(4)).unwrap().probe());
}
fn hard(at: f64) -> Step {
Step::Fail(at, Failure::Hard)
}
fn timeout(at: f64) -> Step {
Step::Fail(at, Failure::Timeout)
}
fn five_hard_at_zero() -> Vec<Step> {
vec![hard(0.0); 5]
}
fn phase(state: &State) -> Phase {
match state {
State::Closed { .. } => Phase::Closed,
State::Open { .. } => Phase::Open,
State::HalfOpen => Phase::HalfOpen,
}
}
fn quiet() -> Breaker {
Breaker::new(Config::default(), Box::new(|_| {}))
}
#[rstest]
#[case::five_hard_failures_open(vec![hard(0.0), hard(1.0), hard(2.0), hard(3.0), hard(4.0)], vec![C, C, C, C, O])]
#[case::a_success_resets_the_hard_failure_count(
vec![hard(0.0), hard(1.0), hard(2.0), hard(3.0), Step::Admit(4.0, true), Step::Succeed(4), hard(5.0), hard(6.0), hard(7.0), hard(8.0)],
vec![C; 10]
)]
#[case::a_timeout_streak_opens_only_after_the_min_duration(
vec![timeout(0.0), timeout(1.0), timeout(2.0), timeout(3.0), timeout(4.9), timeout(5.0)],
vec![C, C, C, C, C, O]
)]
#[case::a_hard_failure_restarts_the_timeout_streak(
vec![timeout(0.0), timeout(1.0), timeout(2.0), timeout(3.0), hard(4.0), timeout(5.0), timeout(6.0), timeout(7.0), timeout(8.0), timeout(9.9), timeout(10.0)],
vec![C, C, C, C, C, C, C, C, C, C, O]
)]
#[case::the_probe_waits_until_strictly_past_the_recovery_timeout(
[five_hard_at_zero(), vec![Step::Admit(60.0, false), Step::Admit(60.5, true)]].concat(),
vec![C, C, C, C, O, O, H]
)]
#[case::a_probe_success_closes(
[five_hard_at_zero(), vec![Step::Admit(61.0, true), Step::Succeed(5)]].concat(),
vec![C, C, C, C, O, H, C]
)]
#[case::a_probe_failure_reopens_with_a_fresh_timer(
[five_hard_at_zero(), vec![Step::Admit(61.0, true), hard(62.0), Step::Admit(122.0, false), Step::Admit(122.5, true)]].concat(),
vec![C, C, C, C, O, H, O, O, H]
)]
#[case::only_one_probe_at_a_time(
[five_hard_at_zero(), vec![Step::Admit(61.0, true), Step::Admit(61.5, false)]].concat(),
vec![C, C, C, C, O, H, H]
)]
#[case::a_late_failure_while_open_pushes_back_the_probe(
[five_hard_at_zero(), vec![hard(30.0), Step::Admit(61.0, false), Step::Admit(90.5, true)]].concat(),
vec![C, C, C, C, O, O, O, H]
)]
#[case::a_success_from_before_the_breaker_moved_is_ignored(
vec![Step::Admit(0.0, true), hard(1.0), hard(1.0), hard(1.0), hard(1.0), hard(1.0), Step::Admit(62.0, true), Step::Succeed(0), Step::Succeed(6)],
vec![C, C, C, C, C, O, H, H, C]
)]
fn matches_the_python_redis_circuit_breaker(
#[case] steps: Vec<Step>,
#[case] expected: Vec<Phase>,
) {
let breaker = quiet();
let t0 = Instant::now();
let mut admissions: HashMap<usize, Admission> = HashMap::new();
let phases: Vec<Phase> = steps
.iter()
.enumerate()
.map(|(index, step)| {
match *step {
Step::Fail(at, failure) => {
breaker.failed(failure, t0 + Duration::from_secs_f64(at))
}
Step::Admit(at, admitted) => {
let admission = breaker.admit(t0 + Duration::from_secs_f64(at));
assert_eq!(admission.is_ok(), admitted, "step {index}");
admissions.extend(admission.ok().map(|admission| (index, admission)));
}
Step::Succeed(admitted_at) => {
admissions.remove(&admitted_at).unwrap().succeeded()
}
}
phase(&breaker.state())
})
.collect();
assert_eq!(phases, expected);
}
#[rstest]
fn a_dropped_probe_reopens_instead_of_wedging_half_open() {
let breaker = quiet();
let t0 = Instant::now();
(0..5).for_each(|_| breaker.failed(Failure::Hard, t0));
let probe = breaker.admit(t0 + Duration::from_secs(61)).unwrap();
assert!(probe.probe());
drop(probe);
assert!(matches!(breaker.state(), State::Open { .. }));
}
#[rstest]
fn a_dropped_regular_call_records_nothing() {
let breaker = quiet();
let t0 = Instant::now();
breaker.failed(Failure::Hard, t0);
drop(breaker.admit(t0).unwrap());
assert_eq!(
breaker.state(),
State::Closed {
hard: 1,
timeouts: 0,
streak_started: None
}
);
}
#[derive(Clone, Copy, Debug)]
enum Action {
Admit,
SucceedOldest,
FailOldest(Failure),
Fail(Failure),
}
fn action() -> impl Strategy<Value = (u64, Action)> {
let failure = prop_oneof![Just(Failure::Hard), Just(Failure::Timeout)];
let action = prop_oneof![
Just(Action::Admit),
Just(Action::SucceedOldest),
failure.clone().prop_map(Action::FailOldest),
failure.prop_map(Action::Fail),
];
(0u64..30_000, action)
}
proptest! {
#[test]
fn open_never_closes_directly_and_each_half_open_visit_has_one_probe(
actions in prop::collection::vec(action(), 1..200)
) {
let transitions: Arc<Mutex<Vec<Transition<State>>>> = Arc::default();
let seen = Arc::clone(&transitions);
let config = Config {
failure_threshold: 3,
recovery_timeout: Duration::from_secs(10),
timeout_min_duration: Duration::from_secs(2),
};
let breaker = Breaker::new(config, Box::new(move |t| seen.lock().unwrap().push(t.clone())));
let t0 = Instant::now();
let mut now = t0;
let mut pending: Vec<Admission> = Vec::new();
let mut probes_per_visit: Vec<u32> = Vec::new();
for (advance_ms, action) in actions {
now += Duration::from_millis(advance_ms);
match action {
Action::Admit => {
if let Ok(admission) = breaker.admit(now) {
if admission.probe() {
probes_per_visit.push(1);
}
pending.push(admission);
}
}
Action::SucceedOldest if !pending.is_empty() => pending.remove(0).succeeded(),
Action::FailOldest(failure) if !pending.is_empty() => pending.remove(0).failed(failure, now),
Action::SucceedOldest | Action::FailOldest(_) => {}
Action::Fail(failure) => breaker.failed(failure, now),
}
}
let transitions = transitions.lock().unwrap();
let open_to_closed = transitions
.iter()
.filter(|t| matches!((&t.from, &t.to), (State::Open { .. }, State::Closed { .. })))
.count();
let half_open_visits = transitions.iter().filter(|t| t.to == State::HalfOpen).count();
prop_assert_eq!(open_to_closed, 0);
prop_assert_eq!(probes_per_visit.len(), half_open_visits);
}
}
}

174
litellm-rust/crates/circuit-breaker/src/fsm.rs Normal file
View file

@ -0,0 +1,174 @@
use std::{
fmt::Debug,
mem,
sync::{Mutex, PoisonError},
};
pub trait Machine {
type State: Clone + Debug;
type Event;
type Output;
fn step(&self, state: &Self::State, event: Self::Event) -> (Self::State, Self::Output);
}
#[derive(Clone, Debug)]
pub struct Transition<S> {
pub from: S,
pub to: S,
pub generation: u64,
}
pub type Observer<S> = Box<dyn Fn(&Transition<S>) + Send + Sync>;
pub struct Shared<M: Machine> {
machine: M,
current: Mutex<(M::State, u64)>,
observer: Observer<M::State>,
}
impl<M: Machine> Shared<M> {
pub fn new(machine: M, initial: M::State, observer: Observer<M::State>) -> Self {
Self {
machine,
current: Mutex::new((initial, 0)),
observer,
}
}
pub fn fire(&self, event: M::Event) -> (M::Output, u64) {
let mut current = self.current.lock().unwrap_or_else(PoisonError::into_inner);
self.apply(&mut current, event)
}
pub fn fire_if(&self, generation: u64, event: M::Event) -> Option<M::Output> {
let mut current = self.current.lock().unwrap_or_else(PoisonError::into_inner);
(current.1 == generation).then(|| self.apply(&mut current, event).0)
}
pub fn state(&self) -> M::State {
self.current
.lock()
.unwrap_or_else(PoisonError::into_inner)
.0
.clone()
}
fn apply(&self, current: &mut (M::State, u64), event: M::Event) -> (M::Output, u64) {
let (to, output) = self.machine.step(&current.0, event);
if mem::discriminant(&to) != mem::discriminant(&current.0) {
let from = mem::replace(&mut current.0, to.clone());
current.1 += 1;
(self.observer)(&Transition {
from,
to,
generation: current.1,
});
} else {
current.0 = to;
}
(output, current.1)
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use rstest::rstest;
use super::{Machine, Shared, Transition};
struct Wrapping(u32);
#[derive(Clone, Debug, PartialEq, Eq)]
enum Counter {
Low(u32),
High(u32),
}
enum Tick {
Up,
Stay,
}
impl Machine for Wrapping {
type State = Counter;
type Event = Tick;
type Output = ();
fn step(&self, state: &Counter, event: Tick) -> (Counter, ()) {
let value = match (state, event) {
(Counter::Low(_) | Counter::High(_), Tick::Stay) => return (state.clone(), ()),
(Counter::Low(value) | Counter::High(value), Tick::Up) => (value + 1) % self.0,
};
match value < self.0 / 2 {
true => (Counter::Low(value), ()),
false => (Counter::High(value), ()),
}
}
}
type Seen = Arc<Mutex<Vec<Transition<Counter>>>>;
fn recorded(wrap: u32) -> (Shared<Wrapping>, Seen) {
let transitions: Seen = Arc::default();
let seen = Arc::clone(&transitions);
let shared = Shared::new(
Wrapping(wrap),
Counter::Low(0),
Box::new(move |transition| seen.lock().unwrap().push(transition.clone())),
);
(shared, transitions)
}
#[rstest]
fn concurrent_events_are_all_applied_and_transitions_are_observed_in_generation_order() {
let (shared, transitions) = recorded(4);
std::thread::scope(|scope| {
for _ in 0..64 {
scope.spawn(|| {
for _ in 0..100 {
shared.fire(Tick::Up);
}
});
}
});
assert_eq!(shared.state(), Counter::Low(6400 % 4));
let generations: Vec<u64> = transitions
.lock()
.unwrap()
.iter()
.map(|t| t.generation)
.collect();
assert_eq!(generations, (1..=3200).collect::<Vec<u64>>());
}
#[rstest]
fn a_step_that_keeps_the_variant_neither_counts_nor_notifies() {
let (shared, transitions) = recorded(4);
let (_, after_inner_change) = shared.fire(Tick::Up);
let (_, after_self_loop) = shared.fire(Tick::Stay);
assert_eq!(shared.state(), Counter::Low(1));
assert_eq!((after_inner_change, after_self_loop), (0, 0));
assert!(transitions.lock().unwrap().is_empty());
}
#[rstest]
fn fire_if_skips_an_event_from_an_older_generation() {
let (shared, _) = recorded(4);
shared.fire(Tick::Up);
let (_, generation) = shared.fire(Tick::Up);
let stale = shared.fire_if(generation - 1, Tick::Up);
let current = shared.fire_if(generation, Tick::Up);
assert!(stale.is_none());
assert!(current.is_some());
assert_eq!(shared.state(), Counter::High(3));
}
}

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

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

4
litellm-rust/crates/config/Cargo.toml
View file

@ -9,11 +9,13 @@ repository.workspace = true
litellm-auth-types.workspace = true
litellm-router-types.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_yaml_ng = "0.10.0"
strum.workspace = true
thiserror.workspace = true
tracing.workspace = true
[dev-dependencies]
rstest.workspace = true
serde_json.workspace = true
tempfile.workspace = true
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }

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

@ -4,4 +4,11 @@ pub enum Error {
Read(#[from] std::io::Error),
#[error("invalid YAML config")]
Parse(#[from] serde_yaml_ng::Error),
#[error("Redis host does not yet support configured setting {0}")]
RedisSetting(String),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
}

5
litellm-rust/crates/config/src/lib.rs
View file

@ -2,6 +2,7 @@ mod error;
mod includes;
mod mcp;
mod model;
mod redis;
mod settings;
mod value;
@ -13,6 +14,10 @@ pub use error::Error;
pub use litellm_router_types::{LitellmParams, Spelled};
pub use mcp::{McpAuth, McpServer, McpTransport};
pub use model::Model;
pub use redis::{
CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings,
ResolvedRedis, environment_has_target,
};
pub use settings::{
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
TracingStoreSettings,

667
litellm-rust/crates/config/src/redis.rs Normal file
View file

@ -0,0 +1,667 @@
use std::{collections::BTreeMap, fmt, time::Duration};
use litellm_auth_types::SecretValue;
use litellm_router_types::Spelled;
use serde::{Deserialize, Deserializer, de::Error as _};
use crate::{AdditionalFields, Error, Object, Value};
pub type Lookup<'a> = &'a dyn Fn(&str) -> Option<String>;
const SUPPORTED: &[&str] = &[
"url",
"host",
"port",
"username",
"password",
"ssl",
"db",
"socket_timeout",
"startup_nodes",
"namespace",
];
const IGNORED: &[&str] = &[
"max_connections",
"single_connection_client",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
"decode_responses",
"encoding",
"encoding_errors",
"charset",
"errors",
"client_name",
"lib_name",
"lib_version",
"sentinel_password",
"gcp_ssl_ca_certs",
"azure_client_id",
"azure_tenant_id",
"azure_client_secret",
"aws_iam_user_name",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
];
const WARN_IGNORED: &[&str] = &[
"max_connections",
"health_check_interval",
"socket_keepalive",
"socket_keepalive_options",
];
const REDIS_PY_KWARGS: &[&str] = &[
"aws_iam_auth",
"aws_iam_cache_name",
"aws_iam_region",
"aws_iam_serverless",
"aws_iam_user_name",
"azure_client_id",
"azure_client_secret",
"azure_redis_ad_token",
"azure_tenant_id",
"cache",
"cache_config",
"charset",
"client_name",
"db",
"decode_responses",
"encoding",
"encoding_errors",
"errors",
"event_dispatcher",
"gcp_service_account",
"gcp_ssl_ca_certs",
"health_check_interval",
"host",
"lib_name",
"lib_version",
"max_connections",
"password",
"port",
"protocol",
"redis_connect_func",
"retry_on_error",
"retry_on_timeout",
"single_connection_client",
"socket_connect_timeout",
"socket_keepalive",
"socket_keepalive_options",
"socket_timeout",
"ssl",
"ssl_ca_certs",
"ssl_ca_data",
"ssl_ca_path",
"ssl_cert_reqs",
"ssl_certfile",
"ssl_check_hostname",
"ssl_ciphers",
"ssl_keyfile",
"ssl_min_version",
"ssl_ocsp_context",
"ssl_ocsp_expected_cert",
"ssl_password",
"ssl_validate_ocsp",
"ssl_validate_ocsp_stapled",
"unix_socket_path",
"url",
"username",
];
const UNSUPPORTED_EXTRA: &[&str] = &[
"credential_provider",
"sentinel_nodes",
"retry",
"connection_pool",
];
const CACHE_ONLY: &[&str] = &[
"type",
"mode",
"ttl",
"default_in_memory_ttl",
"default_in_redis_ttl",
"supported_call_types",
"similarity_threshold",
"max_messages",
"_backend",
];
#[derive(Clone, Default)]
pub struct RedisSettings {
pub url: Option<SecretValue>,
pub host: Option<Spelled<String>>,
pub port: Option<Spelled<u16>>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub ssl: Option<Spelled<bool>>,
pub db: Option<Spelled<i64>>,
pub socket_timeout: Option<Spelled<f64>>,
pub startup_nodes: Option<Value>,
pub namespace: Option<Spelled<String>>,
pub additional_fields: AdditionalFields,
unparsed: DeferredInputs,
}
impl<'de> Deserialize<'de> for RedisSettings {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let fields = BTreeMap::<String, Value>::deserialize(deserializer)?;
let scalar = |key: &str| fields.get(key).and_then(as_text);
Ok(Self {
url: scalar("url").map(SecretValue::new),
host: scalar("host").map(Spelled::Value),
port: spelled(scalar("port")),
username: scalar("username").map(SecretValue::new),
password: scalar("password").map(SecretValue::new),
ssl: spelled(scalar("ssl")),
db: spelled(scalar("db")),
socket_timeout: spelled(scalar("socket_timeout")),
startup_nodes: fields
.get("startup_nodes")
.filter(|value| !value.is_null())
.cloned(),
namespace: scalar("namespace").map(Spelled::Value),
additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)),
unparsed: deferred_inputs(&fields),
})
}
}
impl fmt::Debug for RedisSettings {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("RedisSettings")
.field("url", &self.url)
.field("host", &self.host)
.field("port", &self.port)
.field("username", &self.username)
.field("password", &self.password)
.field("ssl", &self.ssl)
.field("db", &self.db)
.field("socket_timeout", &self.socket_timeout)
.field(
"startup_nodes",
&self.startup_nodes.as_ref().map(|_| "<configured>"),
)
.field("namespace", &self.namespace)
.field("additional_fields", &self.additional_fields.keys())
.field("unparsed", &self.unparsed.keys())
.finish()
}
}
#[derive(Clone, Debug)]
pub enum CacheSettings {
Redis(Box<RedisSettings>),
Other(Object),
}
impl<'de> Deserialize<'de> for CacheSettings {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let fields = BTreeMap::<String, Value>::deserialize(deserializer)?;
let is_redis = matches!(fields.get("type"), Some(Value::String(kind)) if kind == "redis");
if !is_redis {
return Ok(Self::Other(Object::from_fields(fields)));
}
let mapping = fields
.into_iter()
.map(|(key, value)| (Value::String(key), value))
.collect();
serde_yaml_ng::from_value(Value::Mapping(mapping))
.map(|settings| Self::Redis(Box::new(settings)))
.map_err(D::Error::custom)
}
}
#[derive(Clone, PartialEq)]
pub struct ResolvedRedis {
pub endpoint: RedisEndpoint,
pub namespace: Option<String>,
pub timeout: Option<Duration>,
}
impl fmt::Debug for ResolvedRedis {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ResolvedRedis")
.field("endpoint", &self.endpoint)
.field("namespace", &self.namespace)
.field("timeout", &self.timeout)
.finish()
}
}
#[derive(Clone, PartialEq)]
pub enum RedisEndpoint {
Url(SecretValue),
Host(RedisHost),
Cluster(RedisCluster),
}
impl fmt::Debug for RedisEndpoint {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Url(_) => formatter.debug_tuple("Url").field(&"<redacted>").finish(),
Self::Host(host) => host.fmt(formatter),
Self::Cluster(cluster) => cluster.fmt(formatter),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisHost {
pub host: String,
pub port: u16,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
pub db: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisCluster {
pub nodes: Vec<RedisNode>,
pub username: Option<SecretValue>,
pub password: Option<SecretValue>,
pub tls: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone)]
struct Entry {
value: Value,
source: String,
}
type Kwargs = BTreeMap<String, Entry>;
impl RedisSettings {
pub fn resolve(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, false)
}
pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result<ResolvedRedis, Error> {
self.resolve_with(lookup, true)
}
fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result<ResolvedRedis, Error> {
let block = self.kwargs(cache);
let block: Kwargs = block
.into_iter()
.filter_map(|(key, entry)| {
let value = environment_reference(&entry.value, lookup)?;
Some((key, Entry { value, ..entry }))
})
.collect();
let environment = environment_kwargs(lookup);
let names_a_target = ["host", "startup_nodes", "sentinel_nodes"]
.iter()
.any(|key| block.contains_key(*key));
let drop_environment_url = names_a_target && !block.contains_key("url");
let merged: Kwargs = environment
.into_iter()
.filter(|(key, _)| !(key == "url" && drop_environment_url))
.chain(block)
.collect();
reject_unsupported(&merged, lookup)?;
warn_ignored(&merged);
let startup_nodes = match merged.get("startup_nodes") {
Some(entry) => Some(cluster_nodes(&entry.value)?),
None => lookup("REDIS_CLUSTER_NODES")
.map(|nodes| cluster_nodes(&Value::String(nodes)))
.transpose()?,
};
let username = secret_value(&merged, "username");
let password = secret_value(&merged, "password");
let tls = bool_value(&merged, "ssl");
let endpoint = match (startup_nodes, text_value(&merged, "url")) {
(Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster {
nodes,
username,
password,
tls,
}),
(_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)),
_ => RedisEndpoint::Host(RedisHost {
host: text_value(&merged, "host").ok_or(Error::RedisInvalid {
setting: "host",
reason: "either host or url must be specified for redis",
})?,
port: integer(&merged, "port")?.unwrap_or(6379),
username,
password,
tls,
db: integer(&merged, "db")?.unwrap_or(0),
}),
};
let timeout = merged
.get("socket_timeout")
.and_then(|entry| as_text(&entry.value))
.map(|seconds| {
seconds
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: "socket_timeout",
reason: "must be a non-negative number of seconds",
})
})
.transpose()?;
Ok(ResolvedRedis {
endpoint,
namespace: text_value(&merged, "namespace"),
timeout,
})
}
fn kwargs(&self, cache: bool) -> Kwargs {
let known = [
(
"url",
self.url.as_ref().map(|value| value.expose().to_owned()),
),
("host", spelled_text(self.host.as_ref())),
("port", spelled_text(self.port.as_ref())),
(
"username",
self.username
.as_ref()
.map(|value| value.expose().to_owned()),
),
(
"password",
self.password
.as_ref()
.map(|value| value.expose().to_owned()),
),
("ssl", spelled_text(self.ssl.as_ref())),
("db", spelled_text(self.db.as_ref())),
("socket_timeout", spelled_text(self.socket_timeout.as_ref())),
("namespace", spelled_text(self.namespace.as_ref())),
];
let known = known.into_iter().filter_map(|(key, value)| {
value.map(|value| {
(
key.to_owned(),
Entry {
value: Value::String(value),
source: key.to_owned(),
},
)
})
});
let startup_nodes = self.startup_nodes.clone().map(|value| {
(
"startup_nodes".to_owned(),
Entry {
value,
source: "startup_nodes".to_owned(),
},
)
});
let extras = self.additional_fields.iter().filter_map(|(key, value)| {
if cache && CACHE_ONLY.contains(&key.as_str()) {
return None;
}
let key = match (cache, key.as_str()) {
(true, "redis_startup_nodes") => "startup_nodes".to_owned(),
_ => key.clone(),
};
Some((
key.clone(),
Entry {
value: value.clone(),
source: key,
},
))
});
let unparsed = self.unparsed.iter().map(|(key, value)| {
(
key.clone(),
Entry {
value: value.0.clone(),
source: key.clone(),
},
)
});
extras
.chain(unparsed)
.chain(known)
.chain(startup_nodes)
.collect()
}
}
pub fn environment_has_target(lookup: Lookup<'_>) -> bool {
[
"REDIS_HOST",
"REDIS_URL",
"REDIS_CLUSTER_NODES",
"REDIS_SENTINEL_NODES",
]
.iter()
.any(|name| lookup(name).is_some())
}
fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> {
let active_auth = [
merged
.get("gcp_service_account")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")),
text_value(merged, "azure_redis_ad_token")
.or_else(|| lookup("REDIS_AZURE_AD_TOKEN"))
.filter(|value| value.eq_ignore_ascii_case("true"))
.map(|_| {
merged
.get("azure_redis_ad_token")
.map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str())
}),
bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()),
merged
.get("sentinel_nodes")
.map(|entry| entry.source.as_str())
.or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")),
];
if let Some(setting) = active_auth.into_iter().flatten().next() {
return Err(Error::RedisSetting(setting.to_owned()));
}
match merged
.iter()
.find(|(key, _)| is_unsupported(key))
.map(|(_, entry)| entry.source.clone())
{
Some(setting) => Err(Error::RedisSetting(setting)),
None => Ok(()),
}
}
fn warn_ignored(merged: &Kwargs) {
for setting in WARN_IGNORED {
if merged.contains_key(*setting) {
tracing::warn!(setting, "Redis setting is accepted but ignored");
}
}
}
fn is_unsupported(key: &str) -> bool {
let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key);
let handled = SUPPORTED.contains(&key)
|| IGNORED.contains(&key)
|| matches!(
key,
"gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth"
);
known && !handled
}
fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs {
REDIS_PY_KWARGS
.iter()
.filter_map(|key| {
let source = format!("REDIS_{}", key.to_ascii_uppercase());
let value = lookup(&source)?;
Some((
(*key).to_owned(),
Entry {
value: Value::String(value),
source,
},
))
})
.collect()
}
fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option<Value> {
match value {
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(name) => lookup(name).map(Value::String),
None => Some(value.clone()),
},
Value::Null => None,
_ => Some(value.clone()),
}
}
fn as_text(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
}
}
fn text_value(kwargs: &Kwargs, key: &str) -> Option<String> {
kwargs.get(key).and_then(|entry| as_text(&entry.value))
}
fn secret_value(kwargs: &Kwargs, key: &str) -> Option<SecretValue> {
text_value(kwargs, key).map(SecretValue::new)
}
fn bool_value(kwargs: &Kwargs, key: &str) -> bool {
match kwargs.get(key).map(|entry| &entry.value) {
Some(Value::Bool(flag)) => *flag,
Some(value) => as_text(value)
.is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")),
None => false,
}
}
fn integer<T: std::str::FromStr>(kwargs: &Kwargs, key: &'static str) -> Result<Option<T>, Error> {
kwargs
.get(key)
.and_then(|entry| as_text(&entry.value))
.map(|text| {
text.parse().map_err(|_| Error::RedisInvalid {
setting: key,
reason: "must be an integer",
})
})
.transpose()
}
const INVALID_NODES: Error = Error::RedisInvalid {
setting: "startup_nodes",
reason: "must be a list of {host, port} objects",
};
fn cluster_nodes(value: &Value) -> Result<Vec<RedisNode>, Error> {
match value {
Value::String(json) => {
let nodes: Vec<serde_json::Map<String, serde_json::Value>> =
serde_json::from_str(json).map_err(|_| INVALID_NODES)?;
nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(json_text),
node.get("port").and_then(json_text),
)
})
.collect()
}
Value::Sequence(nodes) => nodes
.iter()
.map(|node| {
cluster_node(
node.get("host").and_then(as_text),
node.get("port").and_then(as_text),
)
})
.collect(),
_ => Err(INVALID_NODES),
}
}
fn cluster_node(host: Option<String>, port: Option<String>) -> Result<RedisNode, Error> {
match (host, port.and_then(|port| port.parse().ok())) {
(Some(host), Some(port)) => Ok(RedisNode { host, port }),
_ => Err(INVALID_NODES),
}
}
fn json_text(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::String(text) => Some(text.clone()),
serde_json::Value::Number(number) => Some(number.to_string()),
_ => None,
}
}
fn spelled_text<T: ToString>(value: Option<&Spelled<T>>) -> Option<String> {
value.map(|value| match value {
Spelled::Value(value) => value.to_string(),
Spelled::Text(value) => value.clone(),
})
}
fn spelled<T: std::str::FromStr>(text: Option<String>) -> Option<Spelled<T>> {
text.map(|text| match text.parse() {
Ok(value) => Spelled::Value(value),
Err(_) => Spelled::Text(text),
})
}
fn cloned(
fields: &BTreeMap<String, Value>,
keep: impl Fn(&str, &Value) -> bool,
) -> AdditionalFields {
fields
.iter()
.filter(|(key, value)| keep(key, value))
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
#[derive(Clone)]
struct DeferredInput(Value);
type DeferredInputs = BTreeMap<String, DeferredInput>;
fn deferred_inputs(fields: &BTreeMap<String, Value>) -> DeferredInputs {
fields
.iter()
.filter(|(key, value)| is_unparsed(key, value))
.map(|(key, value)| (key.clone(), DeferredInput(value.clone())))
.collect()
}
fn is_unparsed(key: &str, value: &Value) -> bool {
SUPPORTED.contains(&key)
&& match key {
"startup_nodes" => value.is_null(),
_ => as_text(value).is_none(),
}
}

8
litellm-rust/crates/config/src/settings.rs
View file

@ -3,7 +3,9 @@ use std::fmt;
use litellm_auth_types::SecretValue;
use serde::Deserialize;
use crate::{AdditionalFields, Object, Spelled, Value, value::one_or_many};
use crate::{
AdditionalFields, CacheSettings, Object, RedisSettings, Spelled, Value, value::one_or_many,
};
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "lowercase")]
@ -76,7 +78,7 @@ pub struct GeneralSettings {
pub enable_public_model_hub: bool,
pub dangerously_permit_weak_or_unset_master_key: Option<bool>,
pub plugins: Option<Box<[Object]>>,
pub coordination_redis: Option<Object>,
pub coordination_redis: Option<RedisSettings>,
pub mcp_allowed_hosts: Option<Box<[String]>>,
pub mcp_allowed_origins: Box<[String]>,
#[serde(flatten)]
@ -220,7 +222,7 @@ pub struct LiteLlmSettings {
pub request_timeout: Option<Spelled<f64>>,
pub num_retries: Option<u64>,
pub cache: Option<bool>,
pub cache_params: Option<Object>,
pub cache_params: Option<CacheSettings>,
#[serde(deserialize_with = "one_or_many")]
pub callbacks: Option<Vec<Value>>,
#[serde(deserialize_with = "one_or_many")]

4
litellm-rust/crates/config/src/value.rs
View file

@ -10,6 +10,10 @@ pub type AdditionalFields = BTreeMap<String, Value>;
pub struct Object(BTreeMap<String, Value>);
impl Object {
pub(crate) fn from_fields(fields: BTreeMap<String, Value>) -> Self {
Self(fields)
}
pub fn get(&self, key: &str) -> Option<&Value> {
self.0.get(key)
}

405
litellm-rust/crates/config/tests/redis.rs Normal file
View file

@ -0,0 +1,405 @@
use std::{
collections::HashMap,
fmt::Write,
sync::{Arc, Mutex},
};
use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis};
use rstest::rstest;
use tracing::{Event, Subscriber, field::Visit};
use tracing_subscriber::{Layer, layer::Context, prelude::*};
fn resolve(
yaml: &str,
environment: &[(&str, &str)],
) -> Result<litellm_config::ResolvedRedis, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
let config = Config::from_yaml(yaml).unwrap();
config
.general_settings
.coordination_redis
.as_ref()
.unwrap_or(&RedisSettings::default())
.resolve(&|name| environment.get(name).cloned())
}
fn resolve_cache(yaml: &str) -> Result<ResolvedRedis, Error> {
let config = Config::from_yaml(yaml).unwrap();
match config.litellm_settings.cache_params.unwrap() {
CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None),
CacheSettings::Other(_) => panic!("expected a Redis cache"),
}
}
fn cluster_ports(resolved: &ResolvedRedis) -> Vec<u16> {
match &resolved.endpoint {
RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(),
endpoint => panic!("expected a cluster, got {endpoint:?}"),
}
}
#[rstest]
#[case::host(
"general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5",
&[],
"host:redis:6380:2:true:app:1500"
)]
#[case::environment_url(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379")],
"url:redis://environment:6379"
)]
#[case::explicit_host_drops_environment_url(
"general_settings:\n coordination_redis:\n host: configured",
&[("REDIS_URL", "redis://environment:6379")],
"host:configured:6379:0:false::none"
)]
#[case::yaml_cluster(
"general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"",
&[],
"cluster:node-a:7000,node-b:7001"
)]
#[case::json_environment_cluster(
"model_list: []",
&[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)],
"cluster:node-a:7000,node-b:7001"
)]
#[case::missing_reference_keeps_environment_fallback(
"general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
#[case::configured_values_beat_the_same_environment_values(
"general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3",
&[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")],
"host:configured:6380:3:false::none"
)]
#[case::environment_url_beats_the_environment_host(
"model_list: []",
&[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")],
"url:redis://environment:6379"
)]
#[case::startup_nodes_beat_a_url(
"general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000",
&[],
"cluster:node-a:7000"
)]
#[case::environment_references_resolve_the_host_and_port(
"general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT",
&[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")],
"host:referenced:6400:0:false::none"
)]
#[case::a_null_host_falls_back_to_the_environment(
"general_settings:\n coordination_redis:\n host: null",
&[("REDIS_HOST", "environment")],
"host:environment:6379:0:false::none"
)]
fn resolves_redis_through_the_parsed_config(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let resolved = resolve(yaml, environment).unwrap();
let description = match resolved.endpoint {
RedisEndpoint::Url(url) => format!("url:{}", url.expose()),
RedisEndpoint::Host(host) => format!(
"host:{}:{}:{}:{}:{}:{}",
host.host,
host.port,
host.db,
host.tls,
resolved.namespace.unwrap_or_default(),
resolved
.timeout
.map(|timeout| timeout.as_millis().to_string())
.unwrap_or_else(|| "none".to_owned())
),
RedisEndpoint::Cluster(cluster) => format!(
"cluster:{}",
cluster
.nodes
.iter()
.map(|node| format!("{}:{}", node.host, node.port))
.collect::<Vec<_>>()
.join(",")
),
};
assert_eq!(description, expected);
}
#[rstest]
#[case::custom_tls_file(
"general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem",
&[],
"ssl_ca_certs"
)]
#[case::gcp_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account",
&[],
"gcp_service_account"
)]
#[case::gcp_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")],
"REDIS_GCP_SERVICE_ACCOUNT"
)]
#[case::azure_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true",
&[],
"azure_redis_ad_token"
)]
#[case::azure_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")],
"REDIS_AZURE_AD_TOKEN"
)]
#[case::aws_auth_from_config(
"general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true",
&[],
"aws_iam_auth"
)]
#[case::aws_auth_from_environment(
"model_list: []",
&[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")],
"REDIS_AWS_IAM_AUTH"
)]
#[case::sentinel_from_config(
"general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]",
&[],
"sentinel_nodes"
)]
#[case::sentinel_from_environment(
"model_list: []",
&[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")],
"REDIS_SENTINEL_NODES"
)]
fn rejected_settings_name_their_source(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: &str,
) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(error, Error::RedisSetting(setting) if setting == expected));
}
#[rstest]
fn an_explicit_host_keeps_environment_credentials() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: configured",
&[
("REDIS_PASSWORD", "environment-secret"),
("REDIS_SSL", "True"),
],
)
.unwrap();
let RedisEndpoint::Host(host) = resolved.endpoint else {
panic!("expected a host endpoint");
};
assert_eq!(
host.password.as_ref().map(|password| password.expose()),
Some("environment-secret")
);
assert!(host.tls);
}
#[rstest]
#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])]
#[case::a_structured_host_with_an_environment_target(
"general_settings:\n coordination_redis:\n host: []",
&[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")]
)]
fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) {
let error = resolve(yaml, environment).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "host",
..
}
));
}
#[rstest]
fn the_cache_alias_names_cluster_nodes() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn canonical_startup_nodes_beat_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101",
)
.unwrap();
assert_eq!(cluster_ports(&resolved), vec![7000]);
}
#[rstest]
fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() {
let resolved = resolve_cache(
"litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100",
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
#[rstest]
fn rejected_environment_settings_never_print_the_value() {
let error = resolve(
"model_list: []",
&[
("REDIS_HOST", "redis"),
("REDIS_SSL_CERT_REQS", "private-option-value"),
],
)
.unwrap_err();
let message = error.to_string();
assert!(message.contains("REDIS_SSL_CERT_REQS"));
assert!(!message.contains("private-option-value"));
}
#[rstest]
#[case::negative("-1")]
#[case::not_a_number("slow")]
fn rejects_bad_socket_timeout(#[case] timeout: &str) {
let yaml = format!(
"general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}"
);
let error = resolve(&yaml, &[]).unwrap_err();
assert!(matches!(
error,
Error::RedisInvalid {
setting: "socket_timeout",
..
}
));
}
#[rstest]
fn defaults_socket_timeout_for_the_gateway() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis",
&[],
)
.unwrap();
assert_eq!(resolved.timeout, None);
}
#[rstest]
fn debug_redacts_redis_secrets_and_url() {
let config = Config::from_yaml(
"general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password",
)
.unwrap();
let settings = config.general_settings.coordination_redis.unwrap();
let resolved = settings.resolve(&|_| None).unwrap();
let debug = format!("{settings:?} {resolved:?}");
assert!(!debug.contains("redis://user:password"));
assert!(!debug.contains("credential-user"));
assert!(!debug.contains("credential-password"));
assert!(debug.contains("redacted") || debug.contains("REDACTED"));
}
#[rstest]
fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() {
let config = Config::from_yaml(
"litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Redis(_))
));
}
#[rstest]
fn non_redis_cache_settings_remain_opaque() {
let config = Config::from_yaml(
"litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true",
)
.unwrap();
assert!(matches!(
config.litellm_settings.cache_params,
Some(litellm_config::CacheSettings::Other(_))
));
}
#[rstest]
fn legacy_unknown_settings_do_not_change_resolution() {
let resolved = resolve(
"general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary",
&[],
)
.unwrap();
assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_)));
}
struct WarningCapture(Arc<Mutex<Vec<String>>>);
impl<S: Subscriber> Layer<S> for WarningCapture {
fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) {
let mut fields = Fields::default();
event.record(&mut fields);
self.0.lock().unwrap().push(fields.0);
}
}
#[derive(Default)]
struct Fields(String);
impl Visit for Fields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
write!(self.0, "{}={value:?} ", field.name()).unwrap();
}
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
write!(self.0, "{}={value} ", field.name()).unwrap();
}
}
#[rstest]
fn ignored_setting_warning_names_the_setting_without_its_value() {
let events = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone()));
tracing::subscriber::with_default(subscriber, || {
resolve(
"general_settings:\n coordination_redis:\n host: redis\n max_connections: 47",
&[],
)
.unwrap();
});
let warnings = events.lock().unwrap().join(" ");
assert!(warnings.contains("max_connections"));
assert!(!warnings.contains("47"));
}

65
litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs
View file

@ -51,7 +51,7 @@ mod rules;
mod status;
mod vertex_ai;
pub use original::{ExceptionFamily, OriginalException};
pub use original::{ExceptionFamily, OriginalException, TimeoutPhase};
pub use public::{MappedFailure, PublicError, UpstreamResponse};
use rules::{Rule, contains_any, first_match};
@ -67,7 +67,6 @@ const TIMEOUT_MARKERS: &[&str] = &[
pub struct ExceptionContext {
pub model: String,
pub custom_llm_provider: String,
pub asynchronous: bool,
pub vertex_project: Option<String>,
pub vertex_location: Option<String>,
pub model_group: Option<String>,
@ -105,11 +104,12 @@ pub fn exception_type(
(None, message.clone(), None)
}
OriginalException::Timeout {
phase,
timeout_seconds,
elapsed_seconds,
} => (
None,
timeout_message(context.asynchronous, *timeout_seconds, *elapsed_seconds),
timeout_message(*phase, *timeout_seconds, *elapsed_seconds),
None,
),
};
@ -161,21 +161,16 @@ fn family_rules(family: ExceptionFamily) -> &'static [Rule] {
}
}
/// The text the Python HTTP handler's timeout carries: the sync and async handlers word it
/// differently.
fn timeout_message(
asynchronous: bool,
/// The text the Python HTTP handler's timeout carries: the phase that expired and the limit
/// httpx enforced for it, both rendered as Python would print them.
pub fn timeout_message(
phase: TimeoutPhase,
timeout_seconds: Option<f64>,
elapsed_seconds: Option<f64>,
) -> String {
let timeout = python_float(timeout_seconds);
if asynchronous {
let elapsed =
python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0));
format!("Connection timed out. Timeout passed={timeout}, time taken={elapsed} seconds")
} else {
format!("Connection timed out after {timeout} seconds.")
}
let elapsed = python_float(elapsed_seconds.map(|seconds| (seconds * 1000.0).round() / 1000.0));
format!("{phase} timed out. Timeout passed={timeout}, time taken={elapsed} seconds")
}
fn python_float(value: Option<f64>) -> String {
@ -356,6 +351,7 @@ mod tests {
#[test]
fn a_handler_timeout_is_a_408_without_a_response() {
let original = OriginalException::Timeout {
phase: TimeoutPhase::Read,
timeout_seconds: Some(0.5),
elapsed_seconds: Some(0.5031),
};
@ -363,7 +359,7 @@ mod tests {
mapped("mistral", &original),
MappedFailure {
error: PublicError::Timeout { status: 408 },
message: "MistralException - Connection timed out after 0.5 seconds.".into(),
message: "MistralException - Read timed out. Timeout passed=0.5, time taken=0.503 seconds".into(),
upstream: None,
debug_info: DEBUG.into(),
}
@ -440,39 +436,44 @@ mod tests {
}
#[rstest::rstest]
#[case::sync(
false,
#[case::connect(TimeoutPhase::Connect, "Connection timed out.")]
#[case::read(TimeoutPhase::Read, "Read timed out.")]
#[case::write(TimeoutPhase::Write, "Write timed out.")]
#[case::pool(TimeoutPhase::Pool, "Connection pool timed out.")]
fn timeout_text_names_the_phase_that_expired(
#[case] phase: TimeoutPhase,
#[case] expected: &str,
) {
assert_eq!(
timeout_message(phase, Some(0.5), Some(0.5031)),
format!("{expected} Timeout passed=0.5, time taken=0.503 seconds")
);
}
#[rstest::rstest]
#[case::rounds_the_elapsed_time(
Some(0.5),
Some(0.5031),
"Connection timed out after 0.5 seconds."
)]
#[case::async_rounds_the_elapsed_time(
true,
Some(0.5),
Some(0.5031),
"Connection timed out. Timeout passed=0.5, time taken=0.503 seconds"
"Timeout passed=0.5, time taken=0.503 seconds"
)]
#[case::whole_seconds_keep_a_decimal(
true,
Some(600.0),
Some(2.0),
"Connection timed out. Timeout passed=600.0, time taken=2.0 seconds"
"Timeout passed=600.0, time taken=2.0 seconds"
)]
#[case::unknown_values_render_as_none(
true,
None,
None,
"Connection timed out. Timeout passed=None, time taken=None seconds"
"Timeout passed=None, time taken=None seconds"
)]
fn timeout_text_follows_the_delivery_mode(
#[case] asynchronous: bool,
fn timeout_text_renders_durations_as_python_floats(
#[case] timeout_seconds: Option<f64>,
#[case] elapsed_seconds: Option<f64>,
#[case] expected: &str,
) {
assert_eq!(
timeout_message(asynchronous, timeout_seconds, elapsed_seconds),
expected
timeout_message(TimeoutPhase::Read, timeout_seconds, elapsed_seconds),
format!("Read timed out. {expected}")
);
}

14
litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs
View file

@ -10,6 +10,7 @@ pub enum OriginalException {
message: String,
},
Timeout {
phase: TimeoutPhase,
timeout_seconds: Option<f64>,
elapsed_seconds: Option<f64>,
},
@ -20,6 +21,19 @@ pub enum OriginalException {
},
}
/// The httpx timeout that expired, spelled the way the Python HTTP handler names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::VariantArray)]
pub enum TimeoutPhase {
#[strum(serialize = "Connection")]
Connect,
#[strum(serialize = "Read")]
Read,
#[strum(serialize = "Write")]
Write,
#[strum(serialize = "Connection pool")]
Pool,
}
/// Which provider-specific text rules apply before the shared status table.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExceptionFamily {

2
litellm-rust/crates/db-testing/Cargo.toml
View file

@ -7,7 +7,7 @@ repository.workspace = true
[dependencies]
sqlx = { workspace = true, features = ["migrate"] }
testcontainers-modules = { version = "0.15.0", features = ["postgres"] }
testcontainers-modules = { workspace = true, features = ["postgres"] }
thiserror.workspace = true
tokio.workspace = true

23
litellm-rust/crates/db/.sqlx/query-085edfb9af2a914466c75b447bc7bb2d97532a9cd461cc5d8da1e1090560ca1f.json generated Normal file
View file

@ -0,0 +1,23 @@
{
"db_name": "PostgreSQL",
"query": "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS \"locked!\"\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "locked!",
"type_info": "Bool",
"origin": "Expression"
}
],
"parameters": {
"Left": [
"Text"
]
},
"nullable": [
null
]
},
"hash": "085edfb9af2a914466c75b447bc7bb2d97532a9cd461cc5d8da1e1090560ca1f"
}

40
litellm-rust/crates/db/.sqlx/query-32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b.json generated Normal file
View file

@ -0,0 +1,40 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_SpendLogs\" (\n request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, \"startTime\", \"endTime\", \"completionStartTime\",\n request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base,\n \"user\", metadata, cache_hit, cache_key, request_tags, team_id, organization_id,\n end_user, requester_ip_address, messages, response, proxy_server_request,\n session_id, status, litellm_call_id\n)\nSELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens,\n l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3),\n l.completion_start_time::timestamp(3), l.request_duration_ms, l.model,\n l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id,\n l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id,\n l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb,\n '{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id\nFROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[],\n $7::int4[], $8::timestamp[], $9::timestamp[], $10::timestamp[], $11::int4[], $12::text[],\n $13::text[], $14::text[], $15::text[], $16::text[], $17::text[],\n $18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[],\n $23::text[], $24::text[], $25::text[], $26::text[], $27::text[])\n AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens,\n completion_tokens, start_time, end_time, completion_start_time,\n request_duration_ms, model, model_id, model_group, custom_llm_provider,\n api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id,\n end_user, requester_ip_address, session_id, status, litellm_call_id)\nON CONFLICT DO NOTHING\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"TextArray",
"TextArray",
"Float8Array",
"Int4Array",
"Int4Array",
"Int4Array",
"TimestampArray",
"TimestampArray",
"TimestampArray",
"Int4Array",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"JsonbArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray"
]
},
"nullable": []
},
"hash": "32f3c07ced96324b535bbc703c1acf6d06eb22facdf478f88107c9dccdac326b"
}

15
litellm-rust/crates/db/.sqlx/query-345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a.json generated Normal file
View file

@ -0,0 +1,15 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_EndUserTable\" (user_id, spend, blocked)\nSELECT id, SUM(cost), false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)\nGROUP BY id\nORDER BY id\nON CONFLICT (user_id) DO UPDATE\nSET spend = \"LiteLLM_EndUserTable\".spend + EXCLUDED.spend\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"Float8Array"
]
},
"nullable": []
},
"hash": "345c02595f55a7f0274ef7214b02288e0f2c1f12a2a79ed53f0002826724cd0a"
}

78
litellm-rust/crates/db/.sqlx/query-377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e.json generated Normal file
View file

@ -0,0 +1,78 @@
{
"db_name": "PostgreSQL",
"query": "SELECT e.spend AS \"spend?\",\n e.blocked AS \"blocked?\",\n b.max_budget AS \"max_budget?\",\n b.rpm_limit AS \"rpm_limit?\",\n b.tpm_limit AS \"tpm_limit?\"\nFROM (SELECT 1) AS one\nLEFT JOIN \"LiteLLM_EndUserTable\" AS e ON e.user_id = $1\nLEFT JOIN \"LiteLLM_BudgetTable\" AS assigned ON assigned.budget_id = e.budget_id\nLEFT JOIN \"LiteLLM_BudgetTable\" AS key_budget ON key_budget.budget_id = $3\nLEFT JOIN \"LiteLLM_BudgetTable\" AS b ON b.budget_id = COALESCE(assigned.budget_id, key_budget.budget_id, $2)\n",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "spend?",
"type_info": "Float8",
"origin": {
"Table": {
"table": "\"LiteLLM_EndUserTable\"",
"name": "spend"
}
}
},
{
"ordinal": 1,
"name": "blocked?",
"type_info": "Bool",
"origin": {
"Table": {
"table": "\"LiteLLM_EndUserTable\"",
"name": "blocked"
}
}
},
{
"ordinal": 2,
"name": "max_budget?",
"type_info": "Float8",
"origin": {
"Table": {
"table": "\"LiteLLM_BudgetTable\"",
"name": "max_budget"
}
}
},
{
"ordinal": 3,
"name": "rpm_limit?",
"type_info": "Int8",
"origin": {
"Table": {
"table": "\"LiteLLM_BudgetTable\"",
"name": "rpm_limit"
}
}
},
{
"ordinal": 4,
"name": "tpm_limit?",
"type_info": "Int8",
"origin": {
"Table": {
"table": "\"LiteLLM_BudgetTable\"",
"name": "tpm_limit"
}
}
}
],
"parameters": {
"Left": [
"Text",
"Text",
"Text"
]
},
"nullable": [
true,
true,
true,
true,
true
]
},
"hash": "377d063cd0fc3a265d07d01f5aa64311d25428fd393bac4250891f148597e92e"
}

15
litellm-rust/crates/db/.sqlx/query-3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563.json generated Normal file
View file

@ -0,0 +1,15 @@
{
"db_name": "PostgreSQL",
"query": "WITH added AS (\n SELECT id, SUM(cost) AS cost\n FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)\n GROUP BY id\n), locked AS (\n SELECT t.token FROM \"LiteLLM_VerificationToken\" t\n WHERE t.token IN (SELECT id FROM added)\n ORDER BY t.token FOR UPDATE\n)\nUPDATE \"LiteLLM_VerificationToken\" t\nSET spend = t.spend + added.cost,\n total_spend = t.total_spend + added.cost,\n last_active = (clock_timestamp() AT TIME ZONE 'UTC'),\n updated_at = (clock_timestamp() AT TIME ZONE 'UTC')\nFROM added\nWHERE t.token = added.id AND t.token IN (SELECT token FROM locked)\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"Float8Array"
]
},
"nullable": []
},
"hash": "3829a383d0ae50578f0fd75da3910731a0cacc1ae5b92eec32967f539948f563"
}

15
litellm-rust/crates/db/.sqlx/query-43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54.json generated Normal file
View file

@ -0,0 +1,15 @@
{
"db_name": "PostgreSQL",
"query": "WITH added AS (\n SELECT id, SUM(cost) AS cost\n FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)\n GROUP BY id\n), locked AS (\n SELECT u.user_id FROM \"LiteLLM_UserTable\" u\n WHERE u.user_id IN (SELECT id FROM added)\n ORDER BY u.user_id FOR UPDATE\n)\nUPDATE \"LiteLLM_UserTable\" u\nSET spend = u.spend + added.cost,\n updated_at = (NOW() AT TIME ZONE 'UTC')\nFROM added\nWHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"Float8Array"
]
},
"nullable": []
},
"hash": "43f2e12a3ffb1de29b5b031a1173b87a20a2664b0bef3616631d2c24487c1e54"
}

21
litellm-rust/crates/db/.sqlx/query-5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4.json generated Normal file
View file

@ -0,0 +1,21 @@
{
"db_name": "PostgreSQL",
"query": "SELECT set_config('lock_timeout', '5000ms', true)",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "set_config",
"type_info": "Text",
"origin": "Expression"
}
],
"parameters": {
"Left": []
},
"nullable": [
null
]
},
"hash": "5c74f47a2b84574133c520d235192bf33e16fea1dbb0c32cf89a376a874927f4"
}

31
litellm-rust/crates/db/.sqlx/query-9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416.json generated Normal file
View file

@ -0,0 +1,31 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_DailyTeamSpend\" (\n id, team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,\n completion_tokens, api_requests, successful_requests, failed_requests,\n cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,\n timed_requests, spend, updated_at\n)\nSELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\nFROM UNNEST(\n $1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[], $18::text[]\n) AS r(\n entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend, mcp_namespaced_tool_name\n)\nORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,\n r.mcp_namespaced_tool_name, r.endpoint\nON CONFLICT (\n team_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint\n) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyTeamSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyTeamSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyTeamSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyTeamSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyTeamSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyTeamSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyTeamSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyTeamSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyTeamSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Float8Array",
"TextArray"
]
},
"nullable": []
},
"hash": "9c51401a8907f818e96da28050b74dca981ff954452b35ffbace3e022035b416"
}

14
litellm-rust/crates/db/.sqlx/query-b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147.json generated Normal file
View file

@ -0,0 +1,14 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_TeamMembership\" (user_id, team_id, spend, total_spend)\nSELECT member.user_id, member.team_id, member.cost, member.cost\nFROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)\nWHERE EXISTS (\n SELECT 1 FROM \"LiteLLM_TeamTable\" t\n WHERE t.team_id = member.team_id\n AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))\n)\n OR EXISTS (\n SELECT 1 FROM \"LiteLLM_TeamMembership\" m\n WHERE m.user_id = member.user_id AND m.team_id = member.team_id\n)\nON CONFLICT (user_id, team_id) DO UPDATE\nSET spend = \"LiteLLM_TeamMembership\".spend + EXCLUDED.spend,\n total_spend = \"LiteLLM_TeamMembership\".total_spend + EXCLUDED.total_spend\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Jsonb"
]
},
"nullable": []
},
"hash": "b0eb48b795bdaf728c79e986d2027df60ff71d8fed05b53be4458b335af88147"
}

31
litellm-rust/crates/db/.sqlx/query-bec3b53762af021a118d63b841b8c2ebd7fc9fcbf19c1c25aa131a1053d96cfe.json generated Normal file
View file

@ -0,0 +1,31 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_DailyUserSpend\" (\n id, user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,\n completion_tokens, api_requests, successful_requests, failed_requests,\n cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,\n timed_requests, spend, updated_at\n)\nSELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\nFROM UNNEST(\n $1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[], $18::text[]\n) AS r(\n entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend, mcp_namespaced_tool_name\n)\nORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,\n r.mcp_namespaced_tool_name, r.endpoint\nON CONFLICT (\n user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint\n) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyUserSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyUserSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyUserSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyUserSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyUserSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyUserSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyUserSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyUserSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyUserSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyUserSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Float8Array",
"TextArray"
]
},
"nullable": []
},
"hash": "bec3b53762af021a118d63b841b8c2ebd7fc9fcbf19c1c25aa131a1053d96cfe"
}

15
litellm-rust/crates/db/.sqlx/query-c3e5542906e89995fae2a16eca1edd37a575d641bb68d91e61f5123173c275ce.json generated Normal file
View file

@ -0,0 +1,15 @@
{
"db_name": "PostgreSQL",
"query": "WITH added AS (\n SELECT id, SUM(cost) AS cost\n FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)\n GROUP BY id\n),\nlocked AS (\n SELECT t.team_id FROM \"LiteLLM_TeamTable\" t\n WHERE t.team_id IN (SELECT id FROM added)\n ORDER BY t.team_id FOR UPDATE\n)\nUPDATE \"LiteLLM_TeamTable\" t\nSET spend = t.spend + added.cost,\n updated_at = (NOW() AT TIME ZONE 'UTC')\nFROM added\nWHERE t.team_id = added.id AND t.team_id IN (SELECT team_id FROM locked)\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"Float8Array"
]
},
"nullable": []
},
"hash": "c3e5542906e89995fae2a16eca1edd37a575d641bb68d91e61f5123173c275ce"
}

31
litellm-rust/crates/db/.sqlx/query-f207ff0f5310dea0833361d55993bcab69ff6ad3555a3926537586f5d81d933d.json generated Normal file
View file

@ -0,0 +1,31 @@
{
"db_name": "PostgreSQL",
"query": "INSERT INTO \"LiteLLM_DailyEndUserSpend\" (\n id, end_user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,\n completion_tokens, api_requests, successful_requests, failed_requests,\n cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,\n timed_requests, spend, updated_at\n)\nSELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,\n r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,\n r.completion_tokens, r.api_requests, r.successful_requests,\n r.failed_requests, r.cache_read, r.cache_creation, r.response_time,\n r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')\nFROM UNNEST(\n $1::text[], $2::text[], $3::text[], $4::text[], $5::text[],\n $6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],\n $11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],\n $16::int8[], $17::float8[], $18::text[]\n) AS r(\n entity_id, date, api_key, model, custom_llm_provider, endpoint,\n model_group, prompt_tokens, completion_tokens, api_requests,\n successful_requests, failed_requests, cache_read, cache_creation,\n response_time, timed_requests, spend, mcp_namespaced_tool_name\n)\nORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,\n r.mcp_namespaced_tool_name, r.endpoint\nON CONFLICT (\n end_user_id, date, api_key, model, custom_llm_provider,\n mcp_namespaced_tool_name, endpoint\n) DO UPDATE SET\n prompt_tokens = \"LiteLLM_DailyEndUserSpend\".prompt_tokens + EXCLUDED.prompt_tokens,\n completion_tokens = \"LiteLLM_DailyEndUserSpend\".completion_tokens + EXCLUDED.completion_tokens,\n api_requests = \"LiteLLM_DailyEndUserSpend\".api_requests + EXCLUDED.api_requests,\n successful_requests = \"LiteLLM_DailyEndUserSpend\".successful_requests + EXCLUDED.successful_requests,\n failed_requests = \"LiteLLM_DailyEndUserSpend\".failed_requests + EXCLUDED.failed_requests,\n cache_read_input_tokens = \"LiteLLM_DailyEndUserSpend\".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,\n cache_creation_input_tokens = \"LiteLLM_DailyEndUserSpend\".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,\n total_response_time_ms = \"LiteLLM_DailyEndUserSpend\".total_response_time_ms + EXCLUDED.total_response_time_ms,\n timed_requests = \"LiteLLM_DailyEndUserSpend\".timed_requests + EXCLUDED.timed_requests,\n spend = \"LiteLLM_DailyEndUserSpend\".spend + EXCLUDED.spend,\n updated_at = (NOW() AT TIME ZONE 'UTC')\n",
"describe": {
"columns": [],
"parameters": {
"Left": [
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"TextArray",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Int8Array",
"Float8Array",
"TextArray"
]
},
"nullable": []
},
"hash": "f207ff0f5310dea0833361d55993bcab69ff6ad3555a3926537586f5d81d933d"
}

25
litellm-rust/crates/db/AGENTS.md Normal file
View file

@ -0,0 +1,25 @@
# Database crate
Keep each static query in `queries/` and call it through a typed SQLx query function. Rust does not create or migrate the database
- Own reusable row projections in `src/` and export those consumed by other crates or integration tests
- Apply `#[macro_rules_attribute::apply(crate::db_row)]` to row types for shared derives, without adding serialization or schema derives automatically
- Apply `#[macro_rules_attribute::apply(crate::db_input)]` to write inputs, adding `Default` only where supported
- Keep function tests inline in their owning query module and reuse the strict database fixture in `test_support`
- Transaction-local setup in `begin` is the approved exception to static query files
- Keep write inputs separate from stored rows and keep test-only assertions out of public types
- `FromRow` supports runtime reads, but production static reads still use `query_file_as!` for compile-time checking
Database integration tests are ignored by default, but still compile in workspace checks. Run them explicitly against a local database already migrated by Prisma. Missing or invalid database configuration must fail the selected tests:
```sh
DATABASE_URL=postgresql://... cargo test -p litellm-db --locked -- --include-ignored
```
After changing a query, refresh the checked query metadata from this directory:
```sh
SQLX_OFFLINE=false cargo sqlx prepare
```
Tests use transactions and leave them uncommitted so fixture rows roll back when each test finishes

9
litellm-rust/crates/db/Cargo.toml
View file

@ -6,5 +6,14 @@ license.workspace = true
repository.workspace = true
[dependencies]
chrono.workspace = true
macro_rules_attribute.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
strum.workspace = true
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

10
litellm-rust/crates/db/queries/end_user/find_end_user.sql Normal file
View file

@ -0,0 +1,10 @@
SELECT e.spend AS "spend?",
e.blocked AS "blocked?",
b.max_budget AS "max_budget?",
b.rpm_limit AS "rpm_limit?",
b.tpm_limit AS "tpm_limit?"
FROM (SELECT 1) AS one
LEFT JOIN "LiteLLM_EndUserTable" AS e ON e.user_id = $1
LEFT JOIN "LiteLLM_BudgetTable" AS assigned ON assigned.budget_id = e.budget_id
LEFT JOIN "LiteLLM_BudgetTable" AS key_budget ON key_budget.budget_id = $3
LEFT JOIN "LiteLLM_BudgetTable" AS b ON b.budget_id = COALESCE(assigned.budget_id, key_budget.budget_id, $2)

6
litellm-rust/crates/db/queries/spend/add_end_user_spend.sql Normal file
View file

@ -0,0 +1,6 @@
INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked)
SELECT id, SUM(cost), false FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)
GROUP BY id
ORDER BY id
ON CONFLICT (user_id) DO UPDATE
SET spend = "LiteLLM_EndUserTable".spend + EXCLUDED.spend

16
litellm-rust/crates/db/queries/spend/add_key_spend.sql Normal file
View file

@ -0,0 +1,16 @@
WITH added AS (
SELECT id, SUM(cost) AS cost
FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)
GROUP BY id
), locked AS (
SELECT t.token FROM "LiteLLM_VerificationToken" t
WHERE t.token IN (SELECT id FROM added)
ORDER BY t.token FOR UPDATE
)
UPDATE "LiteLLM_VerificationToken" t
SET spend = t.spend + added.cost,
total_spend = t.total_spend + added.cost,
last_active = (clock_timestamp() AT TIME ZONE 'UTC'),
updated_at = (clock_timestamp() AT TIME ZONE 'UTC')
FROM added
WHERE t.token = added.id AND t.token IN (SELECT token FROM locked)

15
litellm-rust/crates/db/queries/spend/add_team_spend.sql Normal file
View file

@ -0,0 +1,15 @@
WITH added AS (
SELECT id, SUM(cost) AS cost
FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)
GROUP BY id
),
locked AS (
SELECT t.team_id FROM "LiteLLM_TeamTable" t
WHERE t.team_id IN (SELECT id FROM added)
ORDER BY t.team_id FOR UPDATE
)
UPDATE "LiteLLM_TeamTable" t
SET spend = t.spend + added.cost,
updated_at = (NOW() AT TIME ZONE 'UTC')
FROM added
WHERE t.team_id = added.id AND t.team_id IN (SELECT team_id FROM locked)

14
litellm-rust/crates/db/queries/spend/add_user_spend.sql Normal file
View file

@ -0,0 +1,14 @@
WITH added AS (
SELECT id, SUM(cost) AS cost
FROM UNNEST($1::text[], $2::float8[]) AS a(id, cost)
GROUP BY id
), locked AS (
SELECT u.user_id FROM "LiteLLM_UserTable" u
WHERE u.user_id IN (SELECT id FROM added)
ORDER BY u.user_id FOR UPDATE
)
UPDATE "LiteLLM_UserTable" u
SET spend = u.spend + added.cost,
updated_at = (NOW() AT TIME ZONE 'UTC')
FROM added
WHERE u.user_id = added.id AND u.user_id IN (SELECT user_id FROM locked)

26
litellm-rust/crates/db/queries/spend/insert_spend_logs.sql Normal file
View file

@ -0,0 +1,26 @@
INSERT INTO "LiteLLM_SpendLogs" (
request_id, call_type, api_key, spend, total_tokens, prompt_tokens,
completion_tokens, "startTime", "endTime", "completionStartTime",
request_duration_ms, model, model_id, model_group, custom_llm_provider, api_base,
"user", metadata, cache_hit, cache_key, request_tags, team_id, organization_id,
end_user, requester_ip_address, messages, response, proxy_server_request,
session_id, status, litellm_call_id
)
SELECT l.request_id, l.call_type, l.api_key, l.spend, l.total_tokens, l.prompt_tokens,
l.completion_tokens, l.start_time::timestamp(3), l.end_time::timestamp(3),
l.completion_start_time::timestamp(3), l.request_duration_ms, l.model,
l.model_id, l.model_group, l.custom_llm_provider, l.api_base, l.user_id,
l.metadata, l.cache_hit, l.cache_key, '[]'::jsonb, l.team_id,
l.organization_id, l.end_user, l.requester_ip_address, '{}'::jsonb,
'{}'::jsonb, '{}'::jsonb, l.session_id, l.status, l.litellm_call_id
FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::int4[], $6::int4[],
$7::int4[], $8::timestamp[], $9::timestamp[], $10::timestamp[], $11::int4[], $12::text[],
$13::text[], $14::text[], $15::text[], $16::text[], $17::text[],
$18::jsonb[], $19::text[], $20::text[], $21::text[], $22::text[],
$23::text[], $24::text[], $25::text[], $26::text[], $27::text[])
AS l(request_id, call_type, api_key, spend, total_tokens, prompt_tokens,
completion_tokens, start_time, end_time, completion_start_time,
request_duration_ms, model, model_id, model_group, custom_llm_provider,
api_base, user_id, metadata, cache_hit, cache_key, team_id, organization_id,
end_user, requester_ip_address, session_id, status, litellm_call_id)
ON CONFLICT DO NOTHING

1
litellm-rust/crates/db/queries/spend/lock_team.sql Normal file
View file

@ -0,0 +1 @@
SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS "locked!"

40
litellm-rust/crates/db/queries/spend/upsert_daily_end_user_spend.sql Normal file
View file

@ -0,0 +1,40 @@
INSERT INTO "LiteLLM_DailyEndUserSpend" (
id, end_user_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,
completion_tokens, api_requests, successful_requests, failed_requests,
cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, updated_at
)
SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,
r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,
r.completion_tokens, r.api_requests, r.successful_requests,
r.failed_requests, r.cache_read, r.cache_creation, r.response_time,
r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')
FROM UNNEST(
$1::text[], $2::text[], $3::text[], $4::text[], $5::text[],
$6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],
$11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],
$16::int8[], $17::float8[], $18::text[]
) AS r(
entity_id, date, api_key, model, custom_llm_provider, endpoint,
model_group, prompt_tokens, completion_tokens, api_requests,
successful_requests, failed_requests, cache_read, cache_creation,
response_time, timed_requests, spend, mcp_namespaced_tool_name
)
ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,
r.mcp_namespaced_tool_name, r.endpoint
ON CONFLICT (
end_user_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint
) DO UPDATE SET
prompt_tokens = "LiteLLM_DailyEndUserSpend".prompt_tokens + EXCLUDED.prompt_tokens,
completion_tokens = "LiteLLM_DailyEndUserSpend".completion_tokens + EXCLUDED.completion_tokens,
api_requests = "LiteLLM_DailyEndUserSpend".api_requests + EXCLUDED.api_requests,
successful_requests = "LiteLLM_DailyEndUserSpend".successful_requests + EXCLUDED.successful_requests,
failed_requests = "LiteLLM_DailyEndUserSpend".failed_requests + EXCLUDED.failed_requests,
cache_read_input_tokens = "LiteLLM_DailyEndUserSpend".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,
cache_creation_input_tokens = "LiteLLM_DailyEndUserSpend".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,
total_response_time_ms = "LiteLLM_DailyEndUserSpend".total_response_time_ms + EXCLUDED.total_response_time_ms,
timed_requests = "LiteLLM_DailyEndUserSpend".timed_requests + EXCLUDED.timed_requests,
spend = "LiteLLM_DailyEndUserSpend".spend + EXCLUDED.spend,
updated_at = (NOW() AT TIME ZONE 'UTC')

40
litellm-rust/crates/db/queries/spend/upsert_daily_team_spend.sql Normal file
View file

@ -0,0 +1,40 @@
INSERT INTO "LiteLLM_DailyTeamSpend" (
id, team_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,
completion_tokens, api_requests, successful_requests, failed_requests,
cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, updated_at
)
SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,
r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,
r.completion_tokens, r.api_requests, r.successful_requests,
r.failed_requests, r.cache_read, r.cache_creation, r.response_time,
r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')
FROM UNNEST(
$1::text[], $2::text[], $3::text[], $4::text[], $5::text[],
$6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],
$11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],
$16::int8[], $17::float8[], $18::text[]
) AS r(
entity_id, date, api_key, model, custom_llm_provider, endpoint,
model_group, prompt_tokens, completion_tokens, api_requests,
successful_requests, failed_requests, cache_read, cache_creation,
response_time, timed_requests, spend, mcp_namespaced_tool_name
)
ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,
r.mcp_namespaced_tool_name, r.endpoint
ON CONFLICT (
team_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint
) DO UPDATE SET
prompt_tokens = "LiteLLM_DailyTeamSpend".prompt_tokens + EXCLUDED.prompt_tokens,
completion_tokens = "LiteLLM_DailyTeamSpend".completion_tokens + EXCLUDED.completion_tokens,
api_requests = "LiteLLM_DailyTeamSpend".api_requests + EXCLUDED.api_requests,
successful_requests = "LiteLLM_DailyTeamSpend".successful_requests + EXCLUDED.successful_requests,
failed_requests = "LiteLLM_DailyTeamSpend".failed_requests + EXCLUDED.failed_requests,
cache_read_input_tokens = "LiteLLM_DailyTeamSpend".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,
cache_creation_input_tokens = "LiteLLM_DailyTeamSpend".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,
total_response_time_ms = "LiteLLM_DailyTeamSpend".total_response_time_ms + EXCLUDED.total_response_time_ms,
timed_requests = "LiteLLM_DailyTeamSpend".timed_requests + EXCLUDED.timed_requests,
spend = "LiteLLM_DailyTeamSpend".spend + EXCLUDED.spend,
updated_at = (NOW() AT TIME ZONE 'UTC')

40
litellm-rust/crates/db/queries/spend/upsert_daily_user_spend.sql Normal file
View file

@ -0,0 +1,40 @@
INSERT INTO "LiteLLM_DailyUserSpend" (
id, user_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint, model_group, prompt_tokens,
completion_tokens, api_requests, successful_requests, failed_requests,
cache_read_input_tokens, cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, updated_at
)
SELECT gen_random_uuid()::text, r.entity_id, r.date, r.api_key, r.model,
r.custom_llm_provider, r.mcp_namespaced_tool_name, r.endpoint, r.model_group, r.prompt_tokens,
r.completion_tokens, r.api_requests, r.successful_requests,
r.failed_requests, r.cache_read, r.cache_creation, r.response_time,
r.timed_requests, r.spend, (NOW() AT TIME ZONE 'UTC')
FROM UNNEST(
$1::text[], $2::text[], $3::text[], $4::text[], $5::text[],
$6::text[], $7::text[], $8::int8[], $9::int8[], $10::int8[],
$11::int8[], $12::int8[], $13::int8[], $14::int8[], $15::int8[],
$16::int8[], $17::float8[], $18::text[]
) AS r(
entity_id, date, api_key, model, custom_llm_provider, endpoint,
model_group, prompt_tokens, completion_tokens, api_requests,
successful_requests, failed_requests, cache_read, cache_creation,
response_time, timed_requests, spend, mcp_namespaced_tool_name
)
ORDER BY r.date, r.entity_id, r.api_key, r.model, r.custom_llm_provider,
r.mcp_namespaced_tool_name, r.endpoint
ON CONFLICT (
user_id, date, api_key, model, custom_llm_provider,
mcp_namespaced_tool_name, endpoint
) DO UPDATE SET
prompt_tokens = "LiteLLM_DailyUserSpend".prompt_tokens + EXCLUDED.prompt_tokens,
completion_tokens = "LiteLLM_DailyUserSpend".completion_tokens + EXCLUDED.completion_tokens,
api_requests = "LiteLLM_DailyUserSpend".api_requests + EXCLUDED.api_requests,
successful_requests = "LiteLLM_DailyUserSpend".successful_requests + EXCLUDED.successful_requests,
failed_requests = "LiteLLM_DailyUserSpend".failed_requests + EXCLUDED.failed_requests,
cache_read_input_tokens = "LiteLLM_DailyUserSpend".cache_read_input_tokens + EXCLUDED.cache_read_input_tokens,
cache_creation_input_tokens = "LiteLLM_DailyUserSpend".cache_creation_input_tokens + EXCLUDED.cache_creation_input_tokens,
total_response_time_ms = "LiteLLM_DailyUserSpend".total_response_time_ms + EXCLUDED.total_response_time_ms,
timed_requests = "LiteLLM_DailyUserSpend".timed_requests + EXCLUDED.timed_requests,
spend = "LiteLLM_DailyUserSpend".spend + EXCLUDED.spend,
updated_at = (NOW() AT TIME ZONE 'UTC')

15
litellm-rust/crates/db/queries/spend/upsert_team_member_spend.sql Normal file
View file

@ -0,0 +1,15 @@
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
SELECT member.user_id, member.team_id, member.cost, member.cost
FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8)
WHERE EXISTS (
SELECT 1 FROM "LiteLLM_TeamTable" t
WHERE t.team_id = member.team_id
AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))
)
OR EXISTS (
SELECT 1 FROM "LiteLLM_TeamMembership" m
WHERE m.user_id = member.user_id AND m.team_id = member.team_id
)
ON CONFLICT (user_id, team_id) DO UPDATE
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend

76
litellm-rust/crates/db/src/daily_spend.rs Normal file
View file

@ -0,0 +1,76 @@
use sqlx::types::chrono::NaiveDateTime;
#[macro_rules_attribute::apply(crate::db_row)]
pub struct DailyUserRow {
pub id: String,
pub user_id: Option<String>,
pub api_key: String,
pub model: Option<String>,
pub custom_llm_provider: Option<String>,
pub endpoint: Option<String>,
pub model_group: Option<String>,
pub mcp_namespaced_tool_name: Option<String>,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub api_requests: i64,
pub successful_requests: i64,
pub failed_requests: i64,
pub cache_read_input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub total_response_time_ms: i64,
pub timed_requests: i64,
pub spend: f64,
pub date: String,
pub created_at: NaiveDateTime,
pub updated_at: NaiveDateTime,
}
#[macro_rules_attribute::apply(crate::db_row)]
pub struct DailyTeamRow {
pub id: String,
pub team_id: Option<String>,
pub api_key: String,
pub model: Option<String>,
pub custom_llm_provider: Option<String>,
pub endpoint: Option<String>,
pub model_group: Option<String>,
pub mcp_namespaced_tool_name: Option<String>,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub api_requests: i64,
pub successful_requests: i64,
pub failed_requests: i64,
pub cache_read_input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub total_response_time_ms: i64,
pub timed_requests: i64,
pub spend: f64,
pub date: String,
pub created_at: NaiveDateTime,
pub updated_at: NaiveDateTime,
}
#[macro_rules_attribute::apply(crate::db_row)]
pub struct DailyEndUserRow {
pub id: String,
pub end_user_id: Option<String>,
pub api_key: String,
pub model: Option<String>,
pub custom_llm_provider: Option<String>,
pub endpoint: Option<String>,
pub model_group: Option<String>,
pub mcp_namespaced_tool_name: Option<String>,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub api_requests: i64,
pub successful_requests: i64,
pub failed_requests: i64,
pub cache_read_input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub total_response_time_ms: i64,
pub timed_requests: i64,
pub spend: f64,
pub date: String,
pub created_at: NaiveDateTime,
pub updated_at: NaiveDateTime,
}

359
litellm-rust/crates/db/src/end_user.rs Normal file
View file

@ -0,0 +1,359 @@
use sqlx::PgExecutor;
use crate::Error;
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct RateLimits {
pub rpm: Option<i64>,
pub tpm: Option<i64>,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct EndUser {
pub spend: Option<f64>,
pub blocked: bool,
pub max_budget: Option<f64>,
pub limits: RateLimits,
}
#[macro_rules_attribute::apply(crate::db_row)]
struct EndUserRow {
spend: Option<f64>,
blocked: Option<bool>,
max_budget: Option<f64>,
rpm_limit: Option<i64>,
tpm_limit: Option<i64>,
}
pub async fn find_end_user(
executor: impl PgExecutor<'_>,
end_user_id: &str,
default_budget_id: Option<&str>,
key_budget_id: Option<&str>,
) -> Result<EndUser, Error> {
let row = sqlx::query_file_as!(
EndUserRow,
"queries/end_user/find_end_user.sql",
end_user_id,
default_budget_id,
key_budget_id,
)
.fetch_one(executor)
.await?;
Ok(EndUser {
spend: row.spend,
blocked: row.blocked.unwrap_or(false),
max_budget: row.max_budget,
limits: RateLimits {
rpm: row.rpm_limit,
tpm: row.tpm_limit,
},
})
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{EndUser, RateLimits, find_end_user};
use crate::test_support::database_pool;
use rstest::rstest;
use sqlx::{PgPool, Postgres, Transaction};
#[derive(Clone, Copy)]
struct BudgetFixture {
budget_id: &'static str,
max_budget: f64,
rpm_limit: i64,
tpm_limit: i64,
}
async fn insert_budget(
tx: &mut Transaction<'_, Postgres>,
budget: BudgetFixture,
) -> Result<(), sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "Fixture inserts are outside the compile-checked query set"
)]
let result = sqlx::query(
r#"
INSERT INTO "LiteLLM_BudgetTable"
(budget_id, max_budget, rpm_limit, tpm_limit, created_by, updated_by)
VALUES ($1, $2, $3, $4, 'db-query-test', 'db-query-test')
"#,
)
.bind(budget.budget_id)
.bind(budget.max_budget)
.bind(budget.rpm_limit)
.bind(budget.tpm_limit)
.execute(&mut **tx)
.await;
result.map(|_| ())
}
async fn insert_end_user(
tx: &mut Transaction<'_, Postgres>,
user_id: &str,
spend: f64,
blocked: bool,
budget_id: Option<&str>,
) -> Result<(), sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "Fixture inserts are outside the compile-checked query set"
)]
let result = sqlx::query(
r#"
INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked, budget_id)
VALUES ($1, $2, $3, $4)
"#,
)
.bind(user_id)
.bind(spend)
.bind(blocked)
.bind(budget_id)
.execute(&mut **tx)
.await;
result.map(|_| ())
}
#[rstest]
#[case::without_default_budget(None)]
#[case::with_default_budget(Some(BudgetFixture {
budget_id: "missing-user-default-budget",
max_budget: 11.5,
rpm_limit: 12,
tpm_limit: 1300,
}))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn missing_end_user_maps_to_defaults(
#[case] default_budget: Option<BudgetFixture>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
if let Some(budget) = default_budget {
insert_budget(&mut tx, budget).await?;
}
let actual = find_end_user(
&mut *tx,
"missing-end-user",
default_budget.map(|budget| budget.budget_id),
None,
)
.await?;
assert_eq!(
actual,
EndUser {
spend: None,
blocked: false,
max_budget: default_budget.map(|budget| budget.max_budget),
limits: RateLimits {
rpm: default_budget.map(|budget| budget.rpm_limit),
tpm: default_budget.map(|budget| budget.tpm_limit),
},
}
);
Ok(())
}
#[rstest]
#[case::row_budget_precedes_default(
"row-budget-priority-user",
Some(BudgetFixture {
budget_id: "row-budget-priority-row",
max_budget: 10.0,
rpm_limit: 101,
tpm_limit: 1001,
}),
BudgetFixture {
budget_id: "row-budget-priority-default",
max_budget: 20.0,
rpm_limit: 201,
tpm_limit: 2001,
},
None,
Some(BudgetFixture {
budget_id: "row-budget-priority-row",
max_budget: 10.0,
rpm_limit: 101,
tpm_limit: 1001,
}),
false
)]
#[case::default_applies_without_row_budget(
"default-budget-user",
None,
BudgetFixture {
budget_id: "default-budget-fallback",
max_budget: 30.0,
rpm_limit: 301,
tpm_limit: 3001,
},
None,
Some(BudgetFixture {
budget_id: "default-budget-fallback",
max_budget: 30.0,
rpm_limit: 301,
tpm_limit: 3001,
}),
false
)]
#[case::assigned_budget_precedes_key_when_all_ids_differ(
"override-priority-user",
Some(BudgetFixture {
budget_id: "override-priority-row",
max_budget: 40.0,
rpm_limit: 401,
tpm_limit: 4001,
}),
BudgetFixture {
budget_id: "override-priority-default",
max_budget: 50.0,
rpm_limit: 501,
tpm_limit: 5001,
},
Some(BudgetFixture {
budget_id: "override-priority-override",
max_budget: 60.0,
rpm_limit: 601,
tpm_limit: 6001,
}),
Some(BudgetFixture {
budget_id: "override-priority-row",
max_budget: 40.0,
rpm_limit: 401,
tpm_limit: 4001,
}),
true
)]
#[case::missing_key_budget_keeps_assigned_budget(
"dangling-override-user",
Some(BudgetFixture {
budget_id: "dangling-override-row",
max_budget: 70.0,
rpm_limit: 701,
tpm_limit: 7001,
}),
BudgetFixture {
budget_id: "dangling-override-default",
max_budget: 80.0,
rpm_limit: 801,
tpm_limit: 8001,
},
Some(BudgetFixture {
budget_id: "dangling-override-missing",
max_budget: 90.0,
rpm_limit: 901,
tpm_limit: 9001,
}),
Some(BudgetFixture {
budget_id: "dangling-override-row",
max_budget: 70.0,
rpm_limit: 701,
tpm_limit: 7001,
}),
false
)]
#[case::key_budget_precedes_proxy_default(
"key-default-user",
None,
BudgetFixture {
budget_id: "key-default-proxy",
max_budget: 80.0,
rpm_limit: 801,
tpm_limit: 8001,
},
Some(BudgetFixture {
budget_id: "key-default-key",
max_budget: 90.0,
rpm_limit: 901,
tpm_limit: 9001,
}),
Some(BudgetFixture {
budget_id: "key-default-key",
max_budget: 90.0,
rpm_limit: 901,
tpm_limit: 9001,
}),
true
)]
#[case::missing_key_budget_falls_back_to_proxy_default(
"missing-key-default-user",
None,
BudgetFixture {
budget_id: "missing-key-default-proxy",
max_budget: 80.0,
rpm_limit: 801,
tpm_limit: 8001,
},
Some(BudgetFixture {
budget_id: "missing-key-default-key",
max_budget: 90.0,
rpm_limit: 901,
tpm_limit: 9001,
}),
Some(BudgetFixture {
budget_id: "missing-key-default-proxy",
max_budget: 80.0,
rpm_limit: 801,
tpm_limit: 8001,
}),
false
)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn budget_resolution_matches_python_precedence(
#[case] user_id: &'static str,
#[case] user_budget: Option<BudgetFixture>,
#[case] default_budget: BudgetFixture,
#[case] override_budget: Option<BudgetFixture>,
#[case] expected_budget: Option<BudgetFixture>,
#[case] insert_override_budget: bool,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
insert_budget(&mut tx, default_budget).await?;
if let Some(budget) = user_budget {
insert_budget(&mut tx, budget).await?;
}
if let (Some(budget), true) = (override_budget, insert_override_budget) {
insert_budget(&mut tx, budget).await?;
}
insert_end_user(
&mut tx,
user_id,
7.5,
true,
user_budget.map(|budget| budget.budget_id),
)
.await?;
let actual = find_end_user(
&mut *tx,
user_id,
Some(default_budget.budget_id),
override_budget.map(|budget| budget.budget_id),
)
.await?;
assert_eq!(
actual,
EndUser {
spend: Some(7.5),
blocked: true,
max_budget: expected_budget.map(|budget| budget.max_budget),
limits: RateLimits {
rpm: expected_budget.map(|budget| budget.rpm_limit),
tpm: expected_budget.map(|budget| budget.tpm_limit),
},
}
);
Ok(())
}
}

79
litellm-rust/crates/db/src/error.rs Normal file
View file

@ -0,0 +1,79 @@
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("database query failed")]
Query(#[from] sqlx::Error),
#[error("spend batch has {ids} ids and {costs} costs")]
InvalidSpendBatch { ids: usize, costs: usize },
#[error("invalid spend timestamp")]
InvalidTimestamp(#[from] chrono::ParseError),
}
impl Error {
pub fn is_retryable(&self) -> bool {
match self {
Self::Query(sqlx::Error::Database(error)) => {
matches!(
error.code().as_deref(),
Some("40P01" | "55P03" | "40001" | "53300")
)
}
Self::Query(sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut) => true,
_ => false,
}
}
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use rstest::rstest;
use sqlx::PgPool;
use super::Error;
use crate::test_support::database_pool;
#[rstest]
#[case::closed_pool(sqlx::Error::PoolClosed, false)]
#[case::pool_timeout(sqlx::Error::PoolTimedOut, true)]
#[case::io(sqlx::Error::Io(std::io::ErrorKind::ConnectionReset.into()), true)]
#[case::missing_row(sqlx::Error::RowNotFound, false)]
fn classifies_client_errors(#[case] error: sqlx::Error, #[case] retryable: bool) {
assert_eq!(Error::from(error).is_retryable(), retryable);
}
#[rstest]
#[case::connection_capacity("53300", true)]
#[case::deadlock("40P01", true)]
#[case::lock_unavailable("55P03", true)]
#[case::serialization_failure("40001", true)]
#[case::unique_violation("23505", false)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn classifies_server_errors(
#[case] code: &str,
#[case] retryable: bool,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test query must raise the requested SQLSTATE"
)]
let error = sqlx::query(sqlx::AssertSqlSafe(format!(
"DO $$ BEGIN RAISE EXCEPTION USING ERRCODE = '{code}'; END $$"
)))
.execute(&mut *tx)
.await
.expect_err("PostgreSQL must raise the requested error");
assert_eq!(
error.as_database_error().and_then(|error| error.code()),
Some(code.into())
);
assert_eq!(Error::from(error).is_retryable(), retryable);
tx.rollback().await?;
Ok(())
}
}

25
litellm-rust/crates/db/src/lib.rs
View file

@ -1,3 +1,26 @@
mod failure;
macro_rules_attribute::attribute_alias! {
#[apply(db_input)] =
#[derive(Clone, Debug, PartialEq)];
#[apply(db_row)] =
#[derive(Clone, Debug, PartialEq, sqlx::FromRow)];
}
mod daily_spend;
mod end_user;
mod error;
mod failure;
mod spend;
#[cfg(test)]
mod test_support;
pub use daily_spend::{DailyEndUserRow, DailyTeamRow, DailyUserRow};
pub use end_user::{EndUser, RateLimits, find_end_user};
pub use error::Error;
pub use failure::DbFailure;
pub use spend::{
DailySpend, DailyTable, EntitySpend, SpendLog, SpendLogRow, TeamMemberSpend,
add_end_user_spend, add_key_spend, add_team_member_spend, add_team_spend, add_user_spend,
begin, insert_spend_logs, lock_team, upsert_daily_end_user_spend, upsert_daily_spend,
upsert_daily_team_spend, upsert_daily_user_spend,
};

584
litellm-rust/crates/db/src/spend/daily_end_user_spend.rs Normal file
View file

@ -0,0 +1,584 @@
use sqlx::PgExecutor;
use super::{DailySpend, upsert_daily_team_spend, upsert_daily_user_spend};
use crate::Error;
pub async fn upsert_daily_end_user_spend(
executor: impl PgExecutor<'_>,
rows: &[DailySpend],
) -> Result<(), Error> {
if rows.is_empty() {
return Ok(());
}
macro_rules! column {
($field:ident) => {
rows.iter()
.map(|row| row.$field.clone())
.collect::<Vec<_>>()
};
}
sqlx::query_file!(
"queries/spend/upsert_daily_end_user_spend.sql",
&column!(entity_id),
&column!(date),
&column!(api_key),
&column!(model),
&column!(custom_llm_provider),
&column!(endpoint),
&column!(model_group) as &[Option<String>],
&column!(prompt_tokens),
&column!(completion_tokens),
&column!(api_requests),
&column!(successful_requests),
&column!(failed_requests),
&column!(cache_read_input_tokens),
&column!(cache_creation_input_tokens),
&column!(total_response_time_ms),
&column!(timed_requests),
&column!(spend),
&column!(mcp_namespaced_tool_name),
)
.execute(executor)
.await?;
Ok(())
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::VariantArray)]
pub enum DailyTable {
User,
Team,
EndUser,
}
pub async fn upsert_daily_spend(
executor: impl PgExecutor<'_>,
table: DailyTable,
rows: &[DailySpend],
) -> Result<(), Error> {
match table {
DailyTable::User => upsert_daily_user_spend(executor, rows).await,
DailyTable::Team => upsert_daily_team_spend(executor, rows).await,
DailyTable::EndUser => upsert_daily_end_user_spend(executor, rows).await,
}
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{DailySpend, DailyTable, upsert_daily_end_user_spend, upsert_daily_spend};
use crate::{DailyEndUserRow, test_support::database_pool};
use rstest::rstest;
use sqlx::{PgPool, types::chrono::NaiveDateTime};
#[rstest]
#[case::full_keys(
"daily-end-user-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
Some("first-group")
)]
#[case::initial_model_group_none(
"daily-end-user-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
None
)]
#[case::empty_optional_keys("", "", "", "", "", Some(""))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_end_user_spend_upsert_preserves_keys_and_accumulates(
#[case] end_user_id: &'static str,
#[case] api_key: &'static str,
#[case] model: &'static str,
#[case] custom_llm_provider: &'static str,
#[case] endpoint: &'static str,
#[case] first_model_group: Option<&'static str>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test exercises UTC timestamp behavior under a non-UTC session"
)]
sqlx::query("SET LOCAL TIME ZONE 'America/Los_Angeles'")
.execute(&mut *tx)
.await?;
let first_batch = [
DailySpend {
entity_id: end_user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: first_model_group.map(str::to_owned),
prompt_tokens: 1,
completion_tokens: 2,
api_requests: 3,
successful_requests: 4,
failed_requests: 5,
cache_read_input_tokens: 6,
cache_creation_input_tokens: 7,
total_response_time_ms: 8,
timed_requests: 9,
spend: 0.125,
},
DailySpend {
entity_id: end_user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: "endpoint-sentinel".to_owned(),
model_group: Some("sentinel-endpoint-group".to_owned()),
prompt_tokens: 101,
completion_tokens: 102,
api_requests: 103,
successful_requests: 104,
failed_requests: 105,
cache_read_input_tokens: 106,
cache_creation_input_tokens: 107,
total_response_time_ms: 108,
timed_requests: 109,
spend: 1.25,
},
DailySpend {
entity_id: end_user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: "api-key-sentinel".to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-key-group".to_owned()),
prompt_tokens: 201,
completion_tokens: 202,
api_requests: 203,
successful_requests: 204,
failed_requests: 205,
cache_read_input_tokens: 206,
cache_creation_input_tokens: 207,
total_response_time_ms: 208,
timed_requests: 209,
spend: 2.5,
},
DailySpend {
entity_id: "entity-sentinel".to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-entity-group".to_owned()),
prompt_tokens: 301,
completion_tokens: 302,
api_requests: 303,
successful_requests: 304,
failed_requests: 305,
cache_read_input_tokens: 306,
cache_creation_input_tokens: 307,
total_response_time_ms: 308,
timed_requests: 309,
spend: 3.75,
},
];
upsert_daily_end_user_spend(&mut *tx, &first_batch).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the inserted row identity before conflict handling"
)]
let (initial_id, initial_created_at): (String, NaiveDateTime) = sqlx::query_as(
r#"SELECT id, created_at
FROM "LiteLLM_DailyEndUserSpend"
WHERE end_user_id = $1 AND date = $2 AND api_key = $3 AND model = $4
AND custom_llm_provider = $5 AND endpoint = $6"#,
)
.bind(end_user_id)
.bind("2026-01-15")
.bind(api_key)
.bind(model)
.bind(custom_llm_provider)
.bind(endpoint)
.fetch_one(&mut *tx)
.await?;
assert!(!initial_id.is_empty());
#[expect(
clippy::disallowed_methods,
reason = "The integration test seeds updated_at before exercising the conflict update"
)]
let seeded_updated_at: NaiveDateTime = sqlx::query_scalar(
r#"UPDATE "LiteLLM_DailyEndUserSpend"
SET updated_at = TIMESTAMP '2000-01-01'
WHERE id = $1
RETURNING updated_at"#,
)
.bind(&initial_id)
.fetch_one(&mut *tx)
.await?;
let increment = [DailySpend {
entity_id: end_user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("second-group".to_owned()),
prompt_tokens: 10,
completion_tokens: 20,
api_requests: 30,
successful_requests: 40,
failed_requests: 50,
cache_read_input_tokens: 60,
cache_creation_input_tokens: 70,
total_response_time_ms: 80,
timed_requests: 90,
spend: 0.625,
}];
upsert_daily_end_user_spend(&mut *tx, &increment).await?;
upsert_daily_end_user_spend(&mut *tx, &[]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the transaction timestamp for assertions"
)]
let expected_updated_at: NaiveDateTime =
sqlx::query_scalar("SELECT (NOW() AT TIME ZONE 'UTC')::timestamp(3)")
.fetch_one(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads all rows to verify attribution and accumulation"
)]
let rows: Vec<DailyEndUserRow> = sqlx::query_as(
r#"SELECT id, end_user_id, api_key, model, custom_llm_provider, endpoint,
model_group, mcp_namespaced_tool_name, prompt_tokens,
completion_tokens, api_requests, successful_requests,
failed_requests, cache_read_input_tokens,
cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, date, created_at, updated_at
FROM "LiteLLM_DailyEndUserSpend"
WHERE date = $1 AND model = $2 AND custom_llm_provider = $3
AND mcp_namespaced_tool_name = '' AND end_user_id = ANY($4)
ORDER BY CASE
WHEN end_user_id = $5 AND api_key = $6 AND endpoint = $7 THEN 0
WHEN end_user_id = $5 AND api_key = $6 AND endpoint = $9 THEN 1
WHEN end_user_id = $5 AND api_key = $8 AND endpoint = $7 THEN 2
WHEN end_user_id = $10 AND api_key = $6 AND endpoint = $7 THEN 3
ELSE 4
END"#,
)
.bind("2026-01-15")
.bind(model)
.bind(custom_llm_provider)
.bind(vec![end_user_id, "entity-sentinel"])
.bind(end_user_id)
.bind(api_key)
.bind(endpoint)
.bind("api-key-sentinel")
.bind("endpoint-sentinel")
.bind("entity-sentinel")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 4);
assert_eq!(rows[0].id, initial_id);
assert_eq!(rows[0].end_user_id.as_deref(), Some(end_user_id));
assert_eq!(rows[0].api_key, api_key);
assert_eq!(rows[0].model.as_deref(), Some(model));
assert_eq!(
rows[0].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[0].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[0].model_group.as_deref(), first_model_group);
assert_eq!(rows[0].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[0].prompt_tokens,
rows[0].completion_tokens,
rows[0].api_requests,
rows[0].successful_requests,
rows[0].failed_requests,
rows[0].cache_read_input_tokens,
rows[0].cache_creation_input_tokens,
rows[0].total_response_time_ms,
rows[0].timed_requests,
rows[0].spend
),
(11, 22, 33, 44, 55, 66, 77, 88, 99, 0.75)
);
assert_eq!(rows[0].date, "2026-01-15");
assert_eq!(rows[0].created_at, initial_created_at);
assert!(rows[0].updated_at > seeded_updated_at);
assert_eq!(rows[0].updated_at, expected_updated_at);
assert!(!rows[1].id.is_empty());
assert_eq!(rows[1].end_user_id.as_deref(), Some(end_user_id));
assert_eq!(rows[1].api_key, api_key);
assert_eq!(rows[1].model.as_deref(), Some(model));
assert_eq!(
rows[1].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[1].endpoint.as_deref(), Some("endpoint-sentinel"));
assert_eq!(
rows[1].model_group.as_deref(),
Some("sentinel-endpoint-group")
);
assert_eq!(rows[1].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[1].prompt_tokens,
rows[1].completion_tokens,
rows[1].api_requests,
rows[1].successful_requests,
rows[1].failed_requests,
rows[1].cache_read_input_tokens,
rows[1].cache_creation_input_tokens,
rows[1].total_response_time_ms,
rows[1].timed_requests,
rows[1].spend
),
(101, 102, 103, 104, 105, 106, 107, 108, 109, 1.25)
);
assert_eq!(rows[1].date, "2026-01-15");
assert!(rows[1].updated_at > seeded_updated_at);
assert_eq!(rows[1].updated_at, expected_updated_at);
assert!(!rows[2].id.is_empty());
assert_eq!(rows[2].end_user_id.as_deref(), Some(end_user_id));
assert_eq!(rows[2].api_key, "api-key-sentinel");
assert_eq!(rows[2].model.as_deref(), Some(model));
assert_eq!(
rows[2].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[2].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[2].model_group.as_deref(), Some("sentinel-key-group"));
assert_eq!(rows[2].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[2].prompt_tokens,
rows[2].completion_tokens,
rows[2].api_requests,
rows[2].successful_requests,
rows[2].failed_requests,
rows[2].cache_read_input_tokens,
rows[2].cache_creation_input_tokens,
rows[2].total_response_time_ms,
rows[2].timed_requests,
rows[2].spend
),
(201, 202, 203, 204, 205, 206, 207, 208, 209, 2.5)
);
assert_eq!(rows[2].date, "2026-01-15");
assert!(rows[2].updated_at > seeded_updated_at);
assert_eq!(rows[2].updated_at, expected_updated_at);
assert!(!rows[3].id.is_empty());
assert_eq!(rows[3].end_user_id.as_deref(), Some("entity-sentinel"));
assert_eq!(rows[3].api_key, api_key);
assert_eq!(rows[3].model.as_deref(), Some(model));
assert_eq!(
rows[3].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[3].endpoint.as_deref(), Some(endpoint));
assert_eq!(
rows[3].model_group.as_deref(),
Some("sentinel-entity-group")
);
assert_eq!(rows[3].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[3].prompt_tokens,
rows[3].completion_tokens,
rows[3].api_requests,
rows[3].successful_requests,
rows[3].failed_requests,
rows[3].cache_read_input_tokens,
rows[3].cache_creation_input_tokens,
rows[3].total_response_time_ms,
rows[3].timed_requests,
rows[3].spend
),
(301, 302, 303, 304, 305, 306, 307, 308, 309, 3.75)
);
assert_eq!(rows[3].date, "2026-01-15");
assert!(rows[3].updated_at > seeded_updated_at);
assert_eq!(rows[3].updated_at, expected_updated_at);
Ok(())
}
#[rstest]
#[case::distinct_tools("server/tool-a", "server/tool-b")]
#[case::empty_and_named_tool("", "server/tool-b")]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_end_user_spend_preserves_tool_attribution(
#[case] first_tool: &str,
#[case] second_tool: &str,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
let first = DailySpend {
entity_id: "daily-end-user-tools".to_owned(),
date: "2026-02-21".to_owned(),
api_key: "tool-key".to_owned(),
model: "tool-model".to_owned(),
custom_llm_provider: "tool-provider".to_owned(),
endpoint: "/mcp".to_owned(),
mcp_namespaced_tool_name: first_tool.to_owned(),
prompt_tokens: 2,
spend: 0.25,
..Default::default()
};
let second = DailySpend {
mcp_namespaced_tool_name: second_tool.to_owned(),
prompt_tokens: 5,
spend: 0.5,
..first.clone()
};
upsert_daily_end_user_spend(&mut *tx, &[second, first.clone()]).await?;
upsert_daily_end_user_spend(&mut *tx, &[first]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads stored tool attribution and counters"
)]
let rows: Vec<DailyEndUserRow> = sqlx::query_as(
r#"SELECT * FROM "LiteLLM_DailyEndUserSpend"
WHERE end_user_id = $1 AND date = $2 AND api_key = $3
ORDER BY mcp_namespaced_tool_name"#,
)
.bind("daily-end-user-tools")
.bind("2026-02-21")
.bind("tool-key")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 2);
assert_eq!(
rows[0].mcp_namespaced_tool_name.as_deref(),
Some(first_tool)
);
assert_eq!(
rows[1].mcp_namespaced_tool_name.as_deref(),
Some(second_tool)
);
assert_eq!((rows[0].prompt_tokens, rows[0].spend), (4, 0.5));
assert_eq!((rows[1].prompt_tokens, rows[1].spend), (5, 0.5));
assert_ne!(rows[0].id, rows[1].id);
Ok(())
}
#[rstest]
#[case::user(DailyTable::User, "daily-dispatch-user")]
#[case::team(DailyTable::Team, "daily-dispatch-team")]
#[case::end_user(DailyTable::EndUser, "daily-dispatch-end-user")]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_spend_dispatches_to_selected_table(
#[case] table: DailyTable,
#[case] entity_id: &'static str,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
let api_key = format!("{entity_id}-key");
let rows = [DailySpend {
entity_id: entity_id.to_owned(),
date: "2026-02-20".to_owned(),
api_key: api_key.clone(),
model: "dispatcher-model".to_owned(),
custom_llm_provider: "dispatcher-provider".to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: "/dispatcher".to_owned(),
model_group: Some("dispatcher-group".to_owned()),
prompt_tokens: 1,
completion_tokens: 2,
api_requests: 3,
successful_requests: 4,
failed_requests: 5,
cache_read_input_tokens: 6,
cache_creation_input_tokens: 7,
total_response_time_ms: 8,
timed_requests: 9,
spend: 0.125,
}];
upsert_daily_spend(&mut *tx, table, &rows).await?;
let count: i64 = match table {
DailyTable::User =>
{
#[expect(
clippy::disallowed_methods,
reason = "The integration test checks the user-table dispatcher route"
)]
sqlx::query_scalar(
r#"SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"
WHERE user_id = $1 AND date = $2 AND api_key = $3 AND model = $4"#,
)
.bind(entity_id)
.bind("2026-02-20")
.bind(&api_key)
.bind("dispatcher-model")
.fetch_one(&mut *tx)
.await?
}
DailyTable::Team =>
{
#[expect(
clippy::disallowed_methods,
reason = "The integration test checks the team-table dispatcher route"
)]
sqlx::query_scalar(
r#"SELECT COUNT(*) FROM "LiteLLM_DailyTeamSpend"
WHERE team_id = $1 AND date = $2 AND api_key = $3 AND model = $4"#,
)
.bind(entity_id)
.bind("2026-02-20")
.bind(&api_key)
.bind("dispatcher-model")
.fetch_one(&mut *tx)
.await?
}
DailyTable::EndUser =>
{
#[expect(
clippy::disallowed_methods,
reason = "The integration test checks the end-user-table dispatcher route"
)]
sqlx::query_scalar(
r#"SELECT COUNT(*) FROM "LiteLLM_DailyEndUserSpend"
WHERE end_user_id = $1 AND date = $2 AND api_key = $3 AND model = $4"#,
)
.bind(entity_id)
.bind("2026-02-20")
.bind(&api_key)
.bind("dispatcher-model")
.fetch_one(&mut *tx)
.await?
}
};
assert_eq!(count, 1);
Ok(())
}
}

485
litellm-rust/crates/db/src/spend/daily_team_spend.rs Normal file
View file

@ -0,0 +1,485 @@
use sqlx::PgExecutor;
use super::DailySpend;
use crate::Error;
pub async fn upsert_daily_team_spend(
executor: impl PgExecutor<'_>,
rows: &[DailySpend],
) -> Result<(), Error> {
if rows.is_empty() {
return Ok(());
}
macro_rules! column {
($field:ident) => {
rows.iter()
.map(|row| row.$field.clone())
.collect::<Vec<_>>()
};
}
sqlx::query_file!(
"queries/spend/upsert_daily_team_spend.sql",
&column!(entity_id),
&column!(date),
&column!(api_key),
&column!(model),
&column!(custom_llm_provider),
&column!(endpoint),
&column!(model_group) as &[Option<String>],
&column!(prompt_tokens),
&column!(completion_tokens),
&column!(api_requests),
&column!(successful_requests),
&column!(failed_requests),
&column!(cache_read_input_tokens),
&column!(cache_creation_input_tokens),
&column!(total_response_time_ms),
&column!(timed_requests),
&column!(spend),
&column!(mcp_namespaced_tool_name),
)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{DailySpend, upsert_daily_team_spend};
use crate::{DailyTeamRow, test_support::database_pool};
use rstest::rstest;
use sqlx::{PgPool, types::chrono::NaiveDateTime};
#[rstest]
#[case::distinct_tools("server.fetch", "server.search")]
#[case::default_and_named_tool("", "server.search")]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_team_spend_separates_tools_and_accumulates(
#[case] first_tool: &str,
#[case] second_tool: &str,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
let first = DailySpend {
entity_id: "daily-team-tools".to_owned(),
date: "2026-01-15".to_owned(),
api_key: "daily-key-tools".to_owned(),
model: "daily-model".to_owned(),
custom_llm_provider: "daily-provider".to_owned(),
endpoint: "/mcp".to_owned(),
mcp_namespaced_tool_name: first_tool.to_owned(),
api_requests: 1,
spend: 0.125,
..Default::default()
};
let second = DailySpend {
mcp_namespaced_tool_name: second_tool.to_owned(),
api_requests: 2,
spend: 0.25,
..first.clone()
};
upsert_daily_team_spend(&mut *tx, &[second.clone(), first.clone()]).await?;
upsert_daily_team_spend(
&mut *tx,
&[
DailySpend {
api_requests: 3,
spend: 0.5,
..first
},
DailySpend {
api_requests: 5,
spend: 1.0,
..second
},
],
)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads stored tool attribution and counters"
)]
let rows = sqlx::query_as::<_, DailyTeamRow>(
r#"SELECT * FROM "LiteLLM_DailyTeamSpend"
WHERE team_id = $1 AND date = $2
ORDER BY mcp_namespaced_tool_name"#,
)
.bind("daily-team-tools")
.bind("2026-01-15")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 2);
assert_eq!(
rows[0].mcp_namespaced_tool_name.as_deref(),
Some(first_tool)
);
assert_eq!(
rows[1].mcp_namespaced_tool_name.as_deref(),
Some(second_tool)
);
assert_eq!((rows[0].api_requests, rows[0].spend), (4, 0.625));
assert_eq!((rows[1].api_requests, rows[1].spend), (7, 1.25));
Ok(())
}
#[rstest]
#[case::full_keys(
"daily-team-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
Some("first-group")
)]
#[case::initial_model_group_none(
"daily-team-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
None
)]
#[case::empty_optional_keys("", "", "", "", "", Some(""))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_team_spend_upsert_preserves_keys_and_accumulates(
#[case] team_id: &'static str,
#[case] api_key: &'static str,
#[case] model: &'static str,
#[case] custom_llm_provider: &'static str,
#[case] endpoint: &'static str,
#[case] first_model_group: Option<&'static str>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test exercises UTC timestamp behavior under a non-UTC session"
)]
sqlx::query("SET LOCAL TIME ZONE 'America/Los_Angeles'")
.execute(&mut *tx)
.await?;
let first_batch = [
DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: first_model_group.map(str::to_owned),
prompt_tokens: 1,
completion_tokens: 2,
api_requests: 3,
successful_requests: 4,
failed_requests: 5,
cache_read_input_tokens: 6,
cache_creation_input_tokens: 7,
total_response_time_ms: 8,
timed_requests: 9,
spend: 0.125,
},
DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: "endpoint-sentinel".to_owned(),
model_group: Some("sentinel-endpoint-group".to_owned()),
prompt_tokens: 101,
completion_tokens: 102,
api_requests: 103,
successful_requests: 104,
failed_requests: 105,
cache_read_input_tokens: 106,
cache_creation_input_tokens: 107,
total_response_time_ms: 108,
timed_requests: 109,
spend: 1.25,
},
DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: "api-key-sentinel".to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-key-group".to_owned()),
prompt_tokens: 201,
completion_tokens: 202,
api_requests: 203,
successful_requests: 204,
failed_requests: 205,
cache_read_input_tokens: 206,
cache_creation_input_tokens: 207,
total_response_time_ms: 208,
timed_requests: 209,
spend: 2.5,
},
DailySpend {
entity_id: "entity-sentinel".to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-entity-group".to_owned()),
prompt_tokens: 301,
completion_tokens: 302,
api_requests: 303,
successful_requests: 304,
failed_requests: 305,
cache_read_input_tokens: 306,
cache_creation_input_tokens: 307,
total_response_time_ms: 308,
timed_requests: 309,
spend: 3.75,
},
];
upsert_daily_team_spend(&mut *tx, &first_batch).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the inserted row identity before conflict handling"
)]
let (initial_id, initial_created_at): (String, NaiveDateTime) = sqlx::query_as(
r#"SELECT id, created_at
FROM "LiteLLM_DailyTeamSpend"
WHERE team_id = $1 AND date = $2 AND api_key = $3 AND model = $4
AND custom_llm_provider = $5 AND endpoint = $6"#,
)
.bind(team_id)
.bind("2026-01-15")
.bind(api_key)
.bind(model)
.bind(custom_llm_provider)
.bind(endpoint)
.fetch_one(&mut *tx)
.await?;
assert!(!initial_id.is_empty());
#[expect(
clippy::disallowed_methods,
reason = "The integration test seeds updated_at before exercising the conflict update"
)]
let seeded_updated_at: NaiveDateTime = sqlx::query_scalar(
r#"UPDATE "LiteLLM_DailyTeamSpend"
SET updated_at = TIMESTAMP '2000-01-01'
WHERE id = $1
RETURNING updated_at"#,
)
.bind(&initial_id)
.fetch_one(&mut *tx)
.await?;
let increment = [DailySpend {
entity_id: team_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("second-group".to_owned()),
prompt_tokens: 10,
completion_tokens: 20,
api_requests: 30,
successful_requests: 40,
failed_requests: 50,
cache_read_input_tokens: 60,
cache_creation_input_tokens: 70,
total_response_time_ms: 80,
timed_requests: 90,
spend: 0.625,
}];
upsert_daily_team_spend(&mut *tx, &increment).await?;
upsert_daily_team_spend(&mut *tx, &[]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the transaction timestamp for assertions"
)]
let expected_updated_at: NaiveDateTime =
sqlx::query_scalar("SELECT (NOW() AT TIME ZONE 'UTC')::timestamp(3)")
.fetch_one(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads all rows to verify attribution and accumulation"
)]
let rows: Vec<DailyTeamRow> = sqlx::query_as(
r#"SELECT id, team_id, api_key, model, custom_llm_provider, endpoint,
model_group, mcp_namespaced_tool_name, prompt_tokens,
completion_tokens, api_requests, successful_requests,
failed_requests, cache_read_input_tokens,
cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, date, created_at, updated_at
FROM "LiteLLM_DailyTeamSpend"
WHERE date = $1 AND model = $2 AND custom_llm_provider = $3
AND mcp_namespaced_tool_name = '' AND team_id = ANY($4)
ORDER BY CASE
WHEN team_id = $5 AND api_key = $6 AND endpoint = $7 THEN 0
WHEN team_id = $5 AND api_key = $6 AND endpoint = $9 THEN 1
WHEN team_id = $5 AND api_key = $8 AND endpoint = $7 THEN 2
WHEN team_id = $10 AND api_key = $6 AND endpoint = $7 THEN 3
ELSE 4
END"#,
)
.bind("2026-01-15")
.bind(model)
.bind(custom_llm_provider)
.bind(vec![team_id, "entity-sentinel"])
.bind(team_id)
.bind(api_key)
.bind(endpoint)
.bind("api-key-sentinel")
.bind("endpoint-sentinel")
.bind("entity-sentinel")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 4);
assert_eq!(rows[0].id, initial_id);
assert_eq!(rows[0].team_id.as_deref(), Some(team_id));
assert_eq!(rows[0].api_key, api_key);
assert_eq!(rows[0].model.as_deref(), Some(model));
assert_eq!(
rows[0].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[0].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[0].model_group.as_deref(), first_model_group);
assert_eq!(rows[0].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[0].prompt_tokens,
rows[0].completion_tokens,
rows[0].api_requests,
rows[0].successful_requests,
rows[0].failed_requests,
rows[0].cache_read_input_tokens,
rows[0].cache_creation_input_tokens,
rows[0].total_response_time_ms,
rows[0].timed_requests,
rows[0].spend
),
(11, 22, 33, 44, 55, 66, 77, 88, 99, 0.75)
);
assert_eq!(rows[0].date, "2026-01-15");
assert_eq!(rows[0].created_at, initial_created_at);
assert!(rows[0].updated_at > seeded_updated_at);
assert_eq!(rows[0].updated_at, expected_updated_at);
assert!(!rows[1].id.is_empty());
assert_eq!(rows[1].team_id.as_deref(), Some(team_id));
assert_eq!(rows[1].api_key, api_key);
assert_eq!(rows[1].model.as_deref(), Some(model));
assert_eq!(
rows[1].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[1].endpoint.as_deref(), Some("endpoint-sentinel"));
assert_eq!(
rows[1].model_group.as_deref(),
Some("sentinel-endpoint-group")
);
assert_eq!(rows[1].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[1].prompt_tokens,
rows[1].completion_tokens,
rows[1].api_requests,
rows[1].successful_requests,
rows[1].failed_requests,
rows[1].cache_read_input_tokens,
rows[1].cache_creation_input_tokens,
rows[1].total_response_time_ms,
rows[1].timed_requests,
rows[1].spend
),
(101, 102, 103, 104, 105, 106, 107, 108, 109, 1.25)
);
assert_eq!(rows[1].date, "2026-01-15");
assert!(rows[1].updated_at > seeded_updated_at);
assert_eq!(rows[1].updated_at, expected_updated_at);
assert!(!rows[2].id.is_empty());
assert_eq!(rows[2].team_id.as_deref(), Some(team_id));
assert_eq!(rows[2].api_key, "api-key-sentinel");
assert_eq!(rows[2].model.as_deref(), Some(model));
assert_eq!(
rows[2].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[2].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[2].model_group.as_deref(), Some("sentinel-key-group"));
assert_eq!(rows[2].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[2].prompt_tokens,
rows[2].completion_tokens,
rows[2].api_requests,
rows[2].successful_requests,
rows[2].failed_requests,
rows[2].cache_read_input_tokens,
rows[2].cache_creation_input_tokens,
rows[2].total_response_time_ms,
rows[2].timed_requests,
rows[2].spend
),
(201, 202, 203, 204, 205, 206, 207, 208, 209, 2.5)
);
assert_eq!(rows[2].date, "2026-01-15");
assert!(rows[2].updated_at > seeded_updated_at);
assert_eq!(rows[2].updated_at, expected_updated_at);
assert!(!rows[3].id.is_empty());
assert_eq!(rows[3].team_id.as_deref(), Some("entity-sentinel"));
assert_eq!(rows[3].api_key, api_key);
assert_eq!(rows[3].model.as_deref(), Some(model));
assert_eq!(
rows[3].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[3].endpoint.as_deref(), Some(endpoint));
assert_eq!(
rows[3].model_group.as_deref(),
Some("sentinel-entity-group")
);
assert_eq!(rows[3].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[3].prompt_tokens,
rows[3].completion_tokens,
rows[3].api_requests,
rows[3].successful_requests,
rows[3].failed_requests,
rows[3].cache_read_input_tokens,
rows[3].cache_creation_input_tokens,
rows[3].total_response_time_ms,
rows[3].timed_requests,
rows[3].spend
),
(301, 302, 303, 304, 305, 306, 307, 308, 309, 3.75)
);
assert_eq!(rows[3].date, "2026-01-15");
assert!(rows[3].updated_at > seeded_updated_at);
assert_eq!(rows[3].updated_at, expected_updated_at);
Ok(())
}
}

500
litellm-rust/crates/db/src/spend/daily_user_spend.rs Normal file
View file

@ -0,0 +1,500 @@
use sqlx::PgExecutor;
use crate::Error;
#[macro_rules_attribute::apply(crate::db_input)]
#[derive(Default)]
pub struct DailySpend {
pub entity_id: String,
pub date: String,
pub api_key: String,
pub model: String,
pub custom_llm_provider: String,
pub mcp_namespaced_tool_name: String,
pub endpoint: String,
pub model_group: Option<String>,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub api_requests: i64,
pub successful_requests: i64,
pub failed_requests: i64,
pub cache_read_input_tokens: i64,
pub cache_creation_input_tokens: i64,
pub total_response_time_ms: i64,
pub timed_requests: i64,
pub spend: f64,
}
pub async fn upsert_daily_user_spend(
executor: impl PgExecutor<'_>,
rows: &[DailySpend],
) -> Result<(), Error> {
if rows.is_empty() {
return Ok(());
}
macro_rules! column {
($field:ident) => {
rows.iter()
.map(|row| row.$field.clone())
.collect::<Vec<_>>()
};
}
sqlx::query_file!(
"queries/spend/upsert_daily_user_spend.sql",
&column!(entity_id),
&column!(date),
&column!(api_key),
&column!(model),
&column!(custom_llm_provider),
&column!(endpoint),
&column!(model_group) as &[Option<String>],
&column!(prompt_tokens),
&column!(completion_tokens),
&column!(api_requests),
&column!(successful_requests),
&column!(failed_requests),
&column!(cache_read_input_tokens),
&column!(cache_creation_input_tokens),
&column!(total_response_time_ms),
&column!(timed_requests),
&column!(spend),
&column!(mcp_namespaced_tool_name),
)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{DailySpend, upsert_daily_user_spend};
use crate::{DailyUserRow, test_support::database_pool};
use rstest::rstest;
use sqlx::{PgPool, types::chrono::NaiveDateTime};
#[rstest]
#[case::distinct_tools("server/search", "server/write")]
#[case::default_empty_tool("", "server/search")]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_user_spend_separates_tools_and_accumulates(
#[case] first_tool: &str,
#[case] second_tool: &str,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
let first = DailySpend {
entity_id: format!("daily-user-tools-{first_tool}"),
date: "2026-01-15".to_owned(),
api_key: "tool-key".to_owned(),
model: "tool-model".to_owned(),
custom_llm_provider: "tool-provider".to_owned(),
endpoint: "/mcp".to_owned(),
model_group: Some("tool-group".to_owned()),
mcp_namespaced_tool_name: first_tool.to_owned(),
prompt_tokens: 3,
api_requests: 1,
spend: 0.25,
..Default::default()
};
let second = DailySpend {
mcp_namespaced_tool_name: second_tool.to_owned(),
prompt_tokens: 7,
api_requests: 2,
spend: 0.5,
..first.clone()
};
upsert_daily_user_spend(&mut *tx, &[second.clone(), first.clone()]).await?;
upsert_daily_user_spend(&mut *tx, &[first.clone(), second]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads stored rows to verify tool attribution"
)]
let rows: Vec<DailyUserRow> = sqlx::query_as(
r#"SELECT * FROM "LiteLLM_DailyUserSpend"
WHERE user_id = $1
ORDER BY mcp_namespaced_tool_name, endpoint"#,
)
.bind(&first.entity_id)
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 2);
assert_eq!(
rows[0].mcp_namespaced_tool_name.as_deref(),
Some(first_tool)
);
assert_eq!(
rows[1].mcp_namespaced_tool_name.as_deref(),
Some(second_tool)
);
assert_eq!(
(rows[0].prompt_tokens, rows[0].api_requests, rows[0].spend),
(6, 2, 0.5)
);
assert_eq!(
(rows[1].prompt_tokens, rows[1].api_requests, rows[1].spend),
(14, 4, 1.0)
);
Ok(())
}
#[rstest]
#[case::full_keys(
"daily-user-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
Some("first-group")
)]
#[case::initial_model_group_none(
"daily-user-base",
"daily-key-base",
"daily-model",
"daily-provider",
"/chat/completions",
None
)]
#[case::empty_optional_keys("", "", "", "", "", Some(""))]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn daily_user_spend_upsert_preserves_keys_and_accumulates(
#[case] user_id: &'static str,
#[case] api_key: &'static str,
#[case] model: &'static str,
#[case] custom_llm_provider: &'static str,
#[case] endpoint: &'static str,
#[case] first_model_group: Option<&'static str>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test exercises UTC timestamp behavior under a non-UTC session"
)]
sqlx::query("SET LOCAL TIME ZONE 'America/Los_Angeles'")
.execute(&mut *tx)
.await?;
let first_batch = [
DailySpend {
entity_id: user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: first_model_group.map(str::to_owned),
prompt_tokens: 1,
completion_tokens: 2,
api_requests: 3,
successful_requests: 4,
failed_requests: 5,
cache_read_input_tokens: 6,
cache_creation_input_tokens: 7,
total_response_time_ms: 8,
timed_requests: 9,
spend: 0.125,
},
DailySpend {
entity_id: user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: "endpoint-sentinel".to_owned(),
model_group: Some("sentinel-endpoint-group".to_owned()),
prompt_tokens: 101,
completion_tokens: 102,
api_requests: 103,
successful_requests: 104,
failed_requests: 105,
cache_read_input_tokens: 106,
cache_creation_input_tokens: 107,
total_response_time_ms: 108,
timed_requests: 109,
spend: 1.25,
},
DailySpend {
entity_id: user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: "api-key-sentinel".to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-key-group".to_owned()),
prompt_tokens: 201,
completion_tokens: 202,
api_requests: 203,
successful_requests: 204,
failed_requests: 205,
cache_read_input_tokens: 206,
cache_creation_input_tokens: 207,
total_response_time_ms: 208,
timed_requests: 209,
spend: 2.5,
},
DailySpend {
entity_id: "entity-sentinel".to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("sentinel-entity-group".to_owned()),
prompt_tokens: 301,
completion_tokens: 302,
api_requests: 303,
successful_requests: 304,
failed_requests: 305,
cache_read_input_tokens: 306,
cache_creation_input_tokens: 307,
total_response_time_ms: 308,
timed_requests: 309,
spend: 3.75,
},
];
upsert_daily_user_spend(&mut *tx, &first_batch).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the inserted row identity before conflict handling"
)]
let (initial_id, initial_created_at): (String, NaiveDateTime) = sqlx::query_as(
r#"SELECT id, created_at
FROM "LiteLLM_DailyUserSpend"
WHERE user_id = $1 AND date = $2 AND api_key = $3 AND model = $4
AND custom_llm_provider = $5 AND endpoint = $6"#,
)
.bind(user_id)
.bind("2026-01-15")
.bind(api_key)
.bind(model)
.bind(custom_llm_provider)
.bind(endpoint)
.fetch_one(&mut *tx)
.await?;
assert!(!initial_id.is_empty());
#[expect(
clippy::disallowed_methods,
reason = "The integration test seeds updated_at before exercising the conflict update"
)]
let seeded_updated_at: NaiveDateTime = sqlx::query_scalar(
r#"UPDATE "LiteLLM_DailyUserSpend"
SET updated_at = TIMESTAMP '2000-01-01'
WHERE id = $1
RETURNING updated_at"#,
)
.bind(&initial_id)
.fetch_one(&mut *tx)
.await?;
let increment = [DailySpend {
entity_id: user_id.to_owned(),
date: "2026-01-15".to_owned(),
api_key: api_key.to_owned(),
model: model.to_owned(),
custom_llm_provider: custom_llm_provider.to_owned(),
mcp_namespaced_tool_name: String::new(),
endpoint: endpoint.to_owned(),
model_group: Some("second-group".to_owned()),
prompt_tokens: 10,
completion_tokens: 20,
api_requests: 30,
successful_requests: 40,
failed_requests: 50,
cache_read_input_tokens: 60,
cache_creation_input_tokens: 70,
total_response_time_ms: 80,
timed_requests: 90,
spend: 0.625,
}];
upsert_daily_user_spend(&mut *tx, &increment).await?;
upsert_daily_user_spend(&mut *tx, &[]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test captures the transaction timestamp for assertions"
)]
let expected_updated_at: NaiveDateTime =
sqlx::query_scalar("SELECT (NOW() AT TIME ZONE 'UTC')::timestamp(3)")
.fetch_one(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "The integration test reads all rows to verify attribution and accumulation"
)]
let rows: Vec<DailyUserRow> = sqlx::query_as(
r#"SELECT id, user_id, api_key, model, custom_llm_provider, endpoint,
model_group, mcp_namespaced_tool_name, prompt_tokens,
completion_tokens, api_requests, successful_requests,
failed_requests, cache_read_input_tokens,
cache_creation_input_tokens, total_response_time_ms,
timed_requests, spend, date, created_at, updated_at
FROM "LiteLLM_DailyUserSpend"
WHERE date = $1 AND model = $2 AND custom_llm_provider = $3
AND mcp_namespaced_tool_name = '' AND user_id = ANY($4)
ORDER BY CASE
WHEN user_id = $5 AND api_key = $6 AND endpoint = $7 THEN 0
WHEN user_id = $5 AND api_key = $6 AND endpoint = $9 THEN 1
WHEN user_id = $5 AND api_key = $8 AND endpoint = $7 THEN 2
WHEN user_id = $10 AND api_key = $6 AND endpoint = $7 THEN 3
ELSE 4
END"#,
)
.bind("2026-01-15")
.bind(model)
.bind(custom_llm_provider)
.bind(vec![user_id, "entity-sentinel"])
.bind(user_id)
.bind(api_key)
.bind(endpoint)
.bind("api-key-sentinel")
.bind("endpoint-sentinel")
.bind("entity-sentinel")
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows.len(), 4);
assert_eq!(rows[0].id, initial_id);
assert_eq!(rows[0].user_id.as_deref(), Some(user_id));
assert_eq!(rows[0].api_key, api_key);
assert_eq!(rows[0].model.as_deref(), Some(model));
assert_eq!(
rows[0].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[0].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[0].model_group.as_deref(), first_model_group);
assert_eq!(rows[0].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[0].prompt_tokens,
rows[0].completion_tokens,
rows[0].api_requests,
rows[0].successful_requests,
rows[0].failed_requests,
rows[0].cache_read_input_tokens,
rows[0].cache_creation_input_tokens,
rows[0].total_response_time_ms,
rows[0].timed_requests,
rows[0].spend
),
(11, 22, 33, 44, 55, 66, 77, 88, 99, 0.75)
);
assert_eq!(rows[0].date, "2026-01-15");
assert_eq!(rows[0].created_at, initial_created_at);
assert!(rows[0].updated_at > seeded_updated_at);
assert_eq!(rows[0].updated_at, expected_updated_at);
assert!(!rows[1].id.is_empty());
assert_eq!(rows[1].user_id.as_deref(), Some(user_id));
assert_eq!(rows[1].api_key, api_key);
assert_eq!(rows[1].model.as_deref(), Some(model));
assert_eq!(
rows[1].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[1].endpoint.as_deref(), Some("endpoint-sentinel"));
assert_eq!(
rows[1].model_group.as_deref(),
Some("sentinel-endpoint-group")
);
assert_eq!(rows[1].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[1].prompt_tokens,
rows[1].completion_tokens,
rows[1].api_requests,
rows[1].successful_requests,
rows[1].failed_requests,
rows[1].cache_read_input_tokens,
rows[1].cache_creation_input_tokens,
rows[1].total_response_time_ms,
rows[1].timed_requests,
rows[1].spend
),
(101, 102, 103, 104, 105, 106, 107, 108, 109, 1.25)
);
assert_eq!(rows[1].date, "2026-01-15");
assert!(rows[1].updated_at > seeded_updated_at);
assert_eq!(rows[1].updated_at, expected_updated_at);
assert!(!rows[2].id.is_empty());
assert_eq!(rows[2].user_id.as_deref(), Some(user_id));
assert_eq!(rows[2].api_key, "api-key-sentinel");
assert_eq!(rows[2].model.as_deref(), Some(model));
assert_eq!(
rows[2].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[2].endpoint.as_deref(), Some(endpoint));
assert_eq!(rows[2].model_group.as_deref(), Some("sentinel-key-group"));
assert_eq!(rows[2].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[2].prompt_tokens,
rows[2].completion_tokens,
rows[2].api_requests,
rows[2].successful_requests,
rows[2].failed_requests,
rows[2].cache_read_input_tokens,
rows[2].cache_creation_input_tokens,
rows[2].total_response_time_ms,
rows[2].timed_requests,
rows[2].spend
),
(201, 202, 203, 204, 205, 206, 207, 208, 209, 2.5)
);
assert_eq!(rows[2].date, "2026-01-15");
assert!(rows[2].updated_at > seeded_updated_at);
assert_eq!(rows[2].updated_at, expected_updated_at);
assert!(!rows[3].id.is_empty());
assert_eq!(rows[3].user_id.as_deref(), Some("entity-sentinel"));
assert_eq!(rows[3].api_key, api_key);
assert_eq!(rows[3].model.as_deref(), Some(model));
assert_eq!(
rows[3].custom_llm_provider.as_deref(),
Some(custom_llm_provider)
);
assert_eq!(rows[3].endpoint.as_deref(), Some(endpoint));
assert_eq!(
rows[3].model_group.as_deref(),
Some("sentinel-entity-group")
);
assert_eq!(rows[3].mcp_namespaced_tool_name.as_deref(), Some(""));
assert_eq!(
(
rows[3].prompt_tokens,
rows[3].completion_tokens,
rows[3].api_requests,
rows[3].successful_requests,
rows[3].failed_requests,
rows[3].cache_read_input_tokens,
rows[3].cache_creation_input_tokens,
rows[3].total_response_time_ms,
rows[3].timed_requests,
rows[3].spend
),
(301, 302, 303, 304, 305, 306, 307, 308, 309, 3.75)
);
assert_eq!(rows[3].date, "2026-01-15");
assert!(rows[3].updated_at > seeded_updated_at);
assert_eq!(rows[3].updated_at, expected_updated_at);
Ok(())
}
}

197
litellm-rust/crates/db/src/spend/end_user_spend.rs Normal file
View file

@ -0,0 +1,197 @@
use sqlx::PgExecutor;
use super::EntitySpend;
use crate::Error;
pub async fn add_end_user_spend(
executor: impl PgExecutor<'_>,
spend: &EntitySpend,
) -> Result<(), Error> {
spend.validate()?;
if spend.is_empty() {
return Ok(());
}
sqlx::query_file!(
"queries/spend/add_end_user_spend.sql",
&spend.ids,
&spend.costs
)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{EntitySpend, add_end_user_spend};
use crate::test_support::database_pool;
use rstest::rstest;
use sqlx::PgPool;
#[rstest]
#[case::fractional_costs(0.25, 0.75, false, 2.5, 1.5)]
#[case::zero_costs(0.0, 0.0, false, 2.0, 0.0)]
#[case::duplicate_ids(0.25, 0.75, true, 3.0, 3.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn end_user_spend_matches_python_upsert(
#[case] existing_cost: f64,
#[case] new_cost: f64,
#[case] duplicate_ids: bool,
#[case] expected_existing_spend: f64,
#[case] expected_new_spend: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _budget = sqlx::query(
r#"INSERT INTO "LiteLLM_BudgetTable" (budget_id, created_by, updated_by)
VALUES ('end-user-spend-budget', 'db-query-test', 'db-query-test')"#,
)
.execute(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _users = sqlx::query(
r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked, budget_id)
VALUES ('end-user-spend-z', 2, true, 'end-user-spend-budget'),
('end-user-spend-untouched', 100, true, NULL)"#,
)
.execute(&mut *tx)
.await?;
let spend = if duplicate_ids {
EntitySpend {
ids: vec![
"end-user-spend-z".into(),
"end-user-spend-a".into(),
"end-user-spend-z".into(),
"end-user-spend-a".into(),
],
costs: vec![existing_cost, new_cost, existing_cost, new_cost],
}
} else {
EntitySpend {
ids: vec!["end-user-spend-z".into(), "end-user-spend-a".into()],
costs: vec![existing_cost, new_cost],
}
};
add_end_user_spend(&mut *tx, &spend).await?;
add_end_user_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, f64, bool, Option<String>)> = sqlx::query_as(
r#"SELECT user_id, spend, blocked, budget_id FROM "LiteLLM_EndUserTable"
WHERE user_id = ANY($1) ORDER BY user_id"#,
)
.bind(vec![
"end-user-spend-a",
"end-user-spend-z",
"end-user-spend-untouched",
])
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
("end-user-spend-a".into(), expected_new_spend, false, None),
("end-user-spend-untouched".into(), 100.0, true, None),
(
"end-user-spend-z".into(),
expected_existing_spend,
true,
Some("end-user-spend-budget".into()),
),
]
);
Ok(())
}
#[rstest]
#[case::missing_cost(vec!["end-user-invalid-existing".into(), "end-user-invalid-new".into()], vec![1.0])]
#[case::extra_cost(vec!["end-user-invalid-existing".into()], vec![1.0, 2.0])]
#[case::empty_ids(vec![], vec![1.0])]
#[case::empty_costs(vec!["end-user-invalid-existing".into()], vec![])]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn malformed_end_user_spend_is_rejected_before_writes(
#[case] ids: Vec<String>,
#[case] costs: Vec<f64>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _user = sqlx::query(
r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked)
VALUES ('end-user-invalid-existing', 8.5, true)"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend { ids, costs };
let result = add_end_user_spend(&mut *tx, &spend).await;
assert!(matches!(
result,
Err(crate::Error::InvalidSpendBatch { ids, costs })
if ids == spend.ids.len() && costs == spend.costs.len()
));
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, f64, bool)> = sqlx::query_as(
r#"SELECT user_id, spend, blocked FROM "LiteLLM_EndUserTable"
WHERE user_id IN ('end-user-invalid-existing', 'end-user-invalid-new')"#,
)
.fetch_all(&mut *tx)
.await?;
assert_eq!(rows, vec![("end-user-invalid-existing".into(), 8.5, true)]);
Ok(())
}
#[rstest]
#[case::empty_batch()]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn empty_end_user_spend_is_a_noop(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _user = sqlx::query(
r#"INSERT INTO "LiteLLM_EndUserTable" (user_id, spend, blocked)
VALUES ('end-user-empty-spend', 8.5, true)"#,
)
.execute(&mut *tx)
.await?;
add_end_user_spend(&mut *tx, &EntitySpend::default()).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, bool, Option<String>) = sqlx::query_as(
r#"SELECT spend, blocked, budget_id FROM "LiteLLM_EndUserTable"
WHERE user_id = 'end-user-empty-spend'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (8.5, true, None));
Ok(())
}
}

309
litellm-rust/crates/db/src/spend/key_spend.rs Normal file
View file

@ -0,0 +1,309 @@
use sqlx::PgExecutor;
use std::collections::BTreeMap;
use crate::Error;
#[macro_rules_attribute::apply(crate::db_input)]
#[derive(Default)]
pub struct EntitySpend {
pub ids: Vec<String>,
pub costs: Vec<f64>,
}
impl EntitySpend {
pub(crate) fn validate(&self) -> Result<(), Error> {
if self.ids.len() != self.costs.len() {
return Err(Error::InvalidSpendBatch {
ids: self.ids.len(),
costs: self.costs.len(),
});
}
Ok(())
}
pub fn is_empty(&self) -> bool {
self.ids.is_empty()
}
}
impl FromIterator<(String, f64)> for EntitySpend {
fn from_iter<I: IntoIterator<Item = (String, f64)>>(entries: I) -> Self {
let entries = entries
.into_iter()
.fold(BTreeMap::new(), |mut sums, (id, cost)| {
*sums.entry(id).or_insert(0.0) += cost;
sums
});
let (ids, costs) = entries.into_iter().unzip();
Self { ids, costs }
}
}
pub async fn add_key_spend(
executor: impl PgExecutor<'_>,
spend: &EntitySpend,
) -> Result<(), Error> {
spend.validate()?;
if spend.is_empty() {
return Ok(());
}
sqlx::query_file!("queries/spend/add_key_spend.sql", &spend.ids, &spend.costs)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{EntitySpend, add_key_spend};
use crate::{Error, test_support::database_pool};
use rstest::rstest;
use sqlx::PgPool;
#[rstest]
#[case::direct(false)]
#[case::collected(true)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn duplicate_key_spend_accumulates(
#[case] collected: bool,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
sqlx::query(
r#"INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, updated_at)
VALUES ('key-spend-duplicate', 2, 7, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = if collected {
[
("key-spend-duplicate".into(), 0.25),
("key-spend-duplicate".into(), 0.75),
]
.into_iter()
.collect()
} else {
EntitySpend {
ids: vec!["key-spend-duplicate".into(), "key-spend-duplicate".into()],
costs: vec![0.25, 0.75],
}
};
add_key_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, f64) = sqlx::query_as(
r#"SELECT spend, total_spend FROM "LiteLLM_VerificationToken"
WHERE token = 'key-spend-duplicate'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (3.0, 8.0));
Ok(())
}
#[rstest]
#[case::extra_id(vec!["key-spend-invalid".into(), "key-spend-invalid".into()], vec![0.25])]
#[case::extra_cost(vec!["key-spend-invalid".into()], vec![0.25, 0.75])]
#[case::empty_ids(vec![], vec![0.25])]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn malformed_key_spend_leaves_transaction_usable(
#[case] ids: Vec<String>,
#[case] costs: Vec<f64>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
sqlx::query(
r#"INSERT INTO "LiteLLM_VerificationToken"
(token, spend, total_spend, last_active, updated_at)
VALUES ('key-spend-invalid', 2, 7, NULL, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend { ids, costs };
let result = add_key_spend(&mut *tx, &spend).await;
assert!(matches!(
result,
Err(Error::InvalidSpendBatch { ids, costs })
if ids == spend.ids.len() && costs == spend.costs.len()
));
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let unchanged: (f64, f64, bool, bool) = sqlx::query_as(
r#"SELECT spend, total_spend, last_active IS NULL,
updated_at = TIMESTAMP '2000-01-01'
FROM "LiteLLM_VerificationToken" WHERE token = 'key-spend-invalid'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(unchanged, (2.0, 7.0, true, true));
let valid = EntitySpend {
ids: vec!["key-spend-invalid".into()],
costs: vec![0.5],
};
add_key_spend(&mut *tx, &valid).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let updated: (f64, f64) = sqlx::query_as(
r#"SELECT spend, total_spend FROM "LiteLLM_VerificationToken"
WHERE token = 'key-spend-invalid'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(updated, (2.5, 7.5));
Ok(())
}
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn key_activity_uses_update_time_in_old_transaction(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
sqlx::query(
r#"INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, updated_at)
VALUES ('key-spend-old-transaction', 2, 7, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "Transaction timing setup is outside the compile-checked query set"
)]
sqlx::query("SELECT pg_sleep(0.02)")
.execute(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let boundary: chrono::NaiveDateTime = sqlx::query_scalar(
"SELECT date_trunc('milliseconds', clock_timestamp() AT TIME ZONE 'UTC')",
)
.fetch_one(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec!["key-spend-old-transaction".into()],
costs: vec![0.25],
};
add_key_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (chrono::NaiveDateTime, chrono::NaiveDateTime) = sqlx::query_as(
r#"SELECT last_active, updated_at FROM "LiteLLM_VerificationToken"
WHERE token = 'key-spend-old-transaction'"#,
)
.fetch_one(&mut *tx)
.await?;
assert!(row.0 >= boundary, "last_active predates update boundary");
assert!(row.1 >= boundary, "updated_at predates update boundary");
Ok(())
}
#[rstest]
#[case::fractional_costs(0.25, 2.5, 7.5, 4.0, 12.0)]
#[case::zero_costs(0.0, 2.0, 7.0, 3.0, 11.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn key_spend_matches_python_update_many(
#[case] cost: f64,
#[case] expected_a_spend: f64,
#[case] expected_a_total: f64,
#[case] expected_z_spend: f64,
#[case] expected_z_total: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_VerificationToken"
(token, spend, total_spend, last_active, updated_at)
VALUES ('key-spend-a', 2, 7, NULL, '2000-01-01'),
('key-spend-z', 3, 11, NULL, '2000-01-01'),
('key-spend-untouched', 100, 200, NULL, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec![
"key-spend-z".into(),
"key-spend-a".into(),
"key-spend-missing".into(),
],
costs: vec![2.0 * cost, cost, 50.0],
};
add_key_spend(&mut *tx, &EntitySpend::default()).await?;
add_key_spend(&mut *tx, &spend).await?;
add_key_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, f64, f64, bool, bool)> = sqlx::query_as(
r#"SELECT token, spend, total_spend, last_active IS NOT NULL,
updated_at > TIMESTAMP '2000-01-01'
FROM "LiteLLM_VerificationToken"
WHERE token = ANY($1) ORDER BY token"#,
)
.bind(vec![
"key-spend-a",
"key-spend-z",
"key-spend-missing",
"key-spend-untouched",
])
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
(
"key-spend-a".into(),
expected_a_spend,
expected_a_total,
true,
true
),
("key-spend-untouched".into(), 100.0, 200.0, false, false),
(
"key-spend-z".into(),
expected_z_spend,
expected_z_total,
true,
true
),
]
);
Ok(())
}
}

21
litellm-rust/crates/db/src/spend/mod.rs Normal file
View file

@ -0,0 +1,21 @@
mod daily_end_user_spend;
mod daily_team_spend;
mod daily_user_spend;
mod end_user_spend;
mod key_spend;
mod spend_logs;
mod team_lock;
mod team_member_spend;
mod team_spend;
mod user_spend;
pub use daily_end_user_spend::{DailyTable, upsert_daily_end_user_spend, upsert_daily_spend};
pub use daily_team_spend::upsert_daily_team_spend;
pub use daily_user_spend::{DailySpend, upsert_daily_user_spend};
pub use end_user_spend::add_end_user_spend;
pub use key_spend::{EntitySpend, add_key_spend};
pub use spend_logs::{SpendLog, SpendLogRow, insert_spend_logs};
pub use team_lock::{begin, lock_team};
pub use team_member_spend::{TeamMemberSpend, add_team_member_spend};
pub use team_spend::add_team_spend;
pub use user_spend::add_user_spend;

517
litellm-rust/crates/db/src/spend/spend_logs.rs Normal file
View file

@ -0,0 +1,517 @@
use chrono::{DateTime, NaiveDateTime};
use sqlx::PgExecutor;
use crate::Error;
#[macro_rules_attribute::apply(crate::db_input)]
#[derive(Default)]
pub struct SpendLog {
pub request_id: String,
pub call_type: String,
pub api_key: String,
pub spend: f64,
pub total_tokens: i32,
pub prompt_tokens: i32,
pub completion_tokens: i32,
pub start_time: String,
pub end_time: String,
pub completion_start_time: Option<String>,
pub request_duration_ms: Option<i32>,
pub model: String,
pub model_id: String,
pub model_group: String,
pub custom_llm_provider: String,
pub api_base: String,
pub user: String,
pub metadata: serde_json::Value,
pub cache_hit: String,
pub cache_key: String,
pub team_id: String,
pub organization_id: String,
pub end_user: String,
pub requester_ip_address: Option<String>,
pub session_id: Option<String>,
pub status: String,
pub litellm_call_id: Option<String>,
}
#[macro_rules_attribute::apply(crate::db_row)]
pub struct SpendLogRow {
pub request_id: String,
pub call_type: String,
pub api_key: String,
pub spend: f64,
pub total_tokens: i32,
pub prompt_tokens: i32,
pub completion_tokens: i32,
#[sqlx(rename = "startTime")]
pub start_time: NaiveDateTime,
#[sqlx(rename = "endTime")]
pub end_time: NaiveDateTime,
#[sqlx(rename = "completionStartTime")]
pub completion_start_time: Option<NaiveDateTime>,
pub request_duration_ms: Option<i32>,
pub model: String,
pub model_id: Option<String>,
pub model_group: Option<String>,
pub custom_llm_provider: Option<String>,
pub api_base: Option<String>,
pub user: Option<String>,
pub metadata: Option<serde_json::Value>,
pub cache_hit: Option<String>,
pub cache_key: Option<String>,
pub request_tags: Option<serde_json::Value>,
pub team_id: Option<String>,
pub organization_id: Option<String>,
pub end_user: Option<String>,
pub requester_ip_address: Option<String>,
pub messages: Option<serde_json::Value>,
pub response: Option<serde_json::Value>,
pub proxy_server_request: Option<serde_json::Value>,
pub session_id: Option<String>,
pub status: Option<String>,
pub litellm_call_id: Option<String>,
}
fn parse_timestamp(value: &str) -> Result<NaiveDateTime, chrono::ParseError> {
DateTime::parse_from_rfc3339(value)
.map(|timestamp| timestamp.naive_utc())
.or_else(|_| NaiveDateTime::parse_from_str(value, "%Y-%m-%dT%H:%M:%S%.f"))
.or_else(|_| NaiveDateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S%.f"))
}
pub async fn insert_spend_logs(
executor: impl PgExecutor<'_>,
logs: &[SpendLog],
) -> Result<(), Error> {
if logs.is_empty() {
return Ok(());
}
let start_times = logs
.iter()
.map(|log| parse_timestamp(&log.start_time))
.collect::<Result<Vec<_>, _>>()?;
let end_times = logs
.iter()
.map(|log| parse_timestamp(&log.end_time))
.collect::<Result<Vec<_>, _>>()?;
let completion_start_times = logs
.iter()
.map(|log| {
log.completion_start_time
.as_deref()
.map(parse_timestamp)
.transpose()
})
.collect::<Result<Vec<_>, _>>()?;
macro_rules! column {
($field:ident) => {
logs.iter()
.map(|log| log.$field.clone())
.collect::<Vec<_>>()
};
}
sqlx::query_file!(
"queries/spend/insert_spend_logs.sql",
&column!(request_id),
&column!(call_type),
&column!(api_key),
&column!(spend),
&column!(total_tokens),
&column!(prompt_tokens),
&column!(completion_tokens),
&start_times,
&end_times,
&completion_start_times as &[Option<NaiveDateTime>],
&column!(request_duration_ms) as &[Option<i32>],
&column!(model),
&column!(model_id),
&column!(model_group),
&column!(custom_llm_provider),
&column!(api_base),
&column!(user),
&column!(metadata),
&column!(cache_hit),
&column!(cache_key),
&column!(team_id),
&column!(organization_id),
&column!(end_user),
&column!(requester_ip_address) as &[Option<String>],
&column!(session_id) as &[Option<String>],
&column!(status),
&column!(litellm_call_id) as &[Option<String>],
)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{SpendLog, SpendLogRow, insert_spend_logs};
use crate::test_support::database_pool;
use rstest::{fixture, rstest};
use serde_json::json;
use sqlx::{Acquire, Connection, PgExecutor, PgPool, types::chrono::NaiveDateTime};
#[fixture]
fn logs() -> [SpendLog; 2] {
[
SpendLog {
request_id: "spend-logs-a".into(),
call_type: "acompletion".into(),
api_key: "key-a".into(),
spend: 0.125,
total_tokens: 17,
prompt_tokens: 11,
completion_tokens: 6,
start_time: "2026-01-02 03:04:05.123".into(),
end_time: "2026-01-02 03:04:06.789".into(),
completion_start_time: Some("2026-01-02 03:04:05.456".into()),
request_duration_ms: Some(1666),
model: "model-a".into(),
model_id: "deployment-a".into(),
model_group: "group-a".into(),
custom_llm_provider: "provider-a".into(),
api_base: "https://a.example.test".into(),
user: "user-a".into(),
metadata: json!({"tenant": "a", "nested": {"tokens": 11}, "flag": false}),
cache_hit: "False".into(),
cache_key: "cache-a".into(),
team_id: "team-a".into(),
organization_id: "org-a".into(),
end_user: "end-user-a".into(),
requester_ip_address: Some("192.0.2.1".into()),
session_id: Some("session-a".into()),
status: "success".into(),
litellm_call_id: Some("call-a".into()),
},
SpendLog {
request_id: "spend-logs-b".into(),
call_type: "aembedding".into(),
api_key: "key-b".into(),
spend: 2.75,
total_tokens: 43,
prompt_tokens: 29,
completion_tokens: 14,
start_time: "2026-02-03 14:15:16.234".into(),
end_time: "2026-02-03 14:15:18.567".into(),
completion_start_time: None,
request_duration_ms: None,
model: "model-b".into(),
model_id: "deployment-b".into(),
model_group: "group-b".into(),
custom_llm_provider: "provider-b".into(),
api_base: "https://b.example.test".into(),
user: "user-b".into(),
metadata: json!({"tenant": "b", "tags": ["distinct", 7], "nullable": null}),
cache_hit: "True".into(),
cache_key: "cache-b".into(),
team_id: "team-b".into(),
organization_id: "org-b".into(),
end_user: "end-user-b".into(),
requester_ip_address: None,
session_id: None,
status: "failure".into(),
litellm_call_id: None,
},
]
}
async fn read_log(
executor: impl PgExecutor<'_>,
request_id: &str,
) -> Result<SpendLogRow, sqlx::Error> {
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
sqlx::query_as::<_, SpendLogRow>(
r#"SELECT * FROM "LiteLLM_SpendLogs" WHERE request_id = $1"#,
)
.bind(request_id)
.fetch_one(executor)
.await
}
fn expected_log(log: &SpendLog) -> Result<SpendLogRow, chrono::ParseError> {
Ok(SpendLogRow {
request_id: log.request_id.clone(),
call_type: log.call_type.clone(),
api_key: log.api_key.clone(),
spend: log.spend,
total_tokens: log.total_tokens,
prompt_tokens: log.prompt_tokens,
completion_tokens: log.completion_tokens,
start_time: log.start_time.replace(' ', "T").parse()?,
end_time: log.end_time.replace(' ', "T").parse()?,
completion_start_time: log
.completion_start_time
.as_ref()
.map(|value| value.replace(' ', "T").parse())
.transpose()?,
request_duration_ms: log.request_duration_ms,
model: log.model.clone(),
model_id: Some(log.model_id.clone()),
model_group: Some(log.model_group.clone()),
custom_llm_provider: Some(log.custom_llm_provider.clone()),
api_base: Some(log.api_base.clone()),
user: Some(log.user.clone()),
metadata: Some(log.metadata.clone()),
cache_hit: Some(log.cache_hit.clone()),
cache_key: Some(log.cache_key.clone()),
request_tags: Some(json!([])),
team_id: Some(log.team_id.clone()),
organization_id: Some(log.organization_id.clone()),
end_user: Some(log.end_user.clone()),
requester_ip_address: log.requester_ip_address.clone(),
messages: Some(json!({})),
response: Some(json!({})),
proxy_server_request: Some(json!({})),
session_id: log.session_id.clone(),
status: Some(log.status.clone()),
litellm_call_id: log.litellm_call_id.clone(),
})
}
#[rstest]
#[case::all_columns_and_duplicate()]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn inserts_distinct_rows_without_replacing_duplicates(
logs: [SpendLog; 2],
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
insert_spend_logs(&mut *tx, &logs).await?;
let replacement = SpendLog {
request_id: logs[0].request_id.clone(),
..logs[1].clone()
};
insert_spend_logs(&mut *tx, &[replacement]).await?;
assert_eq!(
read_log(&mut *tx, &logs[0].request_id).await?,
expected_log(&logs[0])?
);
assert_eq!(
read_log(&mut *tx, &logs[1].request_id).await?,
expected_log(&logs[1])?
);
tx.rollback().await?;
Ok(())
}
#[rstest]
#[case::partitioned_duplicate()]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn inserts_into_partitioned_logs_without_replacing_duplicates(
logs: [SpendLog; 2],
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "The test creates a temporary partitioned table at runtime"
)]
sqlx::raw_sql(
r#"CREATE TEMP TABLE "LiteLLM_SpendLogs" (
LIKE public."LiteLLM_SpendLogs" INCLUDING DEFAULTS INCLUDING GENERATED,
PRIMARY KEY (request_id, "startTime")
) PARTITION BY RANGE ("startTime");
CREATE TEMP TABLE spend_logs_default
PARTITION OF pg_temp."LiteLLM_SpendLogs" DEFAULT;
SET LOCAL search_path TO pg_temp, public;"#,
)
.execute(&mut *tx)
.await?;
tx.clear_cached_statements().await?;
insert_spend_logs(&mut *tx, &logs).await?;
let replacement = SpendLog {
spend: 99.0,
..logs[0].clone()
};
insert_spend_logs(&mut *tx, &[replacement]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The partitioned test fixture is created at runtime"
)]
let count: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM pg_temp."LiteLLM_SpendLogs""#)
.fetch_one(&mut *tx)
.await?;
assert_eq!(count.0, 2);
assert_eq!(
read_log(&mut *tx, &logs[0].request_id).await?,
expected_log(&logs[0])?
);
assert_eq!(
read_log(&mut *tx, &logs[1].request_id).await?,
expected_log(&logs[1])?
);
let distinct_partition_key = SpendLog {
request_id: logs[0].request_id.clone(),
..logs[1].clone()
};
insert_spend_logs(&mut *tx, &[distinct_partition_key]).await?;
#[expect(
clippy::disallowed_methods,
reason = "The partitioned test fixture is created at runtime"
)]
let count: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM pg_temp."LiteLLM_SpendLogs""#)
.fetch_one(&mut *tx)
.await?;
assert_eq!(count.0, 3);
tx.rollback().await?;
Ok(())
}
#[rstest]
#[case::positive_offset_across_midnight(
["2026-01-02T00:04:05.1236+02:00", "2026-01-02T00:04:06.7894+02:00", "2026-01-02T00:04:05.4567+02:00"],
["2026-01-01T22:04:05.124", "2026-01-01T22:04:06.789", "2026-01-01T22:04:05.457"]
)]
#[case::negative_offset_across_midnight(
["2026-01-01T23:04:05-02:00", "2026-01-01T23:04:06-02:00", "2026-01-01T23:04:05.456-02:00"],
["2026-01-02T01:04:05", "2026-01-02T01:04:06", "2026-01-02T01:04:05.456"]
)]
#[case::utc(
["2026-01-02T03:04:05Z", "2026-01-02T03:04:06Z", "2026-01-02T03:04:05.456Z"],
["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"]
)]
#[case::naive_iso(
["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"],
["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"]
)]
#[case::naive_space(
["2026-01-02 03:04:05", "2026-01-02 03:04:06", "2026-01-02 03:04:05.456"],
["2026-01-02T03:04:05", "2026-01-02T03:04:06", "2026-01-02T03:04:05.456"]
)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn stores_utc_timestamps(
#[case] input: [&str; 3],
#[case] expected: [&str; 3],
logs: [SpendLog; 2],
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "The test sets a transaction-local timezone to verify timestamp handling"
)]
sqlx::raw_sql("SET LOCAL TIME ZONE 'America/Los_Angeles'")
.execute(&mut *tx)
.await?;
let log = SpendLog {
start_time: input[0].into(),
end_time: input[1].into(),
completion_start_time: Some(input[2].into()),
..logs[0].clone()
};
insert_spend_logs(&mut *tx, std::slice::from_ref(&log)).await?;
let row = read_log(&mut *tx, &log.request_id).await?;
assert_eq!(row.start_time, expected[0].parse::<NaiveDateTime>()?);
assert_eq!(row.end_time, expected[1].parse::<NaiveDateTime>()?);
assert_eq!(
row.completion_start_time,
Some(expected[2].parse::<NaiveDateTime>()?)
);
tx.rollback().await?;
Ok(())
}
#[rstest]
#[case::empty_batch()]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn empty_input_leaves_rows_unchanged(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let before: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM "LiteLLM_SpendLogs""#)
.fetch_one(&mut *tx)
.await?;
insert_spend_logs(&mut *tx, &[]).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let after: (i64,) = sqlx::query_as(r#"SELECT count(*) FROM "LiteLLM_SpendLogs""#)
.fetch_one(&mut *tx)
.await?;
assert_eq!(before, after);
tx.rollback().await?;
Ok(())
}
#[rstest]
#[case::invalid_timestamp("not-a-timestamp", "2026-01-02T03:04:06", None)]
#[case::invalid_end("2026-01-02T03:04:05", "not-a-timestamp", None)]
#[case::invalid_completion(
"2026-01-02T03:04:05",
"2026-01-02T03:04:06",
Some("not-a-timestamp")
)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn invalid_timestamp_allows_caller_rollback(
#[case] start_time: &str,
#[case] end_time: &str,
#[case] completion_start_time: Option<&str>,
logs: [SpendLog; 2],
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut outer = pool.begin().await?;
let mut tx = outer.begin().await?;
let valid = SpendLog {
request_id: "spend-logs-rollback-valid".into(),
..logs[0].clone()
};
let invalid = SpendLog {
request_id: "spend-logs-rollback-invalid".into(),
start_time: start_time.into(),
end_time: end_time.into(),
completion_start_time: completion_start_time.map(str::to_owned),
..logs[1].clone()
};
insert_spend_logs(&mut *tx, std::slice::from_ref(&valid)).await?;
let result = insert_spend_logs(&mut *tx, &[logs[0].clone(), invalid]).await;
assert!(matches!(result, Err(crate::Error::InvalidTimestamp(_))));
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let partial_batch: (i64,) = sqlx::query_as(
r#"SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1)"#,
)
.bind([logs[0].request_id.as_str(), "spend-logs-rollback-invalid"].as_slice())
.fetch_one(&mut *tx)
.await?;
assert_eq!(partial_batch.0, 0);
tx.rollback().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let remaining: (i64,) = sqlx::query_as(
r#"SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1)"#,
)
.bind([valid.request_id.as_str(), "spend-logs-rollback-invalid"].as_slice())
.fetch_one(&mut *outer)
.await?;
assert_eq!(remaining.0, 0);
outer.rollback().await?;
Ok(())
}
}

98
litellm-rust/crates/db/src/spend/team_lock.rs Normal file
View file

@ -0,0 +1,98 @@
use sqlx::{PgPool, Postgres, Transaction};
use crate::Error;
pub async fn begin(pool: &PgPool) -> Result<Transaction<'static, Postgres>, Error> {
let mut transaction = pool.begin().await?;
sqlx::query!("SELECT set_config('lock_timeout', '5000ms', true)")
.fetch_one(&mut *transaction)
.await?;
Ok(transaction)
}
pub async fn lock_team(
transaction: &mut Transaction<'_, Postgres>,
team_id: &str,
) -> Result<(), Error> {
sqlx::query_file!("queries/spend/lock_team.sql", team_id)
.fetch_one(&mut **transaction)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{begin, lock_team};
use crate::test_support::database_pool;
use rstest::rstest;
use sqlx::{PgPool, postgres::PgPoolOptions};
#[rstest]
#[case::transaction_scoped("team-lock-scope")]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn team_lock_is_held_until_rollback(
#[case] team_id: &str,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let database_pool = database_pool?;
let pool = PgPoolOptions::new()
.max_connections(2)
.connect_with(database_pool.connect_options().as_ref().clone())
.await?;
let mut first = begin(&pool).await?;
let mut second = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let original_timeout: String = sqlx::query_scalar("SELECT current_setting('lock_timeout')")
.fetch_one(&mut *second)
.await?;
lock_team(&mut first, team_id).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let (backend, timeout): (i32, String) =
sqlx::query_as("SELECT pg_backend_pid(), current_setting('lock_timeout')")
.fetch_one(&mut *first)
.await?;
assert_eq!(timeout, "5s");
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let held: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock(hashtext($1))")
.bind(team_id)
.fetch_one(&mut *second)
.await?;
assert!(!held);
first.rollback().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let released: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock(hashtext($1))")
.bind(team_id)
.fetch_one(&mut *second)
.await?;
assert!(released);
let mut next = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let (next_backend, next_timeout): (i32, String) =
sqlx::query_as("SELECT pg_backend_pid(), current_setting('lock_timeout')")
.fetch_one(&mut *next)
.await?;
assert_eq!(next_backend, backend);
assert_eq!(next_timeout, original_timeout);
next.rollback().await?;
second.rollback().await?;
Ok(())
}
}

162
litellm-rust/crates/db/src/spend/team_member_spend.rs Normal file
View file

@ -0,0 +1,162 @@
use sqlx::{Postgres, Transaction};
use super::lock_team;
use crate::Error;
#[macro_rules_attribute::apply(crate::db_input)]
pub struct TeamMemberSpend {
pub team_id: String,
pub user_id: String,
pub cost: f64,
}
pub async fn add_team_member_spend(
transaction: &mut Transaction<'_, Postgres>,
members: &[TeamMemberSpend],
) -> Result<(), Error> {
if members.is_empty() {
return Ok(());
}
let mut members = members.to_vec();
members.sort_by(|left, right| {
(&left.team_id, &left.user_id).cmp(&(&right.team_id, &right.user_id))
});
let mut teams: Vec<&str> = members
.iter()
.map(|member| member.team_id.as_str())
.collect();
teams.dedup();
for team in teams {
lock_team(transaction, team).await?;
}
let rows = serde_json::Value::Array(
members
.iter()
.map(|member| {
serde_json::json!({
"user_id": member.user_id,
"team_id": member.team_id,
"cost": member.cost,
})
})
.collect(),
);
sqlx::query_file!("queries/spend/upsert_team_member_spend.sql", rows)
.execute(&mut **transaction)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::super::begin;
use super::{TeamMemberSpend, add_team_member_spend};
use crate::test_support::database_pool;
use rstest::rstest;
use sqlx::PgPool;
fn members(cost: f64) -> [TeamMemberSpend; 3] {
[
TeamMemberSpend {
team_id: "member-spend-z".into(),
user_id: "off-roster-existing".into(),
cost: 2.0 * cost,
},
TeamMemberSpend {
team_id: "member-spend-a".into(),
user_id: "on-roster-missing".into(),
cost,
},
TeamMemberSpend {
team_id: "member-spend-a".into(),
user_id: "off-roster-missing".into(),
cost: 50.0 * cost,
},
]
}
#[rstest]
#[case::zero_then_quarter(0.25, 0.5, 3.0, 8.0)]
#[case::zero_then_eighth(0.125, 0.25, 2.5, 7.5)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn membership_spend_preserves_roster_guard_and_fractional_costs(
#[case] cost: f64,
#[case] expected_created: f64,
#[case] expected_spend: f64,
#[case] expected_total: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = begin(&pool).await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _teams = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles)
VALUES ('member-spend-a', '[{"user_id":"on-roster-missing","role":"user"}]'),
('member-spend-z', '[]')"#,
)
.execute(&mut *tx)
.await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _members = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamMembership" (team_id, user_id, spend, total_spend)
VALUES ('member-spend-z', 'off-roster-existing', 2, 7),
('member-spend-a', 'untouched', 100, 200)"#,
)
.execute(&mut *tx)
.await?;
add_team_member_spend(&mut tx, &[]).await?;
add_team_member_spend(&mut tx, &members(0.0)).await?;
add_team_member_spend(&mut tx, &members(cost)).await?;
add_team_member_spend(&mut tx, &members(cost)).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, String, f64, f64)> = sqlx::query_as(
r#"SELECT team_id, user_id, spend, total_spend FROM "LiteLLM_TeamMembership"
WHERE team_id IN ('member-spend-a', 'member-spend-z') ORDER BY team_id, user_id"#,
)
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
(
"member-spend-a".into(),
"on-roster-missing".into(),
expected_created,
expected_created
),
("member-spend-a".into(), "untouched".into(), 100.0, 200.0),
(
"member-spend-z".into(),
"off-roster-existing".into(),
expected_spend,
expected_total
),
]
);
tx.rollback().await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let remaining: i64 = sqlx::query_scalar(
r#"SELECT COUNT(*) FROM "LiteLLM_TeamMembership"
WHERE team_id IN ('member-spend-a', 'member-spend-z')"#,
)
.fetch_one(&pool)
.await?;
assert_eq!(remaining, 0);
Ok(())
}
}

182
litellm-rust/crates/db/src/spend/team_spend.rs Normal file
View file

@ -0,0 +1,182 @@
use sqlx::PgExecutor;
use super::EntitySpend;
use crate::Error;
pub async fn add_team_spend(
executor: impl PgExecutor<'_>,
spend: &EntitySpend,
) -> Result<(), Error> {
spend.validate()?;
if spend.is_empty() {
return Ok(());
}
sqlx::query_file!("queries/spend/add_team_spend.sql", &spend.ids, &spend.costs)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{EntitySpend, add_team_spend};
use crate::Error;
use crate::test_support::database_pool;
use rstest::rstest;
use sqlx::PgPool;
#[rstest]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn duplicate_team_ids_add_every_cost(
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamTable" (team_id, spend, created_at, updated_at)
VALUES ('team-spend-duplicate-a', 2, '2000-01-01', '2000-01-01'),
('team-spend-duplicate-z', 3, '2000-01-01', '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec![
"team-spend-duplicate-z".into(),
"team-spend-duplicate-a".into(),
"team-spend-duplicate-z".into(),
"team-spend-duplicate-a".into(),
],
costs: vec![0.5, 0.25, 1.0, 0.5],
};
add_team_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, f64)> = sqlx::query_as(
r#"SELECT team_id, spend FROM "LiteLLM_TeamTable"
WHERE team_id = ANY($1) ORDER BY team_id"#,
)
.bind(&spend.ids)
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
("team-spend-duplicate-a".into(), 2.75),
("team-spend-duplicate-z".into(), 4.5),
]
);
Ok(())
}
#[rstest]
#[case::missing_cost(vec!["team-spend-malformed".into()], vec![])]
#[case::extra_cost(vec!["team-spend-malformed".into()], vec![0.25, 0.5])]
#[case::empty_ids(vec![], vec![0.25])]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn malformed_team_batch_does_not_write(
#[case] ids: Vec<String>,
#[case] costs: Vec<f64>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamTable" (team_id, spend, created_at, updated_at)
VALUES ('team-spend-malformed', 2, '2000-01-01', '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let result = add_team_spend(&mut *tx, &EntitySpend { ids, costs }).await;
assert!(matches!(result, Err(Error::InvalidSpendBatch { .. })));
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, bool) = sqlx::query_as(
r#"SELECT spend, updated_at = TIMESTAMP '2000-01-01'
FROM "LiteLLM_TeamTable" WHERE team_id = 'team-spend-malformed'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (2.0, true));
Ok(())
}
#[rstest]
#[case::fractional_costs(0.25, 2.5, 4.0)]
#[case::zero_costs(0.0, 2.0, 3.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn team_spend_matches_python_update_many(
#[case] cost: f64,
#[case] expected_a_spend: f64,
#[case] expected_z_spend: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_TeamTable" (team_id, spend, created_at, updated_at)
VALUES ('team-spend-a', 2, '2000-01-01', '2000-01-01'),
('team-spend-z', 3, '2000-01-01', '2000-01-01'),
('team-spend-untouched', 100, '2000-01-01', '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec![
"team-spend-z".into(),
"team-spend-a".into(),
"team-spend-missing".into(),
],
costs: vec![2.0 * cost, cost, 50.0],
};
add_team_spend(&mut *tx, &EntitySpend::default()).await?;
add_team_spend(&mut *tx, &spend).await?;
add_team_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, f64, bool)> = sqlx::query_as(
r#"SELECT team_id, spend, updated_at > TIMESTAMP '2000-01-01'
FROM "LiteLLM_TeamTable"
WHERE team_id = ANY($1) ORDER BY team_id"#,
)
.bind(vec![
"team-spend-a",
"team-spend-z",
"team-spend-missing",
"team-spend-untouched",
])
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
("team-spend-a".into(), expected_a_spend, true),
("team-spend-untouched".into(), 100.0, false),
("team-spend-z".into(), expected_z_spend, true),
]
);
Ok(())
}
}

173
litellm-rust/crates/db/src/spend/user_spend.rs Normal file
View file

@ -0,0 +1,173 @@
use sqlx::PgExecutor;
use super::EntitySpend;
use crate::Error;
pub async fn add_user_spend(
executor: impl PgExecutor<'_>,
spend: &EntitySpend,
) -> Result<(), Error> {
spend.validate()?;
if spend.is_empty() {
return Ok(());
}
sqlx::query_file!("queries/spend/add_user_spend.sql", &spend.ids, &spend.costs)
.execute(executor)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use std::error::Error as StdError;
use super::{EntitySpend, add_user_spend};
use crate::{Error, test_support::database_pool};
use rstest::rstest;
use sqlx::PgPool;
#[rstest]
#[case::fractional_costs(0.25, 2.5, 4.0)]
#[case::zero_costs(0.0, 2.0, 3.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn user_spend_matches_python_update_many(
#[case] cost: f64,
#[case] expected_a_spend: f64,
#[case] expected_z_spend: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_UserTable" (user_id, spend, updated_at)
VALUES ('user-spend-a', 2, '2000-01-01'),
('user-spend-z', 3, '2000-01-01'),
('user-spend-untouched', 100, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec![
"user-spend-z".into(),
"user-spend-a".into(),
"user-spend-missing".into(),
],
costs: vec![2.0 * cost, cost, 50.0],
};
add_user_spend(&mut *tx, &EntitySpend::default()).await?;
add_user_spend(&mut *tx, &spend).await?;
add_user_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let rows: Vec<(String, f64, bool)> = sqlx::query_as(
r#"SELECT user_id, spend, updated_at > TIMESTAMP '2000-01-01'
FROM "LiteLLM_UserTable"
WHERE user_id = ANY($1) ORDER BY user_id"#,
)
.bind(vec![
"user-spend-a",
"user-spend-z",
"user-spend-missing",
"user-spend-untouched",
])
.fetch_all(&mut *tx)
.await?;
assert_eq!(
rows,
vec![
("user-spend-a".into(), expected_a_spend, true),
("user-spend-untouched".into(), 100.0, false),
("user-spend-z".into(), expected_z_spend, true),
]
);
Ok(())
}
#[rstest]
#[case::direct_duplicate_batch(vec![0.25, 0.75], 3.0)]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn user_spend_sums_duplicate_ids(
#[case] costs: Vec<f64>,
#[case] expected_spend: f64,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_UserTable" (user_id, spend, updated_at)
VALUES ('user-spend-duplicate', 2, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec!["user-spend-duplicate".into(), "user-spend-duplicate".into()],
costs,
};
add_user_spend(&mut *tx, &spend).await?;
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, bool) = sqlx::query_as(
r#"SELECT spend, updated_at > TIMESTAMP '2000-01-01'
FROM "LiteLLM_UserTable" WHERE user_id = 'user-spend-duplicate'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (expected_spend, true));
Ok(())
}
#[rstest]
#[case::missing_cost(vec![])]
#[case::extra_cost(vec![0.25, 0.75])]
#[tokio::test]
#[ignore = "requires DATABASE_URL pointing to a Prisma-migrated PostgreSQL database"]
async fn user_spend_rejects_malformed_batch_without_writes(
#[case] costs: Vec<f64>,
#[future(awt)] database_pool: Result<PgPool, Box<dyn StdError>>,
) -> Result<(), Box<dyn StdError>> {
let pool = database_pool?;
let mut tx = pool.begin().await?;
#[expect(
clippy::disallowed_methods,
reason = "Fixture writes are outside the compile-checked query set"
)]
let _inserted = sqlx::query(
r#"INSERT INTO "LiteLLM_UserTable" (user_id, spend, updated_at)
VALUES ('user-spend-malformed', 2, '2000-01-01')"#,
)
.execute(&mut *tx)
.await?;
let spend = EntitySpend {
ids: vec!["user-spend-malformed".into()],
costs,
};
let result = add_user_spend(&mut *tx, &spend).await;
assert!(matches!(result, Err(Error::InvalidSpendBatch { .. })));
#[expect(
clippy::disallowed_methods,
reason = "Test assertions are outside the compile-checked query set"
)]
let row: (f64, bool) = sqlx::query_as(
r#"SELECT spend, updated_at = TIMESTAMP '2000-01-01'
FROM "LiteLLM_UserTable" WHERE user_id = 'user-spend-malformed'"#,
)
.fetch_one(&mut *tx)
.await?;
assert_eq!(row, (2.0, true));
Ok(())
}
}

10
litellm-rust/crates/db/src/test_support.rs Normal file
View file

@ -0,0 +1,10 @@
use std::error::Error as StdError;
use rstest::fixture;
use sqlx::PgPool;
#[fixture]
pub async fn database_pool() -> Result<PgPool, Box<dyn StdError>> {
let database_url = std::env::var("DATABASE_URL")?;
Ok(PgPool::connect(&database_url).await?)
}

7
litellm-rust/crates/gateway/AGENTS.md
View file

@ -3,3 +3,10 @@
- Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates. Provider execution belongs to the `inference-<fmt>` crates and `llms`
- Inject shared state and infrastructure. Avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups
- Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate
## Redis
- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored
- Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands
- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped
- Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python
- `/health/readiness` never touches Redis, matching the proxy

3
litellm-rust/crates/gateway/Cargo.toml
View file

@ -19,6 +19,8 @@ litellm-gateway-auth.workspace = true
litellm-gateway-ui.workspace = true
litellm-auth-types.workspace = true
litellm-config.workspace = true
litellm-circuit-breaker.workspace = true
litellm-redis-client.workspace = true
litellm-http = { workspace = true, features = ["mcp"] }
litellm-llms.workspace = true
litellm-secrets.workspace = true
@ -31,6 +33,7 @@ uuid.workspace = true
tower-sessions-moka-store = "0.15.0"
[dev-dependencies]
litellm-redis-testing.workspace = true
base64.workspace = true
rstest.workspace = true
tempfile.workspace = true

9
litellm-rust/crates/gateway/src/error.rs
View file

@ -8,6 +8,15 @@ pub enum Error {
McpSetting(String),
#[error(transparent)]
Mcp(#[from] litellm_gateway_mcp::ConnectError),
#[error(transparent)]
Config(#[from] litellm_config::Error),
#[error("invalid Redis setting {setting}: {reason}")]
RedisInvalid {
setting: &'static str,
reason: &'static str,
},
#[error(transparent)]
RedisConnect(#[from] litellm_redis_client::ConnectError),
#[error("MCP requires a configured master key")]
Auth(#[from] litellm_gateway_auth::Error),
}

22
litellm-rust/crates/gateway/src/lib.rs
View file

@ -1,7 +1,9 @@
mod auth;
mod error;
mod redis;
mod secrets;
pub use error::Error;
pub use redis::RedisClients;
use std::{sync::Arc, time::Instant};
@ -70,6 +72,17 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
)?))
}
pub async fn build_redis(config: &Config) -> Result<RedisClients, Error> {
let environment = secrets::environment_values(&config.environment_variables)?;
let lookup = |name: &str| {
environment
.get(name)
.map(|value| value.expose().to_owned())
.or_else(|| std::env::var(name).ok())
};
redis::connect(config, &lookup).await
}
pub async fn build_mcp(
config: &Config,
secrets: Arc<dyn litellm_secrets::source::SecretSource>,
@ -128,12 +141,17 @@ pub fn router(
litellm_gateway_auth::authenticate,
))
.layer(axum::middleware::from_fn(log_request));
let health = Router::new().route("/health/readiness", axum::routing::get(readiness));
match ui {
Some(ui) => inference.merge(ui),
None => inference,
Some(ui) => inference.merge(ui).merge(health),
None => inference.merge(health),
}
}
async fn readiness() -> axum::Json<serde_json::Value> {
axum::Json(serde_json::json!({"status": "healthy", "db": "Not connected"}))
}
async fn log_request(request: Request, next: Next) -> Response {
let request_id = Uuid::new_v4().to_string();
let log_body_chunks = tracing::enabled!(tracing::Level::DEBUG);

1
litellm-rust/crates/gateway/src/main.rs
View file

@ -46,6 +46,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
Logger::new(StderrSink { level }).install_global()?;
let config = Config::load(settings.litellm_config)?;
let inference = litellm_gateway::build_inference(&config)?;
let _redis = litellm_gateway::build_redis(&config).await?;
let ui = match std::env::var_os("LITELLM_UI_PATH") {
Some(directory) => {
let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?;

577
litellm-rust/crates/gateway/src/redis.rs Normal file
View file

@ -0,0 +1,577 @@
use std::time::Duration;
use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State};
use litellm_config::{
CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value,
environment_has_target,
};
use litellm_redis_client::{
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient,
RedisConfig,
};
use crate::Error;
const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2);
type Params = ResolvedRedis;
pub struct RedisClients {
pub coordination: Option<RedisClient>,
pub transaction_buffer: Option<RedisClient>,
pub response_cache: Option<RedisClient>,
}
#[derive(Debug, PartialEq)]
pub(crate) struct Plan {
coordination: Option<Target>,
transaction_buffer_fallback: Option<Params>,
response_cache: Option<Params>,
}
#[derive(Debug, PartialEq)]
enum Target {
Configured(Params),
Environment(Params),
}
pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result<Plan, Error> {
let response_cache = redis_cache_params(config)
.map(|settings| settings.resolve_cache(lookup))
.transpose()?;
let coordination = match &config.general_settings.coordination_redis {
Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)),
None => match (&response_cache, config.litellm_settings.cache == Some(true)) {
(Some(cache), _) => Some(Target::Configured(cache.clone())),
(None, true) => environment(lookup)?.map(Target::Configured),
(None, false) => optional_environment(lookup).map(Target::Environment),
},
};
let transaction_buffer_fallback = match (&coordination, use_transaction_buffer(config)) {
(Some(Target::Configured(_)), _) | (_, false) => None,
(Some(Target::Environment(_)) | None, true) => environment(lookup)?,
};
Ok(Plan {
coordination,
transaction_buffer_fallback,
response_cache,
})
}
pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result<RedisClients, Error> {
let plan = plan(config, lookup)?;
let breaker = breaker_config(lookup)?;
let coordination = match plan.coordination {
Some(Target::Configured(params)) => Some(client(params, breaker)?),
Some(Target::Environment(params)) => match client(params, breaker) {
Ok(client) => reachable(client).await,
Err(_) => {
tracing::warn!(
"REDIS_* environment variables did not form a valid Redis client, cross-pod coordination stays in memory"
);
None
}
},
None => None,
};
let transaction_buffer = match (&coordination, plan.transaction_buffer_fallback) {
(Some(coordination), _) => Some(coordination.clone()),
(None, Some(params)) => Some(client(params, breaker)?),
(None, None) => None,
};
let response_cache = plan
.response_cache
.map(|params| client(params, breaker))
.transpose()?;
Ok(RedisClients {
coordination,
transaction_buffer,
response_cache,
})
}
fn redis_cache_params(config: &Config) -> Option<&RedisSettings> {
let settings = &config.litellm_settings;
let params = settings.cache_params.as_ref()?;
match (settings.cache, params) {
(Some(true), CacheSettings::Redis(redis)) => Some(redis),
_ => None,
}
}
fn environment(lookup: Lookup<'_>) -> Result<Option<Params>, Error> {
match environment_has_target(lookup) {
true => RedisSettings::default()
.resolve(lookup)
.map(Some)
.map_err(Into::into),
false => Ok(None),
}
}
fn optional_environment(lookup: Lookup<'_>) -> Option<Params> {
match environment(lookup) {
Ok(params) => params,
Err(error) => {
tracing::warn!(error = %error, "REDIS_* environment variables are invalid, cross-pod coordination stays in memory");
None
}
}
}
fn use_transaction_buffer(config: &Config) -> bool {
match config
.general_settings
.additional_fields
.get("use_redis_transaction_buffer")
{
Some(Value::Bool(enabled)) => *enabled,
Some(Value::String(text)) => matches!(
text.to_ascii_lowercase().as_str(),
"true" | "1" | "yes" | "on"
),
_ => false,
}
}
async fn reachable(client: RedisClient) -> Option<RedisClient> {
match tokio::time::timeout(ENVIRONMENT_PING_TIMEOUT, client.ping()).await {
Ok(Ok(())) => Some(client),
Ok(Err(error)) => {
tracing::warn!(%error, "REDIS_* environment variables named a Redis that is not reachable, cross-pod coordination stays in memory");
None
}
Err(_) => {
tracing::warn!(
"REDIS_* environment variables named a Redis that did not answer PING in 2s, cross-pod coordination stays in memory"
);
None
}
}
}
fn client(params: Params, breaker: Option<BreakerConfig>) -> Result<RedisClient, Error> {
tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured");
let endpoint = match params.endpoint {
RedisEndpoint::Url(url) => Endpoint::Url(url.expose().to_owned()),
RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint {
host: host.host,
port: host.port,
username: host.username.map(|value| value.expose().to_owned()),
password: host.password.map(|value| value.expose().to_owned()),
tls: host.tls,
db: host.db,
}),
RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint {
nodes: cluster
.nodes
.into_iter()
.map(|node| ClusterNode {
host: node.host,
port: node.port,
})
.collect(),
username: cluster.username.map(|value| value.expose().to_owned()),
password: cluster.password.map(|value| value.expose().to_owned()),
tls: cluster.tls,
}),
};
Ok(RedisClient::new(RedisConfig {
endpoint,
namespace: params.namespace.map(Namespace::new).unwrap_or_default(),
timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT),
breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))),
})?)
}
fn log_transition(transition: &litellm_circuit_breaker::Transition<State>) {
match transition.to {
State::Open { .. } => tracing::warn!(
generation = transition.generation,
"redis circuit breaker opened, fast-failing redis calls"
),
State::HalfOpen => tracing::info!(
generation = transition.generation,
"redis circuit breaker half open, probing redis"
),
State::Closed { .. } => tracing::info!(
generation = transition.generation,
"redis circuit breaker closed, redis recovered"
),
}
}
fn breaker_config(lookup: Lookup<'_>) -> Result<Option<BreakerConfig>, Error> {
let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED")
.is_none_or(|value| value.eq_ignore_ascii_case("true"));
if !enabled {
return Ok(None);
}
let defaults = BreakerConfig::default();
let seconds = |name: &'static str, default: Duration| -> Result<Duration, Error> {
lookup(name)
.map(|value| {
value
.parse::<f64>()
.ok()
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
.ok_or(Error::RedisInvalid {
setting: name,
reason: "must be a non-negative number of seconds",
})
})
.unwrap_or(Ok(default))
};
let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD")
.map(|value| {
value.parse().map_err(|_| Error::RedisInvalid {
setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD",
reason: "must be an integer",
})
})
.unwrap_or(Ok(defaults.failure_threshold))?;
Ok(Some(BreakerConfig {
failure_threshold,
recovery_timeout: seconds(
"REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT",
defaults.recovery_timeout,
)?,
timeout_min_duration: seconds(
"REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION",
defaults.timeout_min_duration,
)?,
}))
}
#[cfg(test)]
mod tests {
use std::{collections::HashMap, time::Duration};
use litellm_config::{Config, RedisEndpoint, RedisHost};
use litellm_redis_testing::RedisServer;
use rstest::rstest;
use super::{Params, Plan, Target, connect, plan};
use crate::Error;
fn host(name: &str, port: u16) -> Params {
Params {
endpoint: RedisEndpoint::Host(RedisHost {
host: name.to_owned(),
port,
username: None,
password: None,
tls: false,
db: 0,
}),
namespace: None,
timeout: None,
}
}
fn resolved(yaml: &str, environment: &[(&str, &str)]) -> Result<Plan, Error> {
let environment: HashMap<String, String> = environment
.iter()
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
.collect();
plan(&Config::from_yaml(yaml).unwrap(), &|name| {
environment.get(name).cloned()
})
}
fn configured(coordination: Params, response_cache: Option<Params>) -> Plan {
Plan {
coordination: Some(Target::Configured(coordination)),
transaction_buffer_fallback: None,
response_cache,
}
}
const NO_REDIS: Plan = Plan {
coordination: None,
transaction_buffer_fallback: None,
response_cache: None,
};
#[rstest]
#[case::nothing_configured("model_list: []", &[], NO_REDIS)]
#[case::coordination_block_wins_over_the_cache_redis(
"general_settings:\n coordination_redis:\n host: coord\nlitellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache",
&[],
configured(host("coord", 6379), Some(host("cache", 6379)))
)]
#[case::cache_params_without_cache_true_are_unused(
"litellm_settings:\n cache_params:\n type: redis\n host: cache",
&[],
NO_REDIS
)]
#[case::non_redis_cache_takes_coordination_from_the_environment_without_a_ping(
"litellm_settings:\n cache: true\n cache_params:\n type: local",
&[("REDIS_HOST", "env")],
configured(host("env", 6379), None)
)]
#[case::environment_only_coordination_is_probed(
"model_list: []",
&[("REDIS_HOST", "env")],
Plan { coordination: Some(Target::Environment(host("env", 6379))), transaction_buffer_fallback: None, response_cache: None }
)]
#[case::the_transaction_buffer_falls_back_to_the_environment_only_when_enabled(
"general_settings:\n use_redis_transaction_buffer: \"true\"",
&[("REDIS_HOST", "env")],
Plan { coordination: Some(Target::Environment(host("env", 6379))), transaction_buffer_fallback: Some(host("env", 6379)), response_cache: None }
)]
fn resolves_each_redis_role_like_the_proxy(
#[case] yaml: &str,
#[case] environment: &[(&str, &str)],
#[case] expected: Plan,
) {
assert_eq!(resolved(yaml, environment).unwrap(), expected);
}
#[rstest]
fn redis_migration_regression_preserves_borrowed_cache_namespace() {
let yaml = "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n namespace: responses\n ttl: 600";
let cache = Params {
namespace: Some("responses".to_owned()),
..host("cache", 6379)
};
assert_eq!(
resolved(yaml, &[]).unwrap(),
configured(cache.clone(), Some(cache))
);
}
#[rstest]
#[tokio::test]
async fn borrowed_coordination_and_transaction_buffer_share_the_cache_namespace() {
let server = RedisServer::start().await.unwrap();
let config = Config::from_yaml(&format!(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: 127.0.0.1\n port: {}\n namespace: responses",
port_of(&server)
))
.unwrap();
let clients = connect(&config, &|_| None).await.unwrap();
let coordination = clients.coordination.unwrap();
let buffer = clients.transaction_buffer.unwrap();
let cache = clients.response_cache.unwrap();
let unnamespaced = connect(&coordination_config(&server), &|_| None)
.await
.unwrap()
.coordination
.unwrap();
coordination
.set("rate-limit", "1", Duration::from_secs(60))
.await
.unwrap();
buffer
.set("spend", "2", Duration::from_secs(60))
.await
.unwrap();
assert_eq!(
cache.get::<String>("rate-limit").await.unwrap().as_deref(),
Some("1")
);
assert_eq!(
cache.get::<String>("spend").await.unwrap().as_deref(),
Some("2")
);
assert_eq!(
unnamespaced
.get::<String>("responses:rate-limit")
.await
.unwrap()
.as_deref(),
Some("1")
);
assert_eq!(
unnamespaced
.get::<String>("responses:spend")
.await
.unwrap()
.as_deref(),
Some("2")
);
assert!(
unnamespaced
.get::<String>("rate-limit")
.await
.unwrap()
.is_none()
);
assert!(unnamespaced.get::<String>("spend").await.unwrap().is_none());
}
#[rstest]
#[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])]
#[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])]
fn redis_migration_regression_malformed_optional_environment_is_absent(
#[case] environment: &[(&str, &str)],
) {
assert_eq!(resolved("model_list: []", environment).unwrap(), NO_REDIS);
}
#[rstest]
fn redis_migration_regression_explicit_malformed_setting_still_fails() {
let error = resolved(
"general_settings:\n coordination_redis:\n host: coord\n port: not-a-port",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() {
let error = resolved(
"general_settings:\n use_redis_transaction_buffer: true",
&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
fn redis_migration_regression_invalid_response_cache_still_fails() {
let error = resolved(
"litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port",
&[],
)
.unwrap_err();
assert!(matches!(
error,
Error::Config(litellm_config::Error::RedisInvalid {
setting: "port",
..
})
));
}
#[rstest]
#[tokio::test]
async fn redis_migration_regression_invalid_optional_environment_url_is_absent() {
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| {
(name == "REDIS_URL").then(|| "not a redis URL".to_owned())
})
.await
.unwrap();
assert!(clients.coordination.is_none());
}
fn port_of(server: &RedisServer) -> String {
server.url().rsplit(':').next().unwrap().to_owned()
}
fn coordination_config(server: &RedisServer) -> Config {
Config::from_yaml(&format!(
"general_settings:\n coordination_redis:\n host: 127.0.0.1\n port: {}\n socket_timeout: 0.5",
port_of(server)
))
.unwrap()
}
async fn eventually_pings(client: &litellm_redis_client::RedisClient) -> bool {
let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
while tokio::time::Instant::now() < deadline {
if client.ping().await.is_ok() {
return true;
}
tokio::time::sleep(Duration::from_millis(200)).await;
}
false
}
#[rstest]
#[tokio::test]
async fn boot_does_not_wait_for_a_configured_redis_that_is_down() {
let server = RedisServer::start().await.unwrap();
server.stop().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
let while_down = coordination.ping().await;
server.resume().await.unwrap();
assert!(while_down.is_err());
assert!(eventually_pings(&coordination).await);
}
#[rstest]
#[tokio::test]
async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() {
let server = RedisServer::start().await.unwrap();
let clients = connect(&coordination_config(&server), &|_| None)
.await
.unwrap();
let coordination = clients.coordination.unwrap();
coordination
.set("written", "1", Duration::from_secs(60))
.await
.unwrap();
server.restart().await.unwrap();
assert!(eventually_pings(&coordination).await);
let buffer = clients.transaction_buffer.unwrap();
assert!(buffer.ping().await.is_ok());
}
#[rstest]
#[tokio::test]
async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() {
let server = RedisServer::start().await.unwrap();
let port = port_of(&server);
server.stop().await.unwrap();
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,
})
.await
.unwrap();
assert!(clients.coordination.is_none());
assert!(clients.transaction_buffer.is_none());
}
#[rstest]
#[tokio::test]
async fn a_reachable_environment_only_redis_becomes_coordination() {
let server = RedisServer::start().await.unwrap();
let port = port_of(&server);
let config = Config::from_yaml("model_list: []").unwrap();
let clients = connect(&config, &|name| match name {
"REDIS_HOST" => Some("127.0.0.1".to_owned()),
"REDIS_PORT" => Some(port.clone()),
_ => None,
})
.await
.unwrap();
assert!(clients.coordination.unwrap().ping().await.is_ok());
}
}

30
litellm-rust/crates/gateway/tests/server.rs
View file

@ -287,3 +287,33 @@ async fn ui_routes_are_absent_when_not_mounted(
.unwrap();
assert_eq!(response.status().as_u16(), 404);
}
#[rstest]
#[case::with_master_key("model_list: []\ngeneral_settings:\n master_key: gateway-key\n")]
#[case::without_master_key("model_list: []\n")]
#[tokio::test]
async fn readiness_answers_without_credentials_like_the_proxy_with_no_database(
inference: Arc<Gateway>,
#[case] yaml: &str,
) {
let config = Config::from_yaml(yaml).unwrap();
let app = litellm_gateway::router(inference, &config, None, None);
let response = app
.oneshot(
Request::get("/health/readiness")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status().as_u16(), 200);
let body = axum::body::to_bytes(response.into_body(), 1024)
.await
.unwrap();
assert_eq!(
serde_json::from_slice::<Value>(&body).unwrap(),
json!({"status": "healthy", "db": "Not connected"})
);
}

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

@ -6,7 +6,7 @@ use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::messages::transformation::BaseMessagesConfig,
bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
bedrock::messages::invoke_transformations::anthropic_claude_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG,
edenai::messages::transformation::EDENAI_MESSAGES_CONFIG,
minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG,

0
litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs → litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs
View file

2
litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs
View file

@ -1 +1 @@
pub mod anthropic_claude3_transformation;
pub mod anthropic_claude_transformation;

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

@ -53,7 +53,7 @@ pyo3.workspace = true
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
redis = { workspace = true, features = ["tls-rustls"] }
serde_json.workspace = true
strum.workspace = true
veil.workspace = true

19
litellm-rust/crates/python-bridge/src/secrets/provider.rs
View file

@ -88,20 +88,23 @@ fn aws_context(
}
fn read_timeout(value: Option<&Bound<'_, PyAny>>) -> PyResult<Option<Duration>> {
let Some(value) = value.filter(|value| !value.is_none()) else {
return Ok(None);
};
let seconds = match value.extract::<f64>() {
Ok(value) => Some(value),
Err(_) => value.getattr("read")?.extract::<Option<f64>>()?,
};
seconds
read_timeout_seconds(value)?
.map(|value| {
Duration::try_from_secs_f64(value).map_err(|_| PyValueError::new_err("invalid timeout"))
})
.transpose()
}
pub(super) fn read_timeout_seconds(value: Option<&Bound<'_, PyAny>>) -> PyResult<Option<f64>> {
let Some(value) = value.filter(|value| !value.is_none()) else {
return Ok(None);
};
match value.extract::<f64>() {
Ok(value) => Ok(Some(value)),
Err(_) => value.getattr("read")?.extract::<Option<f64>>(),
}
}
fn vault_context(params: Option<&Bound<'_, PyAny>>) -> PyResult<HashicorpOperationContext> {
let params = params.and_then(|value| value.cast::<PyDict>().ok());
let nested = params

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

@ -2,10 +2,12 @@ mod operation;
pub(super) use operation::{Failure, FailureKind, FailureStage, delete, rotate, write};
use litellm_core_utils::exception_mapping_utils::{TimeoutPhase, timeout_message};
use litellm_secrets::hashicorp::{Error, RawOperationError};
use pyo3::prelude::*;
use super::mutation::{error_value, http_message, json_value};
use super::provider::read_timeout_seconds;
const HTTP_TRANSPORT_MODULE: &str = "litellm.rust_bridge.host.http_transport";
@ -69,21 +71,17 @@ fn cause_message(py: Python<'_>, kind: FailureKind, context: &ErrorContext) -> P
}
FailureKind::Native(RawOperationError::Timeout { method, elapsed }) => {
if method == "POST" {
let elapsed = py
.import("builtins")?
.call_method1("round", (elapsed.as_secs_f64(), 3))?;
let kwargs = pyo3::types::PyDict::new(py);
kwargs.set_item(
"message",
format!(
"Connection timed out. Timeout passed={}, time taken={} seconds",
context.timeout.as_deref().unwrap_or("None"),
elapsed.str()?
timeout_message(
TimeoutPhase::Read,
context.timeout,
Some(elapsed.as_secs_f64()),
),
)?;
kwargs.set_item("model", "default-model-name")?;
kwargs.set_item("llm_provider", "litellm-httpx-handler")?;
kwargs.set_item("headers", pyo3::types::PyDict::new(py))?;
py.import("litellm.exceptions")?
.getattr("Timeout")?
.call((), Some(&kwargs))?
@ -159,7 +157,7 @@ fn response_text(py: Python<'_>, body: &[u8]) -> PyResult<String> {
#[derive(Default)]
pub(super) struct ErrorContext {
timeout: Option<String>,
timeout: Option<f64>,
aiohttp: bool,
}
@ -173,7 +171,7 @@ impl ErrorContext {
return Ok(Self::default());
}
Ok(Self {
timeout: timeout.map(|value| value.str()?.extract()).transpose()?,
timeout: read_timeout_seconds(timeout)?,
aiohttp: py
.import(HTTP_TRANSPORT_MODULE)?
.getattr("uses_aiohttp_transport")?

43
litellm-rust/crates/redis-client/AGENTS.md Normal file
View file

@ -0,0 +1,43 @@
- The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace
- It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default
- Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route
- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`)
- Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values
## Errors
- `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it
- Classify in `From<RedisError> for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error
- A `Server` error never drops the connection
- `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome
- Never swallow an error or turn it into a miss. The caller picks the fallback
- Reads return `Option<T>`: a missing key is `None`, and a value that doesn't decode to `T` (`"15.0"` as `i64`) is `Decode`
## Circuit breaker
- The `litellm-circuit-breaker` `Breaker` is injected through `RedisConfig::breaker` and shared by every clone. `None` turns it off
- Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it
- `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success
- An open breaker returns `Error::BreakerOpen` without touching the socket
- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count
## Namespace
- Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace`
- The rule never changes, because Python and Rust pods share keys in the same Redis
## Commands
- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
- On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test
- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`)
## Scripts and pipelines
- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async`
- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order
- Pipeline ops take one key each. A multi-key op would span slots inside one group
- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once
## Tests
- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes
- Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning)
- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables
- Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline

17
litellm-rust/crates/redis-client/Cargo.toml Normal file
View file

@ -0,0 +1,17 @@
[package]
name = "litellm-redis-client"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-circuit-breaker.workspace = true
futures-util.workspace = true
redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] }
thiserror.workspace = true
[dev-dependencies]
litellm-redis-testing.workspace = true
rstest.workspace = true
tokio.workspace = true

647
litellm-rust/crates/redis-client/src/client.rs Normal file
View file

@ -0,0 +1,647 @@
use std::{
future::Future,
time::{Duration, Instant},
};
use litellm_circuit_breaker::{Admission, Breaker, Failure};
use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value};
use crate::{
ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands,
connection::Connection,
};
#[derive(Clone)]
pub struct RedisClient {
pub(crate) connection: Connection,
namespace: Namespace,
breaker: Option<Breaker>,
}
#[derive(Debug, PartialEq, Eq)]
pub enum Ttl {
Missing,
Persistent,
Expires(Duration),
}
impl RedisClient {
pub fn new(config: RedisConfig) -> Result<Self, ConnectError> {
let connection = match &config.endpoint {
Endpoint::Url(url) => {
Connection::standalone(url.as_str().into_connection_info()?, config.timeout)?
}
Endpoint::Host(host) => {
Connection::standalone(host.connection_info()?, config.timeout)?
}
Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?,
};
Ok(Self {
connection,
namespace: config.namespace,
breaker: config.breaker,
})
}
pub async fn ping(&self) -> Result<(), Error> {
self.query(&commands::ping()).await
}
pub async fn get<T: FromRedisValue>(&self, key: &str) -> Result<Option<T>, Error> {
self.query(&commands::get(self.key(key))).await
}
pub async fn mget<T: FromRedisValue>(&self, keys: &[&str]) -> Result<Vec<Option<T>>, Error> {
if keys.is_empty() {
return Ok(Vec::new());
}
self.query(&commands::mget(self.keys(keys))).await
}
pub async fn set(
&self,
key: &str,
value: impl ToSingleRedisArg,
ttl: Duration,
) -> Result<(), Error> {
self.query(&commands::set(self.key(key), value, ttl)).await
}
pub async fn set_nx(
&self,
key: &str,
value: impl ToSingleRedisArg,
ttl: Duration,
) -> Result<bool, Error> {
let reply: Option<String> = self
.query(&commands::set_nx(self.key(key), value, ttl))
.await?;
Ok(reply.is_some())
}
pub async fn del(&self, keys: &[&str]) -> Result<usize, Error> {
if keys.is_empty() {
return Ok(0);
}
self.query(&commands::del(self.keys(keys))).await
}
pub async fn expire(&self, key: &str, ttl: Duration) -> Result<bool, Error> {
self.query(&commands::expire(self.key(key), ttl)).await
}
pub async fn ttl(&self, key: &str) -> Result<Ttl, Error> {
let reply: i64 = self.query(&commands::ttl(self.key(key))).await?;
Ok(match reply {
-2 => Ttl::Missing,
-1 => Ttl::Persistent,
millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())),
})
}
pub async fn incr_by(&self, key: &str, delta: i64) -> Result<i64, Error> {
self.query(&commands::incr_by(self.key(key), delta)).await
}
pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result<f64, Error> {
self.query(&commands::incr_by_float(self.key(key), delta))
.await
}
pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result<usize, Error> {
self.query(&commands::rpush(self.key(key), values)).await
}
pub async fn eval<T: FromRedisValue>(
&self,
script: &Script,
keys: &[&str],
args: impl ToRedisArgs,
) -> Result<T, Error> {
let keys = self.keys(keys);
let args = args.to_redis_args();
self.guarded(async {
Ok(script
.invocation(&keys, &args)
.invoke_async(&mut self.connection.clone())
.await?)
})
.await
}
pub fn pipeline(&self) -> Pipeline {
Pipeline::new(self.namespace.clone())
}
pub async fn run(&self, pipeline: &Pipeline) -> Vec<Result<Value, Error>> {
let admission = match self.admit() {
Ok(admission) => admission,
Err(_) => {
return (0..pipeline.len())
.map(|_| Err(Error::BreakerOpen))
.collect();
}
};
let results = pipeline.run(&self.connection).await;
let transport = results.iter().find_map(|result| {
result
.as_ref()
.err()
.filter(|error| matches!(error, Error::Transport(_)))
});
settle(admission, transport);
results
}
fn key(&self, key: &str) -> String {
self.namespace.apply(key)
}
fn keys(&self, keys: &[&str]) -> Vec<String> {
keys.iter().map(|key| self.key(key)).collect()
}
async fn query<T: FromRedisValue>(&self, command: &Cmd) -> Result<T, Error> {
self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) })
.await
}
async fn guarded<T>(&self, call: impl Future<Output = Result<T, Error>>) -> Result<T, Error> {
let admission = self.admit()?;
let result = call.await;
settle(admission, result.as_ref().err());
result
}
fn admit(&self) -> Result<Option<Admission>, Error> {
self.breaker
.as_ref()
.map(|breaker| {
breaker
.admit(Instant::now())
.map_err(|_| Error::BreakerOpen)
})
.transpose()
}
}
fn settle(admission: Option<Admission>, error: Option<&Error>) {
let Some(admission) = admission else {
return;
};
match error {
None => admission.succeeded(),
Some(Error::Transport(error)) if error.is_timeout() => {
admission.failed(Failure::Timeout, Instant::now())
}
Some(Error::Transport(_)) => admission.failed(Failure::Hard, Instant::now()),
Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission),
}
}
#[cfg(test)]
mod tests {
use std::time::{Duration, Instant};
use litellm_redis_testing::RedisServer;
use rstest::rstest;
use litellm_circuit_breaker::{Breaker, Config, State};
use super::{RedisClient, Ttl};
use crate::{
DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig,
test_support::{Server, Topology, client_id},
};
#[rstest]
#[tokio::test]
async fn namespaces_isolate_keys_on_the_server(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let a = server.client("a", DEFAULT_TIMEOUT).await;
let b = server.client("b", DEFAULT_TIMEOUT).await;
let bare = server.client("", DEFAULT_TIMEOUT).await;
a.set("x", "1", Duration::from_secs(60)).await.unwrap();
assert_eq!(a.get::<String>("x").await.unwrap().as_deref(), Some("1"));
assert_eq!(a.get::<String>("a:x").await.unwrap().as_deref(), Some("1"));
assert_eq!(
bare.get::<String>("a:x").await.unwrap().as_deref(),
Some("1")
);
assert_eq!(b.get::<String>("x").await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn a_rejected_command_keeps_the_connection(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
client.rpush("list", "item").await.unwrap();
let before = client_id(&client, "list").await;
let rejected = client.incr_by("list", 1).await;
assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}");
assert_eq!(client.rpush("list", "next").await.unwrap(), 2);
assert_eq!(client_id(&client, "list").await, before);
}
#[rstest]
#[tokio::test]
async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
client
.set("float", "15.0", Duration::from_secs(60))
.await
.unwrap();
let decoded = client.get::<i64>("float").await;
assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}");
assert_eq!(client.get::<i64>("missing").await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", Duration::from_millis(500)).await;
client
.set("first", "1", Duration::from_secs(60))
.await
.unwrap();
client
.set("second", "2", Duration::from_secs(60))
.await
.unwrap();
server.freeze().await;
let timed_out = client.get::<String>("first").await;
server.thaw().await;
assert!(
matches!(timed_out, Err(Error::Transport(_))),
"{timed_out:?}"
);
assert_eq!(
client.get::<String>("second").await.unwrap().as_deref(),
Some("2")
);
}
#[rstest]
#[tokio::test]
async fn the_same_client_recovers_after_redis_restarts(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", Duration::from_millis(500)).await;
server.restart().await;
let deadline = Instant::now() + Duration::from_secs(30);
let recovered = loop {
match client.incr_by("after_restart", 1).await {
Ok(value) => break Some(value),
Err(_) if Instant::now() < deadline => {
tokio::time::sleep(Duration::from_millis(200)).await
}
Err(_) => break None,
}
};
assert_eq!(recovered, Some(1));
}
#[rstest]
#[tokio::test]
async fn a_client_created_while_redis_is_down_works_once_redis_is_up(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
server.stop().await;
let client = server.client("", Duration::from_millis(500)).await;
let while_down = client.ping().await;
server.resume().await;
assert!(
matches!(while_down, Err(Error::Transport(_))),
"{while_down:?}"
);
let deadline = Instant::now() + Duration::from_secs(30);
let recovered = loop {
match client.ping().await {
Ok(()) => break true,
Err(_) if Instant::now() < deadline => {
tokio::time::sleep(Duration::from_millis(200)).await
}
Err(_) => break false,
}
};
assert!(recovered);
}
#[rstest]
#[tokio::test]
async fn empty_batches_succeed_without_asking_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
assert_eq!(
client.mget::<String>(&[]).await.unwrap(),
Vec::<Option<String>>::new()
);
assert_eq!(client.del(&[]).await.unwrap(), 0);
}
#[rstest]
#[tokio::test]
async fn mget_and_del_span_slots_in_input_order(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
client.set("a", "1", Duration::from_secs(60)).await.unwrap();
client.set("c", "3", Duration::from_secs(60)).await.unwrap();
let values = client.mget::<String>(&["c", "b", "a"]).await.unwrap();
let deleted = client.del(&["a", "b", "c"]).await.unwrap();
assert_eq!(
values,
vec![Some("3".to_owned()), None, Some("1".to_owned())]
);
assert_eq!(deleted, 2);
assert_eq!(
client.mget::<String>(&["a", "c"]).await.unwrap(),
vec![None, None]
);
}
#[rstest]
#[tokio::test]
async fn set_nx_only_writes_a_missing_key(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let first = client
.set_nx("lock", "owner-1", Duration::from_secs(60))
.await
.unwrap();
let second = client
.set_nx("lock", "owner-2", Duration::from_secs(60))
.await
.unwrap();
assert!(first);
assert!(!second);
assert_eq!(
client.get::<String>("lock").await.unwrap().as_deref(),
Some("owner-1")
);
}
#[rstest]
#[tokio::test]
async fn ttl_reports_missing_persistent_and_remaining_time(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
client.incr_by("persistent", 1).await.unwrap();
client
.set("expiring", "1", Duration::from_secs(60))
.await
.unwrap();
client.incr_by_float("expired_later", 0.5).await.unwrap();
assert!(
client
.expire("expired_later", Duration::from_secs(30))
.await
.unwrap()
);
assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing);
assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent);
assert!(matches!(
client.ttl("expiring").await.unwrap(),
Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60)
));
assert!(matches!(
client.ttl("expired_later").await.unwrap(),
Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30)
));
assert!(
!client
.expire("missing", Duration::from_secs(30))
.await
.unwrap()
);
}
#[rstest]
#[tokio::test]
async fn incr_by_float_accumulates_fractions(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
client.incr_by_float("spend", 0.25).await.unwrap();
assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75);
}
#[rstest]
#[tokio::test]
async fn concurrent_commands_on_one_cloned_client_all_apply(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let tasks: Vec<_> = (0..100)
.map(|_| {
let client = client.clone();
tokio::spawn(async move {
for _ in 0..10 {
client.incr_by("shared", 1).await.unwrap();
}
})
})
.collect();
for task in tasks {
task.await.unwrap();
}
assert_eq!(client.get::<i64>("shared").await.unwrap(), Some(1000));
}
fn breaker(failure_threshold: u32) -> Breaker {
Breaker::new(
Config {
failure_threshold,
recovery_timeout: Duration::from_millis(200),
timeout_min_duration: Duration::ZERO,
},
Box::new(|_| {}),
)
}
#[rstest]
#[tokio::test]
async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let breaker = breaker(2);
let client = server
.guarded_client("", Duration::from_millis(300), Some(breaker.clone()))
.await;
server.freeze().await;
let first = client.get::<String>("k").await;
let second = client.get::<String>("k").await;
let started = Instant::now();
let third = client.get::<String>("k").await;
let rejected_in = started.elapsed();
server.thaw().await;
assert!(matches!(first, Err(Error::Transport(_))), "{first:?}");
assert!(matches!(second, Err(Error::Transport(_))), "{second:?}");
assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}");
assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}");
let deadline = Instant::now() + Duration::from_secs(30);
let recovered = loop {
match client.get::<String>("k").await {
Ok(value) => break Some(value),
Err(_) if Instant::now() < deadline => {
tokio::time::sleep(Duration::from_millis(50)).await
}
Err(_) => break None,
}
};
assert_eq!(recovered, Some(None));
assert!(
matches!(breaker.state(), State::Closed { .. }),
"{:?}",
breaker.state()
);
}
#[rstest]
#[tokio::test]
async fn rejected_commands_never_open_the_breaker(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let breaker = breaker(2);
let client = server
.guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone()))
.await;
client.rpush("list", "item").await.unwrap();
let rejected = [
client.incr_by("list", 1).await,
client.incr_by("list", 1).await,
client.incr_by("list", 1).await,
];
assert!(
matches!(rejected[2], Err(Error::Server(_))),
"{:?}",
rejected[2]
);
assert!(
matches!(breaker.state(), State::Closed { hard: 0, .. }),
"{:?}",
breaker.state()
);
}
#[rstest]
#[tokio::test]
async fn an_open_breaker_fails_every_pipeline_op_without_sending(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server
.guarded_client("", Duration::from_millis(300), Some(breaker(1)))
.await;
let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1);
server.freeze().await;
let failed = client.run(&pipeline).await;
let rejected = client.run(&pipeline).await;
server.thaw().await;
assert!(
matches!(failed[0], Err(Error::Transport(_))),
"{:?}",
failed[0]
);
assert!(
matches!(rejected[0], Err(Error::BreakerOpen)),
"{:?}",
rejected[0]
);
assert!(
matches!(rejected[1], Err(Error::BreakerOpen)),
"{:?}",
rejected[1]
);
}
#[rstest]
#[tokio::test]
async fn a_host_endpoint_connects_to_its_database() {
let server = RedisServer::start().await.unwrap();
let port = server.url().rsplit(':').next().unwrap().parse().unwrap();
let connect = |endpoint| {
RedisClient::new(RedisConfig {
endpoint,
namespace: Namespace::default(),
timeout: DEFAULT_TIMEOUT,
breaker: None,
})
};
let in_db_3 = connect(Endpoint::Host(HostEndpoint {
host: "127.0.0.1".to_owned(),
port,
username: None,
password: None,
tls: false,
db: 3,
}))
.unwrap();
let in_db_0 = connect(Endpoint::Url(server.url().to_owned())).unwrap();
let via_url = connect(Endpoint::Url(format!("{}/3", server.url()))).unwrap();
in_db_3
.set("k", "3", Duration::from_secs(60))
.await
.unwrap();
assert_eq!(in_db_0.get::<String>("k").await.unwrap(), None);
assert_eq!(
via_url.get::<String>("k").await.unwrap().as_deref(),
Some("3")
);
}
}

62
litellm-rust/crates/redis-client/src/commands.rs Normal file
View file

@ -0,0 +1,62 @@
use std::time::Duration;
use redis::{Cmd, ToRedisArgs, ToSingleRedisArg};
pub(crate) fn ping() -> Cmd {
redis::cmd("PING")
}
pub(crate) fn get(key: String) -> Cmd {
redis::cmd("GET").arg(key).to_owned()
}
pub(crate) fn mget(keys: Vec<String>) -> Cmd {
redis::cmd("MGET").arg(keys).to_owned()
}
pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("PX")
.arg(millis(ttl))
.to_owned()
}
pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("NX")
.arg("PX")
.arg(millis(ttl))
.to_owned()
}
pub(crate) fn del(keys: Vec<String>) -> Cmd {
redis::cmd("DEL").arg(keys).to_owned()
}
pub(crate) fn expire(key: String, ttl: Duration) -> Cmd {
redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned()
}
pub(crate) fn ttl(key: String) -> Cmd {
redis::cmd("PTTL").arg(key).to_owned()
}
pub(crate) fn incr_by(key: String, delta: i64) -> Cmd {
redis::cmd("INCRBY").arg(key).arg(delta).to_owned()
}
pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd {
redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned()
}
pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd {
redis::cmd("RPUSH").arg(key).arg(values).to_owned()
}
fn millis(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}

152
litellm-rust/crates/redis-client/src/config.rs Normal file
View file

@ -0,0 +1,152 @@
use std::{fmt, time::Duration};
use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult};
use litellm_circuit_breaker::Breaker;
use crate::Namespace;
pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
pub struct RedisConfig {
pub endpoint: Endpoint,
pub namespace: Namespace,
pub timeout: Duration,
pub breaker: Option<Breaker>,
}
#[derive(Clone, PartialEq)]
pub enum Endpoint {
Url(String),
Host(HostEndpoint),
Cluster(ClusterEndpoint),
}
#[derive(Clone, PartialEq)]
pub struct HostEndpoint {
pub host: String,
pub port: u16,
pub username: Option<String>,
pub password: Option<String>,
pub tls: bool,
pub db: i64,
}
#[derive(Clone, PartialEq)]
pub struct ClusterEndpoint {
pub nodes: Vec<ClusterNode>,
pub username: Option<String>,
pub password: Option<String>,
pub tls: bool,
}
#[derive(Clone, Debug, PartialEq)]
pub struct ClusterNode {
pub host: String,
pub port: u16,
}
impl ClusterEndpoint {
pub(crate) fn addrs(&self) -> impl Iterator<Item = ConnectionAddr> + '_ {
self.nodes
.iter()
.map(|node| addr(&node.host, node.port, self.tls))
}
}
impl HostEndpoint {
pub(crate) fn connection_info(&self) -> RedisResult<ConnectionInfo> {
let settings = RedisConnectionInfo::default().set_db(self.db);
let settings = match &self.username {
Some(username) => settings.set_username(username),
None => settings,
};
let settings = match &self.password {
Some(password) => settings.set_password(password),
None => settings,
};
Ok(addr(&self.host, self.port, self.tls)
.into_connection_info()?
.set_redis_settings(settings))
}
}
fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr {
match tls {
true => ConnectionAddr::TcpTls {
host: host.to_owned(),
port,
insecure: false,
tls_params: None,
},
false => ConnectionAddr::Tcp(host.to_owned(), port),
}
}
impl fmt::Debug for Endpoint {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Url(url) => formatter
.debug_tuple("Url")
.field(&redacted_url(url))
.finish(),
Self::Host(host) => formatter
.debug_struct("Host")
.field("host", &host.host)
.field("port", &host.port)
.field("username", &host.username)
.field("password", &host.password.as_ref().map(|_| "<redacted>"))
.field("tls", &host.tls)
.field("db", &host.db)
.finish(),
Self::Cluster(cluster) => formatter
.debug_struct("Cluster")
.field("nodes", &cluster.nodes)
.field("username", &cluster.username)
.field("password", &cluster.password.as_ref().map(|_| "<redacted>"))
.field("tls", &cluster.tls)
.finish(),
}
}
}
fn redacted_url(url: &str) -> String {
match url
.split_once("://")
.and_then(|(scheme, rest)| Some((scheme, rest.rsplit_once('@')?.1)))
{
Some((scheme, host)) => format!("{scheme}://<redacted>@{host}"),
None => url.to_owned(),
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::{Endpoint, HostEndpoint};
#[rstest]
#[case::url_with_password(Endpoint::Url("redis://user:secret@cache:6379/0".to_owned()), "redis://<redacted>@cache:6379/0")]
#[case::url_with_at_in_password(Endpoint::Url("rediss://:p@ss@cache:6380".to_owned()), "rediss://<redacted>@cache:6380")]
#[case::host_password(
Endpoint::Host(HostEndpoint {
host: "cache".to_owned(),
port: 6379,
username: None,
password: Some("secret".to_owned()),
tls: false,
db: 0,
}),
"<redacted>"
)]
fn debug_output_never_shows_the_password(#[case] endpoint: Endpoint, #[case] shown: &str) {
let debug = format!("{endpoint:?}");
assert!(
!debug.contains("secret") && !debug.contains("p@ss"),
"{debug}"
);
assert!(debug.contains(shown), "{debug}");
}
}

80
litellm-rust/crates/redis-client/src/connection.rs Normal file
View file

@ -0,0 +1,80 @@
use std::time::Duration;
use redis::{
Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value,
aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig},
cluster::{ClusterClientBuilder, ClusterConfig},
cluster_async::ClusterConnection,
};
use crate::ClusterEndpoint;
const RETRIES: u32 = 3;
#[derive(Clone)]
pub(crate) enum Connection {
Standalone(ConnectionManager),
Cluster(ClusterConnection),
}
impl Connection {
pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult<Self> {
let config = ConnectionManagerConfig::new()
.set_connection_timeout(Some(timeout))
.set_response_timeout(Some(timeout))
.set_number_of_retries(RETRIES as usize);
let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?;
Ok(Self::Standalone(manager))
}
pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult<Self> {
let builder = ClusterClientBuilder::new(endpoint.addrs())
.connection_timeout(timeout)
.response_timeout(timeout)
.retries(RETRIES);
let builder = match &endpoint.username {
Some(username) => builder.username(username),
None => builder,
};
let builder = match &endpoint.password {
Some(password) => builder.password(password),
None => builder,
};
let config = ClusterConfig::new()
.set_connection_timeout(timeout)
.set_response_timeout(timeout);
Ok(Self::Cluster(
builder
.build()?
.get_pending_async_connection_with_config(config),
))
}
}
impl ConnectionLike for Connection {
fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> {
match self {
Self::Standalone(connection) => connection.req_packed_command(cmd),
Self::Cluster(connection) => connection.req_packed_command(cmd),
}
}
fn req_packed_commands<'a>(
&'a mut self,
pipeline: &'a Pipeline,
offset: usize,
count: usize,
) -> RedisFuture<'a, Vec<Value>> {
match self {
Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count),
Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count),
}
}
fn get_db(&self) -> i64 {
match self {
Self::Standalone(connection) => connection.get_db(),
Self::Cluster(connection) => connection.get_db(),
}
}
}

69
litellm-rust/crates/redis-client/src/error.rs Normal file
View file

@ -0,0 +1,69 @@
use redis::{ErrorKind, RedisError, ServerErrorKind};
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("Redis did not answer, so the command may or may not have applied")]
Transport(#[source] RedisError),
#[error("Redis rejected the command")]
Server(#[source] RedisError),
#[error("Redis replied with a value of a different type")]
Decode(#[source] RedisError),
#[error("the Redis circuit breaker is open, so the command was not sent")]
BreakerOpen,
}
#[derive(Debug, thiserror::Error)]
#[error("connecting to Redis")]
pub struct ConnectError(#[from] RedisError);
impl From<RedisError> for Error {
fn from(error: RedisError) -> Self {
match error.kind() {
ErrorKind::Server(
ServerErrorKind::BusyLoading
| ServerErrorKind::ClusterDown
| ServerErrorKind::MasterDown
| ServerErrorKind::TryAgain
| ServerErrorKind::ReadOnly,
) => Self::Transport(error),
ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error),
ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error),
_ => Self::Transport(error),
}
}
}
#[cfg(test)]
mod tests {
use redis::{RedisError, Value};
use rstest::rstest;
use super::Error;
fn reply(resp: &str) -> RedisError {
match redis::parse_redis_value(resp.as_bytes()).unwrap() {
Value::ServerError(error) => error.into(),
other => panic!("{resp:?} parsed as {other:?}"),
}
}
#[rstest]
#[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")]
#[case::no_script("-NOSCRIPT No matching script\r\n")]
#[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")]
#[case::generic("-ERR value is not an integer or out of range\r\n")]
#[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")]
fn rejections_are_server_errors(#[case] resp: &str) {
assert!(matches!(Error::from(reply(resp)), Error::Server(_)));
}
#[rstest]
#[case::loading("-LOADING Redis is loading the dataset in memory\r\n")]
#[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")]
#[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")]
#[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")]
#[case::read_only("-READONLY You can't write against a read only replica.\r\n")]
fn unavailable_nodes_are_transport_errors(#[case] resp: &str) {
assert!(matches!(Error::from(reply(resp)), Error::Transport(_)));
}
}

20
litellm-rust/crates/redis-client/src/lib.rs Normal file
View file

@ -0,0 +1,20 @@
mod client;
mod commands;
mod config;
mod connection;
mod error;
mod namespace;
mod pipeline;
mod script;
#[cfg(test)]
mod test_support;
pub use client::{RedisClient, Ttl};
pub use config::{
ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig,
};
pub use error::{ConnectError, Error};
pub use namespace::Namespace;
pub use pipeline::Pipeline;
pub use redis::Value;
pub use script::Script;

57
litellm-rust/crates/redis-client/src/namespace.rs Normal file
View file

@ -0,0 +1,57 @@
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Namespace(Option<String>);
impl Namespace {
pub fn new(namespace: impl Into<String>) -> Self {
let namespace = namespace.into();
Self((!namespace.is_empty()).then_some(namespace))
}
pub fn as_str(&self) -> Option<&str> {
self.0.as_deref()
}
pub(crate) fn apply(&self, key: &str) -> String {
match &self.0 {
Some(namespace) if !is_prefixed(key, namespace) => {
format!("{namespace}:{key}")
}
_ => key.to_owned(),
}
}
}
fn is_prefixed(key: &str, namespace: &str) -> bool {
key.strip_prefix(namespace)
.is_some_and(|rest| rest.starts_with(':'))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::Namespace;
#[rstest]
#[case::plain(Namespace::new("litellm"), "spend:key:abc", "litellm:spend:key:abc")]
#[case::no_namespace(Namespace::default(), "spend:key:abc", "spend:key:abc")]
#[case::empty_namespace(Namespace::new(""), "spend:key:abc", "spend:key:abc")]
#[case::already_prefixed(Namespace::new("litellm"), "litellm:abc", "litellm:abc")]
#[case::shares_prefix_without_delimiter(
Namespace::new("litellm"),
"litellm_spend_update_buffer",
"litellm:litellm_spend_update_buffer"
)]
#[case::hash_tagged_key(
Namespace::new("litellm"),
"{api_key:sk-1}:requests",
"litellm:{api_key:sk-1}:requests"
)]
fn prefixes_keys_like_python_check_and_fix_namespace(
#[case] namespace: Namespace,
#[case] key: &str,
#[case] expected: &str,
) {
assert_eq!(namespace.apply(key), expected);
}
}

292
litellm-rust/crates/redis-client/src/pipeline.rs Normal file
View file

@ -0,0 +1,292 @@
use std::{collections::HashMap, time::Duration};
use futures_util::future::join_all;
use redis::{
Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value,
cluster_routing::Slot,
};
use crate::{Error, Namespace, Script, commands, connection::Connection};
pub struct Pipeline {
namespace: Namespace,
ops: Vec<Op>,
}
struct Op {
command: Cmd,
key: String,
script: Option<ScriptCall>,
}
struct ScriptCall {
script: Script,
keys: Vec<String>,
args: Vec<Vec<u8>>,
}
type Settled = (usize, Result<Value, Error>);
impl Pipeline {
pub(crate) fn new(namespace: Namespace) -> Self {
Self {
namespace,
ops: Vec::new(),
}
}
pub fn get(self, key: &str) -> Self {
let key = self.namespace.apply(key);
self.command(commands::get(key.clone()), key)
}
pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self {
let key = self.namespace.apply(key);
self.command(commands::set(key.clone(), value, ttl), key)
}
pub fn del(self, key: &str) -> Self {
let key = self.namespace.apply(key);
self.command(commands::del(vec![key.clone()]), key)
}
pub fn expire(self, key: &str, ttl: Duration) -> Self {
let key = self.namespace.apply(key);
self.command(commands::expire(key.clone(), ttl), key)
}
pub fn incr_by(self, key: &str, delta: i64) -> Self {
let key = self.namespace.apply(key);
self.command(commands::incr_by(key.clone(), delta), key)
}
pub fn incr_by_float(self, key: &str, delta: f64) -> Self {
let key = self.namespace.apply(key);
self.command(commands::incr_by_float(key.clone(), delta), key)
}
pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self {
let keys: Vec<String> = keys.iter().map(|key| self.namespace.apply(key)).collect();
let args = args.to_redis_args();
let op = Op {
command: script.evalsha(&keys, &args),
key: keys.first().cloned().unwrap_or_default(),
script: Some(ScriptCall {
script: script.clone(),
keys,
args,
}),
};
self.push(op)
}
fn command(self, command: Cmd, key: String) -> Self {
self.push(Op {
command,
key,
script: None,
})
}
fn push(mut self, op: Op) -> Self {
self.ops.push(op);
self
}
pub(crate) fn len(&self) -> usize {
self.ops.len()
}
pub(crate) async fn run(&self, connection: &Connection) -> Vec<Result<Value, Error>> {
let groups = self.groups(connection);
let settled = join_all(
groups
.iter()
.map(|group| self.run_group(connection.clone(), group)),
)
.await;
let mut settled: Vec<Settled> = settled.into_iter().flatten().collect();
settled.sort_by_key(|(index, _)| *index);
settled.into_iter().map(|(_, result)| result).collect()
}
fn groups(&self, connection: &Connection) -> Vec<Vec<usize>> {
match connection {
_ if self.ops.is_empty() => Vec::new(),
Connection::Standalone(_) => vec![(0..self.ops.len()).collect()],
Connection::Cluster(_) => {
let mut by_slot: HashMap<Slot, Vec<usize>> = HashMap::new();
for (index, op) in self.ops.iter().enumerate() {
by_slot
.entry(Slot::for_key(&op.key))
.or_default()
.push(index);
}
by_slot.into_values().collect()
}
}
}
async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec<Settled> {
let mut pipeline = redis::pipe();
pipeline.ignore_errors();
for &index in indices {
pipeline.add_command(self.ops[index].command.clone());
}
let replies: RedisResult<Vec<RedisResult<Value>>> =
pipeline.query_async(&mut connection).await;
let replies = match replies {
Ok(replies) => replies,
Err(error) => {
return indices
.iter()
.map(|&index| (index, Err(Error::from(error.clone()))))
.collect();
}
};
let mut settled = Vec::with_capacity(indices.len());
for (&index, reply) in indices.iter().zip(replies) {
settled.push((index, self.settle(&mut connection, index, reply).await));
}
settled
}
async fn settle(
&self,
connection: &mut Connection,
index: usize,
reply: RedisResult<Value>,
) -> Result<Value, Error> {
match (reply, &self.ops[index].script) {
(Err(error), Some(call))
if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) =>
{
Ok(call
.script
.invocation(&call.keys, &call.args)
.invoke_async(connection)
.await?)
}
(reply, _) => Ok(reply?),
}
}
}
#[cfg(test)]
mod tests {
use std::time::{Duration, Instant};
use redis::Value;
use rstest::rstest;
use crate::{
DEFAULT_TIMEOUT, Error, Script,
test_support::{Server, Topology, flush_scripts},
};
#[rstest]
#[tokio::test]
async fn each_op_settles_on_its_own_in_input_order(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
client
.set("text", "not a number", Duration::from_secs(60))
.await
.unwrap();
let pipeline = client
.pipeline()
.incr_by("a", 1)
.incr_by("text", 1)
.incr_by("c", 5)
.set("b", "2", Duration::from_secs(60))
.get("b")
.incr_by("a", 1)
.get("missing");
let results = client.run(&pipeline).await;
assert_eq!(results.len(), 7);
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1));
assert!(
matches!(results[1], Err(Error::Server(_))),
"{:?}",
results[1]
);
assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5));
assert_eq!(results[3].as_ref().unwrap(), &Value::Okay);
assert_eq!(
results[4].as_ref().unwrap(),
&Value::BulkString(b"2".to_vec())
);
assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2));
assert_eq!(results[6].as_ref().unwrap(), &Value::Nil);
}
#[rstest]
#[tokio::test]
async fn a_flushed_script_in_a_pipeline_reruns_alone(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("ns", DEFAULT_TIMEOUT).await;
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
flush_scripts(&client).await;
let pipeline = client
.pipeline()
.eval(&incr, &["counter"], 2)
.incr_by("other", 7);
let results = client.run(&pipeline).await;
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3));
assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7));
assert_eq!(client.get::<i64>("counter").await.unwrap(), Some(3));
}
#[rstest]
#[tokio::test]
async fn an_unreachable_server_fails_every_op_as_transport(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", Duration::from_millis(500)).await;
let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1);
server.freeze().await;
let results = client.run(&pipeline).await;
server.thaw().await;
assert_eq!(results.len(), 2);
assert!(
matches!(results[0], Err(Error::Transport(_))),
"{:?}",
results[0]
);
assert!(
matches!(results[1], Err(Error::Transport(_))),
"{:?}",
results[1]
);
}
#[rstest]
#[tokio::test]
async fn an_empty_pipeline_returns_without_asking_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", Duration::from_secs(2)).await;
server.freeze().await;
let started = Instant::now();
let results = client.run(&client.pipeline()).await;
let elapsed = started.elapsed();
server.thaw().await;
assert!(results.is_empty());
assert!(elapsed < Duration::from_secs(1), "{elapsed:?}");
}
}

88
litellm-rust/crates/redis-client/src/script.rs Normal file
View file

@ -0,0 +1,88 @@
use redis::{Cmd, ScriptInvocation};
#[derive(Clone)]
pub struct Script(redis::Script);
impl Script {
pub fn new(source: &str) -> Self {
Self(redis::Script::new(source))
}
pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec<u8>]) -> Cmd {
redis::cmd("EVALSHA")
.arg(self.0.get_hash())
.arg(keys.len())
.arg(keys)
.arg(args)
.to_owned()
}
pub(crate) fn invocation(&self, keys: &[String], args: &[Vec<u8>]) -> ScriptInvocation<'_> {
let mut invocation = self.0.prepare_invoke();
invocation.key(keys).arg(args);
invocation
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use rstest::rstest;
use super::Script;
use crate::{
DEFAULT_TIMEOUT, Error,
test_support::{Server, Topology, flush_scripts},
};
#[rstest]
#[tokio::test]
async fn script_keys_are_namespaced(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("a", DEFAULT_TIMEOUT).await;
let bare = server.client("", DEFAULT_TIMEOUT).await;
let read = Script::new("return redis.call('GET', KEYS[1])");
bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap();
let value: Option<String> = client
.eval(&read, &["x"], Vec::<String>::new())
.await
.unwrap();
assert_eq!(value.as_deref(), Some("1"));
}
#[rstest]
#[tokio::test]
async fn a_failing_script_is_a_server_error_and_keeps_the_connection(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.rpush("list", "item").await.unwrap();
let failed = client.eval::<i64>(&incr, &["list"], 1).await;
assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}");
assert_eq!(client.eval::<i64>(&incr, &["counter"], 3).await.unwrap(), 3);
}
#[rstest]
#[tokio::test]
async fn a_flushed_script_is_loaded_again(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = Server::start(topology).await;
let client = server.client("", DEFAULT_TIMEOUT).await;
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
flush_scripts(&client).await;
assert_eq!(client.eval::<i64>(&incr, &["counter"], 1).await.unwrap(), 2);
}
}

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