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

merge: integrate main while preserving silent model guards

Browse source
This commit is contained in:
Jaeyeon Kim 2026-10-28 09:27:20 +00:00
commit 483e70552c
1748 changed files with 75759 additions and 31388 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

1
.circleci/config.yml
View file

@ -1379,7 +1379,6 @@ jobs:
uv export --frozen --all-groups --all-extras --no-emit-project --no-emit-workspace --no-header --no-annotate \
| awk '/^[^ ]/ { keep = $0 ~ /^(annotated-types|coverage|pydantic|pydantic-core|typing-extensions|typing-inspection)==/ } keep' \
> /tmp/fan-in-requirements.txt
test "$(grep -c '^[^ ]' /tmp/fan-in-requirements.txt)" -eq 6
uv venv .venv --python 3.12
uv pip install --no-config --python .venv --require-hashes -r /tmp/fan-in-requirements.txt
- attach_workspace:

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

2
.github/workflows/cost-map-guard.yml vendored
View file

@ -40,4 +40,4 @@ jobs:
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
HEAD_REF: ${{ github.event.pull_request.head.ref }}
run: |
uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"
uv run --no-project python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"

49
.github/workflows/spend-logs-guard.yml vendored Normal file
View file

@ -0,0 +1,49 @@
name: Spend logs guard
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed
pull_request_target:
branches:
- main
pull_request_review:
types:
- submitted
- dismissed
permissions:
contents: read
pull-requests: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
cancel-in-progress: true
jobs:
spend-logs-guard:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.base.ref }}
persist-credentials: false
- name: Fetch the pull request head, its merge base and its reviews
id: revisions
env:
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
PR_NUMBER: ${{ github.event.pull_request.number }}
run: |
merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')"
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
gh api "repos/${GITHUB_REPOSITORY}/pulls/${PR_NUMBER}/reviews" --paginate --slurp | jq 'add // []' > "$RUNNER_TEMP/reviews.json"
echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Run the guard
env:
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
uv run --no-project --with pydantic==2.14.1 python ci_cd/spend_logs_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --approver kerry-berri --reviews "$RUNNER_TEMP/reviews.json"

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

6
.github/workflows/test-linting.yml vendored
View file

@ -524,6 +524,12 @@ jobs:
- name: check_provider_folders_documented
run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
- name: check_provider_endpoints_support_generated
run: uv run --no-sync python scripts/generate_provider_endpoints_support.py --check
- name: test_generate_provider_endpoints_support
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_generate_provider_endpoints_support.py
- name: check_prisma_binary_cache
run: uv run --no-sync python ./tests/code_coverage_tests/check_prisma_binary_cache.py

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

57
.github/workflows/test-unit.yml vendored
View file

@ -302,6 +302,7 @@ jobs:
tests/unit/llms/modelscope
tests/unit/llms/mongodb
tests/unit/llms/moonshot
tests/unit/llms/morph
tests/unit/llms/nadir
tests/unit/llms/nebius
tests/unit/llms/neosantara
@ -579,13 +580,13 @@ jobs:
- shard: mcp-elicitation
test-path: |-
tests/unit/experimental_mcp_client/test_mcp_client.py
tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
tests/unit/proxy/_experimental/mcp_server/test_interactions.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
tests/unit/proxy/_experimental/mcp_server/test_operations.py
tests/unit/mcp_client/test_mcp_client.py
tests/unit/proxy/mcp_server/test_capabilities.py
tests/unit/proxy/mcp_server/test_interactions.py
tests/unit/proxy/mcp_server/test_mcp_elicitation_handler.py
tests/unit/proxy/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py
tests/unit/proxy/mcp_server/test_operations.py
python-version: "3.12"
workers: 4
reruns: 0
@ -603,7 +604,7 @@ jobs:
tests/unit/proxy/spend_tracking
--ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py
tests/unit/proxy/pass_through_endpoints
tests/unit/proxy/experimental
tests/unit/proxy/mcp_server/test_tool_registry.py
tests/unit/proxy/common_utils
--ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py
--ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py
@ -811,14 +812,14 @@ jobs:
- shard: mcp-oauth
test-path: |-
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
tests/unit/proxy/_experimental/mcp_server/outbound_credentials
tests/unit/proxy/mcp_server/test_discoverable_endpoints.py
tests/unit/proxy/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/mcp_server/test_db_credentials.py
tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py
tests/unit/proxy/mcp_server/test_mcp_oauth_passthrough_tools.py
tests/unit/proxy/mcp_server/test_mcp_max_concurrent_requests.py
tests/unit/proxy/mcp_server/faults/test_list_outcomes.py
tests/unit/proxy/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4
reruns: 2
@ -828,19 +829,21 @@ jobs:
- shard: mcp-server
test-path: |-
tests/unit/mcp_client
tests/unit/experimental_mcp_client
tests/unit/proxy/mcp_server
tests/unit/proxy/_experimental/mcp_server
tests/unit/responses/mcp
--ignore=tests/unit/experimental_mcp_client/test_mcp_client.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_interactions.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_operations.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
--ignore=tests/unit/proxy/_experimental/mcp_server/outbound_credentials
--ignore=tests/unit/mcp_client/test_mcp_client.py
--ignore=tests/unit/proxy/mcp_server/test_capabilities.py
--ignore=tests/unit/proxy/mcp_server/test_interactions.py
--ignore=tests/unit/proxy/mcp_server/test_mcp_elicitation_handler.py
--ignore=tests/unit/proxy/mcp_server/test_mcp_server_manager.py
--ignore=tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py
--ignore=tests/unit/proxy/mcp_server/test_operations.py
--ignore=tests/unit/proxy/mcp_server/test_discoverable_endpoints.py
--ignore=tests/unit/proxy/mcp_server/test_db_credentials.py
--ignore=tests/unit/proxy/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4
reruns: 0
@ -950,7 +953,7 @@ jobs:
dist: loadscope
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
timeout-minutes: 3
timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
with:
persist-credentials: false
fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}

2
.gitignore vendored
View file

