scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

chore: merge main into anthropic stream fix

Browse source
This commit is contained in:
yyobject 2026-10-27 13:02:30 +08:00
commit e2ff54054d
2574 changed files with 138338 additions and 37081 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

2
.circleci/config.yml
View file

@ -1291,7 +1291,7 @@ jobs:
name: Install MCP SDK1 peer
command: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- when:
condition:

2
.circleci/scripts/classify_changes.sh
View file

@ -33,7 +33,7 @@ while IFS= read -r file || [ -n "$file" ]; do
*) has_backend=true ;;
esac
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json)
has_cost_map=true ;;
tests/unit/proxy/*) : ;;
*) outside_cost_map_set=true ;;

4
.github/CODEOWNERS vendored
View file

@ -1,4 +1,4 @@
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
/helm/ @yassin-berriai
/terraform/ @yassin-berriai

9
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -4,6 +4,13 @@
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
"## Affected release" with nothing under it must not appear in the final description -->
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
<details>
<summary>Details</summary>
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
@ -164,3 +171,5 @@ Example checklists:
## Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
</details>

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

91
.github/workflows/interactions-spec-canary.yml vendored Normal file
View file

@ -0,0 +1,91 @@
name: Interactions spec canary
on:
schedule:
- cron: "17 7 * * *"
pull_request:
paths:
- .github/workflows/interactions-spec-canary.yml
- tests/unit/interactions/**
- litellm/types/interactions/**
workflow_dispatch:
permissions: {}
concurrency: interactions-spec-canary
jobs:
interactions-spec-canary:
name: interactions-spec-canary
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 15
permissions:
contents: read
outputs:
compliance: ${{ steps.compliance.outcome }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project
- name: Run the compliance tests against Google's live spec
id: compliance
run: |
uv run --no-sync pytest tests/unit/interactions/test_openapi_compliance.py \
--interactions-spec-url=https://ai.google.dev/static/api/interactions.openapi.json \
--tb=short -vv -p no:rerunfailures -p no:pytest-retry \
--allow-hosts=127.0.0.1,::1,localhost,ai.google.dev
tracking-issue:
name: tracking-issue
needs: interactions-spec-canary
if: always() && github.event_name != 'pull_request' && github.ref_name == github.event.repository.default_branch && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
issues: write
steps:
- name: File the tracking issue unless the compliance step passed, close it when it did
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
env:
COMPLIANCE_OUTCOME: ${{ needs.interactions-spec-canary.outputs.compliance }}
RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
with:
script: |
const title = "Interactions spec canary is red";
const outcome = process.env.COMPLIANCE_OUTCOME;
const runUrl = process.env.RUN_URL;
const query = `repo:${context.repo.owner}/${context.repo.repo} is:issue is:open in:title "${title}"`;
const { data: found } = await github.rest.search.issuesAndPullRequests({ q: query, per_page: 20 });
const open = found.items.find((issue) => issue.title === title);
if (outcome === "success") {
if (open === undefined) return;
await github.rest.issues.createComment({ ...context.repo, issue_number: open.number, body: `Green again: ${runUrl}` });
await github.rest.issues.update({ ...context.repo, issue_number: open.number, state: "closed" });
return;
}
const state = outcome === "failure" ? "red" : `not run (compliance step outcome: ${outcome || "none"})`;
if (open !== undefined) {
await github.rest.issues.createComment({ ...context.repo, issue_number: open.number, body: `Still ${state}: ${runUrl}` });
return;
}
await github.rest.issues.create({
...context.repo,
title,
assignees: ["mateo-berri"],
body: `The scheduled run of tests/unit/interactions/test_openapi_compliance.py against Google's live Interactions spec is ${state}: ${runUrl}\n\nlitellm/types/interactions/README.md says how to refresh the snapshot and regenerate the types.`,
});

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

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

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

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

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

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

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

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

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

@ -139,6 +139,7 @@ jobs:
test-path: |-
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
@ -215,10 +216,12 @@ jobs:
- shard: llms-providers
test-path: |-
tests/unit/llms/test_api_base_trailing_slash_routes.py
tests/unit/llms/test_cache_control_and_reasoning.py
tests/unit/llms/test_custom_llm.py
tests/unit/llms/test_file_content_block.py
tests/unit/llms/test_file_search_responses.py
tests/unit/llms/test_get_complete_url_trailing_slash.py
tests/unit/llms/test_lifecycle_fix.py
tests/unit/llms/test_oss_decision.py
tests/unit/llms/test_polling_url_origin_match.py
@ -422,6 +425,7 @@ jobs:
tests/unit/realtime_api
tests/unit/rerank_api
tests/unit/search
tests/unit/text_completions
tests/unit/vector_stores
tests/unit/videos
python-version: "3.12"
@ -433,6 +437,13 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py
tests/test_litellm_rust/embeddings/test_dispatch.py
tests/test_litellm_rust/messages/test_dispatch.py
tests/test_litellm_rust/ocr/test_dispatch.py
tests/test_litellm_rust/responses/test_dispatch.py
tests/unit/compression
tests/unit/harness
tests/unit/rust_bridge

6
ARCHITECTURE.md
View file

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

11
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:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
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
@ -61,6 +61,7 @@ RUN apk add --no-cache \
python-3.13-dev \
rust \
openssl \
libssl3 \
openssl-dev \
nodejs \
npm \
@ -120,8 +121,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder
COPY --from=uvbin /uv /usr/local/bin/uv
RUN apk add --no-cache python-3.13
ADD --checksum=sha256:2f7ae5cdd9d91731c0990e74a58239dc3e3fd2bf28dab23b55eafcdc47aaf87e \
https://github.com/BerriAI/litellm-admin-agent/archive/ef501e94bc9fbacb9233b922abf71427f030408c.tar.gz /tmp/liteadmin.tar.gz
ADD --checksum=sha256:4fb1847e28e2be30ba76dac94f02c4903291326516a0f4ba772f6ffb8b25d05c \
https://github.com/BerriAI/litellm-admin-agent/archive/5661c805ce479a0cf11b842ea133cf8d22cd7e34.tar.gz /tmp/liteadmin.tar.gz
RUN mkdir /tmp/liteadmin && tar xzf /tmp/liteadmin.tar.gz --strip-components=1 -C /tmp/liteadmin && \
uv venv /opt/liteadmin --python python3.13 && \
uv pip install --python /opt/liteadmin/bin/python --require-hashes -r /tmp/liteadmin/requirements.txt && \
@ -141,7 +142,7 @@ USER root
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app

6
README.md
View file

@ -616,15 +616,15 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws
### Verify Docker Image Signatures
All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Every release is signed with the same key introduced in [commit `0112e53`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0).
All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Releases are signed with the key introduced in [commit `e82a3c6`](https://github.com/BerriAI/litellm/commit/e82a3c646908ba0a6ee0beddc4310d7a6cac4065).
**Verify using the pinned commit hash (recommended):**
A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the original signing key:
A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the published signing key:
```bash
cosign verify \
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
--key https://raw.githubusercontent.com/BerriAI/litellm/e82a3c646908ba0a6ee0beddc4310d7a6cac4065/cosign.pub \
ghcr.io/berriai/litellm:<release-tag>
```

2
security.md → SECURITY.md
View file

@ -79,7 +79,7 @@ To qualify for a bounty, reports must include clear reproduction steps, a reprod
- Audit Logs with retention policy
- Control Allowed IP Addresses that can access your Cloud LiteLLM Instance
For security inquiries, please contact us at support@berri.ai
For security inquiries, please contact us at security@berri.ai
#### Supported data regions for LiteLLM Cloud

15
backend/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
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
@ -43,7 +43,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -109,7 +109,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -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 && \

6
ci_cd/cost_map_guard.py
View file

@ -4,8 +4,7 @@ Every pull request whose diff against its merge base touches one of the two cost
checks: the files parse and the backup copy matches the root file. A pull request that leaves both untouched
skips them, since merging it keeps the base branch's copies and its head tree only carries whatever state the
branch was cut from. Pull requests from the cost map sync bot (branches named litellm_cost_map_sync_*) always
get the file checks and additionally may only touch those two files or the legacy schema
path and may only add or update models.
get the file checks and additionally may only touch those two files and may only add or update models.
"""
from __future__ import annotations
@ -21,7 +20,6 @@ from typing import Final
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
LEGACY_SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
@ -110,7 +108,7 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS and path != LEGACY_SCHEMA_PATH
if path not in GUARDED_PATHS
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),

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.

41
cookbook/litellm_proxy_server/batch_api/bedrock/README.md Normal file
View file

@ -0,0 +1,41 @@
# Bedrock batch endpoints
Set `aws_bedrock_control_plane_endpoint` on a deployment when batch create, retrieve, and cancel requests must use a custom Bedrock control-plane endpoint. The proxy signs these requests for the configured URL. Without this setting, batch requests use the standard regional Bedrock endpoint
```yaml
model_list:
- model_name: bedrock-batch
litellm_params:
model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0
aws_region_name: us-east-1
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN
s3_output_bucket_name: os.environ/AWS_S3_OUTPUT_BUCKET_NAME
aws_bedrock_control_plane_endpoint: os.environ/BEDROCK_CONTROL_PLANE_ENDPOINT
```
Set `BEDROCK_CONTROL_PLANE_ENDPOINT` to your Bedrock control-plane VPC endpoint or compatible gateway base URL, without `/model-invocation-job`. The proxy adds the operation path. The endpoint must be reachable from the proxy, and the configured AWS region must match the endpoint's signing region
`api_base` and `aws_bedrock_runtime_endpoint` configure runtime inference separately. They do not override the batch control-plane endpoint. `s3_endpoint_url` configures S3 file access separately
The control-plane setting belongs in deployment configuration. Sending `aws_bedrock_control_plane_endpoint` in an API request body is rejected. There is no automatic `AWS_BEDROCK_CONTROL_PLANE_ENDPOINT` environment fallback, so reference the desired environment variable explicitly as shown above
Start the proxy with this configuration and set `LITELLM_API_BASE` to its URL. Set `BEDROCK_BATCH_INPUT_URI` to a valid input JSONL file in S3 that the batch service role can read
```bash
curl -sS "$LITELLM_API_BASE/v1/batches" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
-H 'Content-Type: application/json' \
-d "{\"model\":\"bedrock-batch\",\"input_file_id\":\"$BEDROCK_BATCH_INPUT_URI\",\"endpoint\":\"/v1/chat/completions\",\"completion_window\":\"24h\"}"
```
Use the returned `id` as `BATCH_ID` to retrieve or cancel the same job through the deployment's configured endpoint
```bash
curl -sS "$LITELLM_API_BASE/v1/batches/$BATCH_ID?model=bedrock-batch" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
curl -sS -X POST "$LITELLM_API_BASE/v1/batches/$BATCH_ID/cancel?model=bedrock-batch" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
```

4
cosign.pub
View file

@ -1,4 +1,4 @@
-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEKi4ivqGpE231OGH50PKbqy1Y1Kkb
POJC8+i2Wko82gBOUCe3M0Vw86H/4rhUhfoYEti4gdJ9wZbYmK0I2EE96g==
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEG06dpKM+vEDnmp2zcjyXgVCXxzRg
JA5H2LnmwBKQafZGka6OOrWujCEWATQewiAgyD4WQdJ0pvU9RC4mvNAjvQ==
-----END PUBLIC KEY-----

7
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:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:360896efb0753a5701c48f860e17130c205eae5a264797aa3e7ea126ee2923a7
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
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
@ -59,6 +59,7 @@ RUN apk add --no-cache \
python-3.13 \
python-3.13-dev \
openssl \
libssl3 \
openssl-dev \
nodejs \
npm \
@ -122,7 +123,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
USER root
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app

7
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:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
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
@ -60,6 +60,7 @@ RUN for i in 1 2 3; do \
coreutils \
curl \
openssl \
libssl3 \
libsndfile \
nodejs \
npm && break || sleep 5; \
@ -134,7 +135,7 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
apk add --no-cache python-3.13 bash openssl libssl3 tzdata libsndfile nodejs libevent && break || sleep 5; \
done
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer

2
docker/build_from_pip/Dockerfile.build_from_pip
View file

@ -30,7 +30,7 @@ RUN uv venv --python python && \
"google-cloud-aiplatform==1.133.0" \
"google-genai==1.37.0" \
"anthropic[vertex]==0.84.0" \
"grpcio==1.78.0" \
"grpcio==1.84.1" \
"prometheus-client==0.20.0" \
"langfuse==2.59.7" \
"opentelemetry-api==1.28.0" \

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

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

@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
import asyncio
import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -13,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
row_metadata,
)
import litellm
@ -39,7 +41,12 @@ if TYPE_CHECKING:
from litellm.types.utils import LiteLLMBatch
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -84,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _row_metadata(row: object) -> dict[str, object]:
metadata: Final = getattr(row, "metadata", None)
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
class CheckBatchCost:
def __init__(
self,
@ -109,6 +111,11 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
self._line_item_store_turn: Final = asyncio.Lock()
self._poll_idle: Final = asyncio.Event()
@ -161,29 +168,9 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
return await self._find_jobs(
where={
"file_purpose": "batch",
"status": {
@ -197,10 +184,59 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
**where,
"created_at": {
"lte": before,
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
},
**(
{
"OR": [
{"created_at": {"gt": cursor[0]}},
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
]
}
if cursor is not None
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
if arrivals:
self._arrival_cursor = next_cursor
else:
self._poll_cursor = next_cursor
return page
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
@ -208,24 +244,19 @@ class CheckBatchCost:
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
return
if retired == 0:
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
@ -312,13 +343,13 @@ class CheckBatchCost:
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
404 must not retire a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
@ -333,7 +364,7 @@ class CheckBatchCost:
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -355,18 +386,18 @@ class CheckBatchCost:
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
if finalized == 0:
return False
verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
return True
except Exception as db_err:
verbose_proxy_logger.error(
@ -817,20 +848,26 @@ class CheckBatchCost:
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job, original_status, original_file_object)
raise
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
finally:
if accounting.retryable_failure:
await self._release_job_claim(job, original_status, original_file_object)
if accounting.retryable_failure:
return None
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:
@ -1003,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 = (
@ -1022,7 +1059,9 @@ class CheckBatchCost:
if (has_key_row and key_row is None) or (team_id and team_row is None):
return True
settings: Final = get_dynamic_logging_metadata(
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
UserAPIKeyAuth(
metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id
),
proxy_config,
)
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
@ -1047,6 +1086,7 @@ class CheckBatchCost:
"""
try:
from litellm.integrations.prometheus import PrometheusLogger
prom_logger = PrometheusLogger.get_instance()
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
@ -1071,7 +1111,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await _managed_object_table(self.prisma_client).find_many(
jobs = await self._find_jobs(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1084,8 +1124,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:
@ -1152,7 +1190,7 @@ class CheckBatchCost:
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
f"(job {job.id}): {tracking_err}"
)
self._record_error(prom_logger, "cost_tracking_error")
continue

117
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
The spend event reaches the creating key's or team's own logging callbacks the way that creator's
requests do. Its message content is redacted when the creator's key or team logging turns message
logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before
the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or
whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only
and with its message content redacted.
A read that finds no row, because the create could not store one, stores it under the reading key
and bills it the same way, so a response the create left untracked is still billed once. A streamed
background create is free on its own; its stream stores the row under the creating key and bills
@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
ResponseCreator,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -54,12 +61,29 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
class _RoutedModelInfo(BaseModel):
model_config = ConfigDict(frozen=True)
id: str | None = None
class _RoutedDeployment(BaseModel):
model_config = ConfigDict(frozen=True)
model_info: _RoutedModelInfo = _RoutedModelInfo()
_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"})
_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True})
@dataclass(frozen=True, slots=True)
@ -76,6 +100,7 @@ class BilledBackgroundResponse:
custom_llm_provider: str | None
litellm_params: Mapping[str, object]
metadata: Mapping[str, object]
creator_logging: Mapping[str, object]
@dataclass(frozen=True, slots=True)
@ -94,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
return unpriced
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
def _string_callbacks(callbacks: object) -> tuple[str, ...]:
if not isinstance(callbacks, list):
return ()
return tuple(
callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str)
)
logging_obj: Final = LiteLLMLogging(
def _running_proxy_config() -> "ProxyConfig":
from litellm.proxy.proxy_server import proxy_config
return proxy_config
def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields
if not creator.settings_read:
return _REJECTED_CREATOR_LOGGING
try:
return key_and_team_logging_fields(
user_api_key_dict=UserAPIKeyAuth(
metadata=dict(creator.key_metadata),
team_id=creator.team_id,
team_metadata=dict(creator.team_metadata),
),
proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(),
)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _REJECTED_CREATOR_LOGGING
def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging":
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.utils import split_dynamic_success_callbacks
sync_callbacks, async_callbacks = split_dynamic_success_callbacks(
_string_callbacks(creator_logging.get("success_callback"))
)
return LiteLLMLogging(
model=billed.model,
messages=[{"role": "user", "content": "<retrieve_response>"}],
stream=False,
@ -106,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
start_time=datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id=str(uuid.uuid4()),
dynamic_success_callbacks=list(sync_callbacks),
dynamic_async_success_callbacks=list(async_callbacks) or None,
kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS},
)
def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging":
try:
return _billed_logging_obj(billed, billed.creator_logging)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING)
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj: Final = _creator_or_proxy_logging_obj(billed)
logging_obj.update_environment_variables(
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
optional_params={},
@ -174,6 +259,7 @@ class CheckResponsesCost:
prisma_client: "PrismaClient",
llm_router: "Router",
log_billed_response: BilledResponseLogger = log_billed_background_response,
proxy_config: "ProxyConfig | None" = None,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -183,6 +269,7 @@ class CheckResponsesCost:
self.llm_router: Router = llm_router
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
self._log_billed_response: Final = log_billed_response
self._proxy_config: Final = proxy_config
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
@ -294,11 +381,12 @@ class CheckResponsesCost:
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
)
async def _billed_response(
def _billed_response(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
creator: ResponseCreator,
) -> BilledBackgroundResponse:
deployment: Final = target.deployment
return BilledBackgroundResponse(
@ -310,12 +398,26 @@ class CheckResponsesCost:
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
metadata=MappingProxyType(
{
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
**creator.metadata,
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
**self._requested_model_group(target, job.team_id),
}
),
creator_logging=creator_logging_fields(creator, self._proxy_config),
)
def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]:
requested: Final = target.model_metadata.get("model_group")
if target.deployment is None or not isinstance(requested, str):
return MappingProxyType({})
routed: Final = _ROUTED_DEPLOYMENTS.validate_python(
self.llm_router.get_model_list(model_name=requested, team_id=team_id) or ()
)
serving_id: Final = target.deployment.model_info.id
return MappingProxyType(
{"model_group": requested} if any(row.model_info.id == serving_id for row in routed) else {}
)
async def _bill_once(
@ -328,7 +430,8 @@ class CheckResponsesCost:
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
return None
try:
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id)
charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator))
except BaseException:
await self._release_billing_claim(job)
raise

207
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -9,12 +9,17 @@ from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, cast
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLI_SESSION_KEY_PREFIX
from litellm.repositories.table_repositories import (
DeletedVerificationTokenRepository,
DeprecatedVerificationTokenRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -27,6 +32,12 @@ if TYPE_CHECKING:
class ManagedObjectRow(Protocol):
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def id(self) -> str: ...
@ -63,6 +74,7 @@ class ManagedObjectRow(Protocol):
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
_MAX_ROTATION_HOPS: Final = 16
def _string_tags(request_tags: object) -> list[str] | None:
@ -89,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL
return TeamRepository(prisma_client).table
def _rotated_token_table(
prisma_client: PrismaClient,
) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]:
return DeprecatedVerificationTokenRepository(prisma_client).table
def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]:
return DeletedVerificationTokenRepository(prisma_client).table
def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]:
return TeamRepository(prisma_client).deleted_table
@dataclass(frozen=True, slots=True)
class _CreatorRows:
key: prisma_models.LiteLLM_VerificationToken | None
@ -96,6 +122,15 @@ class _CreatorRows:
user: prisma_models.LiteLLM_UserTable | None
@dataclass(frozen=True, slots=True)
class ResponseCreator:
metadata: Mapping[str, object]
key_metadata: Mapping[str, object]
team_id: str | None
team_metadata: Mapping[str, object]
settings_read: bool = True
class ManagedObjectCreatorAttribution:
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
job's spend row is attributed like the creating request: key hash, team, organization,
@ -107,8 +142,15 @@ class ManagedObjectCreatorAttribution:
user's alias, because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
to increment organization spend, so leaving it out drops the cost from org accounting.
A background response also carries its end user and the key, team, and user budgets, the
way the read that billed it before did, and never names a key-scoped creator as its user.
A background response also carries its end user, the creating key's project and its alias,
and the key, team, and user budgets and per-model budgets, the way the read that billed it
before did, and never names a key-scoped creator as its user. A key that sets no budget of
its own takes its budget tier's, the same rule auth applies. Its creator's key and team
metadata come along too, so the billed event reaches the key's and team's logging callbacks
like that read did. A rotated key's metadata is read from the live key its rotations lead
to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the
creator's settings unread and the attribution to the rows that still load, so the
response still bills and still counts toward their budgets.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
@ -119,17 +161,70 @@ class ManagedObjectCreatorAttribution:
rows: Final = await self._rows(job, object_id)
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
rows: Final = await self._rows(job, object_id)
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
}
async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
try:
return await self._read_response_creator(job, object_id)
except Exception as e:
verbose_proxy_logger.error(
f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their "
f"settings: {e}"
)
return ResponseCreator(
metadata=_response_creator_metadata(job, await self._rows(job, object_id)),
key_metadata={},
team_id=getattr(job, "team_id", None),
team_metadata={},
settings_read=False,
)
async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
metadata=_response_creator_metadata(job, rows),
key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key),
team_id=team_id,
team_metadata=(
row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id)
),
)
async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]:
if not _names_a_key_row(api_key):
return {}
current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS)
current_key: Final = (
await _token_table(self._prisma_client).find_unique(where={"token": current_token})
if current_token != api_key
else None
)
if current_key is not None:
return row_metadata(current_key)
return row_metadata(
await _deleted_token_table(self._prisma_client).find_first(
where={"token": current_token}, order={"deleted_at": "desc"}
)
)
async def _token_rotated_to(self, token: str, hops_left: int) -> str:
rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token})
if rotation is None or hops_left == 0:
return token
return await self._token_rotated_to(rotation.active_token_id, hops_left - 1)
async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]:
if not team_id:
return {}
return row_metadata(
await _deleted_team_table(self._prisma_client).find_first(
where={"team_id": team_id}, order={"deleted_at": "desc"}
)
)
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
return _CreatorRows(
@ -152,23 +247,57 @@ class ManagedObjectCreatorAttribution:
async def _key_row(
self, object_id: str, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
try:
return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
return await self._live_key_row(api_key, created_by)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
return None
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
async def _live_key_row(
self, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
return await _token_table(self._prisma_client).find_unique(
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
)
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
try:
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
return await self._live_team_row(team_id)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
"user_api_key_team_spend",
"user_api_key_team_max_budget",
"user_api_key_team_model_max_budget",
),
**_budget_fields(
rows.user,
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_user_model_max_budget",
),
}
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
api_key: Final = getattr(job, "api_key", None)
@ -196,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str,
}
def row_metadata(row: object) -> Mapping[str, object]:
metadata: Final = getattr(row, "metadata", None)
if not isinstance(metadata, dict):
return {}
return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column
def _names_a_key_row(api_key: object) -> TypeGuard[str]:
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
return isinstance(api_key, str) and is_valid_sha256_hash(api_key)
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
@ -212,7 +354,28 @@ def _organization_id(row: object) -> str | None:
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
if row is None:
return {}
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}
return {
spend_key: getattr(row, "spend", None),
max_budget_key: getattr(row, "max_budget", None),
model_max_budget_key: getattr(row, "model_max_budget", None),
}
def _key_budget_fields(key: object) -> Mapping[str, object]:
if key is None:
return {}
tier: Final = getattr(key, "litellm_budget_table", None)
return {
"user_api_key_spend": getattr(key, "spend", None),
"user_api_key_max_budget": _own_or_tier_budget(key, tier, "max_budget"),
"user_api_key_model_max_budget": _own_or_tier_budget(key, tier, "model_max_budget"),
}
def _own_or_tier_budget(key: object, tier: object, field: str) -> object:
own: Final = getattr(key, field, None)
tier_value: Final = getattr(tier, field, None)
return own if tier_value is None or own not in (None, {}) else tier_value

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:
"""

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(

4
enterprise/pyproject.toml
View file

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

15
gateway/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
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
@ -59,7 +59,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl libssl3 openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -125,7 +125,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic libevent && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -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

4
gateway/routes/allowlist.py
View file

@ -1,7 +1,7 @@
"""Path allowlist for the gateway component.
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image,
audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/decisions",
"/v1/systemone",
"/systemone",
"/v1/extraction",
"/extraction",
"/v1/ocr",
"/ocr",
"/v1/rag/",

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.110"
version = "0.4.111"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.110"
version = "0.4.111"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

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.9"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b"
checksum = "e49d0e65abc00c142e8f9428da991da9f1a250c39be97459901eb12b9b49f62a"
dependencies = [
"rustls-pki-types",
]
[[package]]
name = "webpki-roots"
version = "1.0.9"
version = "1.0.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a"
checksum = "f5f280bbb2b287bc8e7de7e3852945a0ff50d1bbcb5ed42b6c7a568c546df171"
dependencies = [
"rustls-pki-types",
]

9
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"
@ -124,7 +131,7 @@ moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
percent-encoding = "2.3"
webpki-roots = "1"
webpki-roots = "1.0.10"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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