scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 360

fix(chat): merge current main into signed-in team access (#35187)

Browse source
This commit is contained in:
ishaan-j-berri 2026-10-27 14:06:45 -07:00
commit 46fc5fd1ac
1410 changed files with 63258 additions and 23112 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

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

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

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

@ -579,13 +579,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 +603,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 +811,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 +828,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 +952,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

10
ARCHITECTURE.md
View file

@ -332,13 +332,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 +346,7 @@ Each translation is isolated in its own file, making it easy to test and modify
If `/v1/messages` → Bedrock Converse prompt caching isn't working but Bedrock Invoke works:
1. **Bedrock Converse translation**: `llms/bedrock/chat/converse_transformation.py`
2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py`
2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py`
3. Compare how each handles `cache_control` in `transform_request()`
### How translations work

4
Dockerfile
View file

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

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

4
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(

11
backend/Dockerfile
View file

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

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.

4
docker/Dockerfile.database
View file

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

4
docker/Dockerfile.non_root
View file

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

13
enterprise/enterprise_hooks/aporia_ai.py
View file

@ -15,9 +15,10 @@ sys.path.insert(
) # Adds the parent directory to the system path
import json
import sys
from typing import Any, List, Literal, Optional
from typing import Final, List, Literal, Optional
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -29,11 +30,13 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.llms.openai.common_utils import decoded_json_object
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.types.guardrails import GuardrailEventHooks
GUARDRAIL_NAME = "aporia"
_DECODED_JSON: Final = TypeAdapter(object)
class AporiaGuardrail(CustomGuardrail):
@ -68,7 +71,7 @@ class AporiaGuardrail(CustomGuardrail):
async def prepare_aporia_request(
self, new_messages: List[dict], response_string: Optional[str] = None
) -> dict:
data: dict[str, Any] = {}
data: dict[str, object] = {}
if new_messages is not None:
data["messages"] = new_messages
if response_string is not None:
@ -122,8 +125,10 @@ class AporiaGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Aporia AI response: %s", response.text)
if response.status_code == 200:
# check if the response was flagged
_json_response = response.json()
action: str = _json_response.get(
_json_response: Final = decoded_json_object(
_DECODED_JSON.validate_python(response.json())
)
action: Final = _json_response.get(
"action"
) # possible values are modify, passthrough, block, rephrase
if action == "block":

17
enterprise/enterprise_hooks/blocked_user_list.py
View file

@ -7,9 +7,10 @@
## This accepts a list of user id's for whom calls will be rejected
from typing import Literal, Optional
from typing import Final, Literal, Optional
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._internal_context import with_service_target
@ -19,6 +20,9 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import EndUserRepository
_REQUEST_BODY_VALUE: Final = TypeAdapter(object)
class ENTERPRISE_BlockedUserList(CustomLogger):
@ -94,15 +98,14 @@ class ENTERPRISE_BlockedUserList(CustomLogger):
)
if end_user_cache_obj is None and self.prisma_client is not None:
# check db
end_user_obj = (
await self.prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": user}
)
)
end_user_id: Final = _REQUEST_BODY_VALUE.validate_python(user)
end_user_obj = await EndUserRepository(
self.prisma_client
).table.find_unique(where={"user_id": end_user_id})
if end_user_obj is None: # user not in db - assume not blocked
end_user_obj = LiteLLM_EndUserTable(user_id=user, blocked=False)
cache.set_cache(key=cache_key, value=end_user_obj, ttl=60)
if end_user_obj is not None and end_user_obj.blocked is True:
if end_user_obj.blocked is True:
raise HTTPException(
status_code=400,
detail={

13
enterprise/enterprise_hooks/google_text_moderation.py
View file

@ -6,7 +6,10 @@
# +-----------------------------------------------+
# Thank you users! We ❤️ you! - Krrish & Ishaan
from typing import Final
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -15,6 +18,8 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import iter_message_text
from litellm.types.utils import CallTypesLiteral
_THRESHOLD_VALUE: Final = TypeAdapter(object)
class ENTERPRISE_GoogleTextModeration(CustomLogger):
user_api_key_cache = None
@ -69,9 +74,11 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger):
f"{category}_confidence_threshold",
default_confidence_threshold,
)
set_confidence_value = getattr(
self,
f"{category}_confidence_threshold",
set_confidence_value = _THRESHOLD_VALUE.validate_python(
getattr(
self,
f"{category}_confidence_threshold",
)
)
verbose_proxy_logger.info(
f"Google Text Moderation: {category}_confidence_threshold: {set_confidence_value}"

7
enterprise/litellm_enterprise/enterprise_callbacks/example_logging_api.py
View file

@ -1,15 +1,20 @@
# this is an example endpoint to receive data from litellm
from typing import Final
from fastapi import FastAPI, HTTPException, Request
from pydantic import TypeAdapter
app = FastAPI()
_LOGGED_EVENT: Final = TypeAdapter(object)
@app.post("/log-event")
async def log_event(request: Request):
try:
print("Received /log-event request") # noqa
# Assuming the incoming request has JSON data
data = await request.json()
data: Final = _LOGGED_EVENT.validate_python(await request.json())
print("Received request data:") # noqa
print(data) # noqa

18
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View file

@ -8,25 +8,30 @@
## This provides an LLM Guard Integration for content moderation on the proxy
import asyncio
from typing import Final, Optional
from collections.abc import Mapping
from typing import Final
import aiohttp
from fastapi import HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.openai.common_utils import decoded_json_object
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.utils import CallTypesLiteral
_DECODED_JSON: Final = TypeAdapter(object)
class _ENTERPRISE_LLMGuard(CustomLogger):
# Class variables or attributes
def __init__(
self,
mock_testing: bool = False,
mock_redacted_text: Optional[dict] = None,
mock_redacted_text: Mapping[str, object] | None = None,
):
self.mock_redacted_text = mock_redacted_text
self.llm_guard_mode = litellm.llm_guard_mode
@ -66,7 +71,9 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
async with session.post(
analyze_url, json={"prompt": text}
) as response:
redacted_text = await response.json()
redacted_text = _DECODED_JSON.validate_python(
await response.json()
)
verbose_proxy_logger.debug(
f"LLM Guard: Received response - {redacted_text}"
)
@ -77,12 +84,13 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
"error": f"Invalid content moderation response: {redacted_text}"
},
)
if redacted_text.get("is_valid", None) is False:
redacted_response: Final = decoded_json_object(redacted_text)
if redacted_response.get("is_valid", None) is False:
raise HTTPException(
status_code=400,
detail={"error": "Violated content safety policy"},
)
sanitized_prompt = redacted_text.get("sanitized_prompt")
sanitized_prompt: Final = redacted_response.get("sanitized_prompt")
return sanitized_prompt if isinstance(sanitized_prompt, str) else text
except Exception as e:
verbose_proxy_logger.exception(

4
enterprise/litellm_enterprise/proxy/auth/route_checks.py
View file

@ -2,6 +2,8 @@ import os
from fastapi import HTTPException, status
from litellm.proxy._types import MODEL_LIST_ROUTES
class EnterpriseRouteChecks:
@staticmethod
@ -43,7 +45,7 @@ class EnterpriseRouteChecks:
# Routes that should remain accessible even when LLM API endpoints are disabled.
# These are read-only model listing routes needed by the Admin UI.
LLM_API_EXEMPT_ROUTES = ["/models", "/v1/models"]
LLM_API_EXEMPT_ROUTES = MODEL_LIST_ROUTES
@staticmethod
def should_call_route(route: str):

4
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -1040,8 +1040,8 @@ class CheckBatchCost:
from litellm.proxy.proxy_server import proxy_config
from litellm.secret_managers.main import str_to_bool
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
api_key: Final = job.api_key if hasattr(job, "api_key") else None
team_id: Final = job.team_id if hasattr(job, "team_id") else None
has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
try:
key_row: Final = (

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

@ -44,6 +44,7 @@ from litellm.constants import (
)
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.internal_call_metadata import requests_background
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -133,7 +134,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
Span = Union[_Span, Any]
Span = _Span
InternalUsageCache = _InternalUsageCache
PrismaClient = _PrismaClient
else:
@ -490,6 +491,14 @@ _PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDE
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
class _SupportsKeyedGet(Protocol):
def get(self, key: str, default: object, /) -> object: ...
def _keyed_get(source: object, key: str, default: object = None) -> object:
return require_attributes(source, _SupportsKeyedGet, "get").get(key, default)
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(
@ -1611,7 +1620,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return holders if holders or "input_file_id" in request else healthy_deployments
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
self, kwargs: dict[str, object], call_type: Optional[CallTypes]
) -> Optional[dict]:
"""
Allow modifying the request just before it's sent to the deployment.
@ -1623,11 +1632,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
# model_info may be at top-level or nested under litellm_metadata
# (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
model_id = cast(Optional[str], _keyed_get(kwargs.get("model_info", {}), "id"))
if model_id is None:
model_id = cast(
Optional[str],
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
_keyed_get(_keyed_get(kwargs.get("litellm_metadata", {}), "model_info", {}), "id"),
)
for accessor_key in accessor_keys:
input_file_id = cast(Optional[str], kwargs.get(accessor_key))
@ -1824,7 +1833,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
)
async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span) -> dict:
async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span | None) -> dict:
"""
Get model-specific file IDs for a list of proxy file IDs.
Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id
@ -1866,7 +1875,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
llm_router: Optional[Router],
_create_file_request: CreateFileRequest,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
) -> List[OpenAIFileObject]:
if llm_router is None:
raise Exception("LLM Router not initialized. Ensure models added to proxy.")
@ -1882,7 +1891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
create_file_request: CreateFileRequest,
llm_router: Router,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
) -> OpenAIFileObject:
responses = await self.create_file_for_each_model(
@ -1949,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_objects: List[OpenAIFileObject],
create_file_request: CreateFileRequest,
internal_usage_cache: InternalUsageCache,
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
target_model_names_list: List[str],
) -> OpenAIFileObject:
## GET THE FILE TYPE FROM THE CREATE FILE REQUEST
@ -2652,7 +2661,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Continue with other files even if one fails
continue
def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str:
def _get_content_type_from_file_object(self, file_object: object) -> str:
"""
Determine content type from file object.

6
enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
View file

@ -29,7 +29,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
Span = Union[_Span, Any]
Span = _Span
InternalUsageCache = _InternalUsageCache
PrismaClient = _PrismaClient
else:
@ -107,7 +107,7 @@ class _PROXY_LiteLLMManagedVectorStores(
llm_router: Router,
model: str,
request_data: dict[str, object] | VectorStoreCreateOptionalRequestParams,
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
) -> VectorStoreCreateResponse:
"""
Create a vector store for a specific model.
@ -134,7 +134,7 @@ class _PROXY_LiteLLMManagedVectorStores(
create_request: VectorStoreCreateOptionalRequestParams,
llm_router: Router,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
) -> VectorStoreCreateResponse:
"""

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,
)

13
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -36,6 +36,7 @@ from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import PrismaTableRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -73,6 +74,16 @@ def _verification_token_table(
return VerificationTokenRepository(prisma_client).table
class _WriterVerificationTokenRepository(PrismaTableRepository["prisma_models.LiteLLM_VerificationToken"]):
table_name = "litellm_verificationtoken"
def _writer_verification_token_table(
prisma_client: PrismaClient,
) -> TableActions["prisma_models.LiteLLM_VerificationToken"]:
return _WriterVerificationTokenRepository(prisma_client, use_writer=True).table
def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]:
return BudgetRepository(prisma_client).table
@ -849,7 +860,7 @@ async def update_project(
else None
)
if current_project is not None and data.team_id != current_project.team_id:
mismatched_key_count: Final = await prisma_client.writer_db.litellm_verificationtoken.count(
mismatched_key_count: Final = await _writer_verification_token_table(prisma_client).count(
where={
"project_id": data.project_id,
"OR": [{"team_id": {"not": data.team_id}}, {"team_id": None}],

8
enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py
View file

@ -14,6 +14,7 @@ from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, List, Optional, Protocol
from fastapi import APIRouter, Depends, HTTPException
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -37,6 +38,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
router = APIRouter()
_DECODED_JSON: Final = TypeAdapter(object)
class ManagedVectorStoreRow(Protocol):
@ -348,11 +350,11 @@ async def update_vector_store(
raise HTTPException(status_code=500, detail="Database not connected")
try:
update_data = data.model_dump(exclude_unset=True)
vector_store_id: Final[str] = update_data.pop("vector_store_id")
update_data = data.model_dump(exclude_unset=True, exclude={"vector_store_id"})
vector_store_id: Final = data.vector_store_id
if update_data.get("vector_store_metadata") is not None:
update_data["vector_store_metadata"] = safe_dumps(
update_data["vector_store_metadata"]
_DECODED_JSON.validate_python(update_data["vector_store_metadata"])
)
updated = await managed_vector_store_table(prisma_client).update(

11
gateway/Dockerfile
View file

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

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

@ -0,0 +1,639 @@
"""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:"
LEGACY_SPEND_LOG_KEY_REMEDY: Final = (
"Rows written before v1.99.0 hold the raw key. Hash them in place (one-way, idempotent):\n"
' UPDATE "LiteLLM_SpendLogs" SET api_key = encode(sha256(convert_to('
"regexp_replace(api_key, '^[Bb][Ee][Aa][Rr][Ee][Rr] ', ''), 'UTF8')), 'hex')\n"
" WHERE api_key <> '' AND api_key !~ '^[0-9a-fA-F]{64}$' AND api_key !~ '^hashed-jwt-[0-9A-Fa-f]{64}$';"
)
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:
if check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "api_key" in check.columns:
return LEGACY_SPEND_LOG_KEY_REMEDY
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

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