@ -105,7 +105,7 @@ litellm_config.yaml
.cursor
litellm/proxy/to_delete_loadtest_work/*
update_model_cost_map.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/mcp_server/test_mcp_server_manager.py
scripts/test_vertex_ai_search.py
LAZY_LOADING_IMPROVEMENTS.md
STABILIZATION_TODO.md

12
ARCHITECTURE.md
View file

@ -239,6 +239,8 @@ graph LR
7. `DBSpendUpdateWriter.update_database()` queues spend increments to Redis
8. Background job `update_spend` flushes queued spend to PostgreSQL every 60s
Cost lookup falls back to the response model when the provider-reported name cannot be priced. Explicitly registered zero rates remain zero, and custom deployment pricing or an explicit `base_model` retain precedence
### Data Access Layer (Models & Repositories)
Database entities and the operations on them live in two packages at the root of `litellm/` so both the gateway (`proxy/`) and the SDK can use them without importing proxy internals:
@ -332,13 +334,13 @@ 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) | Vertex AI | `llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py` |
| `/v1/messages` (passthrough) | Anthropic | `llms/anthropic/messages/native/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/messages/transformation.py` |
| Passthrough endpoints | All | `proxy/pass_through_endpoints/llm_provider_handlers/` |
### Example: Debugging prompt caching
@ -346,7 +348,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

2
CONTRIBUTING.md
View file

@ -94,6 +94,8 @@ Add your tests to the [`tests/unit/` directory](https://github.com/BerriAI/litel
- **Only add mocked tests** - no real LLM API calls in this directory
- For integration tests with real APIs, use the appropriate test directories
Tests that call `litellm.register_model` or change `litellm.model_cost` should use the `local_model_cost_map` fixture from `tests/unit/conftest.py`. It provides a fresh catalog and restores the previous map, so pricing and capability overrides cannot leak into later tests
### File Naming Convention
The `tests/unit/` directory follows the same structure as `litellm/`:

4
Makefile
View file

@ -313,7 +313,7 @@ test-unit-proxy-core: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/auth tests/unit/proxy/client tests/unit/proxy/db tests/unit/proxy/hooks tests/unit/proxy/policy_engine --ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py --ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py --tb=short -vv -n 4 --durations=20
test-unit-proxy-misc: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/agent_endpoints tests/unit/proxy/anthropic_endpoints tests/unit/proxy/common_utils --ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py --ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py --ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py --ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py --ignore=tests/unit/proxy/common_utils/test_realtime_cache.py tests/unit/proxy/discovery_endpoints tests/unit/proxy/experimental tests/unit/proxy/google_endpoints tests/unit/proxy/health_endpoints tests/unit/proxy/image_endpoints tests/unit/proxy/middleware --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py tests/unit/proxy/openai_files_endpoint tests/unit/proxy/pass_through_endpoints tests/unit/proxy/prompts tests/unit/proxy/public_endpoints tests/unit/proxy/response_api_endpoints tests/unit/proxy/shutdown tests/unit/proxy/spend_tracking --ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py tests/unit/proxy/ui_crud_endpoints tests/unit/proxy/vector_store_endpoints tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py --ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py --ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py --ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/proxy/agent_endpoints tests/unit/proxy/anthropic_endpoints tests/unit/proxy/common_utils --ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py --ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py --ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py --ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py --ignore=tests/unit/proxy/common_utils/test_realtime_cache.py tests/unit/proxy/discovery_endpoints tests/unit/proxy/mcp_server/test_tool_registry.py tests/unit/proxy/google_endpoints tests/unit/proxy/health_endpoints tests/unit/proxy/image_endpoints tests/unit/proxy/middleware --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py tests/unit/proxy/openai_files_endpoint tests/unit/proxy/pass_through_endpoints tests/unit/proxy/prompts tests/unit/proxy/public_endpoints tests/unit/proxy/response_api_endpoints tests/unit/proxy/shutdown tests/unit/proxy/spend_tracking --ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py tests/unit/proxy/ui_crud_endpoints tests/unit/proxy/vector_store_endpoints tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py --ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py --ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py --ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py --ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py --tb=short -vv -n 4 --durations=20
test-unit-proxy-root: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/test_*.py --ignore=tests/unit/proxy/test_aproxy_startup.py --ignore=tests/unit/proxy/test_credential_slot_registry.py --ignore=tests/unit/proxy/test_custom_callback_input.py --ignore=tests/unit/proxy/test_custom_logger_s3_gcs.py --ignore=tests/unit/proxy/test_custom_tokenizer_bug.py --ignore=tests/unit/proxy/test_db_schema_changes.py --ignore=tests/unit/proxy/test_deprecated_key_grace_period.py --ignore=tests/unit/proxy/test_get_favicon.py --ignore=tests/unit/proxy/test_get_image.py --ignore=tests/unit/proxy/test_prisma_client_backoff_retry.py --ignore=tests/unit/proxy/test_prompt_test_endpoint.py --ignore=tests/unit/proxy/test_proxy_config_unit_test.py --ignore=tests/unit/proxy/test_proxy_custom_auth.py --ignore=tests/unit/proxy/test_proxy_reject_logging.py --ignore=tests/unit/proxy/test_proxy_server.py --ignore=tests/unit/proxy/test_proxy_setting_guardrails.py --ignore=tests/unit/proxy/test_proxy_token_counter.py --ignore=tests/unit/proxy/test_proxy_utils.py --ignore=tests/unit/proxy/test_reducto_ocr_route.py --ignore=tests/unit/proxy/test_response_polling_pre_call_checks.py --ignore=tests/unit/proxy/test_server_root_path.py --ignore=tests/unit/proxy/test_ui_path_detection.py --ignore=tests/unit/proxy/test_unit_test_proxy_hooks.py --ignore=tests/unit/proxy/test_update_spend.py --ignore=tests/unit/proxy/test_zero_cost_model_budget_bypass.py --tb=short -vv -n 4 --durations=20
@ -325,7 +325,7 @@ test-unit-core-utils: install-test-deps
$(UV_RUN) pytest tests/unit/litellm_core_utils --tb=short -vv -n 2 --durations=20
test-unit-other: install-test-deps
$(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
test-unit-root: install-test-deps
$(UV_RUN) pytest tests/unit/test_*.py --tb=short -vv -n 4 --durations=20

6
README.md
View file

@ -223,7 +223,7 @@ async with httpx.AsyncClient(headers=headers, timeout=60.0) as http_client:
```python
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from litellm import experimental_mcp_client
from litellm import mcp_client
import litellm
server_params = StdioServerParameters(command="python", args=["mcp_server.py"])
@ -233,7 +233,7 @@ async with stdio_client(server_params) as (read, write):
await session.initialize()
# Load MCP tools in OpenAI format
tools = await experimental_mcp_client.load_mcp_tools(session=session, format="openai")
tools = await mcp_client.load_mcp_tools(session=session, format="openai")
# Use with any LiteLLM model
response = await litellm.acompletion(
@ -660,6 +660,8 @@ This covers:
We welcome contributions to LiteLLM! Whether you're fixing bugs, adding features, or improving documentation, we appreciate your help.
If you want to work closely with our team, check out [LiteLLM OSS Champions](https://www.litellm.ai/champions), our program for trusted OSS contributors with their own Slack channel and a voice in the roadmap
## Quick Start for Contributors
This requires uv to be installed.

119
ci_cd/spend_logs_guard.py Normal file
View file

@ -0,0 +1,119 @@
"""Required check: a pull request that touches LiteLLM_SpendLogs in schema.prisma or a migration
needs an approving review from the spend-log owner before it can merge.
CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies,
so a path rule would put every migration on the owner. This gate reads the diff instead: when no
added line in the watched files mentions the table it passes at once, otherwise it passes only once
the owner's latest review on the pull request is an approval. The workflow re-runs on every review.
"""
from __future__ import annotations
import argparse
import re
import subprocess
import sys
from collections.abc import Sequence
from pathlib import Path
from typing import Final
from pydantic import BaseModel, ConfigDict, TypeAdapter
WATCHED: Final = (
"schema.prisma",
"litellm/proxy/schema.prisma",
"litellm-proxy-extras/litellm_proxy_extras/schema.prisma",
)
MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/"
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs", re.IGNORECASE)
class Reviewer(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
login: str
class Review(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
state: str
user: Reviewer
_REVIEWS: Final = TypeAdapter(list[Review])
def watched(path: str) -> bool:
return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql"))
def _added_lines(section: str) -> tuple[str, ...]:
body: Final = section.split("\n@@", 1)
return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else ()
def spend_log_hits(diff: str) -> tuple[str, ...]:
"""`path: added line` for every added line in a watched file that names the table."""
sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE)
return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line))
def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]:
pairs: Final = tuple((_path_of(section), section) for section in sections)
return tuple(
(path, line) for path, section in pairs if watched(path) for line in _added_lines(section)
) # comprehension-ok: path then its lines
def _path_of(section: str) -> str:
match: Final = re.search(r"^\+\+\+ b/(.+)$", section, flags=re.MULTILINE)
return match.group(1) if match else ""
def approved_by(reviews: Sequence[Review], login: str) -> bool:
"""True when `login`'s most recent review on the pull request is an approval."""
states: Final = tuple(review.state for review in reviews if review.user.login == login)
return bool(states) and states[-1] == "APPROVED"
def added_diff(base: str, head: str) -> str:
result: Final = subprocess.run(
("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX),
check=True,
capture_output=True,
text=True,
)
return result.stdout
class Arguments(BaseModel):
model_config = ConfigDict(frozen=True)
base: str
head: str
approver: str
reviews: Path
def main(argv: Sequence[str]) -> int:
parser: Final = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base", required=True, help="merge base of the pull request")
parser.add_argument("--head", required=True, help="head sha of the pull request")
parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge")
parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews")
args: Final = Arguments.model_validate(vars(parser.parse_args(argv)))
hits: Final = spend_log_hits(added_diff(args.base, args.head))
if not hits:
sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n")
return 0
if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver):
sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n")
return 0
listed: Final = "".join(f" {hit}\n" for hit in hits)
sys.stdout.write(
f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n"
f"{listed}The check re-runs when a review is submitted.\n"
)
return 1
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))

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.

120
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -13,13 +13,19 @@ sys.path.insert(
) # Adds the parent directory to the system path
import configparser
import contextlib
import hashlib
import itertools
import json
import re
import tempfile
from collections import OrderedDict
from collections.abc import Generator, Iterator, Sequence
from contextvars import ContextVar
from functools import partial
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
from pydantic import JsonValue, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import (
@ -27,7 +33,7 @@ from litellm.integrations.custom_guardrail import (
log_guardrail_information,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import walk_user_text
from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs
@ -40,7 +46,7 @@ GUARDRAIL_PROVIDER = "hide-secrets"
# Per-invocation tally of redacted secrets by detect-secrets plugin type; None
# means the guardrail did not run, so _process_response records nothing.
_masked_entity_count: ContextVar[Optional[dict]] = ContextVar(
_masked_entity_count: ContextVar[Optional[dict[str, int]]] = ContextVar(
"hide_secrets_masked_entity_count", default=None
)
@ -449,6 +455,19 @@ _default_detect_secrets_config = {
_CONFIG_SECTION: Final = "litellm-prompt"
_JSON_STRING: Final = TypeAdapter(str)
_KEY_CONTEXT: Final = 64
# Longest JSON line scanned for a multi-line tool string. Escaping puts the whole string on one line, and
# detect-secrets regexes backtrack across it, so a longer multi-line string is cut into runs of whole lines,
# each scanned as its own JSON line beside the key. Every run holds at least two lines, so the high-entropy
# detectors, which match only a quoted string made wholly of their charset, never judge a lone line.
_JSON_LINE_LIMIT: Final = 4096
# A detector can stop a match inside an escape of a JSON-escaped line, as at the backslash of \".
_CUT_ESCAPE: Final = re.compile(r"(?<=.)\\(?:u[0-9A-Fa-f]{0,3})?\Z", re.DOTALL)
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
@ -459,6 +478,10 @@ _SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
_SCAN_SUFFIX: Final = ".py"
# Digests of scan inputs that held no secret, so the history an agent resends every turn is scanned once. 65536
# digests take about 10 MB and hold a few hundred interleaved agent sessions per worker.
_CLEAN_SCAN_CACHE_SIZE: Final = 65536
@contextlib.contextmanager
def _temp_file(text: str) -> Generator[str, None, None]:
@ -532,6 +555,57 @@ def _quoted_assignments(text: str) -> tuple[str, ...]:
)
def _json_unescaped(value: str) -> str:
for candidate in (value, _CUT_ESCAPE.sub("", value)):
try:
return _JSON_STRING.validate_json(f'"{candidate}"')
except ValueError:
continue
return value
def _key_context(key: str) -> str:
return json.dumps(" ".join(key.split()), ensure_ascii=False)[1:-1][-_KEY_CONTEXT:]
def _json_line(context: str | None, text: str) -> str:
value: Final = json.dumps(text, ensure_ascii=False)
return value if context is None else f'"{context}": {value}'
def _json_lines(context: str | None, text: str) -> Iterator[str]:
"""Yield ``text`` as one JSON line, or a long multi-line ``text`` as JSON lines of two or more whole lines."""
line: Final = _json_line(context, text)
lines: Final = text.splitlines(keepends=True)
if len(line) <= _JSON_LINE_LIMIT or len(lines) < 2:
yield line
return
budget: Final = _JSON_LINE_LIMIT - len(_json_line(context, ""))
sizes: Final = (len(json.dumps(part, ensure_ascii=False)) - 2 for part in lines)
states: Final = itertools.accumulate(sizes, partial(_next_run_state, budget), initial=(0, 0, 0))
run_ids: Final = (state[0] for state in itertools.islice(states, 1, None))
runs: Final = tuple(
tuple(part for part, _ in run) for _, run in itertools.groupby(zip(lines, run_ids), key=lambda pair: pair[1])
)
whole_runs: Final = (*runs[:-2], runs[-2] + runs[-1]) if len(runs[-1]) < 2 else runs
for run in whole_runs:
yield _json_line(context, "".join(run[:-1]) + run[-1].splitlines()[0])
def _next_run_state(budget: int, state: tuple[int, int, int], size: int) -> tuple[int, int, int]:
"""Advance ``(run id, lines in run, escaped size of run)`` by one line of escaped ``size``."""
run_id, count, total = state
if count >= 2 and total + size > budget:
return run_id + 1, 1, size
return run_id, count + 1, total + size
def _value_in_texts(value: str, texts: Sequence[str]) -> str | None:
"""Return the detected value as it appears in ``texts``, decoding it when it came from a JSON-escaped key line."""
candidates: Final = (value, _json_unescaped(value))
return next((candidate for candidate in candidates if any(candidate in text for text in texts)), None)
class _ENTERPRISE_SecretDetection(CustomGuardrail):
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
# path skips should_run_check and never sees data["prompt"]).
@ -539,11 +613,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs):
self.user_defined_detect_secrets_config = detect_secrets_config
self._clean_scans: OrderedDict[bytes, None] = OrderedDict()
super().__init__(**kwargs)
def scan_message_for_secrets(self, message_content: str):
def scan_message_for_secrets(self, message_content: str) -> list[dict[str, str]]:
from detect_secrets.settings import transient_settings
digest: Final = hashlib.sha256(message_content.encode()).digest()
if digest in self._clean_scans:
self._clean_scans.move_to_end(digest)
return []
detect_secrets_config = (
self.user_defined_detect_secrets_config or _default_detect_secrets_config
)
@ -552,6 +631,10 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
(*message_content.splitlines(), *_quoted_assignments(message_content))
)
if not found:
self._clean_scans[digest] = None
if len(self._clean_scans) > _CLEAN_SCAN_CACHE_SIZE:
self._clean_scans.popitem(last=False)
return [
{"type": secret_type, "value": value}
for value, secret_type in sorted(
@ -580,6 +663,35 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
)
return pattern.sub("[REDACTED]", text)
def redact_tool_payload(self, payload: JsonValue) -> JsonValue:
"""Redact secrets in a tool payload's string values, scanning each value alone and as its JSON lines."""
fields: Final = tuple(iter_tool_json_fields(payload))
texts: Final = tuple(text for _, text in fields)
contexts: Final = {key: _key_context(key) for key in {key for key, _ in fields if key is not None}}
keyed: Final = tuple(
itertools.chain.from_iterable(
_json_lines(None if key is None else contexts[key], text) for key, text in fields
)
)
detected: Final = tuple(
{**secret, "value": value}
for secret in self.scan_message_for_secrets("\n".join((*texts, *keyed)))
if (value := _value_in_texts(secret["value"], texts)) is not None
)
if not detected:
return payload
counts: Final = _masked_entity_count.get()
if counts is not None:
for secret in detected:
counts[secret["type"]] = counts.get(secret["type"], 0) + 1
verbose_proxy_logger.warning(
"Detected and redacted secrets in tool payload: %s",
sorted(dict.fromkeys(secret["type"] for secret in detected)),
)
values: Final = sorted({secret["value"] for secret in detected}, key=len, reverse=True)
pattern: Final = re.compile("|".join(re.escape(value) for value in values))
return map_tool_json_text(payload, lambda text: pattern.sub("[REDACTED]", text))
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.permissions is not None:
if GUARDRAIL_NAME in user_api_key_dict.permissions:
@ -641,7 +753,7 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
_masked_entity_count.set({})
# Covers multimodal list content + Responses-API input.
inspected = walk_user_text(data, self.redact_text) + self._redact_prompt(data)
inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data)
if inspected == 0:
# Image-only, empty-text, and unsupported payloads inspected

4
enterprise/litellm_enterprise/proxy/liteadmin.py
View file

@ -18,8 +18,8 @@ from fastapi.responses import HTMLResponse, RedirectResponse, Response
from pydantic import ConfigDict, Field, SecretStr, TypeAdapter, ValidationError
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._experimental.mcp_server.oauth_utils import get_request_base_url
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.mcp_server.oauth_utils import get_request_base_url
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
@ -122,7 +122,7 @@ def _cookie_name(token: str) -> str:
async def _session_user(request: Request) -> str | None:
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
from litellm.proxy.mcp_server.byok_oauth_endpoints import (
get_authenticated_browser_user_id,
)

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.79"
version = "0.1.80"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.79"
version = "0.1.80"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

20
helm/litellm/README.md
View file

@ -23,3 +23,23 @@ ui:
backend:
gatewayUrl: https://gateway.example.com
```
## Routes to another Service
For a shared hostname, `ingress.extraPaths` can route to a Service in the same namespace using `serviceName` and `servicePort`. The port can be a number or a named Service port. Omit `service` when using these fields
```yaml
ingress:
enabled: true
className: alb
host: gateway.example.com
extraPaths:
- path: /moyai
pathType: Prefix
serviceName: moyai-ingress
servicePort: 8080
```
The Service and application must already support the path prefix. The chart adds the path to the existing Ingress and leaves its built-in routes in place. It does not create the target Service or rewrite requests
For AWS ALB, configure the target Service's `alb.ingress.kubernetes.io/healthcheck-path` annotation when the application needs a different health check from LiteLLM. Existing `service` aliases (`gateway`, `backend`, `ui`) remain supported

29
helm/litellm/templates/ingress.yaml
View file

@ -277,11 +277,34 @@ spec:
{{- if not $extra.path }}
{{- fail (printf "ingress.extraPaths[%d]: 'path' is required" $idx) }}
{{- end }}
{{- $target := dict }}
{{- $port := dict }}
{{- if or (hasKey $extra "serviceName") (hasKey $extra "servicePort") }}
{{- if hasKey $extra "service" }}
{{- fail (printf "ingress.extraPaths[%d]: service cannot be combined with serviceName or servicePort" $idx) }}
{{- end }}
{{- if not (and (kindIs "string" $extra.serviceName) $extra.serviceName $extra.servicePort) }}
{{- fail (printf "ingress.extraPaths[%d]: a non-empty serviceName and servicePort are required for a custom Service" $idx) }}
{{- end }}
{{- $target = dict "name" $extra.serviceName }}
{{- if kindIs "string" $extra.servicePort }}
{{- $port = dict "name" $extra.servicePort }}
{{- else if or (kindIs "float64" $extra.servicePort) (kindIs "int64" $extra.servicePort) (kindIs "int" $extra.servicePort) }}
{{- if or (lt (float64 $extra.servicePort) 1.0) (gt (float64 $extra.servicePort) 65535.0) (ne (float64 $extra.servicePort) (float64 (int $extra.servicePort))) }}
{{- fail (printf "ingress.extraPaths[%d]: numeric servicePort must be an integer from 1 to 65535" $idx) }}
{{- end }}
{{- $port = dict "number" (int $extra.servicePort) }}
{{- else }}
{{- fail (printf "ingress.extraPaths[%d]: servicePort must be a port number or name" $idx) }}
{{- end }}
{{- else }}
{{- $service := $extra.service | default "gateway" }}
{{- $target := get $extraPathBackends $service }}
{{- $target = get $extraPathBackends $service }}
{{- if not $target }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown service %q, expected one of backend, gateway, ui" $idx $extra.path $service) }}
{{- end }}
{{- $port = dict "number" $target.port }}
{{- end }}
{{- $requestedPathType := $extra.pathType | default "Prefix" }}
{{- if not (has $requestedPathType (list "Prefix" "Exact" "ImplementationSpecific")) }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $requestedPathType) }}
@ -297,9 +320,9 @@ spec:
pathType: {{ $pathType }}
backend:
service:
name: {{ $target.name }}
name: {{ $target.name | quote }}
port:
number: {{ $target.port }}
{{- toYaml $port | nindent 18 }}
{{- end }}
# --- Catch-all → backend (management API: /key/*, /user/*, /team/*, ...) ---
- path: /

134
helm/litellm/tests/ingress_extra_paths_tests.yaml
View file

@ -326,3 +326,137 @@ tests:
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: each entry must be a mapping with a 'path' key"
- it: routes a custom Service by number without replacing the existing ingress
set:
ingress.enabled: true
ingress.className: alb
ingress.host: gateway.example.com
ingress.extraPaths:
- path: /moyai
serviceName: moyai-ingress
servicePort: 8080
asserts:
- hasDocuments:
count: 1
- equal:
path: metadata.name
value: RELEASE-NAME-litellm
- equal:
path: spec.ingressClassName
value: alb
- equal:
path: spec.rules[0].host
value: gateway.example.com
- equal:
path: spec.rules[0].http.paths[-2]
value:
path: /moyai
pathType: Prefix
backend:
service:
name: moyai-ingress
port:
number: 8080
- equal:
path: spec.rules[0].http.paths[-1].backend.service.name
value: RELEASE-NAME-litellm-backend
- equal:
path: spec.rules[0].http.paths[0].backend.service.name
value: RELEASE-NAME-litellm-ui
- it: routes a custom Service by named port alongside a component alias
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
serviceName: moyai-ingress
servicePort: http
- path: /watsonx
service: gateway
asserts:
- equal:
path: spec.rules[0].http.paths[-3].backend.service
value:
name: moyai-ingress
port:
name: http
- equal:
path: spec.rules[0].http.paths[-2].backend.service.name
value: RELEASE-NAME-litellm-gateway
- it: rejects ambiguous component and custom Service selectors
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
service: gateway
serviceName: moyai-ingress
servicePort: 8080
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: service cannot be combined with serviceName or servicePort'
- it: rejects a custom Service without its port
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
serviceName: moyai-ingress
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: a non-empty serviceName and servicePort are required for a custom Service'
- it: rejects a custom Service port without its name
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
servicePort: 8080
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: a non-empty serviceName and servicePort are required for a custom Service'
- it: rejects a non scalar port
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
serviceName: moyai-ingress
servicePort: {number: 8080}
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: servicePort must be a port number or name'
- it: rejects a fractional port
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
serviceName: moyai-ingress
servicePort: 8080.5
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: numeric servicePort must be an integer from 1 to 65535'
- it: rejects an out of range port
set:
ingress.enabled: true
ingress.extraPaths:
- path: /moyai
serviceName: moyai-ingress
servicePort: 65536
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: numeric servicePort must be an integer from 1 to 65535'
- it: rejects a custom Service taking over a built in path
set:
ingress.enabled: true
ingress.extraPaths:
- path: /v1/chat
serviceName: moyai-ingress
servicePort: 8080
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0]: path /v1/chat with pathType Prefix is already routed by this chart, and a duplicate would take it over rather than add to it'

631
litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py Normal file
View file

@ -0,0 +1,631 @@
"""Read-only checks that the rows already in the database can take the pending schema
changes, run before `prisma migrate deploy` and before each request-log index build.
Postgres has a short list of schema statements that fail because of existing rows: a
btree index whose entry for some row would not fit in a third of a page (2704 bytes on
the default 8 KB page, counting every indexed column), UNIQUE (duplicate values), SET NOT NULL and ADD COLUMN NOT NULL with no default (nulls, any row at all),
FOREIGN KEY (orphans) and CHECK (rows failing the expression). Each pending statement
of one of those kinds becomes one count query derived from the statement text and the
column types in the catalog. A count above zero names the table, the column, the
statement it blocks and the SQL that clears it, and the migration is refused before
anything is applied. Every other statement a migration can contain cannot fail on
existing data and gets no query. A statement the derivation cannot read gets none
either, so the migration keeps today's behavior for it.
"""
import os
import re
from collections.abc import Callable, Iterator
from dataclasses import dataclass
from functools import reduce
from itertools import accumulate
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
from litellm_proxy_extras._logging import logger
if TYPE_CHECKING:
import psycopg
from psycopg import sql
DEFAULT_BLOCK_SIZE: Final = 8192
INDEX_TUPLE_HEADER_BYTES: Final = 8
NULL_BITMAP_BYTES: Final = 8
VARLENA_HEADER_BYTES: Final = 4
SHORT_VARLENA_MAX_BYTES: Final = 126
_PAGE_HEADER_WITH_THREE_LINE_POINTERS_BYTES: Final = 40
_BTREE_PAGE_OPAQUE_BYTES: Final = 16
_HEAP_TID_BYTES: Final = 8
_ALIGNMENT: Final = MappingProxyType({"c": 1, "s": 2, "i": 4, "d": 8})
_OCTET_LENGTH_TYPES: Final = frozenset({"text", "varchar", "bpchar", "bytea", "name"})
MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIMEOUT"
DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0
REFUSAL: Final = "Existing rows would make a pending schema change fail:"
Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"]
@dataclass(frozen=True, slots=True)
class Check:
kind: Kind
table: str
columns: tuple[str, ...]
statement: str
source: str
index: str | None = None
referenced_table: str | None = None
referenced_columns: tuple[str, ...] = ()
expression: str | None = None
@dataclass(frozen=True, slots=True)
class BlockedChange:
check: Check
rows: int
detail: str
@dataclass(frozen=True, slots=True)
class CouldNotVerify:
check: Check
timeout_seconds: float
Finding = BlockedChange | CouldNotVerify
@dataclass(frozen=True, slots=True)
class ColumnType:
"""One column as `pg_attribute` and `pg_type` describe it: `length` is `attlen` (-1 for a
variable-length value), `align` is `attalign` (c, s, i or d), `typmod` is `atttypmod`."""
type_name: str
length: int
align: str
typmod: int
@property
def is_variable_length(self) -> bool:
return self.length == -1
@property
def can_overflow_btree(self) -> bool:
"""Whether one value of this column alone can push a btree entry past the limit: a
variable-length type with no declared cap, or a cap of more bytes than the limit."""
if not self.is_variable_length:
return False
return self.typmod == -1 or (self.typmod - VARLENA_HEADER_BYTES) * 4 > btree_max_entry_bytes(DEFAULT_BLOCK_SIZE)
ColumnTypeLookup = Callable[[str, str], ColumnType | None]
_IDENT: Final = r'"?(\w+)"?'
_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?'
_CREATE_INDEX: Final = re.compile(
rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?"
rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)",
re.IGNORECASE | re.DOTALL,
)
_ALTER_TABLE: Final = re.compile(
rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL
)
_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE)
_ADD_COLUMN: Final = re.compile(
rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$",
re.IGNORECASE | re.DOTALL,
)
_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?"
_ADD_FOREIGN_KEY: Final = re.compile(
rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$",
re.IGNORECASE | re.DOTALL,
)
_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL)
_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL)
_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE)
_COMMENT: Final = re.compile(r"--[^\n]*")
_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE)
_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE)
def migration_data_check_timeout_seconds() -> float:
"""Seconds one count query may run before the preflight gives up on verifying it."""
raw: Final = os.getenv(MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR)
if raw is None:
return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
try:
return float(raw)
except ValueError:
logger.warning(
"%s=%r is not a number, using %s",
MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR,
raw,
DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS,
)
return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
def statements(migration_sql: str) -> tuple[str, ...]:
"""The statements of one migration file, comments stripped, empty ones dropped."""
parts: Final = _COMMENT.sub("", migration_sql).split(";")
return tuple(part.strip() for part in parts if part.strip())
def _split_top_level(text: str) -> tuple[str, ...]:
"""Split on the commas outside parentheses."""
depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text)
cuts: Final = tuple(
position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0
)
starts: Final = (0, *(cut + 1 for cut in cuts))
stops: Final = (*cuts, len(text))
pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True))
return tuple(piece for piece in pieces if piece)
def _plain_columns(column_list: str) -> tuple[str, ...] | None:
"""The column names of an index or constraint column list, None when one entry is an
expression the derivation cannot turn into a count."""
matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list))
if any(match is None for match in matches):
return None
return tuple(match.group(1) for match in matches if match is not None)
def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]:
unique, name, table, method, column_list = match.groups()
columns: Final = _plain_columns(column_list)
if columns is None:
return
label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"'
if method is None or method.lower() == "btree":
yield Check("too_long", table, columns, label, source, index=name)
if unique:
yield Check("duplicates", table, columns, label, source, index=name)
def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]:
if (set_not_null := _SET_NOT_NULL.match(action)) is not None:
yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source)
return
if (add_column := _ADD_COLUMN.match(action)) is not None:
column, body = add_column.groups()
if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None:
yield Check(
"any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source
)
return
if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None:
constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups()
own: Final = _plain_columns(columns)
theirs: Final = _plain_columns(referenced_columns)
if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper():
return
yield Check(
"orphans",
table,
own,
_constraint_label(table, constraint, "FOREIGN KEY"),
source,
referenced_table=referenced_table,
referenced_columns=theirs,
)
return
if (unique := _ADD_UNIQUE.match(action)) is not None:
constraint, columns = unique.groups()
unique_columns: Final = _plain_columns(columns)
if unique_columns is not None:
yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source)
return
if (check := _ADD_CHECK.match(action)) is not None:
constraint, expression, not_valid = check.groups()
if not_valid is None:
yield Check(
"check_violations",
table,
(),
_constraint_label(table, constraint, "CHECK"),
source,
expression=expression,
)
def _constraint_label(table: str, constraint: str | None, kind: str) -> str:
named: Final = f'CONSTRAINT "{constraint}" ' if constraint else ""
return f'ALTER TABLE "{table}" ADD {named}{kind}'
def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]:
"""The count queries one SQL statement needs before it can be applied, empty for a
statement that cannot fail on existing rows or that the derivation cannot read."""
if (index := _CREATE_INDEX.match(statement_text)) is not None:
return tuple(_index_checks(index, statement_text, source))
if (alter := _ALTER_TABLE.match(statement_text)) is not None:
table, actions = alter.groups()
return tuple(
check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source)
) # comprehension-ok: flattening the per-action checks of one statement
return ()
def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]:
return tuple(
check
for statement_text in statements(migration_sql)
for check in checks_for_statement(statement_text, f"migration {name}")
) # comprehension-ok: flattening the per-statement checks of one migration
def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]:
"""The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`."""
return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),)
def _applied_migrations(connection: "psycopg.Connection[tuple[object, ...]]", schema: str) -> frozenset[str]:
import psycopg
from psycopg import sql
try:
rows: Final = connection.execute(
sql.SQL("SELECT migration_name FROM {} WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL").format(
sql.Identifier(schema, "_prisma_migrations")
)
).fetchall()
except psycopg.errors.UndefinedTable:
return frozenset()
return frozenset(str(row[0]) for row in rows)
def pending_migrations(
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, migrations_dir: Path
) -> tuple[tuple[str, str], ...]:
"""(name, sql) of every shipped migration the ledger has not finished, oldest first.
Without a ledger every migration is pending, and the checks skip tables that do not exist yet."""
applied: Final = _applied_migrations(connection, schema)
folders: Final = tuple(folder for folder in sorted(migrations_dir.iterdir()) if folder.is_dir())
return tuple(
(folder.name, (folder / "migration.sql").read_text())
for folder in folders
if folder.name not in applied and (folder / "migration.sql").is_file()
)
def _int(value: object, default: int) -> int:
return value if isinstance(value, int) else default
def column_type_lookup(connection: "psycopg.Connection[tuple[object, ...]]", schema: str) -> ColumnTypeLookup:
def lookup(table: str, column: str) -> ColumnType | None:
row: Final = connection.execute(
"SELECT t.typname, a.attlen, a.attalign, a.atttypmod FROM pg_attribute a "
"JOIN pg_class c ON c.oid = a.attrelid JOIN pg_namespace n ON n.oid = c.relnamespace "
"JOIN pg_type t ON t.oid = a.atttypid "
"WHERE n.nspname = %s AND c.relname = %s AND a.attname = %s AND a.attnum > 0 AND NOT a.attisdropped",
(schema, table, column),
).fetchone()
if row is None:
return None
return ColumnType(str(row[0]), _int(row[1], -1), str(row[2]), _int(row[3], -1))
return lookup
def _column_types(check: Check, column_type: ColumnTypeLookup) -> tuple[ColumnType, ...] | None:
"""The catalog types of the check's columns, None when one of them is not there yet."""
types: Final = tuple(column_type(check.table, column) for column in check.columns)
if any(column is None for column in types):
return None
return tuple(column for column in types if column is not None)
def _table_exists(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str) -> bool:
row: Final = connection.execute(
"SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name = %s", (schema, table)
).fetchone()
return row is not None
def _index_is_valid(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: str) -> bool:
row: Final = connection.execute(
"SELECT x.indisvalid FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid "
"JOIN pg_namespace n ON n.oid = i.relnamespace WHERE n.nspname = %s AND i.relname = %s",
(schema, index),
).fetchone()
return row is not None and bool(row[0])
def btree_max_entry_bytes(block_size: int) -> int:
"""`BTMaxItemSize` from nbtree.h for btree version 4: a third of the page left after the page
header, three line pointers and the btree opaque area, rounded down to 8, minus the heap TID a
pivot tuple carries. 2704 on the default 8 KB page, the number in Postgres's error."""
usable: Final = block_size - _PAGE_HEADER_WITH_THREE_LINE_POINTERS_BYTES - _BTREE_PAGE_OPAQUE_BYTES
return (usable // 3) // 8 * 8 - _HEAP_TID_BYTES
def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
from psycopg import sql
if column.type_name in _OCTET_LENGTH_TYPES:
return sql.SQL("octet_length({})").format(identifier)
return sql.SQL("octet_length({}::text)").format(identifier)
def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":
from psycopg import sql
unit: Final = _ALIGNMENT.get(align, 8)
return sql.SQL("(({off} + {pad}) / {unit}) * {unit}").format(
off=offset, pad=sql.Literal(unit - 1), unit=sql.Literal(unit)
)
def _value_end(column: ColumnType, identifier: "sql.Identifier", offset: "sql.Composable") -> "sql.Composed":
"""Where this column's value ends when laid down from `offset` the way `index_form_tuple`
does it: NULL takes no room, a fixed-length value is aligned to its type, a variable-length
value of up to 126 bytes gets a 1-byte header and no alignment, a longer one a 4-byte header."""
from psycopg import sql
length: Final = _value_bytes(column, identifier)
placed: Final = (
sql.SQL("{} + {}").format(_aligned(offset, column.align), sql.Literal(column.length))
if not column.is_variable_length
else sql.SQL("CASE WHEN {len} <= {short} THEN {off} + 1 + {len} ELSE {aligned} + {hdr} + {len} END").format(
len=length,
short=sql.Literal(SHORT_VARLENA_MAX_BYTES),
off=offset,
aligned=_aligned(offset, column.align),
hdr=sql.Literal(VARLENA_HEADER_BYTES),
)
)
return sql.SQL("CASE WHEN {col} IS NULL THEN {off} ELSE {placed} END").format(
col=identifier, off=offset, placed=placed
)
def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
"""count(*) and the largest entry over the rows whose btree entry would not fit, plus the
longest value of each variable-length column: an 8-byte tuple header, 8 more for the null
bitmap when a column is NULL, each value laid down in index order, the total rounded up to 8."""
from psycopg import sql
identifiers: Final = tuple(sql.Identifier(column) for column in check.columns)
any_null: Final = sql.SQL(" OR ").join(sql.SQL("{} IS NULL").format(identifier) for identifier in identifiers)
header: Final = sql.SQL("{} + CASE WHEN {} THEN {} ELSE 0 END").format(
sql.Literal(INDEX_TUPLE_HEADER_BYTES), any_null, sql.Literal(NULL_BITMAP_BYTES)
)
innermost: Final = sql.SQL("SELECT {cols}, {hdr} AS off0 FROM {t}").format(
cols=sql.SQL(", ").join(identifiers), hdr=header, t=sql.Identifier(schema, check.table)
)
def lay_down(inner: "sql.Composed", step: tuple[int, tuple[ColumnType, "sql.Identifier"]]) -> "sql.Composed":
position, (column, identifier) = step
return sql.SQL("SELECT *, {end} AS {next} FROM ({inner}) AS {alias}").format(
end=_value_end(column, identifier, sql.Identifier(f"off{position}")),
next=sql.Identifier(f"off{position + 1}"),
inner=inner,
alias=sql.Identifier(f"s{position}"),
)
laid_out: Final = reduce(lay_down, enumerate(zip(types, identifiers, strict=True)), innermost)
entry: Final = _aligned(sql.Identifier(f"off{len(types)}"), "d")
longest: Final = sql.SQL("").join(
sql.SQL(", max({})").format(_value_bytes(column, identifier))
for column, identifier in zip(types, identifiers, strict=True)
if column.is_variable_length
)
return sql.SQL("SELECT count(*), max({entry}){longest} FROM ({rows}) AS entries WHERE {entry} > {limit}").format(
entry=entry, longest=longest, rows=laid_out, limit=sql.Literal(limit_bytes)
)
def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
from psycopg import sql
table: Final = sql.Identifier(schema, check.table)
columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns)
match check.kind:
case "too_long":
return _entry_size_query(check, schema, types, limit_bytes)
case "duplicates":
not_null: Final = sql.SQL(" AND ").join(
sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d"
).format(t=table, nn=not_null, cols=columns)
case "nulls":
return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format(
t=table, col=sql.Identifier(check.columns[0])
)
case "any_row":
return sql.SQL("SELECT count(*) FROM {t}").format(t=table)
case "orphans":
pairs: Final = sql.SQL(" AND ").join(
sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own))
for own, theirs in zip(check.columns, check.referenced_columns, strict=True)
)
present: Final = sql.SQL(" AND ").join(
sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})"
).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs)
case "check_violations":
return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format(
t=table,
expr=sql.SQL(check.expression or "true"), # pyright: ignore[reportArgumentType] # copied from our own migration file
)
def _applies(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
check: Check,
types: tuple[ColumnType, ...] | None,
) -> bool:
"""False when the check has nothing to count yet: the table or a column is not there
(a pending migration adds it), the index is already valid, or no column can overflow."""
if not _table_exists(connection, schema, check.table):
return False
if check.index is not None and _index_is_valid(connection, schema, check.index):
return False
if check.kind == "any_row":
return types is None
if types is None:
return False
if check.kind == "too_long":
return any(column.can_overflow_btree for column in types)
if check.kind == "orphans":
return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table)
return True
def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
variable: Final = tuple(
column for column, column_type in zip(check.columns, types, strict=True) if column_type.is_variable_length
)
longest: Final = ", ".join(
f"{column} up to {length} bytes"
for column, length in zip(variable, values[2:], strict=True)
if length is not None
)
return f"with an index entry over {limit_bytes} bytes, largest {values[1]} bytes ({longest})"
def _run_check(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
check: Check,
types: tuple[ColumnType, ...],
limit_bytes: int,
timeout_seconds: float,
) -> Finding | None:
import psycopg
try:
row: Final = connection.execute(_count_query(check, schema, types, limit_bytes)).fetchone()
except psycopg.errors.QueryCanceled:
return CouldNotVerify(check, timeout_seconds)
values: Final[tuple[object, ...]] = row if row is not None else ()
rows: Final = _int(values[0], 0) if values else 0
if rows == 0:
return None
match check.kind:
case "too_long":
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
case "duplicates":
return BlockedChange(check, rows, "duplicated value groups")
case "nulls":
return BlockedChange(check, rows, "NULL")
case "any_row":
return BlockedChange(
check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows"
)
case "orphans":
return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"')
case "check_violations":
return BlockedChange(check, rows, f"failing CHECK ({check.expression})")
def find_blockers(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
checks: tuple[Check, ...],
timeout_seconds: float = DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS,
) -> tuple[Finding, ...]:
"""Run every applicable check read-only, each under `timeout_seconds`. Never writes
and never takes a table lock; a check that times out is a finding, not a pass."""
from psycopg import sql
previous_timeout: Final = connection.execute("SHOW statement_timeout").fetchone()
previous_read_only: Final = connection.execute("SHOW default_transaction_read_only").fetchone()
block_size: Final = connection.execute("SHOW block_size").fetchone()
limit_bytes: Final = btree_max_entry_bytes(int(str(block_size[0])) if block_size else DEFAULT_BLOCK_SIZE)
connection.execute("SET default_transaction_read_only = on")
connection.execute(sql.SQL("SET statement_timeout = {}").format(sql.Literal(int(timeout_seconds * 1000))))
column_type: Final = column_type_lookup(connection, schema)
typed: Final = tuple((check, _column_types(check, column_type)) for check in checks)
applicable: Final = tuple(
(check, types if types is not None else ())
for check, types in typed
if _applies(connection, schema, check, types)
)
results: Final = tuple(
_run_check(connection, schema, check, types, limit_bytes, timeout_seconds) for check, types in applicable
)
connection.execute(
sql.SQL("SET statement_timeout = {}").format(sql.Literal(str(previous_timeout[0]) if previous_timeout else "0"))
)
connection.execute(
sql.SQL("SET default_transaction_read_only = {}").format(
sql.Literal(str(previous_read_only[0]) if previous_read_only else "off")
)
)
return tuple(result for result in results if result is not None)
def _remedy(check: Check) -> str:
joined: Final = ", ".join(check.columns)
match check.kind:
case "too_long":
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
case "duplicates":
return f'Remove or merge the duplicate "{check.table}" ({joined}) values.'
case "nulls":
return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.'
case "any_row":
return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards."
case "orphans":
return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.'
case "check_violations":
return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.'
def _describe(finding: Finding) -> str:
check: Final = finding.check
where: Final = f'"{check.table}"' + (f".{', '.join(check.columns)}" if check.columns else "")
match finding:
case CouldNotVerify(timeout_seconds=seconds):
return (
f" {where}: could not be checked within {seconds:g} s, refusing rather than guessing\n"
f" blocks: {check.statement} ({check.source})\n"
f" raise {MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR} to give the scan more time"
)
case BlockedChange(rows=rows, detail=detail):
return (
f" {where}: {rows} row(s) {detail}\n"
f" blocks: {check.statement} ({check.source})\n"
f" fix: {_remedy(check).replace(chr(10), chr(10) + ' ' * 8)}"
)
def report(findings: tuple[Finding, ...]) -> str:
"""The operator-facing refusal: one block per finding, then what to do next."""
lines: Final = "\n".join(_describe(finding) for finding in findings)
return (
f"{REFUSAL}\n{lines}\n"
"No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again."
)
def pending_change_checks(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
migrations_dir: Path | None,
indexes: tuple[tuple[str, str, tuple[str, ...]], ...],
) -> tuple[Check, ...]:
"""Every check the pending migrations and the request-log indexes need. `migrations_dir`
is None when the schema is pushed rather than migrated, so only the indexes are checked."""
migration_checks: Final = (
tuple(
check
for name, migration_sql in pending_migrations(connection, schema, migrations_dir)
for check in checks_for_migration(name, migration_sql)
) # comprehension-ok: flattening the per-migration checks
if migrations_dir is not None
else ()
)
index_checks: Final = tuple(
check for table, name, columns in indexes for check in checks_for_index(table, name, columns)
) # comprehension-ok: flattening the per-index checks
return migration_checks + index_checks

12
litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py
View file

@ -18,6 +18,12 @@ from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.migration_data_checks import (
checks_for_index,
find_blockers,
migration_data_check_timeout_seconds,
report,
)
from litellm_proxy_extras.migration_lock import held_migration_lock
if TYPE_CHECKING:
@ -355,6 +361,12 @@ def _build_leaf_index(
connection.execute(sql.SQL("DROP INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)))
elif _adopt_equivalent_index(connection, schema, table, name, index):
return True
blockers: Final = find_blockers(
connection, schema, checks_for_index(table, name, index.columns), migration_data_check_timeout_seconds()
)
if blockers:
logger.warning("Not building index %s on %s. %s", name, table, report(blockers))
return False
logger.info("Building index %s on %s concurrently", name, table)
prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)

226
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -16,6 +16,13 @@ from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.migration_data_checks import (
Finding,
find_blockers,
migration_data_check_timeout_seconds,
pending_change_checks,
report,
)
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@ -28,7 +35,11 @@ from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes, filter_request_log_index_diff
from litellm_proxy_extras.request_log_indexes import (
REQUEST_LOG_INDEXES,
ensure_request_log_indexes,
filter_request_log_index_diff,
)
if TYPE_CHECKING:
import psycopg
@ -53,9 +64,7 @@ def _get_prisma_env() -> dict:
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
# These env vars prevent Prisma from attempting downloads
prisma_env["NPM_CONFIG_PREFER_OFFLINE"] = "true"
prisma_env["NPM_CONFIG_CACHE"] = os.getenv(
"NPM_CONFIG_CACHE", "/app/.cache/npm"
)
prisma_env["NPM_CONFIG_CACHE"] = os.getenv("NPM_CONFIG_CACHE", "/app/.cache/npm")
return prisma_env
@ -86,6 +95,7 @@ class _InvalidIndex:
name: str
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
LIBPQ_URL_PARAMS: Final = frozenset(
{
@ -135,9 +145,7 @@ class _MigrateAttemptBudget:
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE)
_SPEND_LOGS_PK_CLAUSE_RE = re.compile(
r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
@ -146,9 +154,9 @@ _SPEND_LOGS_PK_CLAUSE_RE = re.compile(
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
'so its primary key must include the partition key ("startTime"). `prisma db push` '
"reconciles the database against schema.prisma, which declares the unpartitioned "
"primary key (\"request_id\"), and Postgres rejects that rewrite with: unique "
'primary key ("request_id"), and Postgres rejects that rewrite with: unique '
"constraint on partitioned table must include all partitioning columns. Start the "
"proxy without --use_prisma_db_push so it uses `prisma migrate deploy`, which only "
"applies shipped migrations and leaves the partitioned primary key alone."
@ -157,9 +165,7 @@ PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
def _without_sql_comments(statement: str) -> str:
return "\n".join(
line
for line in statement.splitlines()
if line.strip() and not line.strip().startswith("--")
line for line in statement.splitlines() if line.strip() and not line.strip().startswith("--")
).strip()
@ -169,7 +175,7 @@ def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
return statement
kept = tuple(
clause.strip()
for clause in statement[prefix_match.end():].split(",\n")
for clause in statement[prefix_match.end() :].split(",\n")
if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
)
if not kept:
@ -243,8 +249,7 @@ def _get_prisma_command() -> str:
# If not found, log warning and fall back
logger.warning(
f"Prisma CLI not found at {default_cli_path}. "
"Falling back to Python wrapper (may attempt downloads)"
f"Prisma CLI not found at {default_cli_path}. Falling back to Python wrapper (may attempt downloads)"
)
# Fall back to the Python wrapper (will work in online mode)
@ -331,9 +336,7 @@ class ProxyExtrasDBManager:
return True
except subprocess.TimeoutExpired:
logger.warning(
"Migration timed out - the database might be under heavy load."
)
logger.warning("Migration timed out - the database might be under heavy load.")
return False
except subprocess.CalledProcessError as e:
logger.warning(
@ -386,10 +389,12 @@ class ProxyExtrasDBManager:
except ImportError:
return None
from psycopg import sql
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
ledger_table: Final = psycopg.sql.SQL("{}.{}").format(
psycopg.sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"),
psycopg.sql.Identifier("_prisma_migrations"),
ledger_table: Final = sql.SQL("{}.{}").format(
sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"),
sql.Identifier("_prisma_migrations"),
)
try:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
@ -500,9 +505,7 @@ class ProxyExtrasDBManager:
return filter_partitioned_spend_logs_diff(without_indexes)
@staticmethod
def _resolve_all_migrations(
migrations_dir: str, schema_path: str, mark_all_applied: bool = True
):
def _resolve_all_migrations(migrations_dir: str, schema_path: str, mark_all_applied: bool = True):
"""
1. Compare the current database state to schema.prisma and generate a migration for the diff.
2. Run prisma migrate deploy to apply any pending migrations.
@ -571,9 +574,7 @@ class ProxyExtrasDBManager:
)
logger.info(f"Applied migration: {mig_file.parent.name}")
except subprocess.CalledProcessError as e:
logger.warning(
f"Failed to apply migration {mig_file.parent.name}: {e.stderr}"
)
logger.warning(f"Failed to apply migration {mig_file.parent.name}: {e.stderr}")
except subprocess.TimeoutExpired:
logger.warning(f"Migration {mig_file.parent.name} timed out.")
return
@ -645,9 +646,7 @@ class ProxyExtrasDBManager:
logger.debug(f"Resolved migration: {migration_name}")
except subprocess.CalledProcessError as e:
if "is already recorded as applied in the database." not in e.stderr:
logger.warning(
f"Failed to resolve migration {migration_name}: {e.stderr}"
)
logger.warning(f"Failed to resolve migration {migration_name}: {e.stderr}")
@staticmethod
def raise_if_lens_rename_pending(connect: LensCheckConnect | None = None) -> None:
@ -707,9 +706,7 @@ class ProxyExtrasDBManager:
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
try:
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
row = conn.execute(
"SELECT 1 "
"FROM pg_partitioned_table pt "
@ -717,10 +714,7 @@ class ProxyExtrasDBManager:
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"WHERE c.relname = 'LiteLLM_SpendLogs' "
" AND n.nspname = %s",
(
ProxyExtrasDBManager._prisma_schema_param(database_url)
or "public",
),
(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",),
).fetchone()
except (psycopg.OperationalError, psycopg.DatabaseError):
return False
@ -815,9 +809,7 @@ class ProxyExtrasDBManager:
# it, psycopg3's `with conn` calls COMMIT on clean exit — which
# fails after `UndefinedTable` (fresh DB) leaves the transaction
# in an aborted state.
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
try:
rows = conn.execute(
"SELECT migration_name FROM _prisma_migrations "
@ -838,9 +830,7 @@ class ProxyExtrasDBManager:
return
head_newest_ts = _max_migration_timestamp(known)
hostile = {
name for name in unknown if _migration_timestamp(name) > head_newest_ts
}
hostile = {name for name in unknown if _migration_timestamp(name) > head_newest_ts}
if not hostile:
return
@ -931,8 +921,7 @@ class ProxyExtrasDBManager:
if not found:
return True
logger.warning(
"Found %d invalid index(es) left by an interrupted CREATE INDEX "
"CONCURRENTLY, rebuilding: %s",
"Found %d invalid index(es) left by an interrupted CREATE INDEX CONCURRENTLY, rebuilding: %s",
len(found),
", ".join(f"{index.name} (table size {index.table_size})" for index in found),
)
@ -1403,8 +1392,7 @@ class ProxyExtrasDBManager:
prisma_env = _get_prisma_env()
except OSError as e:
logger.error(
"Could not resolve the migrations directory for the REPLICA "
"IDENTITY FULL step, skipping it. Error: %s",
"Could not resolve the migrations directory for the REPLICA IDENTITY FULL step, skipping it. Error: %s",
e,
)
return False
@ -1435,18 +1423,57 @@ class ProxyExtrasDBManager:
Returns:
bool: True if setup was successful, False otherwise
"""
ensure_prisma_toolchain(
prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env()
)
migrated = ProxyExtrasDBManager._run_migrations(
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env())
ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate)
migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver)
if not migrated:
return False
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return True
@staticmethod
def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]:
"""What in the existing rows would make a pending migration or request-log index
fail, read with count queries only: no write, no table lock. Empty when nothing
blocks, when no database is configured, or when it cannot be reached, so the
migration step that follows reports the connection problem the way it does today."""
import psycopg
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return ()
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
os.environ.get("DIRECT_URL") or database_url
)
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None
indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES)
try:
with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection:
checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes)
return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds())
except psycopg.OperationalError as exc:
logger.warning(
"Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc))
)
return ()
@staticmethod
def migration_data_check_report(use_migrate: bool = True) -> Optional[str]:
"""The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks."""
findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate)
return report(findings) if findings else None
@staticmethod
def raise_if_rows_block_pending_changes(use_migrate: bool) -> None:
"""Refuse to run migrations that the existing rows would make fail, before any is
applied. The RuntimeError carries the report and reaches the operator through the
proxy's "Database migration cannot proceed" exit."""
refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)
if refusal is not None:
raise RuntimeError(refusal)
@staticmethod
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
"""Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the
@ -1513,9 +1540,7 @@ class ProxyExtrasDBManager:
# Skip sanity check when deploy reports no pending migrations —
# DB already matches schema, no drift to correct.
if "No pending migrations to apply" in result.stdout:
logger.info(
"No pending migrations — skipping post-migration sanity check"
)
logger.info("No pending migrations — skipping post-migration sanity check")
return True
# Run sanity check to ensure DB matches schema
@ -1530,9 +1555,7 @@ class ProxyExtrasDBManager:
logger.info(f"prisma db error: {stderr}, e: {e.stdout}")
if "P3009" in stderr:
# Extract the failed migration name from the error message
migration_match = re.search(
r"`(\d+_.*)` migration", stderr
)
migration_match = re.search(r"`(\d+_.*)` migration", stderr)
if migration_match:
failed_migration = migration_match.group(1)
if ProxyExtrasDBManager._is_idempotent_error(stderr):
@ -1540,9 +1563,7 @@ class ProxyExtrasDBManager:
f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied"
)
try:
ProxyExtrasDBManager._roll_back_migration(
failed_migration
)
ProxyExtrasDBManager._roll_back_migration(failed_migration)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
@ -1552,9 +1573,7 @@ class ProxyExtrasDBManager:
f"It may already be in a rolled-back state."
)
try:
ProxyExtrasDBManager._resolve_specific_migration(
failed_migration
)
ProxyExtrasDBManager._resolve_specific_migration(failed_migration)
logger.info(
f"✅ Migration {failed_migration} resolved, retrying to apply remaining migrations"
)
@ -1562,9 +1581,7 @@ class ProxyExtrasDBManager:
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
logger.warning(
f"Failed to resolve migration {failed_migration}: {resolve_err}"
)
logger.warning(f"Failed to resolve migration {failed_migration}: {resolve_err}")
# Apply any schema drift not covered by the marked-as-applied migration
ProxyExtrasDBManager._resolve_all_migrations(
migrations_dir,
@ -1572,9 +1589,7 @@ class ProxyExtrasDBManager:
mark_all_applied=False,
)
else:
logger.info(
f"Found failed migration: {failed_migration}, marking as rolled back"
)
logger.info(f"Found failed migration: {failed_migration}, marking as rolled back")
# Mark the failed migration as rolled back
prisma_toolchain.run_prisma(
[
@ -1587,23 +1602,14 @@ class ProxyExtrasDBManager:
timeout=prisma_command_timeout(),
env=_get_prisma_env(),
)
logger.info(
f"✅ Migration {failed_migration} marked as rolled back... retrying"
)
elif (
"P3005" in stderr
and "database schema is not empty" in stderr
):
logger.info(f"✅ Migration {failed_migration} marked as rolled back... retrying")
elif "P3005" in stderr and "database schema is not empty" in stderr:
logger.info(
"Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
)
ProxyExtrasDBManager._create_baseline_migration(schema_path)
logger.info(
"Baseline migration created, resolving all migrations"
)
ProxyExtrasDBManager._resolve_all_migrations(
migrations_dir, schema_path
)
logger.info("Baseline migration created, resolving all migrations")
ProxyExtrasDBManager._resolve_all_migrations(migrations_dir, schema_path)
logger.info("✅ All migrations resolved.")
return True
elif "P3018" in stderr:
@ -1611,14 +1617,8 @@ class ProxyExtrasDBManager:
if ProxyExtrasDBManager._is_permission_error(stderr):
# Permission errors should NOT be marked as applied
# Extract migration name for logging
migration_match = re.search(
r"Migration name: (\d+_.*)", stderr
)
migration_name = (
migration_match.group(1)
if migration_match
else "unknown"
)
migration_match = re.search(r"Migration name: (\d+_.*)", stderr)
migration_name = migration_match.group(1) if migration_match else "unknown"
logger.error(
f"❌ Migration {migration_name} failed due to insufficient permissions. "
@ -1628,16 +1628,10 @@ class ProxyExtrasDBManager:
# Mark as rolled back and exit with error
if migration_match:
try:
ProxyExtrasDBManager._roll_back_migration(
migration_name
)
logger.info(
f"Migration {migration_name} marked as rolled back"
)
ProxyExtrasDBManager._roll_back_migration(migration_name)
logger.info(f"Migration {migration_name} marked as rolled back")
except Exception as rollback_error:
logger.warning(
f"Failed to mark migration as rolled back: {rollback_error}"
)
logger.warning(f"Failed to mark migration as rolled back: {rollback_error}")
# Re-raise the error to prevent silent failures
raise RuntimeError(
@ -1652,18 +1646,12 @@ class ProxyExtrasDBManager:
"resolving as applied"
)
# Extract the migration name from the error message
migration_match = re.search(
r"Migration name: (\d+_.*)", stderr
)
migration_match = re.search(r"Migration name: (\d+_.*)", stderr)
if migration_match:
migration_name = migration_match.group(1)
try:
logger.info(
f"Rolling back migration {migration_name}"
)
ProxyExtrasDBManager._roll_back_migration(
migration_name
)
logger.info(f"Rolling back migration {migration_name}")
ProxyExtrasDBManager._roll_back_migration(migration_name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
@ -1677,9 +1665,7 @@ class ProxyExtrasDBManager:
f"Resolving migration {migration_name} that failed "
f"due to existing schema objects"
)
ProxyExtrasDBManager._resolve_specific_migration(
migration_name
)
ProxyExtrasDBManager._resolve_specific_migration(migration_name)
logger.info(
f"✅ Migration {migration_name} resolved, "
f"retrying to apply remaining migrations"
@ -1688,9 +1674,7 @@ class ProxyExtrasDBManager:
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
logger.warning(
f"Failed to resolve migration {migration_name}: {resolve_err}"
)
logger.warning(f"Failed to resolve migration {migration_name}: {resolve_err}")
# Apply any schema drift not covered by the marked-as-applied migration
ProxyExtrasDBManager._resolve_all_migrations(
migrations_dir,
@ -1732,14 +1716,8 @@ class ProxyExtrasDBManager:
time.sleep(random.randrange(5, 15))
except subprocess.CalledProcessError as e:
attempts_left = 3 - attempt
retry_msg = (
f" Retrying... ({attempts_left} attempts left)"
if attempts_left > 0
else ""
)
stderr_detail: Final = (
f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else ""
)
retry_msg = f" Retrying... ({attempts_left} attempts left)" if attempts_left > 0 else ""
stderr_detail: Final = f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else ""
logger.error(
f"The process failed to execute. Details: {_redact_command_error(e)}.{stderr_detail}{retry_msg}"
)

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

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

@ -235,7 +235,7 @@ python_callbacks! {
"async_log_stream_event",
"async_post_mcp_tool_call_hook",
],
Dispatch::Python("litellm.llms.anthropic.pass_through.messages.handler") => [
Dispatch::Python("litellm.llms.anthropic.messages.native.handler") => [
"async_pre_request_hook",
],
Dispatch::Python("litellm.litellm_core_utils.streaming_handler") => [

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

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