scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

Merge remote-tracking branch 'origin/main' into litellm_speech_dynamic_api_key

Browse source
This commit is contained in:
mateo-berri 2026-10-28 13:52:14 -07:00
commit cbeb81ddda
3395 changed files with 180755 additions and 52708 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

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

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

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

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

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

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

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

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

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

@ -0,0 +1,94 @@
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: 30
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: Set up the Rust bridge
uses: ./.github/actions/rust-bridge
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- 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.`,
});

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

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

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

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

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

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

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

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

15
.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"
@ -126,7 +132,7 @@ jobs:
windows-sdk:
name: windows-sdk
runs-on: windows-2025
timeout-minutes: 10
timeout-minutes: 30
env:
LITELLM_LOCAL_MODEL_COST_MAP: "True"
defaults:
@ -148,8 +154,11 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
uses: ./.github/actions/rust-bridge
- name: Install dependencies
run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11
run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11
- name: Run the Windows tests
run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests

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

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

@ -216,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
@ -300,6 +302,7 @@ jobs:
tests/unit/llms/modelscope
tests/unit/llms/mongodb
tests/unit/llms/moonshot
tests/unit/llms/morph
tests/unit/llms/nadir
tests/unit/llms/nebius
tests/unit/llms/neosantara
@ -423,6 +426,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"
@ -576,13 +580,13 @@ jobs:
- shard: mcp-elicitation
test-path: |-
tests/unit/experimental_mcp_client/test_mcp_client.py
tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
tests/unit/proxy/_experimental/mcp_server/test_interactions.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
tests/unit/proxy/_experimental/mcp_server/test_operations.py
tests/unit/mcp_client/test_mcp_client.py
tests/unit/proxy/mcp_server/test_capabilities.py
tests/unit/proxy/mcp_server/test_interactions.py
tests/unit/proxy/mcp_server/test_mcp_elicitation_handler.py
tests/unit/proxy/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py
tests/unit/proxy/mcp_server/test_operations.py
python-version: "3.12"
workers: 4
reruns: 0
@ -600,7 +604,7 @@ jobs:
tests/unit/proxy/spend_tracking
--ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py
tests/unit/proxy/pass_through_endpoints
tests/unit/proxy/experimental
tests/unit/proxy/mcp_server/test_tool_registry.py
tests/unit/proxy/common_utils
--ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py
--ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py
@ -808,14 +812,14 @@ jobs:
- shard: mcp-oauth
test-path: |-
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
tests/unit/proxy/_experimental/mcp_server/outbound_credentials
tests/unit/proxy/mcp_server/test_discoverable_endpoints.py
tests/unit/proxy/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/mcp_server/test_db_credentials.py
tests/unit/proxy/mcp_server/test_gateway_dcr_flow.py
tests/unit/proxy/mcp_server/test_mcp_oauth_passthrough_tools.py
tests/unit/proxy/mcp_server/test_mcp_max_concurrent_requests.py
tests/unit/proxy/mcp_server/faults/test_list_outcomes.py
tests/unit/proxy/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4
reruns: 2
@ -825,19 +829,21 @@ jobs:
- shard: mcp-server
test-path: |-
tests/unit/mcp_client
tests/unit/experimental_mcp_client
tests/unit/proxy/mcp_server
tests/unit/proxy/_experimental/mcp_server
tests/unit/responses/mcp
--ignore=tests/unit/experimental_mcp_client/test_mcp_client.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_interactions.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_operations.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
--ignore=tests/unit/proxy/_experimental/mcp_server/outbound_credentials
--ignore=tests/unit/mcp_client/test_mcp_client.py
--ignore=tests/unit/proxy/mcp_server/test_capabilities.py
--ignore=tests/unit/proxy/mcp_server/test_interactions.py
--ignore=tests/unit/proxy/mcp_server/test_mcp_elicitation_handler.py
--ignore=tests/unit/proxy/mcp_server/test_mcp_server_manager.py
--ignore=tests/unit/proxy/mcp_server/test_mcp_server_tool_calls_and_headers.py
--ignore=tests/unit/proxy/mcp_server/test_operations.py
--ignore=tests/unit/proxy/mcp_server/test_discoverable_endpoints.py
--ignore=tests/unit/proxy/mcp_server/test_db_credentials.py
--ignore=tests/unit/proxy/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4
reruns: 0
@ -947,7 +953,7 @@ jobs:
dist: loadscope
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
timeout-minutes: 3
timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
with:
persist-credentials: false
fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
@ -1300,6 +1306,8 @@ jobs:
done
openai-sdk-compat:
name: OpenAI ${{ matrix.openai-version }} compatibility
needs: rust-bridge
if: ${{ !cancelled() }}
permissions:
contents: read
pull-requests: read
@ -1336,10 +1344,18 @@ jobs:
with:
version: "0.10.9"
- name: Set up the Rust bridge
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
uses: ./.github/actions/rust-bridge
with:
artifact: ${{ needs.rust-bridge.outputs.artifact }}
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace
cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/
uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z
uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__'

2
.gitignore vendored
View file

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

12
ARCHITECTURE.md
View file

@ -239,6 +239,8 @@ graph LR
7. `DBSpendUpdateWriter.update_database()` queues spend increments to Redis
8. Background job `update_spend` flushes queued spend to PostgreSQL every 60s
Cost lookup falls back to the response model when the provider-reported name cannot be priced. Explicitly registered zero rates remain zero, and custom deployment pricing or an explicit `base_model` retain precedence
### Data Access Layer (Models & Repositories)
Database entities and the operations on them live in two packages at the root of `litellm/` so both the gateway (`proxy/`) and the SDK can use them without importing proxy internals:
@ -332,13 +334,13 @@ Each translation is isolated in its own file, making it easy to test and modify
|--------------|----------|------------------|
| `/v1/chat/completions` | Anthropic | `llms/anthropic/chat/transformation.py` |
| `/v1/chat/completions` | Bedrock Converse | `llms/bedrock/chat/converse_transformation.py` |
| `/v1/chat/completions` | Bedrock Invoke | `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py` |
| `/v1/chat/completions` | Bedrock Invoke | `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py` |
| `/v1/chat/completions` | Gemini | `llms/gemini/chat/transformation.py` |
| `/v1/chat/completions` | Vertex AI | `llms/vertex_ai/gemini/transformation.py` |
| `/v1/chat/completions` | OpenAI | `llms/openai/chat/gpt_transformation.py` |
| `/v1/messages` (passthrough) | Anthropic | `llms/anthropic/pass_through/messages/transformation.py` |
| `/v1/messages` (passthrough) | Bedrock | `llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py` |
| `/v1/messages` (passthrough) | Vertex AI | `llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py` |
| `/v1/messages` (passthrough) | Anthropic | `llms/anthropic/messages/native/transformation.py` |
| `/v1/messages` (passthrough) | Bedrock | `llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py` |
| `/v1/messages` (passthrough) | Vertex AI | `llms/vertex_ai/vertex_ai_partner_models/anthropic/messages/transformation.py` |
| Passthrough endpoints | All | `proxy/pass_through_endpoints/llm_provider_handlers/` |
### Example: Debugging prompt caching
@ -346,7 +348,7 @@ Each translation is isolated in its own file, making it easy to test and modify
If `/v1/messages` → Bedrock Converse prompt caching isn't working but Bedrock Invoke works:
1. **Bedrock Converse translation**: `llms/bedrock/chat/converse_transformation.py`
2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py`
2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py`
3. Compare how each handles `cache_control` in `transform_request()`
### How translations work

2
CONTRIBUTING.md
View file

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

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

4
Makefile
View file

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

13
README.md
View file

@ -223,7 +223,7 @@ async with httpx.AsyncClient(headers=headers, timeout=60.0) as http_client:
```python
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from litellm import experimental_mcp_client
from litellm import mcp_client
import litellm
server_params = StdioServerParameters(command="python", args=["mcp_server.py"])
@ -233,7 +233,7 @@ async with stdio_client(server_params) as (read, write):
await session.initialize()
# Load MCP tools in OpenAI format
tools = await experimental_mcp_client.load_mcp_tools(session=session, format="openai")
tools = await mcp_client.load_mcp_tools(session=session, format="openai")
# Use with any LiteLLM model
response = await litellm.acompletion(
@ -427,6 +427,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Watsonx Text (`watsonx_text`)](https://docs.litellm.ai/docs/providers/watsonx) | ✅ | ✅ | ✅ | | | | | | | |
| [xAI (`xai`)](https://docs.litellm.ai/docs/providers/xai) | ✅ | ✅ | ✅ | | | | | | | |
| [Xinference (`xinference`)](https://docs.litellm.ai/docs/providers/xinference) | | | | ✅ | | | | | | |
| [ZeroGPU (`zerogpu`)](https://docs.litellm.ai/docs/providers/zerogpu) | ✅ | ✅ | ✅ | | | | | | | |
[**Read the Docs**](https://docs.litellm.ai/docs/)
@ -616,15 +617,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>
```
@ -660,6 +661,8 @@ This covers:
We welcome contributions to LiteLLM! Whether you're fixing bugs, adding features, or improving documentation, we appreciate your help.
If you want to work closely with our team, check out [LiteLLM OSS Champions](https://www.litellm.ai/champions), our program for trusted OSS contributors with their own Slack channel and a voice in the roadmap
## Quick Start for Contributors
This requires uv to be installed.

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

119
ci_cd/spend_logs_guard.py Normal file
View file

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

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

@ -0,0 +1,78 @@
# 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.
The `lite` CLI can set up this provider each time it starts Codex, with no `config.toml` edit. Export `LITELLM_PROXY_URL` and `LITELLM_PROXY_API_KEY`, then run `lite codex --subscription`. Run `lite config set codex_subscription true` to make it the default.
For a managed workspace rollout, use OpenAI's [gateway configuration guide](https://learn.chatgpt.com/docs/enterprise/sign-in-with-chatgpt-through-a-gateway).
## Native compatibility and Pi
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
### Configure Pi
Sign into the built-in `openai-codex` provider in Pi. Merge this override into `~/.pi/agent/models.json` and export your LiteLLM virtual key as `LITELLM_API_KEY` before launching Pi:
```json
{
"providers": {
"openai-codex": {
"baseUrl": "https://litellm.example.com/backend-api/codex",
"headers": {
"x-litellm-api-key": "LITELLM_API_KEY"
}
}
}
}
```
Keep Pi's subscription login as the provider credential. The header value names an environment variable, not a key to paste into the file.
## Verify account usage
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.
This flow does not connect accounts through the dashboard or pool subscription allowances. Use a Codex build that supports ChatGPT authentication with a custom model provider. The integration forwards HTTP and SSE requests; it does not add native WebSocket or other Codex backend routes.

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

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

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

4
enterprise/litellm_enterprise/proxy/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):

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

@ -5,9 +5,9 @@ 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 collections.abc import Callable, Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
@ -28,6 +28,7 @@ from litellm.constants import (
from litellm.repositories.table_repositories import ManagedObjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.utils import get_utc_datetime
if TYPE_CHECKING:
from prisma import models as prisma_models
@ -42,11 +43,14 @@ if TYPE_CHECKING:
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
tokens: Final[list[str]] = re.findall(r"""[^\s"'`<>(]+""", str(error))
return any(token.rstrip(".,:;)") == resource_id for token in tokens)
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -91,6 +95,10 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _unique_jobs(jobs: Sequence[ManagedObjectRow]) -> tuple[ManagedObjectRow, ...]:
return tuple({job.id: job for job in jobs}.values())
class CheckBatchCost:
def __init__(
self,
@ -98,6 +106,7 @@ class CheckBatchCost:
prisma_client: "PrismaClient",
llm_router: "Router",
track_unmanaged_batch_cost: bool = False,
clock: Callable[[], datetime] = get_utc_datetime,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -106,13 +115,14 @@ class CheckBatchCost:
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
self._clock: Final = clock
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
# Cached after the first poll cycle. Once we know the column is absent we skip
# 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._poll_before: datetime = clock()
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
@ -147,27 +157,6 @@ class CheckBatchCost:
return
self.batch_processed_support_confirmed = True
async def _cleanup_stale_managed_objects(self) -> None:
"""
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
in non-terminal states as 'stale_expired'. These will never complete and
should not be polled.
"""
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
result: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
)
if result > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: marked {result} stale managed objects "
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await self._find_jobs(
@ -187,20 +176,43 @@ class CheckBatchCost:
)
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)
lead_arrivals: Final = self._poll_arrivals
self._poll_arrivals = not lead_arrivals
lead_share: Final = MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2
lead: Final = await self._find_job_page(where, arrivals=lead_arrivals, take=lead_share)
trail: Final = await self._find_job_page(
where, arrivals=not lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead)
)
spill: Final = (
await self._find_job_page(
where, arrivals=lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) - len(trail)
)
if len(lead) == lead_share
else ()
)
return _unique_jobs((*lead, *trail, *spill))
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
async def _find_job_page(
self, where: Mapping[str, object], *, arrivals: bool, take: int
) -> Sequence[ManagedObjectRow]:
if take == 0:
return ()
resumed_mid_lane: Final = (self._arrival_cursor if arrivals else self._poll_cursor) is not None
page: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take)
if len(page) == take or not resumed_mid_lane:
return page
wrapped: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take - len(page))
return _unique_jobs((*page, *wrapped))
async def _find_job_page_from_cursor(
self, where: Mapping[str, object], *, arrivals: bool, take: int
) -> 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)
self._arrival_before = self._clock()
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
self._poll_before = self._clock() - 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={
@ -220,10 +232,10 @@ class CheckBatchCost:
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
take=take,
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
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == take else None
if arrivals:
self._arrival_cursor = next_cursor
else:
@ -231,11 +243,47 @@ class CheckBatchCost:
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 {}),
}
if self._has_batch_processed_column:
return {"id": job.id, "batch_processed": False}
return {"id": job.id, "updated_at": job.updated_at}
def _is_past_give_up_age(self, job: ManagedObjectRow) -> bool:
return job.created_at < self._clock() - GIVE_UP_POLLING_AFTER
async def _expire_if_old_and_unfinished(self, job: ManagedObjectRow) -> None:
if job.status in TERMINAL_MANAGED_OBJECT_STATUSES or not self._is_past_give_up_age(job):
return
await self._expire_unfinished(job)
@staticmethod
def _is_completed_with_output_id_pending(response: "LiteLLMBatch") -> bool:
from litellm.proxy.openai_files_endpoints.common_utils import (
_completed_batch_safe_to_retire,
)
return (
response.status in ("completed", "complete")
and response.output_file_id is None
and not _completed_batch_safe_to_retire(response)
)
async def _expire_unfinished(self, job: ManagedObjectRow) -> None:
try:
expired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
**self._unprocessed_snapshot(job),
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
},
data={"status": "stale_expired"},
)
except Exception as db_err:
verbose_proxy_logger.error(f"CheckBatchCost: failed to mark stale job {job.id} as stale_expired: {db_err}")
return
if expired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} is older than {GIVE_UP_POLLING_AFTER.days} days "
"and still not finished, so it is marked stale_expired and no longer polled"
)
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
@ -356,15 +404,21 @@ class CheckBatchCost:
def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
"""A 404 naming the output file means there is nothing to fetch on this or any
later poll: providers like Vertex AI advertise an output path for every batch,
including terminal ones that never wrote it. Any other failure may be
transient, so it keeps retrying until the staleness sweep bounds it."""
including terminal ones that never wrote it. GCS names a gs:// object without its scheme,
and S3 names an s3:// object by its key alone. Any other failure keeps retrying."""
import openai
from litellm.exceptions import NotFoundError
from litellm.llms.base_llm.chat.transformation import BaseLLMException
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
is_not_found: Final = isinstance(error, (NotFoundError, openai.NotFoundError)) or (
isinstance(error, BaseLLMException) and error.status_code == 404
)
object_path: Final = output_file_id.split("://", 1)[-1]
names: Final = (output_file_id, object_path, object_path.split("/", 1)[-1])
return is_not_found and any(_error_names_resource(error, name) for name in names)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -1040,8 +1094,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 = (
@ -1094,13 +1148,6 @@ class CheckBatchCost:
processed_models: List[Tuple[Optional[str], Optional[str]]] = []
try:
await self._cleanup_stale_managed_objects()
except Exception as cleanup_err:
verbose_proxy_logger.warning(
f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
)
# Look for all batches that have not yet been processed by CheckBatchCost.
# self._has_batch_processed_column is cached after the first probe so that
# older schemas don't pay a guaranteed-failing primary query + warning on
@ -1142,6 +1189,8 @@ class CheckBatchCost:
if routing is None:
if self._has_unified_id_without_model(job):
await self._retire_job(job, "unified object id has no model id")
else:
await self._expire_if_old_and_unfinished(job)
continue
model_id, batch_id = routing
@ -1166,6 +1215,14 @@ class CheckBatchCost:
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
else:
await self._expire_if_old_and_unfinished(job)
continue
if self._is_completed_with_output_id_pending(response):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
)
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE
@ -1217,17 +1274,9 @@ class CheckBatchCost:
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
from litellm.proxy.openai_files_endpoints.common_utils import (
_completed_batch_safe_to_retire,
)
if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
)
continue
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
else:
await self._expire_if_old_and_unfinished(job)
# Record polling run metrics (always, even if nothing was processed)
if prom_logger:

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

@ -6,6 +6,7 @@ import base64
import json
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from datetime import datetime, timedelta
from itertools import chain
from types import MappingProxyType
from typing import (
@ -27,7 +28,7 @@ import httpx
from fastapi import HTTPException
from openai import APIConnectionError, AsyncOpenAI
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -39,18 +40,20 @@ from litellm.constants import (
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_FILE_LIST_LIMIT,
PENDING_RESPONSE_STATUSES,
)
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 (
extract_file_metadata,
)
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
@ -112,6 +115,7 @@ from litellm.types.utils import (
LLMResponseTypes,
SpecialEnums,
)
from litellm.utils import get_utc_datetime
class _ManagedFileRetrieve(Protocol):
@ -133,7 +137,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:
@ -305,6 +309,7 @@ def _is_xai_batch_output(
class _ManagedFileRow(Protocol):
unified_file_id: str
file_object: OpenAIFileObject
model_mappings: Mapping[str, str]
flat_model_file_ids: Sequence[str]
storage_backend: Optional[str]
storage_url: Optional[str]
@ -358,6 +363,8 @@ class _ManagedResourceDatabase(Protocol):
@property
def litellm_managedfiletable(self) -> _ManagedFileTableActions: ...
async def query_raw(self, query: str, *args: object) -> object: ...
@property
def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ...
@ -381,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False):
custom_llm_provider: ReadOnly[str | None]
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
return prisma_client.db
def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions:
return prisma_client.db.litellm_managedfiletable
class _ListedFileRow(Protocol):
@property
def unified_file_id(self) -> str: ...
@property
def file_object(self) -> object: ...
@property
def model_mappings(self) -> Mapping[str, object]: ...
@property
def storage_backend(self) -> str | None: ...
class _ScopedFileRow(BaseModel):
model_config = ConfigDict(frozen=True)
unified_file_id: str
file_object: JsonValue = None
model_mappings: Mapping[str, JsonValue]
storage_backend: str | None = None
_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...])
def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]:
match owner_filter:
case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}:
return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id)
case {"created_by": str() as user_id}:
return '"created_by" = $1', (user_id,)
case {"team_id": str() as team_id}:
return '"team_id" = $1', (team_id,)
case _:
return ("FALSE", ()) if owner_filter else ("TRUE", ())
async def _scoped_file_rows(
database: _ManagedResourceDatabase,
owner_filter: Mapping[str, object],
model_mapping_keys: frozenset[str],
cursor_id: str | None,
direction: Literal["asc", "desc"],
take: int,
) -> tuple[_ScopedFileRow, ...]:
owner_condition, owner_params = _owner_condition(owner_filter)
keys_param: Final = len(owner_params) + 1
comparison: Final = "<" if direction == "desc" else ">"
cursor_condition: Final = (
f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" '
f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) '
if cursor_id
else ""
)
order: Final = direction.upper()
rows: Final = await database.query_raw(
'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" '
f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) '
f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} '
f"LIMIT ${keys_param + (2 if cursor_id else 1)}",
*owner_params,
sorted(model_mapping_keys),
*((cursor_id,) if cursor_id else ()),
take,
)
return _SCOPED_FILE_ROWS.validate_python(rows)
def _iter_provider_file_id_pairs(
rows: Sequence[_ManagedFileRow],
requested_provider_file_ids: frozenset[str],
@ -490,6 +570,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__(
@ -498,10 +586,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prisma_client: PrismaClient,
*,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
clock: Callable[[], datetime] = get_utc_datetime,
):
self.internal_usage_cache = internal_usage_cache
self.prisma_client = prisma_client
self._sleep = sleep
self._clock = clock
@staticmethod
def _get_prometheus_logger():
@ -1611,7 +1701,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 +1713,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 +1914,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 +1956,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 +1972,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 +2039,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
@ -2251,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
limit: Optional[int] = None,
after: Optional[str] = None,
model_scope: Optional[FileListModelScope] = None,
before: Optional[str] = None,
**data: Dict,
) -> FileListPage:
"""List the managed files the caller owns, newest first.
Pagination is keyset based on ``unified_file_id`` so a key that owns
every file on the proxy still reads one bounded page at a time.
``purpose`` is applied after parsing, because the managed file table
keeps it inside the ``file_object`` blob instead of a column, and rows
``after`` pages toward older files and ``before`` toward newer ones,
each page still ordered newest first.
``model_scope`` is applied in the query, so only rows whose
``model_mappings`` name one of its deployment ids or model names, and
storage-backed rows, leave the database; a storage-backed row then
stays only when its model name is one the router sends to the listed
provider. ``purpose`` is applied after parsing, because the managed
file table keeps it inside the ``file_object`` blob instead of a
column. Rows
whose blob will not parse drop out there too, so a chunk of rows can
yield fewer matches than the page holds. Successive chunks are read
until the page is full or the caller's rows run out, which keeps
@ -2267,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
scan has to continue past it, the chunk widens to
``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per
that many rows instead of one per page. That bound is per query, not
per request: the work is still linear in the rows the caller owns, and
a filter matching nothing reads every one of them, with no index
covering either the owner filter or the sort.
per request: a ``purpose`` filter matching nothing still reads every
row that passes the query.
"""
validate_file_list_limit(limit)
validate_file_list_purpose(purpose)
@ -2278,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if owner_filter is None:
return FileListPage.model_validate(build_list_page([]))
if after:
cursor_row = await _managed_file_table(self.prisma_client).find_first(
where={**owner_filter, "unified_file_id": after}
if after and before:
raise ProxyException(
message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.",
type="invalid_request_error",
param="before_id",
code=400,
openai_code="invalid_value",
)
cursor_param: Final = "before_id" if before else "after"
cursor: Final = before or after
database: Final = _managed_database(self.prisma_client)
if cursor:
cursor_row = await database.litellm_managedfiletable.find_first(
where={**owner_filter, "unified_file_id": cursor}
)
if cursor_row is None:
raise ProxyException(
message=f"Invalid 'after' cursor: no file found with id '{after}'.",
message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.",
type="invalid_request_error",
param="after",
param=cursor_param,
code=400,
openai_code="invalid_value",
)
if model_scope is not None and not model_scope.model_mapping_keys:
return FileListPage.model_validate(build_list_page([]))
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
matches: Final[List[OpenAIFileObject]] = []
cursor_id = after
cursor_id = cursor
chunk_size = page_size + 1
direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc"
while len(matches) <= page_size:
cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_file_table(self.prisma_client).find_many(
where=owner_filter,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
**cursor_args,
chunk: Sequence[_ListedFileRow] = (
await _scoped_file_rows(
database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size
)
if model_scope is not None
else await database.litellm_managedfiletable.find_many(
where=owner_filter,
take=chunk_size,
order=[{"created_at": direction}, {"unified_file_id": direction}],
**cursor_args,
)
)
matches.extend(
_public_file_object(parsed_file_object, row.unified_file_id)
for row in chunk
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
if (
model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None)
)
and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
and (purpose is None or parsed_file_object.purpose == purpose)
)
if len(chunk) < chunk_size:
@ -2315,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
cursor_id = chunk[-1].unified_file_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size))
page: Final = matches[:page_size]
return FileListPage.model_validate(
build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size)
)
def _is_batch_polling_enabled(self) -> bool:
"""
@ -2378,12 +2503,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"not_in": ["failed", "expired", "cancelled"]},
"status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
},
take=MAX_MATCHES_TO_RETURN,
order={"created_at": "desc"},
)
input_file_cutoff: Final = self._clock() - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
referencing_batches: Final[list[dict[str, object]]] = []
for batch in batches:
try:
@ -2396,7 +2522,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Extract file IDs from batch
# Batches typically reference the unified file ID in input_file_id
# Output and error files are generated by the provider
input_file_id = batch_data.get("input_file_id")
input_file_id = batch_data.get("input_file_id") if batch.created_at >= input_file_cutoff else None
output_file_id = batch_data.get("output_file_id")
error_file_id = batch_data.get("error_file_id")
@ -2652,7 +2778,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Continue with other files even if one fails
continue
def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str:
def _get_content_type_from_file_object(self, file_object: object) -> str:
"""
Determine content type from file object.

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

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

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

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

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

@ -15,7 +15,11 @@ from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import TypeAdapter
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
describe_missing_quota,
missing_project_model_quota,
)
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -27,6 +31,15 @@ from litellm.proxy.management.teams.dependencies import get_team_access
from litellm.proxy.management.users.service import org_admin_org_ids
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
from litellm.proxy.management_helpers.model_quota_policy import (
ENFORCE_PROJECT_MODEL_QUOTA_SETTING,
REQUIRE_BOTH_LIMITS,
ModelQuotaPolicy,
RouteOf,
no_route,
parse_model_quota_setting,
router_model_routes,
)
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -36,6 +49,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 +87,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
@ -265,26 +289,6 @@ def _check_team_project_limits(
)
def _project_models_missing_positive_quota(
models: list[str] | None,
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
) -> list[str]:
"""Return the models that lack a positive `rpm` AND `tpm` quota.
A valid quota is a positive integer; null, zero, and negative are rejected
because downstream rate limiters treat a non-positive limit as immediately
exhausted (every request blocked).
"""
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
rpm = rpm_limits or {}
tpm = tpm_limits or {}
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
@ -311,64 +315,79 @@ def _raise_on_project_models_expanding_at_request_time(
raise HTTPException(
status_code=400,
detail={
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead."
},
)
def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None:
try:
return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING))
except ValidationError as e:
raise HTTPException(
status_code=500,
detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"},
) from e
def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
data: NewProjectRequest | UpdateProjectRequest,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> None:
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE.
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
model's `set_model_info` validator, so they are read from there.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing = _project_models_missing_positive_quota(
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
missing: Final = missing_project_model_quota(
data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, route_of
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
"error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
},
)
def _raise_on_missing_project_model_quota_on_update(
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
data: UpdateProjectRequest,
existing_project: object,
access_group_names: frozenset[str] = frozenset(),
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
route_of: RouteOf = no_route,
) -> None:
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE.
`/project/update` replaces `models` and `metadata` when they are provided, so the
check runs on what the project WILL look like: a partial update that doesn't touch
models/quota keeps the existing values, while one that adds a model or clears a
model's quota must leave every resulting model with a positive limit.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
model's quota must leave every resulting model with the limits it needs.
"""
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
resulting_metadata = (
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
)
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
missing = _project_models_missing_positive_quota(
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
missing: Final = missing_project_model_quota(
resulting_models,
resulting_metadata.get("model_rpm_limit"),
resulting_metadata.get("model_tpm_limit"),
policy,
route_of,
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
"error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit."
},
)
@ -519,6 +538,7 @@ async def new_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -567,10 +587,6 @@ async def new_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model added to the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -589,6 +605,15 @@ async def new_project(
},
)
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota(
data,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, data.team_id),
)
# Generate project_id if not provided
if data.project_id is None:
data.project_id = str(uuid.uuid4())
@ -710,6 +735,7 @@ async def update_project(
"""
from litellm.proxy.proxy_server import (
general_settings,
general_settings_view,
litellm_proxy_admin_name,
llm_router,
premium_user,
@ -817,10 +843,14 @@ async def update_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
if general_settings.get("enforce_project_model_quota", False):
quota_policy: Final = _project_model_quota_policy(general_settings_view())
if quota_policy is not None:
_raise_on_missing_project_model_quota_on_update(
data, existing_project, _router_access_group_names(llm_router)
data,
existing_project,
_router_access_group_names(llm_router),
quota_policy,
router_model_routes(llm_router, target_team_id),
)
# Prepare update data
@ -849,7 +879,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}],

49
enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py Normal file
View file

@ -0,0 +1,49 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from litellm.proxy.management_helpers.model_quota_policy import (
ModelQuotaPolicy,
QuotaLimit,
RouteOf,
required_limits,
unmet_limits,
)
_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class MissingModelQuota:
model: str
limits: tuple[QuotaLimit, ...]
def _limits_set_for(
model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object]
) -> Mapping[QuotaLimit, object]:
return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)})
def missing_project_model_quota(
models: Sequence[str],
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
policy: ModelQuotaPolicy,
route_of: RouteOf,
) -> tuple[MissingModelQuota, ...]:
rpm_set: Final = rpm_limits or _NO_LIMITS_SET
tpm_set: Final = tpm_limits or _NO_LIMITS_SET
gaps: Final = tuple(
MissingModelQuota(
model,
unmet_limits(required_limits(policy, model, route_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
)
for model in models
)
return tuple(gap for gap in gaps if gap.limits)
def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str:
return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing)

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

17
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
@ -148,7 +153,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__"
USER nonroot

20
helm/litellm/README.md
View file

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

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

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

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

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

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

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

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql Normal file
View file

@ -0,0 +1,2 @@
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx"
ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC);

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

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

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

@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
@@index([unified_file_id])
@@index([flat_model_file_ids], type: Gin)
@@index([team_id, created_at(sort: Desc)])
@@index([created_by, created_at(sort: Desc)])
}
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the

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

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

36
litellm-rust/.agents/skills/rust-error/SKILL.md Normal file
View file

@ -0,0 +1,36 @@
---
name: rust-error
description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings
---
# Rust errors
`crates/http/src/error.rs` is the reference implementation. Read the reference for the part you are touching before editing, each one carries the reasoning and worked examples behind its rules
- [shape](references/shape.md)
- One top-level `Error` per crate in `src/error.rs`. A public enum by default, an opaque struct over a private `Repr` enum when the variants must stay private. No `#[non_exhaustive]`
- A variant is a failure mode a caller handles differently. Provider names and guidance are data, never variants
- [messages and causes](references/messages-and-causes.md)
- Every message is written once, in `#[error(...)]`, and describes its own layer only. Call sites pass typed facts, never a string
- The cause is the `source()`, never part of the message. Keep it typed, never stringify or drop it
- `reqwest::Error` enters our types only through `litellm_http::Error`, which strips its URL
- Anything shown to someone else is rendered with `litellm_error::Report`, never `to_string()`
- [conversions](references/conversions.md)
- One `From` per source and target pair, only when the conversion needs no context. It classifies and sanitizes once
- Nest a lower crate's error in one variant. Translate when the meaning changes, never mirror its variants
- [classification](references/classification.md)
- Callers ask predicates named after their decision, they do not match nested variants. Predicates match exhaustively
- Classification states facts. Retry and fallback policy lives with the caller that knows idempotency
- [boundaries](references/boundaries.md)
- An error becomes text for someone else only in one mapping function per error type, in the gateway crates, the Python bridge and host-python. The inference gateway and the bridge decide by the same predicates
- A failure before the provider was called is a decline. Both ask `declined()` first and hand it back to the host unchanged
- [gateway](references/gateway.md)
- One `Error` with `IntoResponse` per gateway crate that serves a router. Every failure flows through the same `From` chain, is classified once in `status()` and rendered by surface and framing
- [pyo3](references/pyo3.md)
- Two `create_exception!` classes with the payload in `args`, `ValueError` for the caller's fault and `RuntimeError` otherwise, decoded in `failures.py`
- [testing](references/testing.md)
- Test predicates, chain shape and redaction on errors built the way production builds them
- [recipes](references/recipes.md)
- Common tasks walked through end to end
The workspace uses thiserror 2 for every error type, and no library crate uses `anyhow` or `Box<dyn Error>` in a public signature. The [derive docs](https://docs.rs/thiserror/latest/thiserror/) are the authority on macro behavior, and the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html) on what `Display` and `source()` mean

45
litellm-rust/.agents/skills/rust-error/references/boundaries.md Normal file
View file

@ -0,0 +1,45 @@
# Boundaries
A boundary is where an error becomes something another program reads: an HTTP response, a Python exception, a stored or logged message. Each boundary has one mapping function per error type it receives, and nothing outside those functions formats an error for someone else
| Boundary | Mapping | Reference |
|---|---|---|
| Inference HTTP response | `Error::status()` and its renderers in `crates/gateway-inference/src/error.rs` | [gateway](gateway.md) |
| Other gateway surfaces | the `IntoResponse` of each `crates/gateway-*` crate's `Error` | [gateway](gateway.md) |
| Python bridge, inference and OCR routes | `route_failure` in `crates/python-bridge/src/errors.rs` and in `crates/python-bridge/src/routes/ocr/errors.rs` | [pyo3](pyo3.md) |
| Python bridge, everything else | one `From<Error> for PyErr` or mapping function per crate error, all in `crates/python-bridge` | [pyo3](pyo3.md) |
| Python host runtime | `From<Error> for PyErr` in `crates/host-python/src/error.rs` | [pyo3](pyo3.md) |
## Asking the nested error
The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates in the same order, so a new failure lands correctly on both once its predicates answer right
| Predicate | Gateway | Bridge |
|---|---|---|
| `declined()` | 501 | `RustBridgeDeclined`, without the route's public mapping |
| `upstream_response()` | the provider's status and body | `RustUpstreamError` |
| `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false |
| `is_request()` | 400 | `ValueError` |
| none of these | 500 | `RuntimeError` |
If you find yourself adding an arm for your variant at a boundary, check whether a predicate is missing first
The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does
## Declines
A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway answers 501 and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side
Anything after the send is terminal, never a decline. `declined()` answers `None` for a response fault and for a post-call hook failure even when the wrapped error was a decline, and every route that transforms the provider's answer collapses its errors with `Error::after_send`, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a decline can never describe a call that was billed. Ask `declined()`, do not match the variant behind it, because that after-send cut lives in the predicate
## What reaches the client
A boundary renders anything that is not a provider answer with `litellm_error::Report`, so the reader sees every layer of the chain once. That is why messages are written for a reader, each layer leaves its cause to `source()`, and reqwest URLs are stripped on entry. A boundary that must not show the chain, such as the bridge's TLS settings error that names the setting but not the file, renders fixed text of its own and says so in its test
Provider bodies are passed through intact for gateway and Python error mapping. `litellm_http::Error::upstream(status, body)` retains the complete body, but its `Display`, `Debug` and `Report` diagnostics show only the first 256 characters and mark the cut. Do not log the raw body or `RustUpstreamError` args. OCR's `Provider` error also keeps the whole body for Python's `get_error_class`, bounded by the connection's response size limit
Never format `Debug` of an error into a response or a Python exception. A type whose sources may hold secrets derives `veil::Redact`, so its `Debug` is safe to log
## Logging
Propagate with context, and let the terminal handler own the one log or trace event for a failed call. Do not log at every `?`, and do not log an error you also return. Diagnostic tracing excludes raw error strings and bodies, see [`rust-tracing`](../../rust-tracing/SKILL.md)

66
litellm-rust/.agents/skills/rust-error/references/classification.md Normal file
View file

@ -0,0 +1,66 @@
# Classification
## Ask, don't match
Downstream crates ask a lower error questions instead of matching its nested variants. `litellm_http::Error` answers these
| Predicate | Caller decides |
|---|---|
| `dispatched()` | whether retrying elsewhere could repeat a billed call |
| `is_request()` | the caller's input is at fault, so 400 and no reroute |
| `is_response()` | the provider answered with something we cannot use |
| `upstream_response()` | the provider's own status and body to pass through |
| `status()` | an HTTP status to report |
| `is_connect()` | the connection was never made, for a host's own connect-error mapping |
| `tls_source()` | which TLS setting to blame in a config error |
| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into a decline |
Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()`, and `litellm_llms::Error` and `litellm_inference::Error` answer `declined()`, which the boundaries ask first instead of reaching into `Auth` or `Declined`, see [boundaries](boundaries.md)
```rust
// reaches through two crates' variants and breaks when either changes
Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => StatusCode::UNAUTHORIZED,
// asks the question it cares about
Self::Route(error) if error.is_missing_credential() => StatusCode::UNAUTHORIZED,
```
A predicate is an exhaustive `match` with no `_` arm, never `matches!`, so a new variant does not compile until it has answered every question on purpose
```rust
pub fn is_request(&self) -> bool {
match &*self.0 {
Repr::Request(_) | Repr::BlockedUrl | Repr::DownloadTooLarge => true,
Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => false,
// ... every other variant, listed by name
}
}
```
Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()`
A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant
## Facts, not policy
A predicate reports what happened. The decision to retry, reroute or fall back depends on things the error cannot know, such as whether the request is idempotent, whether its body can be replayed, the deadline and the attempt budget. That decision lives with the caller. `dispatched()` exists for the router's reroute decision, so every transport failure answers it even before that caller is wired
A timeout is the main trap. A request that timed out may already have reached the provider and been billed, so `Timeout` answers `dispatched() == true` and is not safe to retry as if it were a connect failure. Only a failure that provably happened before any byte reached the provider answers `false`: connecting, building or signing the request, and media downloads, which fetch the caller's document before the provider call is made
```rust
pub fn dispatched(&self) -> bool {
match &*self.0 {
Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => true,
Repr::Connect(_) | Repr::RequestBody(_) | Repr::Download(_) | Repr::BlockedUrl => false,
// ... every other variant, listed by name
}
}
```
When the transport cannot tell whether a request reached the provider, it must answer `true`. Wrongly answering `false` risks a duplicate billed call, wrongly answering `true` only costs a reroute
Keep one predicate per question. `status()` currently mixes the provider's status, the media download's status and a synthesized 408 for timeouts, so a caller that wants "did the provider answer" must use `upstream_response()` instead. Do not add more meanings to an existing predicate, add a new one
## Equality
`Arc<Repr>` errors compare with `Arc::ptr_eq`, so a clone equals its original and two separately built errors never compare equal. That exists only so a wrapping enum such as `litellm_inference::Error` can keep its `PartialEq` derive. Do not give a new type pointer equality for any other reason, and do not rely on it in tests. Assert predicates instead, see [testing](testing.md)

92
litellm-rust/.agents/skills/rust-error/references/conversions.md Normal file
View file

@ -0,0 +1,92 @@
# Conversions
## One From per source and target pair
A `From<Foreign> for Ours` exists when the conversion is obvious, infallible and needs no context from the call site. It does all classification and sanitizing for that pair, so every call site is `?` or `.map_err(Ours::from)`
```rust
impl From<reqwest::Error> for Error {
fn from(error: reqwest::Error) -> Self {
let error = error.without_url();
if error.is_timeout() {
Repr::Timeout(error).into()
} else if error.is_connect() || error.is_builder() {
Repr::Connect(error).into()
} else {
Repr::Network(error).into()
}
}
}
```
A call site never re-derives that classification
```rust
// every route used to repeat this, and each drifted a little
.map_err(|err| {
if err.is_connect() || err.is_builder() {
Error::Transport(Connect(err.to_string()))
} else {
Error::Transport(Network(err.to_string()))
}
})?
// now
.map_err(litellm_http::Error::from)?
```
The same foreign type may convert into several of our types, one `From` each. What is not allowed is two places converting the same pair differently
## When the call site knows something the From cannot
One `io::Error` can mean "reading the CA bundle failed" or "resolving the media host failed". A global `From<io::Error>` cannot know which, so there is none. Name the operation at the call site instead
```rust
.map_err(Repr::Resolve)?
.map_err(|source| Repr::Read { path, tls_source, source })?
```
Do not drop meaningful context just to use `?`
## Reclassify once, in the shared From
When a lower error needs a different meaning for every caller, do it in the upper type's `From`, not per route
```rust
impl From<litellm_http::Error> for Error {
fn from(error: litellm_http::Error) -> Self {
if error.is_computed_header() {
return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes"));
}
Self::Http(error)
}
}
```
All four inference routes classify a header the signer computes the same way, because none of them does it themselves
## Nest, translate or add context, never mirror
| The upper layer | Do | Example |
|---|---|---|
| has nothing to add | nest in one transparent variant and delegate predicates | `Self::Http(error) => error.is_request()` |
| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Declined` |
| knows something the lower one does not | a variant with the extra fields plus the cause as its source | `Read { path, tls_source, source }` |
Never re-declare a lower crate's variants in the upper type with an upper-layer name and a 1:1 mapping function. OCR once did that for every media download failure, `MediaError::BlockedUrl => Error::BlockedDocumentUrl` and so on, which doubled every change and let the two lists drift.
## Third-party and host faults
Converting a fault from a crate we do not own into a string variant is still flattening. Prefer a variant that carries the fault, and when an existing `ErrorDetail` constructor fits, use it
```rust
// flattens
.map_err(|error| Error::InvalidRequest(error.to_string().into()))
// keeps the cause typed
.map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("header", error)))
```
## Through a foreign error type
A trait we implement for another crate sometimes forces its error type on us, such as `io::Error` from a body stream or `sqlx::Error` from a driver. Wrap ours as that type's source (`io::Error::other(error)`, `azure_core::Error::new(kind, error)`) and recover it with `downcast_ref` on the other side, never `with_message(error.to_string())` or `from_msg`. The error must come back out as the same value it went in as, so one failure classifies the same way on every path

56
litellm-rust/.agents/skills/rust-error/references/gateway.md Normal file
View file

@ -0,0 +1,56 @@
# Gateway responses
The gateway is where a Rust error becomes an HTTP response. This page is about the path an error takes to get there and the one place it is classified. The gateway crates' own conventions, such as which envelope each API surface uses, which extractors exist and how the router is composed, live in the [`rust-gateway`](../../rust-gateway/SKILL.md) skill
## The flow
```text
reqwest / provider failure
-> litellm_http::Error From<reqwest::Error>, strips the URL, classifies once
-> litellm_inference::Error::Http(_) From, one transparent variant, predicates delegate
-> gateway_inference::Error::Route(_) From, one variant per nested crate
-> handler returns Result<_, Error> every `?` ends here
-> IntoResponse status() picks the code, a renderer builds the body
```
Each hop is one `From` that nests the lower error in one variant, see [conversions](conversions.md). No hop re-derives what the lower crate already decided, and nothing between the provider and `IntoResponse` formats an error for the client. That is what makes "one `Error` with `IntoResponse` per gateway crate that serves a router" work: a handler can `?` any failure from any depth and the response comes out right
## Classify once
`status()` is pure and runs once per error. Variants the gateway declares itself, such as a bad body, an unknown model or an oversized body, are matched at the top of its own enum. A nested route error goes to `route_status` and an OCR error to `ocr_status`, which ask the predicates [boundaries](boundaries.md) lists
| Answer | Status |
|---|---|
| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider |
| `upstream_response()` | the provider's own status, 502 when it is not a valid code |
| `is_missing_credential()` | 401 |
| `is_request()` | 400, for a route error |
| `http_status_code()` | the transport's status, for an OCR error, such as 408 for a timeout |
| otherwise | 500 |
Do not reach into a nested enum's variants from the gateway. When a new nested failure needs a different status, add a predicate on the nested error, see [classification](classification.md)
## Render by surface and framing
Rendering is separate from classification. The surface picks the envelope and the moment picks the framing, a JSON body before the headers went out and an `event: error` SSE frame after. Both framings of one surface render the same status and `Report` message, so a streaming failure and a JSON failure of the same error agree on everything but framing
| Surface | JSON | SSE |
|---|---|---|
| Anthropic, `/v1/messages` | `body(request_id)` | `sse_frame()` |
| OpenAI, chat, responses and the rest | `openai_response()` | `responses_sse_frame()`, the Responses API `error` event |
The message of anything that is not a provider answer is the `litellm_error::Report` of the error, and in the OpenAI and Anthropic JSON envelopes the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Responses `error` event has the fixed type `error` and the status as its `code`. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else. The OpenAI envelope always wraps, because that is what the Python proxy returns there
## Every way in uses the same Error
A request can fail before the handler runs, in the handler, or after the headers went out. All three end in the same `Error`, so they get the same classification and envelope
- An extractor sets `type Rejection = Error`, so a bad body is the surface's envelope rather than axum's plain text
- A path the gateway knows but does not serve returns `Error::Unsupported(path)` from one handler, so the client gets a 501 in the surface's envelope rather than a bare 404
- A stream that fails after its headers went out hands the surface's SSE frame of `Error::from(error)` to `litellm_host_http::Sse::new`. `litellm_host_http::Error<E>` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the `From` maps the second to `Internal`
The axum hooks behind each line are in [`rust-gateway`](../../rust-gateway/SKILL.md)
## Testing
`status()` and `body()` are pure, so unit tests call them on errors built the way production builds them, one `#[case]` per fault, as `status_follows_who_is_at_fault`, `an_anthropic_provider_body_passes_through` and `other_failures_get_the_anthropic_envelope_with_their_message` in `crates/gateway-inference/src/error.rs` do. Assert the status and the fields the surface promises, not the whole body, except where passing the body through untouched is the contract. Driving a whole router belongs to the gateway crate's own tests, see [`rust-gateway`](../../rust-gateway/references/testing.md)

112
litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md Normal file
View file

@ -0,0 +1,112 @@
# Messages and causes
## Messages belong to the variant
Every message template is written once, in the variant's `#[error(...)]`. A caller passes the typed facts the template needs and nothing else
```rust
// none of these may appear at a call site
Error::Network("WebSocket handshake returned a wrong Sec-WebSocket-Accept".into())
Error::Request(format!("bad header {name}"))
Error::Signature(error.to_string())
// the call site names what happened, the variant owns the words
Repr::Handshake(HandshakeFault::Accept).into()
Error::signature(error)
Error::NonAsciiHeader { name }
```
When a check fails for a distinct reason, that reason is a variant or a field, never a string inside a generic variant. A `Message(String)` variant, or a `From<String>` / `From<&str>` into an error, is the hole every stringly call site goes through, so neither exists
Format arguments can call functions and reach into fields, so a hand-written `Display` is rarely needed
```rust
#[error("could not read {}", path.display())]
#[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))]
```
Keep messages short and lowercase, with no trailing period, unless Python parity pins the exact text. A reporter appends `: {cause}` after them. Never format `Debug` (`{:?}`) of an error, or of a type with no `Display`, into a message
## The cause is the source, never the message
This is the rule from the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html): an error's `Display` describes that layer only, and the error underneath is returned by `source()`. thiserror wires `source()` from a field marked `#[source]`, a `#[from]` field, or a field named `source`
```rust
// prints the cause twice under any reporter
#[error("could not read {}: {source}", path.display())]
Read { path: PathBuf, source: io::Error },
#[error("invalid JSON: {0}")]
Json(#[from] serde_json::Error),
// this layer's words, the cause is the source
#[error("could not read {}", path.display())]
Read { path: PathBuf, source: io::Error },
#[error("invalid JSON")]
Json(#[from] serde_json::Error),
```
| Field | thiserror does | Use for |
|---|---|---|
| `source: E` or `#[source] cause: E` | exposes it as `source()` | a cause with context fields next to it |
| `#[from] E` | implies `#[source]` and adds `From` | context-free wrapping, see [conversions](conversions.md) |
| `#[error(transparent)]` | forwards both `Display` and `source()` to the inner value | nesting another of our errors that adds no text |
| no attribute, any other name | nothing, the field is plain data | facts the message renders, such as a path or a status |
`#[error(transparent)]` is an invisible layer, not an extra chain node, so a transparent wrapper's `source()` is the inner value's `source()`
A field that holds an error but is neither rendered nor a source hides that error from every reader. Every error-typed field is a source. A fieldless reason enum, such as `HandshakeFault` or `PagesFault`, is data and not a cause, so the variant renders it in its message
`source()` is one chain, so an error with several causes, such as `CredentialChain` trying each credential in turn, makes the most relevant one its source and renders the others in its message, each with `Report`
## Keep causes typed
Store the lower-level error itself, never its string, and never drop it with `map_err(|_| ...)`
| Flattening or dropping loses | Because |
|---|---|
| the chain | a `String` has no `source()` |
| downcasting | the type is gone |
| redaction | `reqwest::Error`'s `Display` includes the request URL, which can carry credentials in the path or query |
If the cause is not `Clone` and the error must be, hold it as `#[source] litellm_error::Source`, a cloneable `Arc<dyn Error + Send + Sync>` that is transparent to `Display` and `source()`, or make the crate's error opaque. `litellm_error::find::<T>(&error)` walks the chain and looks through `Source`, so it is how anyone reaches a typed cause. A third-party `Box<dyn Error + Send + Sync>` converts with `Source::from`, and an `anyhow::Error` a dependency returns goes through `Box` first, `Source::from(Box::<dyn Error + Send + Sync>::from(error))`, and a dependency's own boxed alias such as `tokenizers::Error` is held as it is. A source boxed to keep the variant small is a `Box<E>` node of the chain, so it is reached as `find::<Box<E>>`. Never put `#[error(transparent)]` on a `Source`, the erased error would then be no node of the chain and `find` could not reach it. Do not stringify or drop a cause to get `Clone` or `PartialEq`
```rust
// loses the type and leaks the URL
.map_err(|error| Error::Network(error.to_string()))
// loses everything
.map_err(|_| Error::Transport)
// keeps the cause, the From impl strips the URL
.map_err(litellm_http::Error::from)
```
A cause that must not be shown at all, such as a secret manager's response, is still a source. The type derives [`veil::Redact`](https://docs.rs/veil) with `#[redact]` on the field, so `Debug` hides it, and the boundary decides how much of the chain to show, see [boundaries](boundaries.md)
## Strip reqwest URLs on entry
Every `reqwest::Error` enters our types through `From<reqwest::Error> for litellm_http::Error`, which calls `without_url()`. A crate that talks HTTP holds `litellm_http::Error`, never `reqwest::Error`. When a dependency's own error wraps a `reqwest::Error`, as vaultrs and rustify do, one `From` takes the reqwest error out and converts it, as `secrets-hashicorp` does. That way so every later `Display`, `Debug` or chain walk is safe from that URL
`without_url()` only removes the URL reqwest attached. It does not scrub headers, response bodies, provider messages, file paths or anything a variant stores next to it. Those need their own policy, see [boundaries](boundaries.md)
## Rendering
Anything that shows an error to someone else renders the whole chain with `litellm_error::Report`, which writes each layer once joined by `: `, the single-line form of the unstable [`std::error::Report`](https://doc.rust-lang.org/std/error/struct.Report.html). `to_string()` of an error is only its top layer and is never what a client, a Python exception or a log line shows
```rust
let message = litellm_error::Report::new(&error).to_string();
// "could not reach the provider: error sending request: client error (Connect): tcp connect error: Connection refused"
```
## #[from] limits
`#[from]` only works on a variant whose only fields are the source and an optional backtrace. A variant that needs context such as a path, a phase or a `TlsSource` takes its cause through an explicit `map_err` or a named constructor
```rust
std::fs::read(&path).map_err(|source| Repr::Read {
path: path.clone(),
tls_source: TlsSource::CaBundle,
source,
})?;
```
`#[backtrace]` and automatic `provide()` need a nightly compiler and are not used here

76
litellm-rust/.agents/skills/rust-error/references/pyo3.md Normal file
View file

@ -0,0 +1,76 @@
# Python bridge
`crates/python-bridge` turns a Rust error into a Python exception through one mapping function per error type. A `#[pyfunction]` returns `Err(PyErr)` and pyo3 raises it, so those functions are the whole boundary
| Error | Mapping |
|---|---|
| `litellm_inference::Error` from a route | `errors::route_failure`, with `route_error_to_pyerr` as its shorthand for routes that set no marker |
| the OCR error | `route_failure` and `to_pyerr` in `routes/ocr/errors.rs` |
| token counter, cache, spend, secrets, argument checks | one `From<Error> for PyErr` or mapping function next to the code that registers the functions |
A route host never builds a `PyErr` from a route error itself. It calls the mapping, and the mapping is the only place that knows the exception classes and their attributes
## Exception classes
| Class | Meaning | Payload |
|---|---|---|
| `RustUpstreamError` | the provider call was issued and the provider answered with a failure | `args == (status, body)` and a `headers` attribute of `(name, value)` pairs, always set |
| `RustBridgeDeclined` | the route stopped before calling the provider, so the host may run its own path | `args == (message,)` |
| a standard class from `pyo3::exceptions` | whichever Python class means the same thing, such as `ValueError` or `TypeError` for the caller's input, `FileNotFoundError` or `OSError` for a file, `KeyError` or `OverflowError` where Python raises them, and `RuntimeError` for anything else | the rendered chain, or a parity text where Python pins one |
The two native classes come from `create_exception!` in `crates/python-bridge/src/errors.rs` and are exported with `#[pymodule_export]`, which is how the [pyo3 guide](https://github.com/PyO3/pyo3/blob/main/guide/src/exception.md) makes an exception importable. Their docstrings are the contract Python reads. `errors::upstream_error(py, status, body, headers)` is the one constructor of `RustUpstreamError`, so the payload cannot differ between routes
Python decodes a `RustUpstreamError` in `litellm/rust_bridge/failures.py`, which validates `args` and `headers` with a `TypeAdapter` and rebuilds the provider's `httpx.Response`. `litellm/rust_bridge/runtime.py` catches `RustBridgeDeclined` and reads its reason. `_raise_upstream` in `runtime.py` still reads `args` a second time by hand for the `attempt` path, which is the one decode to fold into `failures.py`. A new fact goes into `args` or `headers`, and the decoder learns to read it in the same change
Two routes also flag a request fault with a boolean attribute, read by their Python route module: `ocr_request_format_error` in `host/routes/ocr.py` and `messages_request_error` in `host/routes/messages.py`. OCR additionally sets `status_code` and `message` for Python parity, and `native_settings_error`, which only the OCR host reads back to tell a settings failure from the caller's exception. Do not add another attribute, nothing on the Python side types them
## Why args, not a typed subclass
The bridge builds with pyo3's `abi3-py310` feature. Under abi3, `#[pyclass(extends = PyException)]` is only allowed on Python 3.12 and newer, so a class with typed fields is not available. Even where it is allowed, pyo3 passes an empty tuple to the native base's `tp_new`, so `.args` and `str()` would need manual overrides. The payload therefore travels in `args`, which a `create_exception!` class fills from `new_err`
## Mapping order
A core decline is returned first, as `RustBridgeDeclined` with the decline's reason as its only arg, without calling the route's public mapping. Then a Python exception raised inside the call, such as a secret manager callback's, is returned as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that
OCR puts Python parity arms in front of that order, each matching a top-level variant of its own public enum
| Arm | Exception |
|---|---|
| the provider answered | `RustUpstreamError` with the provider's headers |
| `RequestFormat` | `ValueError` flagged with `ocr_request_format_error` |
| `FileRead` whose cause is `NotFound` | `FileNotFoundError` naming the path |
| other `FileRead` | `OSError` with the io message |
Both routes ask `litellm_auth::Error::is_request()` through their own `is_request()`, so a missing key or API base is a `RuntimeError` on both, and a missing setting the caller can pass, such as a region, is a `ValueError`
Use the standard class whose Python meaning matches. `PyFileNotFoundError` for a missing file is right, a `RuntimeError` whose text says the file is missing is not. `import_exception!` exists for an exception defined in Python, none is used today because `failures.py` maps to LiteLLM's public exceptions through `litellm.exception_type`
## Declining
`RustBridgeDeclined` means no provider call was made and the host may run its own path for the same call. The tokenizer, embeddings and token counter routes raise it from their own pre-checks, and `route_failure` raises it for a route error whose `declined()` answers, so a decline from deep inside a transformation reaches Python the same way. Never raise it for a failure after the request went out, the host would repeat a billed call. `a_core_decline_hands_the_call_back_to_python_before_the_public_mapping` and `route_failures_are_terminal` in `errors.rs` pin both sides
Route hosts pass a terminal failure through the Python route module's `map_failure`, which runs `litellm.exception_type` and returns the public exception a Python caller expects. A decline must skip that step. Once wrapped it looks like a provider failure to the dispatcher and the host no longer falls back, so `route_failure` returns it before it calls `public`
## Rules
- Never format `Debug` of an error into a `PyErr`
- Never build a `PyErr` from a string at a call site. A failed argument check is a variant of a typed error, as `ProjectionError` in `coercion.rs` is, mapped once
- A Python exception raised inside the call stays that exception. Return it as is, never wrap it in a Rust message or stringify it into a new `ValueError`. The one policy exception is a caller-supplied callable, whose failure `callable::wrap_failure` reports as a `RuntimeError` with the original as `__cause__`
- A `panic!`, `unwrap` or `expect` inside a `#[pyfunction]` or `#[getter]` reaches Python as `pyo3_runtime.PanicException`, which is not an `Exception` subclass. Model the failure as a value instead
- When the mapping itself calls Python, as `setattr` does, propagate its failure with `?` rather than discarding it with `.ok()`
## Testing
Build the Rust error the way production builds it, map it, and assert the contract Python sees. `Python::initialize()` once, then `Python::attach`
```rust
Python::initialize();
Python::attach(|py| {
let failure = route_error_to_pyerr(litellm_http::Error::upstream(429, "slow down".into()).into());
assert!(failure.is_instance_of::<RustUpstreamError>(py));
let args: (u16, String) = failure.value(py).getattr("args").unwrap().extract().unwrap();
assert_eq!(args, (429, "slow down".into()));
});
```
One `#[case]` per fault class, asserting `is_instance_of` for the class and `extract` for the payload. The Python side of the contract is tested in `tests/unit/rust_bridge`, through `failures.py`

64
litellm-rust/.agents/skills/rust-error/references/recipes.md Normal file
View file

@ -0,0 +1,64 @@
# Recipes
## A new failure in litellm_http
Say media downloads must reject a `Content-Type` outside an allow-list
1. Add a `Repr` variant with typed facts and the message, `#[error("media type {0:?} is not allowed")] DisallowedMediaType(String)`
2. Return it where the check fails, `return Err(Repr::DisallowedMediaType(content_type.to_owned()).into())`
3. Add an arm to every predicate's exhaustive `match`, the compiler lists them, the way its neighbor `DownloadTooLarge` answers. A media download happens before the provider call, so it is not in `dispatched()`. The caller chose the URL, so it is in `is_request()`
4. Leave both boundaries alone. With those answers the gateway returns 400 and Python raises `ValueError`
5. Test it with a mock server answering `Content-Type: application/x-msdownload` and assert `error.is_request()`
## A provider-specific failure
A second OCR provider needs a missing-key error whose text names its environment variable and how to set it
1. Look for an existing provider-neutral variant whose template fits. `litellm_auth::Error::MissingApiKey { provider, environment_variable }` renders `Missing {provider} API Key - Set ...`
2. If none fits, find what varies between providers with the same failure, here the environment variable and the guidance after the dash, and add one neutral variant with those as fields, `#[error("Missing {environment_variable} - {guidance}")] MissingCredential { environment_variable: &'static str, guidance: &'static str }`
3. Pass the provider's values from its transformation. The next provider passes its own and adds no variant, and `is_missing_credential()` lists `MissingCredential { .. }` once
## Wrapping a new lower crate in litellm_inference
Say routes start calling `litellm_cache` and its errors must surface
1. Add one variant, `#[error(transparent)] Cache(#[from] litellm_cache::Error)`
2. Extend every predicate's `match` with an arm that calls through to the cache error's own predicates, adding a predicate there if it has none
3. Do not copy the cache error's variants into `litellm_inference::Error`
## A foreign error that needs context
Reading a file fails with `io::Error`, and the message must name the file
1. A variant with the context and the cause, `Read { path: PathBuf, source: io::Error }` with `#[error("could not read {}", path.display())]`
2. Build it with `map_err(|source| Repr::Read { path: path.clone(), source })`, not a `From<io::Error>`
3. If the error must be `Clone`, the crate's error is opaque or the field is `litellm_error::Source`
## Migrating a stringly error
Call sites look like `Error::Network(error.to_string())`, `Error::Request(format!(...))` or `.map_err(|_| Error::Unavailable)`
1. List every call site of the string variant and group them by what actually failed
2. Each group becomes a variant holding the typed cause or facts. Groups no caller tells apart share one
3. Where a foreign error was stringified with classification logic around it, move that logic into the `From` and turn the call sites into `?`
4. Delete the string variant. If it cannot be deleted yet, nothing new may construct it
5. Re-run the boundary tests, since messages a client sees may change
## Moving a crate to source chains
A variant reads `#[error("invalid JSON: {0}")] Json(#[from] serde_json::Error)` or stores `cause` and renders `{cause}`
1. Drop the cause from the template, `#[error("invalid JSON")]`
2. Make the cause the source, `#[from]`, `#[source]` or a field named `source`
3. Find what renders this error for someone else and make it use `litellm_error::Report`
4. Update tests that asserted the cause in `to_string()` to assert it in the report
## Making a public enum opaque
Downstream crates match nested variants, or a variant needs a non-`Clone` cause
1. Rename the enum to `pub(crate) enum Repr` and add `pub struct Error(Arc<Repr>)` with `#[error(transparent)]` and `From<Repr>`
2. For each downstream `match`, find the question it asks and add a predicate. Replace the `match` with the predicate
3. For each downstream construction, add a constructor taking typed data, or move the failure into this crate
4. Swap test assertions to predicates and add the `#[cfg(test)] repr()` accessor for in-crate tests
5. Implement `PartialEq` with `Arc::ptr_eq` only if a wrapping enum's derive needs it

105
litellm-rust/.agents/skills/rust-error/references/shape.md Normal file
View file

@ -0,0 +1,105 @@
# Error shape
## One top-level Error per crate
A crate's errors live in `src/error.rs` and are re-exported from `lib.rs` as `Error`. A caller of any public function in the crate should only ever see `Result<_, crate::Error>`, so one `?` and one set of predicates cover the whole crate
This deliberately departs from the advice to put an error type next to each unit of fallibility ([Sabrina Jewson, "Modular Errors"](https://sabrinajewson.org/blog/errors)). Our callers do not branch on which function failed, they ask the predicates in [classification](classification.md), and the boundaries map one type per crate. Splitting by function would multiply those predicates and mappings without giving any caller a decision it lacks. Layering, the other half of that advice, is kept: each crate nests the one below it, see [conversions](conversions.md)
That one type takes one of two forms
| Form | Use when | Example |
|---|---|---|
| Public enum | Callers outside the crate match variants, and every variant holds our own data or another crate's `Error` that is already `Clone` and comparable | `litellm_inference::Error`, `litellm_auth::Error` |
| Opaque struct over a private `Repr` | The variants hold third-party causes, or the variants must be free to change without breaking callers | `litellm_http::Error` |
Both forms are still one top-level `Error`. The opaque form just moves the enum behind `Repr`, where it is `pub(crate)`
```rust
// public enum
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error(transparent)]
Http(litellm_http::Error),
}
// opaque struct
#[derive(Clone, Debug, thiserror::Error)]
#[error(transparent)]
pub struct Error(Arc<Repr>);
#[derive(Debug, thiserror::Error)]
pub(crate) enum Repr {
#[error("could not reach the provider")]
Connect(#[source] reqwest::Error),
#[error("upstream request failed with status {status}: {body}")]
Status { status: u16, body: String },
}
```
Start with the public enum. Move to opaque when a variant needs a non-`Clone` third-party cause, or when downstream crates keep matching variants in ways that duplicate classification. `reqwest::Error`, `std::io::Error` and `serde_json::Error` are built the opaque way
In the opaque form other crates cannot name a variant, so they cannot build an error with their own message. Public constructors exist only for errors another crate genuinely produces, take typed data and never a message, as in `Error::upstream(status, body)`, `Error::computed_header(name)`, `Error::unsignable_body()` and `Error::signature(cause)`
`Arc` makes cloning a refcount bump, so `Repr` never has to be `Clone`. `Box<Repr>` is enough when the goal is only a small `Result`. `Arc` is for when callers clone
Every crate here is internal to the workspace, so do not mark errors or variants `#[non_exhaustive]`. It only forces a wildcard arm on the other workspace crates and turns off the exhaustiveness check that makes a new variant visit every `match`. When the variants must be free to change, use the opaque form
## What else may sit next to Error
These are not a second top-level error and need no folding in
- Leaf data types whose fields are typed and whose message is in `#[error(...)]`, such as `litellm_http::request::HeaderError`. A wrapping crate nests it as one variant
- Plain data used by predicates, such as `litellm_http::TlsSource`
- A private fault enum one variant carries, such as `PemFault` inside `Repr::InvalidPem`
## Variants are failure modes
A public variant is something a caller handles differently: a phase, a status code, retry versus no retry, fault on the request versus the environment, or a message Python parity pins exactly. Failures no caller tells apart share one variant and differ in their data
```rust
// one variant per provider, so every new provider adds a variant and a match arm
#[error("Missing Azure API Base - set AZURE_API_BASE")]
MissingAzureApiBase,
#[error("Missing Vertex API Base - set VERTEXAI_API_BASE")]
MissingVertexApiBase,
// one failure mode, the provider and its guidance are data
#[error("Missing {provider} API Base - {guidance}")]
MissingApiBase { provider: &'static str, guidance: &'static str },
```
An exact parity message does not justify a provider-specific variant when caller-supplied data can produce the same text
Inside a private `Repr` a variant costs no API, so one per distinct reason is fine when each one carries a different message. When several only differ in a word and share every predicate answer, prefer one variant with a fieldless reason enum, so predicates list one arm instead of four, as `Repr::Handshake(HandshakeFault)` does
```rust
#[error("WebSocket handshake returned {0}")]
Handshake(HandshakeFault),
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandshakeFault {
#[error("a wrong Sec-WebSocket-Accept")]
Accept,
#[error("an invalid Upgrade header")]
Upgrade,
#[error("an invalid Connection header")]
Connection,
#[error("an invalid subprotocol")]
Subprotocol,
}
```
## When to split
Split into a separate type only when different functions fail in disjoint ways and different callers see them, so no caller matches variants its function can never return. A type per function overdoes it
Name a split type after what went wrong, not after the function that returns it. `ConnectError` reads better than `OpenSessionError`, and a unit struct is fine for a single failure mode
A split type still converts into the crate's `Error` when a caller further up needs one `Result` type
`litellm-llms` has one recorded exception, the OCR error in `src/base_llm/ocr/error.rs`, kept apart until OCR folds into `litellm_llms::Error`, see `crates/inference-ocr/AGENTS.md`. Do not add another
A utility crate whose modules are unrelated, such as `core-utils`, has no crate-wide failure, so each module that can fail owns one error next to its code instead of a `src/error.rs`. The rest of this skill applies to each of them

37
litellm-rust/.agents/skills/rust-error/references/testing.md Normal file
View file

@ -0,0 +1,37 @@
# Testing errors
Assert behavior through predicates and accessors, on errors produced the way production produces them
```rust
// guesses the message and the variant, and cannot be written once the type is opaque
assert_eq!(error, Error::Transport(TransportError::Http { status: 429, body: "slow down".into() }));
// asserts what callers rely on
assert_eq!(error.upstream_response(), Some((429, "slow down")));
assert!(error.dispatched());
assert!(!error.is_request());
```
## Producing the error
In order of preference
1. A real failing call. Bind a `TcpListener`, drop it, then connect to its address for a connect failure. Accept and never answer for a timeout
2. A real cause, such as `tungstenite::Error::ConnectionClosed` or a `serde_json::from_str("{")` error, passed through our `From`
3. A public constructor, such as `Error::upstream(500, "boom".into())`
Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, reach the cause the test expects with `litellm_error::find::<T>(&error)`. A transparent layer is no node of the chain, so `find` reaches foreign causes underneath it but never the nested error of ours itself. Ask that error's predicates through the outer type instead, which is what callers do too
## What to cover
| Property | Test shape |
|---|---|
| classification | one `#[case]` per input class, asserting the predicate |
| redaction | build from a URL holding `api_key=secret`, assert the message lacks `secret` and the path |
| cause kept | `litellm_error::find::<T>(&error)` returns the expected cause |
| rendered once | the cause's text appears exactly once in `litellm_error::Report::new(&error).to_string()` and not in `to_string()` |
| boundary mapping | the gateway status and body, or the Python exception type and its `status` attribute |
Assert exact message text only where it is a contract, such as a message Python parity pins. Everywhere else assert that the message contains the facts it must carry. Matching a fieldless variant with `matches!` is fine, full `assert_eq!` on an error that holds a cause is not, since causes compare by pointer
A test that returns `Result` reports only `Debug` of the error on failure. Prefer `expect` with a short reason, or `unwrap_err` followed by assertions on the predicates

20
litellm-rust/.agents/skills/rust-gateway/SKILL.md Normal file
View file

@ -0,0 +1,20 @@
---
name: rust-gateway
description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test
---
# Rust gateway
`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching
- [surfaces](references/surfaces.md)
- Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose
- The error type string is derived from the status and provider bodies pass through bounded
- [requests](references/requests.md)
- Extractors reject with the crate's `Error`, handlers take the body as a `Result` when the error needs the request id
- Every alias of a path is registered together, a known but unserved path goes to one handler
- Streams relay through `litellm_host_http::Sse`, with the error frame rendered by the crate's `Error`
- [testing](references/testing.md)
- Build the router through `tests/support` and drive it with `oneshot`, assert the fields the surface promises
How an error is classified before it reaches a response is the [`rust-error`](../rust-error/SKILL.md) skill, in particular its [gateway](../rust-error/references/gateway.md) page. Axum's own conventions are in its [docs](https://docs.rs/axum/latest/axum/)

34
litellm-rust/.agents/skills/rust-gateway/references/requests.md Normal file
View file

@ -0,0 +1,34 @@
# Requests
## Extractors reject with the crate's Error
A custom extractor sets `type Rejection = Error`, so a bad body gets the surface's envelope rather than axum's plain-text rejection. Never take axum's own `Json`, `Query` or `Path` in a handler, their rejection is plain text. The rejection is the variant's source, never its string. `JsonObject` and `InferenceBody` in `crates/gateway-inference/src/request.rs` are the pattern
```rust
impl<S: Send + Sync> FromRequest<S> for JsonObject {
type Rejection = Error;
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
let body = Bytes::from_request(request, state).await.map_err(|error| {
if error.status() == StatusCode::PAYLOAD_TOO_LARGE {
Error::BodyTooLarge
} else {
BodyError::Read(error).into()
}
})?;
Ok(Self(object(&body)?))
}
}
```
When the handler needs something from the request in the error body, such as the request id on `/v1/messages`, take the extractor as `Result<JsonObject, Error>` so the handler maps the rejection itself, as `messages::create` does. `Infallible` is the right rejection for an extractor that only reads an optional header, as `RequestId` does
## Paths
Every alias the Python proxy registers for a route is registered together in `crates/gateway-inference/src/lib.rs`, such as the bare path, its `/v1` form and the `/engines/{model}` and `/openai/deployments/{model}` forms, otherwise one alias silently behaves differently from the others. A path the gateway knows but does not serve, such as `/v1/completions`, is routed to `request::unsupported`, which returns `Error::Unsupported(path)` so the client gets a 501 in the surface's envelope rather than a bare 404
`crates/gateway` merges the mounted routers and adds the shared inbound middleware as `route_layer`s, so authentication runs only for matched routes. Endpoint paths and parsing stay in the mounted crate, see the two crates' `AGENTS.md`
## Streaming
Once the response headers are sent the status cannot change. A streaming handler builds `litellm_host_http::Sse::new(Json, |error| Bytes::from(Error::from(error).sse_frame()))`, where the closure renders a mid-stream failure as an `event: error` frame with the same status and message a JSON failure on that surface carries, `sse_frame()` for Anthropic and `responses_sse_frame()` for the Responses API's `error` event. `litellm_host_http::Error<E>` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the crate's `From` maps the second to `Internal`

20
litellm-rust/.agents/skills/rust-gateway/references/surfaces.md Normal file
View file

@ -0,0 +1,20 @@
# Envelopes follow the API surface
Each surface mirrors what the Python proxy returns there, so the shapes differ on purpose
| Surface | Envelope | Where |
|---|---|---|
| `/v1/messages` | Anthropic, `{"type": "error", "error": {"type", "message"}, "request_id"}` | `Error::body()` and `sse_frame()` in `crates/gateway-inference/src/error.rs` |
| OpenAI-shaped routes | `{"error": {"message", "type", "param", "code"}}` | `Error::openai_response()` in the same file |
| `/v1/responses` stream | the Responses `error` event, `{"type": "error", "code", "message", "param"}` | `Error::responses_sse_frame()` in the same file |
| MCP | `{"detail": message}` | `crates/gateway-mcp/src/error.rs` |
| UI | `{"error": {"message"}}`, plus `Retry-After` on 429 | `crates/gateway-ui/src/error.rs` |
| auth rejections | plain text, the top-level message only | `crates/gateway-auth/src/error.rs` |
Every surface takes its status from one `status()` on its crate's `Error` and renders the message of anything that is not a provider answer with `litellm_error::Report`. In `gateway-inference` the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else, and the OpenAI envelope always wraps, as the Python proxy does. A provider body is bounded by `litellm_http::Error::upstream` when the error is built
Two surfaces show less than the chain on purpose, and a test in each pins it. An auth rejection goes to a caller who has not authenticated, so it never shows what a secret manager or key store said. An MCP upstream or authorization failure shows only its own layer, because rmcp renders the transport error, URL included, into its own `Display`, and an authorizer's cause is the same backend detail auth rejections hide. The inference surfaces show only the top layer of an auth rejection for the same reason
`/v1/messages` adds the request id to the envelope when the provider did not set one, which is why its handler takes the body as `Result<JsonObject, Error>` and renders the rejection itself, see [requests](requests.md)
A new surface gets its own renderer next to the others, named after the shape it produces, and its status still comes from `Error::status()`. Picking the status is the [`rust-error`](../../rust-error/references/gateway.md) skill's concern, not the surface's

9
litellm-rust/.agents/skills/rust-gateway/references/testing.md Normal file
View file

@ -0,0 +1,9 @@
# Testing
Test HTTP contracts in the crate that owns the endpoint: status codes, forwarded headers, error envelopes and streaming behavior. Inference and provider behavior is tested in the `inference-<fmt>` and `llms` crates, and mounting and server lifecycle in `crates/gateway`
`crates/gateway-inference/tests/support/mod.rs` builds a `Router` with injected deployments, pools and settings through `app()` and its variants, and `post()` drives it with `tower::ServiceExt::oneshot`. `json()` reads the response body back. Use these instead of building a router or a client by hand in a test
Assert the status and the fields the surface promises, not the whole body. A test that pins a whole envelope breaks on every unrelated field, one that checks `status`, `error.type` and the presence of `request_id` catches the regressions that matter
`status()` and `body()` of the error type are pure, so a classification case belongs in the error module's unit tests, see the [`rust-error`](../../rust-error/references/gateway.md) skill. Router tests cover what only the full path shows, such as an extractor rejection arriving in the right envelope or an SSE error frame after the headers went out

12
litellm-rust/AGENTS.md
View file

@ -4,6 +4,8 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a
For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md)
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`
@ -24,13 +26,3 @@ Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new
- Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body. Give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product
- Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries
## Error definitions
- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs`
- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish. Different wording or providers do not justify new variants
- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly). Failures no caller tells apart share one variant and differ only in its message
- Keep shared error enums minimal and provider-neutral. Provider names, credential types, configuration fields, and setup guidance belong in caller-supplied data, not dedicated variants or hardcoded shared messages. Reuse a variant for the same failure mode across providers, such as `MissingApiBase { provider: "Azure", guidance: "..." }`. An exact parity message does not justify a provider-specific variant when caller-supplied context can preserve it
- Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string
- Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return
- Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it

261
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]]
@ -2329,7 +2342,7 @@ dependencies = [
"hmac 0.13.0",
"http 1.4.2",
"jiff",
"reqwest 0.13.5",
"reqwest",
"rustc_version",
"rustls 0.23.42",
"rustls-pki-types",
@ -2387,7 +2400,7 @@ dependencies = [
"pin-project",
"prost",
"prost-types",
"reqwest 0.13.5",
"reqwest",
"rustc_version",
"serde",
"serde_json",
@ -2883,7 +2896,6 @@ dependencies = [
"tokio",
"tokio-rustls 0.26.4",
"tower-service",
"webpki-roots",
]
[[package]]
@ -3373,10 +3385,13 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"bytes",
"futures-util",
"litellm-auth-types",
"litellm-error",
"litellm-http",
"moka",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -3393,6 +3408,7 @@ dependencies = [
"azure_core",
"azure_identity",
"litellm-auth-types",
"litellm-error",
"moka",
"rstest",
"serde_json",
@ -3411,6 +3427,7 @@ dependencies = [
"google-cloud-auth",
"http 1.4.2",
"litellm-auth-types",
"litellm-error",
"moka",
"rstest",
"serde",
@ -3425,6 +3442,7 @@ dependencies = [
name = "litellm-auth-types"
version = "0.1.0"
dependencies = [
"litellm-error",
"rstest",
"serde",
"serde_json",
@ -3438,6 +3456,7 @@ dependencies = [
name = "litellm-cache"
version = "0.1.0"
dependencies = [
"litellm-error",
"rstest",
"serde",
"serde_json",
@ -3460,9 +3479,10 @@ dependencies = [
"litellm-cache-response",
"litellm-cache-testing",
"litellm-http",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
"wiremock",
@ -3481,6 +3501,7 @@ dependencies = [
"serde-pickle",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
]
@ -3493,11 +3514,13 @@ dependencies = [
"litellm-auth-types",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"percent-encoding",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"wiremock",
]
@ -3520,9 +3543,10 @@ dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"qdrant-client",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -3540,11 +3564,13 @@ version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -3597,8 +3623,9 @@ dependencies = [
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"tokio",
@ -3619,11 +3646,13 @@ dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"litellm-error",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"uuid",
]
@ -3633,6 +3662,7 @@ name = "litellm-callbacks-legacy-python"
version = "0.1.0"
dependencies = [
"litellm-auth",
"litellm-error",
"litellm-host",
"litellm-host-python",
"proptest",
@ -3640,6 +3670,15 @@ dependencies = [
"rstest",
"serde_json",
"strum",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-circuit-breaker"
version = "0.1.0"
dependencies = [
"proptest",
"rstest",
]
[[package]]
@ -3647,6 +3686,7 @@ name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"litellm-error",
"litellm-router-types",
"rstest",
"serde",
@ -3655,6 +3695,8 @@ dependencies = [
"strum",
"tempfile",
"thiserror 2.0.19",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -3662,6 +3704,7 @@ name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"litellm-error",
"litellm-llms-types",
"litellm-tracing",
"rstest",
@ -3689,14 +3732,22 @@ dependencies = [
"criterion",
"proptest",
"rstest",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-db"
version = "0.1.0"
dependencies = [
"chrono",
"macro_rules_attribute",
"rstest",
"serde",
"serde_json",
"sqlx",
"strum",
"thiserror 2.0.19",
"tokio",
]
[[package]]
@ -3711,6 +3762,14 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-error"
version = "0.1.0"
dependencies = [
"rstest",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-framer"
version = "0.1.0"
@ -3736,7 +3795,9 @@ dependencies = [
"futures-util",
"http-body-util",
"litellm-auth-types",
"litellm-circuit-breaker",
"litellm-config",
"litellm-error",
"litellm-gateway-auth",
"litellm-gateway-inference",
"litellm-gateway-mcp",
@ -3744,6 +3805,8 @@ dependencies = [
"litellm-http",
"litellm-inference",
"litellm-llms",
"litellm-redis-client",
"litellm-redis-testing",
"litellm-secrets",
"litellm-tracing",
"rstest",
@ -3768,6 +3831,7 @@ dependencies = [
"futures-util",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secrets",
"rstest",
"serde",
@ -3791,6 +3855,7 @@ dependencies = [
"litellm-auth",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-error",
"litellm-gateway-auth",
"litellm-host",
"litellm-host-http",
@ -3810,6 +3875,7 @@ dependencies = [
"serde_json",
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
"tower",
"wiremock",
]
@ -3820,6 +3886,7 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-auth-types",
"litellm-error",
"litellm-gateway-auth",
"rand 0.8.7",
"rstest",
@ -3837,6 +3904,7 @@ dependencies = [
"http 1.4.2",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secrets",
"moka",
"rmcp",
@ -3851,6 +3919,7 @@ dependencies = [
"tower",
"url",
"uuid",
"veil",
]
[[package]]
@ -3860,9 +3929,11 @@ dependencies = [
"axum",
"axum-login",
"base64 0.22.1",
"futures-util",
"governor",
"jsonwebtoken",
"litellm-auth-types",
"litellm-error",
"litellm-gateway-auth",
"rand 0.8.7",
"rstest",
@ -3887,6 +3958,7 @@ dependencies = [
"litellm-coroutine",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -3924,6 +3996,7 @@ dependencies = [
"bytes",
"futures-util",
"litellm-coroutine",
"litellm-error",
"litellm-host",
"parking_lot",
"pyo3",
@ -3944,10 +4017,9 @@ dependencies = [
"http 1.4.2",
"hyper-util",
"litellm-core-utils",
"litellm-error",
"rcgen",
"reqwest 0.12.28",
"reqwest 0.13.5",
"rmcp",
"reqwest",
"rstest",
"rustls 0.23.42",
"rustls-native-certs",
@ -3957,8 +4029,10 @@ dependencies = [
"tempfile",
"thiserror 2.0.19",
"tokio",
"tokio-rustls 0.26.4",
"tokio-tungstenite",
"tracing",
"url",
"veil",
"webpki-roots",
]
@ -3974,13 +4048,14 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-http",
"litellm-llms",
"litellm-secrets",
"litellm-tracing",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4001,6 +4076,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4028,6 +4104,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-host-native",
@ -4039,7 +4116,7 @@ dependencies = [
"litellm-router-types",
"litellm-secrets",
"litellm-tracing",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4059,6 +4136,7 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4076,6 +4154,7 @@ dependencies = [
"time",
"tokio",
"tracing",
"url",
"wiremock",
]
@ -4091,6 +4170,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4100,7 +4180,7 @@ dependencies = [
"litellm-llms-types",
"litellm-secrets",
"litellm-tracing",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde_json",
"tokio",
@ -4128,6 +4208,7 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-inference",
"litellm-inference-testing",
@ -4156,6 +4237,7 @@ dependencies = [
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-http",
@ -4163,7 +4245,7 @@ dependencies = [
"litellm-python-compat",
"litellm-router-types",
"litellm-secrets",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4196,6 +4278,7 @@ version = "0.1.0"
dependencies = [
"indexmap 2.14.0",
"jsonschema",
"litellm-error",
"litellm-llms-types",
"rstest",
"schemars 1.2.2",
@ -4222,6 +4305,7 @@ dependencies = [
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-python",
"litellm-http",
@ -4245,13 +4329,14 @@ dependencies = [
"pyo3",
"pyo3-async-runtimes",
"redis",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"serde_with",
"sqlx",
"strum",
"tempfile",
"thiserror 2.0.19",
"tokio",
"tokio-tungstenite",
@ -4275,6 +4360,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"
@ -4307,6 +4414,7 @@ dependencies = [
"google-cloud-auth",
"google-cloud-kms-v1",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-python-compat",
"litellm-secrets-aws",
@ -4316,7 +4424,7 @@ dependencies = [
"litellm-secrets-hashicorp",
"litellm-secrets-types",
"moka",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
@ -4358,13 +4466,14 @@ dependencies = [
"litellm-http",
"litellm-secrets-types",
"percent-encoding",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4381,13 +4490,14 @@ dependencies = [
"moka",
"percent-encoding",
"rcgen",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4408,12 +4518,13 @@ dependencies = [
"litellm-secrets-types",
"moka",
"percent-encoding",
"reqwest 0.12.28",
"reqwest",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4423,7 +4534,10 @@ name = "litellm-secrets-hashicorp"
version = "0.1.0"
dependencies = [
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-secrets-types",
"reqwest",
"rstest",
"rustify",
"rustify_derive",
@ -4457,6 +4571,7 @@ name = "litellm-spend-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-error",
"litellm-http",
"litellm-storage-clickhouse",
"rstest",
@ -4477,6 +4592,7 @@ name = "litellm-storage-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-error",
"litellm-http",
"rstest",
"serde",
@ -4496,7 +4612,8 @@ version = "0.1.0"
dependencies = [
"flate2",
"futures-util",
"reqwest 0.12.28",
"litellm-http",
"reqwest",
"rstest",
"semver",
"serde",
@ -4519,6 +4636,7 @@ dependencies = [
"criterion",
"indexmap 2.14.0",
"itoa",
"litellm-error",
"litellm-token-counter-fast",
"litellm-token-counter-huggingface",
"litellm-token-counter-tiktoken",
@ -4535,6 +4653,7 @@ name = "litellm-token-counter-fast"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-error",
"rand 0.8.7",
"rstest",
"rustc-hash",
@ -4560,6 +4679,7 @@ name = "litellm-token-counter-tiktoken"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-error",
"once_cell",
"rstest",
"rustc-hash",
@ -5750,22 +5870,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",
]
@ -5869,49 +6001,6 @@ version = "1.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
[[package]]
name = "reqwest"
version = "0.12.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-core",
"futures-util",
"h2 0.4.15",
"http 1.4.2",
"http-body 1.1.0",
"http-body-util",
"hyper 1.10.1",
"hyper-rustls 0.27.9",
"hyper-util",
"js-sys",
"log",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
"rustls 0.23.42",
"rustls-pki-types",
"serde",
"serde_json",
"serde_urlencoded",
"sync_wrapper",
"tokio",
"tokio-rustls 0.26.4",
"tokio-util",
"tower",
"tower-http",
"tower-service",
"url",
"wasm-bindgen",
"wasm-bindgen-futures",
"wasm-streams 0.4.2",
"web-sys",
"webpki-roots",
]
[[package]]
name = "reqwest"
version = "0.13.5"
@ -5933,6 +6022,7 @@ dependencies = [
"js-sys",
"log",
"mime",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -5952,7 +6042,7 @@ dependencies = [
"url",
"wasm-bindgen",
"wasm-bindgen-futures",
"wasm-streams 0.5.0",
"wasm-streams",
"web-sys",
]
@ -5999,7 +6089,7 @@ dependencies = [
"pin-project-lite",
"process-wrap",
"rand 0.10.2",
"reqwest 0.13.5",
"reqwest",
"schemars 1.2.2",
"serde",
"serde_json",
@ -6122,7 +6212,7 @@ dependencies = [
"async-trait",
"bytes",
"http 1.4.2",
"reqwest 0.13.5",
"reqwest",
"rustify_derive",
"serde",
"serde_json",
@ -7154,7 +7244,7 @@ dependencies = [
"memchr",
"parse-display",
"pin-project-lite",
"reqwest 0.13.5",
"reqwest",
"serde",
"serde_json",
"serde_with",
@ -7403,11 +7493,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9"
dependencies = [
"futures-util",
"log",
"rustls 0.23.42",
"rustls-native-certs",
"rustls-pki-types",
"tokio",
"tokio-rustls 0.26.4",
"tungstenite",
]
@ -7752,8 +7838,6 @@ dependencies = [
"httparse",
"log",
"rand 0.8.7",
"rustls 0.23.42",
"rustls-pki-types",
"sha1 0.10.7",
"thiserror 1.0.69",
"utf-8",
@ -7793,7 +7877,7 @@ dependencies = [
"futures",
"pin-project",
"rand 0.10.2",
"reqwest 0.13.5",
"reqwest",
"serde",
"serde_json",
"time",
@ -8007,7 +8091,7 @@ dependencies = [
"async-trait",
"derive_builder",
"http 1.4.2",
"reqwest 0.13.5",
"reqwest",
"rustify",
"rustify_derive",
"serde",
@ -8154,19 +8238,6 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "wasm-streams"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
dependencies = [
"futures-util",
"js-sys",
"wasm-bindgen",
"wasm-bindgen-futures",
"web-sys",
]
[[package]]
name = "wasm-streams"
version = "0.5.0"
@ -8202,18 +8273,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",
]

16
litellm-rust/Cargo.toml
View file

@ -10,6 +10,7 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-config = { path = "crates/config" }
litellm-error = { path = "crates/error" }
litellm-router = { path = "crates/router" }
litellm-router-types = { path = "crates/router-types" }
litellm-tracing = { path = "crates/tracing" }
@ -69,6 +70,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"
@ -87,23 +91,27 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
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"] }
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
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"
wiremock = "0.6.5"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
rustify = "=0.7.0"
rustify_derive = "=0.5.5"
vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] }
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"
@ -111,7 +119,7 @@ thiserror = "2.0"
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
tiktoken-rs = "0.12.0"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] }
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
tokio-tungstenite = { version = "0.24", default-features = false, features = ["handshake"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
flate2 = "1"
@ -124,7 +132,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"

3
litellm-rust/clippy.toml
View file

@ -10,8 +10,11 @@ disallowed-methods = [
{ path = "reqwest::Client::new", reason = "take litellm_http::Client from HttpClientPool" },
{ path = "reqwest::Client::builder", reason = "HttpClientConfig owns client construction" },
{ path = "reqwest::ClientBuilder::danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" },
{ path = "reqwest::ClientBuilder::tls_danger_accept_invalid_certs", reason = "set HttpClientConfig::verify instead" },
{ path = "reqwest::ClientBuilder::tls_danger_accept_invalid_hostnames", reason = "set HttpClientConfig::verify instead" },
{ path = "reqwest::ClientBuilder::identity", reason = "set HttpClientConfig::client_certificate instead" },
{ path = "reqwest::ClientBuilder::use_preconfigured_tls", reason = "HttpClientConfig owns the TLS configuration" },
{ path = "reqwest::ClientBuilder::tls_backend_preconfigured", reason = "HttpClientConfig owns the TLS configuration" },
{ path = "sqlx::query", reason = "use sqlx::query! or query_file! so the SQL is checked against the migrated schema" },
{ path = "sqlx::query_as", reason = "use sqlx::query_as! or query_file_as! so the SQL is checked against the migrated schema" },
{ path = "sqlx::query_scalar", reason = "use sqlx::query_scalar! so the SQL is checked against the migrated schema" },

5
litellm-rust/crates/auth-aws/Cargo.toml
View file

@ -7,7 +7,9 @@ repository.workspace = true
[dependencies]
litellm-auth-types.workspace = true
litellm-error.workspace = true
litellm-http.workspace = true
reqwest.workspace = true
moka = { workspace = true, features = ["sync"] }
serde.workspace = true
@ -26,5 +28,6 @@ aws-smithy-runtime-api = "1.13.0"
[dev-dependencies]
rstest.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
reqwest.workspace = true
bytes.workspace = true
futures-util.workspace = true
tokio.workspace = true

246
litellm-rust/crates/auth-aws/src/aws.rs
View file

@ -1,4 +1,3 @@
use std::collections::BTreeMap;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
@ -12,16 +11,17 @@ use aws_sigv4::http_request::{
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use litellm_auth_types::AwsParams;
use super::Error;
use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES,
};
use crate::error::{CredentialFlow, Error, SigningStage};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
@ -271,7 +271,7 @@ impl AwsAuthService {
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsProfile(error.to_string()))
.map_err(Error::credentials(CredentialFlow::Profile))
}
AwsAuthFlow::AssumeRole { role, session_name } => {
if is_already_running_as_role(&role, &resolved).await? {
@ -287,7 +287,7 @@ impl AwsAuthService {
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
self.set_cached_credentials(
key,
credentials.clone(),
@ -327,7 +327,7 @@ impl AwsAuthService {
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
.map_err(Error::credentials(CredentialFlow::AssumeRole))
}
AwsAuthFlow::WebIdentity {
token,
@ -350,12 +350,14 @@ impl AwsAuthService {
.web_identity_token(token)
.send()
.await
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::WebIdentity))?;
let credentials = response
.credentials()
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
.ok_or(Error::MissingWebIdentityCredentials)?;
let expiration =
SystemTime::try_from(*credentials.expiration()).map_err(|source| {
Error::WebIdentityExpiration(litellm_error::Source::new(source))
})?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
@ -376,7 +378,7 @@ impl AwsAuthService {
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
self.set_cached_credentials(
key,
credentials.clone(),
@ -435,12 +437,12 @@ fn default_session_name() -> String {
/// client header cannot change the canonical request and invalidate the
/// signature. Signing everything instead makes the request 403 on a header the
/// caller supplied, on a deployment that works on the Python path.
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
pub fn aws_signature_headers(headers: &HeaderMap) -> HeaderMap {
headers
.iter()
.filter(|(name, _)| {
let name = name.to_ascii_lowercase();
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
let name = name.as_str();
AWS_SIGNED_HEADER_NAMES.contains(&name)
|| name.starts_with("x-amz-")
|| name.starts_with("x-amzn-")
})
@ -456,45 +458,62 @@ pub fn is_sigv4_computed_header(name: &str) -> bool {
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
}
pub fn sign_post(
#[derive(Clone, Copy, Debug)]
pub struct SigningScope<'a> {
pub region: &'a str,
pub service: &'a str,
pub credentials: &'a Credentials,
pub time: SystemTime,
}
pub fn sign_headers(
method: &str,
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
let identity: Identity = credentials.clone().into();
headers: &HeaderMap,
scope: SigningScope<'_>,
) -> Result<HeaderMap, Error> {
let identity: Identity = scope.credentials.clone().into();
let params = v4::SigningParams::builder()
.identity(&identity)
.region(region)
.name(service)
.time(signing_time)
.region(scope.region)
.name(scope.service)
.time(scope.time)
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
.map_err(Error::signing(SigningStage::Parameters))?;
let header_refs = headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.as_str(), value))
.map_err(|_| Error::NonAsciiHeader {
name: name.to_string(),
})
})
.collect::<Result<Vec<_>, _>>()?;
let request = SignableRequest::new(
method,
url,
header_refs.into_iter(),
SignableBody::Bytes(body),
)
.map_err(Error::signing(SigningStage::Request))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| Error::AwsSigning(error.to_string()))?
.map_err(Error::signing(SigningStage::Signature))?
.into_parts();
Ok(instructions
instructions
.headers()
.map(|(name, value)| {
let normalized_name = match name {
"authorization" => "Authorization",
"x-amz-date" => "X-Amz-Date",
"x-amz-security-token" => "X-Amz-Security-Token",
_ => name,
};
(normalized_name.to_string(), value.to_string())
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(Error::signing(SigningStage::Output))?;
let value =
HeaderValue::from_str(value).map_err(Error::signing(SigningStage::Output))?;
Ok((name, value))
})
.collect())
.collect()
}
/// Model-id and region parsing shared by every Bedrock route.
@ -639,6 +658,7 @@ pub fn host_supplied_credentials(params: &AwsParams) -> Option<Credentials> {
mod tests {
use super::*;
use crate::constants::{AWS_REGION, BEDROCK_SERVICE};
use reqwest::header::CONTENT_TYPE;
fn no_env(_: &str) -> Option<String> {
None
@ -731,12 +751,12 @@ mod tests {
);
}
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
fn parity_inputs() -> (String, Vec<u8>, HeaderMap) {
(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
.to_string(),
br#"{"input":"hello"}"#.to_vec(),
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]),
)
}
@ -961,9 +981,9 @@ mod tests {
// cannot change the canonical request. Signing it instead makes the
// request 403 the moment anything on the wire rewrites or drops it.
let (url, body, mut headers) = parity_inputs();
headers.insert("x-request-id".to_string(), "abc-123".to_string());
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
headers.insert("x-request-id", HeaderValue::from_static("abc-123"));
headers.insert("Accept-Encoding", HeaderValue::from_static("gzip"));
headers.insert("x-amzn-trace-id", HeaderValue::from_static("Root=1-abc"));
let signable = aws_signature_headers(&headers);
assert!(!signable.contains_key("x-request-id"));
@ -979,18 +999,21 @@ mod tests {
None,
"test",
);
let signed = sign_post(
let signed = sign_headers(
"POST",
&url,
&body,
&signable,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
SystemTime::UNIX_EPOCH,
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: SystemTime::UNIX_EPOCH,
},
)
.expect("signs");
let authorization = signed
.get("Authorization")
let authorization = signed["Authorization"]
.to_str()
.expect("carries an authorization header");
assert!(
!authorization.contains("x-request-id"),
@ -1012,29 +1035,24 @@ mod tests {
None,
"test",
);
let signed = sign_post(
let signed = sign_headers(
"POST",
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
},
)
.expect("golden signature");
assert_eq!(signed["X-Amz-Date"], "20240102T030405Z");
assert_eq!(signed["X-Amz-Security-Token"], "session-token");
assert_eq!(
signed.get("X-Amz-Date").map(String::as_str),
Some("20240102T030405Z")
);
assert_eq!(
signed.get("X-Amz-Security-Token").map(String::as_str),
Some("session-token")
);
assert_eq!(
signed.get("Authorization").map(String::as_str),
Some(
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
)
signed["Authorization"],
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
);
}
@ -1048,19 +1066,71 @@ mod tests {
None,
"test",
);
let signed = sign_post(
let signed = sign_headers(
"POST",
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
},
)
.expect("signature");
assert!(!signed.contains_key("X-Amz-Security-Token"));
}
#[test]
fn a_header_value_outside_visible_ascii_is_named_instead_of_signed() {
let (url, body, mut headers) = parity_inputs();
headers.insert("x-amz-meta-note", HeaderValue::from_str("café").unwrap());
let credentials = Credentials::new("AKIDEXAMPLE", "secret", None, None, "test");
let error = sign_headers(
"POST",
&url,
&body,
&headers,
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH,
},
)
.unwrap_err();
assert_eq!(
error,
Error::NonAsciiHeader {
name: "x-amz-meta-note".into()
}
);
}
#[test]
fn the_shared_auth_error_nests_the_aws_failure_and_its_cause() {
let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "profile not found");
let error =
litellm_auth_types::Error::from(Error::credentials(CredentialFlow::Profile)(cause));
assert!(matches!(
litellm_error::find::<Error>(&error),
Some(Error::Credentials {
flow: CredentialFlow::Profile,
..
})
));
let report = litellm_error::Report::new(&error).to_string();
assert!(
report.contains("could not load AWS profile credentials"),
"{report}"
);
assert_eq!(report.matches("profile not found").count(), 1, "{report}");
}
#[ignore]
#[tokio::test]
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
@ -1068,7 +1138,7 @@ mod tests {
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
let headers =
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]);
let credentials = AwsAuthService::default()
.resolve_credentials(
AwsAuthConfig {
@ -1080,30 +1150,32 @@ mod tests {
&no_env,
)
.await?;
let client = litellm_http::Client::plain_for_test();
let client = litellm_http::Client::for_test();
let mut failures = Vec::new();
for region in ["us-west-2", "us-east-1"] {
let url = format!(
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
);
let signed_headers = sign_post(
let signed_headers = sign_headers(
"POST",
&url,
&body,
&headers,
region,
BEDROCK_SERVICE,
&credentials,
SystemTime::now(),
SigningScope {
region,
service: BEDROCK_SERVICE,
credentials: &credentials,
time: SystemTime::now(),
},
)?;
let mut request = client.post(&url).body(body.clone());
for (name, value) in &headers {
request = request.header(name, value);
}
for (name, value) in signed_headers {
request = request.header(name, value);
}
let response = request.send().await?;
let response = client
.post(&url)
.headers(headers.clone())
.headers(signed_headers)
.body(body.clone())
.send()
.await?;
let status = response.status();
let response_body = response.text().await?;
let snippet: String = response_body.chars().take(240).collect();

101
litellm-rust/crates/auth-aws/src/error.rs
View file

@ -1,46 +1,73 @@
use thiserror::Error as ThisError;
use litellm_error::Source;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
AwsProfile(String),
#[error("AWS default credentials failed: {0}")]
AwsDefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AwsAssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
AwsWebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
AwsWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
AwsSigningParameters(String),
#[error("AWS signable request failed: {0}")]
AwsSignableRequest(String),
#[error("AWS request signing failed: {0}")]
AwsSigning(String),
#[error("could not load AWS {flow} credentials")]
Credentials {
flow: CredentialFlow,
#[source]
source: Source,
},
#[error("AWS web identity response had no credentials")]
AwsMissingWebIdentityCredentials,
MissingWebIdentityCredentials,
#[error("AWS web identity response had an invalid expiration")]
WebIdentityExpiration(#[source] Source),
#[error("header {name} has a value that is not visible ASCII")]
NonAsciiHeader { name: String },
#[error("{stage}")]
Signing {
stage: SigningStage,
#[source]
source: Source,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum CredentialFlow {
#[error("profile")]
Profile,
#[error("default chain")]
DefaultChain,
#[error("assumed role")]
AssumeRole,
#[error("web identity")]
WebIdentity,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum SigningStage {
#[error("SigV4 signing parameters are invalid")]
Parameters,
#[error("SigV4 cannot sign this request")]
Request,
#[error("SigV4 could not compute the signature")]
Signature,
#[error("SigV4 produced an invalid header")]
Output,
}
impl Error {
pub(crate) fn credentials<E: std::error::Error + Send + Sync + 'static>(
flow: CredentialFlow,
) -> impl FnOnce(E) -> Self {
move |source| Self::Credentials {
flow,
source: Source::new(source),
}
}
pub(crate) fn signing<E: std::error::Error + Send + Sync + 'static>(
stage: SigningStage,
) -> impl FnOnce(E) -> Self {
move |source| Self::Signing {
stage,
source: Source::new(source),
}
}
}
impl From<Error> for litellm_auth_types::Error {
fn from(error: Error) -> Self {
Self::ProviderAuthentication(error.to_string())
}
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn converts_to_shared_auth_error_without_losing_context() {
let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into()));
assert_eq!(
error,
litellm_auth_types::Error::ProviderAuthentication(
"AWS profile credentials failed: profile not found".into()
)
);
Self::provider(error)
}
}

2
litellm-rust/crates/auth-aws/src/lib.rs
View file

@ -5,5 +5,5 @@ mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::Error;
pub use error::{CredentialFlow, Error, SigningStage};
pub use signer::SigV4Signer;

186
litellm-rust/crates/auth-aws/src/signer.rs
View file

@ -1,11 +1,13 @@
use std::{collections::BTreeMap, time::SystemTime};
use std::time::SystemTime;
use aws_credential_types::Credentials;
use reqwest::{Request, header::HeaderMap};
use crate::{
AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header,
sign_post,
AwsAuthService, AwsCredentialSource, Error, SigningScope, aws_signature_headers,
is_sigv4_computed_header, sign_headers,
};
use aws_credential_types::Credentials;
use litellm_http::outbound::{RequestSigner, UnsignedRequest};
use litellm_http::outbound::RequestSigner;
#[derive(Clone, Debug)]
pub struct SigV4Signer {
@ -45,41 +47,54 @@ impl SigV4Signer {
}
impl RequestSigner for SigV4Signer {
fn sign(
&self,
request: UnsignedRequest<'_>,
) -> Result<Vec<(String, String)>, litellm_http::Error> {
if let Some((name, _)) = request
.headers
.iter()
.find(|(name, _)| is_sigv4_computed_header(name))
fn sign(&self, request: &Request) -> Result<HeaderMap, litellm_http::Error> {
if let Some(name) = request
.headers()
.keys()
.find(|name| is_sigv4_computed_header(name.as_str()))
{
return Err(litellm_http::Error::ComputedHeader(name.clone()));
return Err(litellm_http::Error::computed_header(name.as_str()));
}
let headers: BTreeMap<String, String> = request.headers.iter().cloned().collect();
sign_post(
request.url,
request.body,
&aws_signature_headers(&headers),
&self.region,
self.service,
&self.credentials,
(self.clock)(),
let body = match request.body() {
None => &[],
Some(body) => body
.as_bytes()
.ok_or_else(litellm_http::Error::unsignable_body)?,
};
sign_headers(
request.method().as_str(),
request.url().as_str(),
body,
&aws_signature_headers(request.headers()),
SigningScope {
region: &self.region,
service: self.service,
credentials: &self.credentials,
time: (self.clock)(),
},
)
.map(|signature| signature.into_iter().collect())
.map_err(|error| litellm_http::Error::Signature(error.to_string()))
.map_err(|error| match error {
Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
Error::Credentials { .. }
| Error::MissingWebIdentityCredentials
| Error::WebIdentityExpiration(_)
| Error::Signing { .. } => litellm_http::Error::signature(error),
})
}
}
#[cfg(test)]
mod tests {
use reqwest::{Body, Method, Url};
use serde_json::{Value, json as json_value};
use std::time::{Duration, UNIX_EPOCH};
use litellm_http::outbound::OutboundRequest;
use serde_json::{Value, json};
use litellm_http::outbound::{json, signed};
use super::*;
const URL: &str = "https://textract.us-east-1.amazonaws.com/";
fn fixed_clock() -> SystemTime {
UNIX_EPOCH + Duration::from_secs(1_700_000_000)
}
@ -93,82 +108,121 @@ mod tests {
.with_clock(fixed_clock)
}
fn target() -> (String, String) {
("X-Amz-Target".into(), "Textract.DetectDocumentText".into())
}
fn authorization(body: &Value, service: &'static str) -> String {
OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
body,
None,
signed(
json(URL, &[target()], body, None).unwrap(),
&signer(service),
)
.unwrap()
.header("Authorization")
.unwrap()
.to_string()
.headers()["authorization"]
.to_str()
.unwrap()
.to_string()
}
#[test]
fn the_signature_verifies_against_the_bytes_that_are_sent() {
let sent = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
&json!({"Document": {"Bytes": "aGk="}}),
None,
let sent = signed(
json(
URL,
&[target()],
&json_value!({"Document": {"Bytes": "aGk="}}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap();
let unsigned: BTreeMap<String, String> = sent
let unsigned: HeaderMap = sent
.headers()
.iter()
.filter(|(name, _)| !is_sigv4_computed_header(name))
.cloned()
.filter(|(name, _)| !is_sigv4_computed_header(name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect();
let recomputed = sign_post(
sent.url(),
sent.body(),
let recomputed = sign_headers(
"POST",
sent.url().as_str(),
sent.body().unwrap().as_bytes().unwrap(),
&aws_signature_headers(&unsigned),
"us-east-1",
"textract",
&Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
fixed_clock(),
SigningScope {
region: "us-east-1",
service: "textract",
credentials: &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
time: fixed_clock(),
},
)
.unwrap();
assert_eq!(
sent.header("Authorization"),
Some(recomputed["Authorization"].as_str())
);
assert_eq!(sent.headers()["authorization"], recomputed["authorization"]);
}
#[test]
fn the_signature_depends_on_the_body_and_the_service() {
let original = authorization(&json!({"text": "card 4111"}), "textract");
let original = authorization(&json_value!({"text": "card 4111"}), "textract");
assert_ne!(
original,
authorization(&json!({"text": "card [REDACTED]"}), "textract")
authorization(&json_value!({"text": "card [REDACTED]"}), "textract")
);
assert_ne!(
original,
authorization(&json!({"text": "card 4111"}), "bedrock")
authorization(&json_value!({"text": "card 4111"}), "bedrock")
);
assert!(original.contains("/us-east-1/textract/aws4_request"));
}
#[test]
fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() {
let error = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("authorization".into(), "Bearer caller".into())],
&json!({}),
None,
let error = signed(
json(
URL,
&[("authorization".into(), "Bearer caller".into())],
&json_value!({}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap_err();
assert_eq!(
error,
litellm_http::Error::ComputedHeader("authorization".into())
);
assert!(error.is_computed_header());
assert!(error.to_string().contains("authorization"), "{error}");
}
#[test]
fn a_forwarded_header_that_cannot_be_signed_is_the_callers_fault() {
let error = signed(
json(
URL,
&[("x-amz-meta-note".into(), "café".into())],
&json_value!({}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap_err();
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains("x-amz-meta-note"), "{report}");
assert!(error.is_request());
assert!(!error.dispatched());
}
#[test]
fn a_streaming_body_cannot_be_signed() {
let mut request = Request::new(Method::POST, Url::parse(URL).unwrap());
*request.body_mut() = Some(Body::wrap_stream(futures_util::stream::once(async {
Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{}"))
})));
assert!(matches!(
signed(request, &signer("textract")),
Err(error) if !error.dispatched() && !error.is_request()
));
}
}

1
litellm-rust/crates/auth-azure/Cargo.toml
View file

@ -18,5 +18,6 @@ azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
tokio.workspace = true

43
litellm-rust/crates/auth-azure/src/native.rs
View file

@ -12,7 +12,7 @@ use azure_identity::{
};
use sha2::{Digest, Sha256};
use litellm_auth_types::Error;
use litellm_auth_types::{Error, ErrorDetail};
use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
use super::credential_provider_cache::{
@ -133,7 +133,9 @@ impl NativeAzureTokenAcquirer {
let token = credential
.get_token(&[scope.as_str()], None)
.await
.map_err(|error| Error::CredentialAcquisition(error.to_string().into()))?;
.map_err(|error| {
Error::CredentialAcquisition(ErrorDetail::failed("Azure token request", error))
})?;
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
.ok()
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
@ -250,11 +252,8 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
let Some(authority) = authority else {
return Ok(());
};
let url = url::Url::parse(authority.value()).map_err(|_| {
Error::InvalidConfiguration(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into(),
)
let url = url::Url::parse(authority.value()).map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed("Azure authority parsing", error))
})?;
if url.scheme() != "https"
|| url.host_str().is_none()
@ -264,10 +263,10 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
|| url.fragment().is_some()
|| !matches!(url.path(), "" | "/")
{
return Err(Error::InvalidConfiguration(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into(),
));
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
field: "Azure authority".into(),
expected: "an HTTPS origin without credentials, query, or fragment",
}));
}
Ok(())
}
@ -377,7 +376,7 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
fn mixed_sources<T>() -> Result<T, Error> {
Err(Error::InvalidConfiguration(
"request-controlled Azure auth inputs cannot be combined with host credentials".into(),
ErrorDetail::MixedInputSources { provider: "Azure" },
))
}
@ -444,7 +443,7 @@ fn build_credential(
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed(
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential initialization",
error,
))
@ -662,12 +661,11 @@ mod tests {
))
.unwrap_err();
assert_eq!(
error,
litellm_auth_types::Error::InvalidConfiguration(
assert!(
litellm_error::Report::new(&error).to_string().contains(
"request-controlled Azure auth inputs cannot be combined with host credentials"
.into()
)
),
"{error:?}"
);
}
@ -695,12 +693,11 @@ mod tests {
authority,
))
.unwrap_err();
assert_eq!(
error,
litellm_auth_types::Error::InvalidConfiguration(
assert!(
litellm_error::Report::new(&error).to_string().contains(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into()
)
),
"{error:?}"
);
}
}

119
litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -1,8 +1,8 @@
use litellm_auth_types::Error;
use litellm_auth_types::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
};
use litellm_auth_types::{Error, ErrorDetail};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@ -91,9 +91,9 @@ impl AzureAuthService {
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(Error::EmptyCallerCredential(
"Azure AD token provider returned an empty token",
));
return Err(Error::EmptyCallerCredential {
caller: "Azure AD token provider",
});
}
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
}
@ -106,11 +106,9 @@ impl AzureAuthService {
} => {
let assertion = resolve_reference(inputs, env_lookup, reference.value())
.await?
.ok_or_else(|| {
Error::CredentialAcquisition(
"Azure OIDC reference did not resolve to a value".into(),
)
})?;
.ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "Azure OIDC reference",
}))?;
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
@ -142,7 +140,13 @@ impl AzureAuthService {
Err(error) => failures.push(error),
}
}
Err(Error::CredentialChain(failures))
match failures.pop() {
Some(last) => Err(Error::CredentialChain {
last: Box::new(last),
earlier: failures,
}),
None => Ok(None),
}
}
AzureCredentialPlan::Missing => Ok(None),
}
@ -173,7 +177,12 @@ pub(crate) fn select_auth_plan(
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|_| Error::InvalidConfiguration("invalid Azure credential selector".into()))?;
.map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential selector parsing",
error,
))
})?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
@ -263,9 +272,12 @@ fn select_native_plan(
let selection_source = selected.source();
match selected.into_value() {
AzureCredentialType::ClientSecretCredential => Err(Error::InvalidConfiguration(
"ClientSecretCredential requires tenant_id, client_id, and client_secret".into(),
)),
AzureCredentialType::ClientSecretCredential => {
Err(Error::InvalidConfiguration(ErrorDetail::MissingSettings {
credential: "ClientSecretCredential",
settings: "tenant_id, client_id, and client_secret",
}))
}
AzureCredentialType::WorkloadIdentityCredential => {
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
@ -349,22 +361,21 @@ fn workload_request(
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, Error> {
Ok(NativeAzureRequest::WorkloadIdentity {
tenant_id: tenant_id.ok_or_else(|| {
Error::InvalidConfiguration("WorkloadIdentityCredential requires tenant_id".into())
})?,
client_id: client_id.ok_or_else(|| {
Error::InvalidConfiguration("WorkloadIdentityCredential requires client_id".into())
})?,
token_file_path: token_file_path.ok_or_else(|| {
Error::InvalidConfiguration(
"WorkloadIdentityCredential requires azure_federated_token_file".into(),
)
})?,
tenant_id: tenant_id.ok_or(workload_requires("tenant_id"))?,
client_id: client_id.ok_or(workload_requires("client_id"))?,
token_file_path: token_file_path.ok_or(workload_requires("azure_federated_token_file"))?,
scope,
authority,
})
}
fn workload_requires(settings: &'static str) -> Error {
Error::InvalidConfiguration(ErrorDetail::MissingSettings {
credential: "WorkloadIdentityCredential",
settings,
})
}
fn configured_string(
configured: &ConfigValue<String>,
environment_name: &str,
@ -410,11 +421,10 @@ async fn resolve_reference(
.map_or(CredentialLookup::Missing, CredentialLookup::Found),
CredentialRef::None => return Ok(None),
CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => {
let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| {
Error::InvalidConfiguration(
"credential reference requires a host credential resolver".into(),
)
})?;
let resolver = inputs
.credential_resolver
.as_ref()
.ok_or(Error::InvalidConfiguration(ErrorDetail::MissingResolver))?;
resolver.resolve(reference).await?
}
};
@ -433,7 +443,7 @@ fn oidc_reference(
let value = token.value().expose();
if token.source() == InputSource::Request && value.starts_with("oidc/") {
return Err(Error::InvalidConfiguration(
"request-controlled Azure credential references are not allowed".into(),
ErrorDetail::RequestControlledReference { provider: "Azure" },
));
}
if let Some(name) = value.strip_prefix("oidc/env/") {
@ -457,19 +467,17 @@ fn oidc_reference(
}
if value.starts_with("oidc/") {
return Err(Error::InvalidConfiguration(
"unsupported OIDC reference".into(),
ErrorDetail::UnsupportedReference,
));
}
Ok(None)
}
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
fn non_empty_reference(value: &str, kind: &'static str) -> Result<String, Error> {
if value.is_empty() {
return Err(Error::InvalidConfiguration(
litellm_auth_types::ErrorDetail::Empty {
subject: kind.into(),
},
));
return Err(Error::InvalidConfiguration(ErrorDetail::Empty {
subject: kind,
}));
}
Ok(value.to_string())
}
@ -518,7 +526,9 @@ mod tests {
expires_on: None,
})
} else {
Err(Error::CredentialAcquisition(kind.into()))
Err(Error::CredentialAcquisition(
litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
))
}
})
}
@ -624,7 +634,11 @@ mod tests {
)))
.expect_err("unsupported backend must fail validation");
assert!(error.to_string().contains("unsupported OIDC reference"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("unsupported OIDC reference")
);
}
#[rstest::rstest]
@ -649,11 +663,11 @@ mod tests {
})
.unwrap_err();
assert_eq!(
error,
Error::InvalidConfiguration(
"request-controlled Azure credential references are not allowed".into()
)
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("request-controlled Azure credential references are not allowed"),
"{error:?}"
);
}
@ -715,7 +729,16 @@ mod tests {
.await
.unwrap_err();
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
assert!(matches!(&error, Error::CredentialChain { earlier, .. } if earlier.len() == 1));
let report = litellm_error::Report::new(&error).to_string();
assert!(
report.contains("managed-identity did not resolve"),
"{report}"
);
assert!(
report.contains("developer-tools did not resolve"),
"{report}"
);
}
#[derive(Debug)]
@ -763,7 +786,9 @@ mod tests {
assert_eq!(
error,
Error::EmptyCallerCredential("Azure AD token provider returned an empty token")
Error::EmptyCallerCredential {
caller: "Azure AD token provider"
}
);
}
}

1
litellm-rust/crates/auth-gcp/Cargo.toml
View file

@ -24,5 +24,6 @@ google-cloud-auth = { workspace = true, optional = true }
http = { workspace = true, optional = true }
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
tempfile.workspace = true

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

@ -73,13 +73,17 @@ impl VertexAuth {
.or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV));
let project_id = get_vertex_ai_project(config, env_lookup);
if !has_authorization && static_token.is_none() {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
let static_token = match (has_authorization, static_token) {
(true, _) => None,
(false, Some(token)) => Some(token),
(false, None) => {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
};
let project_id = match project_id {
Some(project_id) => project_id,
@ -90,14 +94,9 @@ impl VertexAuth {
.await?
}
};
let headers = if has_authorization {
headers
} else {
apply_credential(
headers,
static_token.as_deref().expect("static token was checked"),
CredentialPlacement::Bearer,
)?
let headers = match static_token {
Some(token) => apply_credential(headers, &token, CredentialPlacement::Bearer)?,
None => headers,
};
Ok(VertexEnvironment {
headers,
@ -234,14 +233,13 @@ enum CredentialCacheKey {
Adc,
}
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
Set it to a readable file path, or to the credentials JSON itself.";
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
BASE64_CREDENTIALS.into(),
)));
return Err(Error::InvalidConfiguration(
ErrorDetail::Base64EncodedJson {
field: "vertex_credentials",
},
));
}
if names_a_credentials_file(configured) {
CustomServiceAccount::from_file(configured)
@ -299,7 +297,8 @@ mod tests {
else {
panic!("{configured:?} should not load");
};
assert!(error.to_string().contains(expected), "{error}");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains(expected), "{report}");
}
#[rstest::rstest]
@ -321,7 +320,7 @@ mod tests {
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => error.to_string(),
Err(error) => litellm_error::Report::new(&error).to_string(),
}
}

6
litellm-rust/crates/auth-gcp/src/sdk.rs
View file

@ -55,7 +55,7 @@ impl CredentialsProvider for GoogleCredentials {
entity_tag: EntityTag::new(),
data,
})
.map_err(|_| CredentialsError::from_msg(false, "Google authentication failed"))
.map_err(|error| CredentialsError::from_source(false, error))
}
async fn universe_domain(&self) -> Option<String> {
@ -102,5 +102,9 @@ mod tests {
);
let error = credentials.headers(Extensions::new()).await.unwrap_err();
assert!(!format!("{error:?}").contains("private"));
assert_eq!(
std::error::Error::source(&error).and_then(|cause| cause.downcast_ref::<Error>()),
Some(&Error::InvalidHeader)
);
}
}

1
litellm-rust/crates/auth-types/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-error.workspace = true
serde.workspace = true
serde_json.workspace = true
subtle.workspace = true

26
litellm-rust/crates/auth-types/src/credential.rs
View file

@ -86,9 +86,9 @@ impl CredentialPlan {
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(Error::EmptyCallerCredential(
"credential caller returned an empty credential",
));
return Err(Error::EmptyCallerCredential {
caller: "credential caller",
});
}
Ok(CredentialPlanResolution::Resolved(credential))
}
@ -104,8 +104,8 @@ mod tests {
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::Error;
use crate::SecretValue;
use crate::{Error, ErrorDetail};
#[derive(Debug)]
struct HostResolver;
@ -150,9 +150,10 @@ mod tests {
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async {
Err(Error::CredentialAcquisition(
"host credential lookup failed".into(),
))
Err(Error::CredentialAcquisition(ErrorDetail::failed(
"host credential lookup",
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
)))
})
}
}
@ -168,9 +169,14 @@ mod tests {
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(
error,
Error::CredentialAcquisition("host credential lookup failed".into())
assert!(
matches!(error, Error::CredentialAcquisition(_)),
"{error:?}"
);
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("host credential lookup failed")
);
}
}

326
litellm-rust/crates/auth-types/src/error.rs
View file

@ -1,17 +1,23 @@
use thiserror::Error as ThisError;
use litellm_error::{Report, Source};
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
use crate::{ParamKind, ParamSpec};
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("invalid authentication configuration: {0}")]
#[error("invalid authentication configuration")]
InvalidConfiguration(#[source] ErrorDetail),
#[error("credential acquisition failed: {0}")]
#[error("credential acquisition failed")]
CredentialAcquisition(#[source] ErrorDetail),
#[error("credential caller failed: {0}")]
EmptyCallerCredential(&'static str),
#[error("{0}")]
ProviderAuthentication(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<Error>),
#[error("{caller} returned an empty credential")]
EmptyCallerCredential { caller: &'static str },
#[error("provider authentication failed")]
Provider(#[source] Source),
#[error("no credential in the chain succeeded{}", earlier_attempts(earlier))]
CredentialChain {
#[source]
last: Box<Error>,
earlier: Vec<Error>,
},
#[error(
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
)]
@ -22,7 +28,7 @@ pub enum Error {
#[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())]
MissingParam {
provider: &'static str,
spec: &'static crate::ParamSpec,
spec: &'static ParamSpec,
},
#[error("Missing {provider} API Base - {guidance}")]
MissingApiBase {
@ -33,78 +39,282 @@ pub enum Error {
InvalidHeader,
}
fn earlier_attempts(earlier: &[Error]) -> String {
if earlier.is_empty() {
return String::new();
}
let attempts = earlier
.iter()
.map(|failure| Report::new(failure).to_string())
.collect::<Vec<_>>()
.join("; ");
format!(" (earlier attempts: {attempts})")
}
impl Error {
pub fn provider(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
Self::Provider(Source::new(cause))
}
pub fn is_missing_credential(&self) -> bool {
match self {
Self::MissingApiKey { .. } => true,
Self::MissingParam { spec, .. } => match spec.kind {
ParamKind::Credential => true,
ParamKind::Setting => false,
},
Self::InvalidConfiguration(_)
| Self::CredentialAcquisition(_)
| Self::EmptyCallerCredential { .. }
| Self::Provider(_)
| Self::CredentialChain { .. }
| Self::MissingApiBase { .. }
| Self::InvalidHeader => false,
}
}
pub fn is_request(&self) -> bool {
match self {
Self::InvalidConfiguration(_)
| Self::EmptyCallerCredential { .. }
| Self::InvalidHeader => true,
Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain { .. } => {
false
}
Self::MissingParam { spec, .. } => match spec.kind {
ParamKind::Setting => true,
ParamKind::Credential => false,
},
Self::MissingApiKey { .. } | Self::MissingApiBase { .. } => false,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ErrorDetail {
#[error("{0}")]
Message(String),
#[error("{field} must be {expected}")]
InvalidType {
field: String,
expected: &'static str,
},
#[error("{subject} cannot be empty")]
Empty { subject: String },
Empty { subject: &'static str },
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{operation} failed: {source}")]
#[error("credential plan is not allowed by the provider auth policy")]
PlanNotAllowed,
#[error("{credential} requires {settings}")]
MissingSettings {
credential: &'static str,
settings: &'static str,
},
#[error("credentials cannot be sent to a request-controlled {provider} endpoint")]
RequestControlledEndpoint { provider: &'static str },
#[error("request-controlled {provider} credential references are not allowed")]
RequestControlledReference { provider: &'static str },
#[error("request-controlled {provider} auth inputs cannot be combined with host credentials")]
MixedInputSources { provider: &'static str },
#[error("unsupported OIDC reference")]
UnsupportedReference,
#[error("credential reference requires a host credential resolver")]
MissingResolver,
#[error("{subject} did not resolve to a value")]
Unresolved { subject: &'static str },
#[error(
"The `{field}` value is base64-encoded JSON, not the JSON itself. Set it to a readable file path, or to the credentials JSON itself."
)]
Base64EncodedJson { field: &'static str },
#[error("{operation} failed")]
Failed {
operation: &'static str,
#[source]
source: ErrorSource,
cause: Source,
},
}
impl ErrorDetail {
pub fn failed(
operation: &'static str,
source: impl std::error::Error + Send + Sync + 'static,
cause: impl std::error::Error + Send + Sync + 'static,
) -> Self {
Self::Failed {
operation,
source: ErrorSource::new(source),
cause: Source::new(cause),
}
}
}
impl From<String> for ErrorDetail {
fn from(message: String) -> Self {
Self::Message(message)
#[cfg(test)]
mod tests {
use std::io;
use litellm_error::{Report, Source};
use rstest::rstest;
use super::{Error, ErrorDetail};
use crate::{AwsParams, VertexParams};
#[rstest]
#[case::api_key(
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
)]
#[case::another_api_key(
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
)]
#[case::api_base(
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
"Missing Example API Base - Pass api_base"
)]
#[case::another_api_base(
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
)]
fn missing_inputs_keep_the_python_text(#[case] error: Error, #[case] expected: &str) {
assert_eq!(error.to_string(), expected);
}
#[rstest]
#[case::configuration(
Error::InvalidConfiguration(ErrorDetail::Empty { subject: "credential selector" }),
"invalid authentication configuration: credential selector cannot be empty"
)]
#[case::acquisition(
Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }),
"credential acquisition failed: OIDC reference did not resolve to a value"
)]
#[case::empty_caller(
Error::EmptyCallerCredential { caller: "token provider" },
"token provider returned an empty credential"
)]
fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) {
assert!(
Report::new(&error).to_string().contains(expected),
"{}",
Report::new(&error)
);
}
#[rstest]
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
"test credential",
denied()
)))]
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
"test credential",
denied()
)))]
fn a_failure_keeps_its_cause_as_the_source_and_renders_it_once(#[case] error: Error) {
let cause =
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
.find_map(|cause| cause.downcast_ref::<Source>()?.downcast_ref::<io::Error>())
.expect("the io cause is in the source chain");
assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied);
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
}
#[rstest]
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
"test credential",
denied()
)))]
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
"test credential",
denied()
)))]
fn a_contextual_failure_leaves_its_cause_out_of_its_own_message(#[case] error: Error) {
assert!(!error.to_string().contains("access denied"), "{error}");
assert!(
Report::new(&error)
.to_string()
.contains("test credential failed: access denied")
);
}
#[test]
fn a_provider_failure_keeps_the_provider_error_as_its_source() {
let error = Error::provider(denied());
assert_eq!(
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
Some(io::ErrorKind::PermissionDenied)
);
assert!(!error.to_string().contains("access denied"), "{error}");
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
}
#[test]
fn a_chain_keeps_its_last_failure_as_the_source_and_renders_every_attempt_once() {
let error = Error::CredentialChain {
last: Box::new(Error::provider(denied())),
earlier: vec![
Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "OIDC reference",
}),
Error::EmptyCallerCredential {
caller: "token provider",
},
],
};
assert_eq!(
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
Some(io::ErrorKind::PermissionDenied)
);
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
assert_eq!(
report.matches("OIDC reference did not resolve").count(),
1,
"{report}"
);
assert_eq!(
report
.matches("token provider returned an empty credential")
.count(),
1,
"{report}"
);
}
fn denied() -> io::Error {
io::Error::new(io::ErrorKind::PermissionDenied, "access denied")
}
#[rstest]
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, true)]
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), true)]
#[case::aws_secret_access_key(AwsParams::SECRET_ACCESS_KEY.missing("AWS"), true)]
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), false)]
#[case::aws_region(AwsParams::REGION.missing("AWS"), false)]
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)]
fn only_an_absent_credential_is_a_missing_credential(
#[case] error: Error,
#[case] expected: bool,
) {
assert_eq!(error.is_missing_credential(), expected);
}
#[rstest]
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), true)]
#[case::aws_region(AwsParams::REGION.missing("AWS"), true)]
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), false)]
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, false)]
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), true)]
#[case::invalid_header(Error::InvalidHeader, true)]
#[case::identity_provider_outage(
Error::CredentialAcquisition(ErrorDetail::MissingResolver),
false
)]
#[case::provider_sdk_failure(Error::provider(denied()), false)]
fn a_setting_the_caller_can_pass_is_theirs_to_fix_but_credentials_and_api_bases_are_not(
#[case] error: Error,
#[case] expected: bool,
) {
assert_eq!(error.is_request(), expected);
}
}
impl From<&str> for ErrorDetail {
fn from(message: &str) -> Self {
Self::Message(message.into())
}
}
#[derive(Clone, Debug)]
pub struct ErrorSource(std::sync::Arc<dyn std::error::Error + Send + Sync>);
impl std::ops::Deref for ErrorSource {
type Target = dyn std::error::Error + Send + Sync;
fn deref(&self) -> &Self::Target {
self.0.as_ref()
}
}
impl std::fmt::Display for ErrorSource {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.0, formatter)
}
}
impl ErrorSource {
pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self {
Self(std::sync::Arc::new(error))
}
}
impl PartialEq for ErrorSource {
fn eq(&self, other: &Self) -> bool {
std::sync::Arc::ptr_eq(&self.0, &other.0)
}
}
impl Eq for ErrorSource {}

12
litellm-rust/crates/auth-types/src/http.rs
View file

@ -21,9 +21,9 @@ pub fn apply_credential(
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, Error> {
if credential.trim().is_empty() {
return Err(Error::InvalidConfiguration(
"credential cannot be empty".into(),
));
return Err(Error::InvalidConfiguration(crate::ErrorDetail::Empty {
subject: "credential",
}));
}
if headers
.iter()
@ -71,6 +71,10 @@ mod tests {
)
.expect_err("provider policy must handle existing credentials");
assert!(error.to_string().contains("already exists"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("already exists")
);
}
}

4
litellm-rust/crates/auth-types/src/lib.rs
View file

@ -51,9 +51,9 @@ pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
};
pub use error::{Error, ErrorDetail, ErrorSource};
pub use error::{Error, ErrorDetail};
pub use http::CredentialPlacement;
pub use params::{AwsParams, ParamSpec, VertexParams};
pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

49
litellm-rust/crates/auth-types/src/params/aws.rs
View file

@ -4,15 +4,17 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use veil::Redact;
use super::ParamSpec;
use super::{ParamKind, ParamSpec};
const fn spec(
setting: &'static str,
kind: ParamKind,
wire: &'static [&'static str],
env: &'static [&'static str],
) -> ParamSpec {
ParamSpec {
setting,
kind,
wire,
module_global: None,
env,
@ -53,40 +55,67 @@ pub struct AwsParams {
impl AwsParams {
pub const ACCESS_KEY_ID: ParamSpec = spec(
"access key id",
ParamKind::Credential,
&["aws_access_key_id"],
&["AWS_ACCESS_KEY_ID"],
);
pub const SECRET_ACCESS_KEY: ParamSpec = spec(
"secret access key",
ParamKind::Credential,
&["aws_secret_access_key"],
&["AWS_SECRET_ACCESS_KEY"],
);
pub const SESSION_TOKEN: ParamSpec = spec(
"session token",
ParamKind::Credential,
&["aws_session_token"],
&["AWS_SESSION_TOKEN"],
);
pub const REGION: ParamSpec = spec(
"region",
ParamKind::Setting,
&["aws_region_name"],
&["AWS_REGION_NAME", "AWS_REGION"],
);
pub const SESSION_NAME: ParamSpec =
spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]);
pub const PROFILE_NAME: ParamSpec =
spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]);
pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]);
pub const SESSION_NAME: ParamSpec = spec(
"session name",
ParamKind::Setting,
&["aws_session_name"],
&["AWS_SESSION_NAME"],
);
pub const PROFILE_NAME: ParamSpec = spec(
"profile name",
ParamKind::Setting,
&["aws_profile_name"],
&["AWS_PROFILE_NAME"],
);
pub const ROLE_NAME: ParamSpec = spec(
"role name",
ParamKind::Setting,
&["aws_role_name"],
&["AWS_ROLE_NAME"],
);
pub const WEB_IDENTITY_TOKEN: ParamSpec = spec(
"web identity token",
ParamKind::Credential,
&["aws_web_identity_token"],
&["AWS_WEB_IDENTITY_TOKEN"],
);
pub const STS_ENDPOINT: ParamSpec =
spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]);
pub const EXTERNAL_ID: ParamSpec =
spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]);
pub const STS_ENDPOINT: ParamSpec = spec(
"STS endpoint",
ParamKind::Setting,
&["aws_sts_endpoint"],
&["AWS_STS_ENDPOINT"],
);
pub const EXTERNAL_ID: ParamSpec = spec(
"external id",
ParamKind::Setting,
&["aws_external_id"],
&["AWS_EXTERNAL_ID"],
);
pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec(
"Bedrock runtime endpoint",
ParamKind::Setting,
&["aws_bedrock_runtime_endpoint"],
&["AWS_BEDROCK_RUNTIME_ENDPOINT"],
);

7
litellm-rust/crates/auth-types/src/params/mod.rs
View file

@ -10,11 +10,18 @@ pub use vertex::VertexParams;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ParamSpec {
pub setting: &'static str,
pub kind: ParamKind,
pub wire: &'static [&'static str],
pub module_global: Option<&'static str>,
pub env: &'static [&'static str],
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ParamKind {
Credential,
Setting,
}
impl ParamSpec {
/// The first non-blank value: each wire name through `params`, then each environment
/// name through `env`.

5
litellm-rust/crates/auth-types/src/params/vertex.rs
View file

@ -2,7 +2,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::Unexpected};
use serde_json::Value;
use veil::Redact;
use super::ParamSpec;
use super::{ParamKind, ParamSpec};
/// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the
/// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read.
@ -27,18 +27,21 @@ pub struct VertexParams {
impl VertexParams {
pub const CREDENTIALS: ParamSpec = ParamSpec {
setting: "credentials",
kind: ParamKind::Credential,
wire: &["vertex_credentials", "vertex_ai_credentials"],
module_global: None,
env: &["VERTEXAI_CREDENTIALS"],
};
pub const PROJECT: ParamSpec = ParamSpec {
setting: "project",
kind: ParamKind::Setting,
wire: &["vertex_project", "vertex_ai_project"],
module_global: Some("vertex_project"),
env: &["VERTEXAI_PROJECT"],
};
pub const LOCATION: ParamSpec = ParamSpec {
setting: "location",
kind: ParamKind::Setting,
wire: &["vertex_location", "vertex_ai_location"],
module_global: Some("vertex_location"),
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],

35
litellm-rust/crates/auth-types/src/policy.rs
View file

@ -1,5 +1,5 @@
use crate::http::apply_credential;
use crate::{CredentialPlacement, Error, ResolvedCredential};
use crate::{CredentialPlacement, Error, ErrorDetail, ResolvedCredential};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlanKind {
@ -31,11 +31,18 @@ pub struct ProviderAuthPolicy {
impl ProviderAuthPolicy {
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, _)| {
self.accepted_existing_headers
.iter()
.any(|accepted| name.eq_ignore_ascii_case(accepted))
})
self.existing_credential_header(headers).is_some()
}
fn existing_credential_header(&self, headers: &[(String, String)]) -> Option<&'static str> {
self.accepted_existing_headers
.iter()
.copied()
.find(|accepted| {
headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(accepted))
})
}
pub fn apply(
@ -44,11 +51,11 @@ impl ProviderAuthPolicy {
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, Error> {
if self.has_existing_credential(&headers) {
if let Some(existing) = self.existing_credential_header(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(Error::InvalidConfiguration(
"credential header already exists".into(),
ErrorDetail::DuplicateHeader(existing),
)),
};
}
@ -56,11 +63,7 @@ impl ProviderAuthPolicy {
.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or_else(|| {
Error::InvalidConfiguration(
"credential plan is not allowed by the provider auth policy".into(),
)
})?;
.ok_or(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed))?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
@ -108,6 +111,10 @@ mod tests {
)
.unwrap_err();
assert!(error.to_string().contains("not allowed"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("not allowed")
);
}
}

74
litellm-rust/crates/auth-types/tests/error.rs
View file

@ -1,74 +0,0 @@
use litellm_auth_types::Error;
use rstest::rstest;
#[rstest]
#[case::api_key(
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
)]
#[case::another_api_key(
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
)]
#[case::api_base(
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
"Missing Example API Base - Pass api_base"
)]
#[case::another_api_base(
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
)]
#[case::configuration(
Error::InvalidConfiguration("credential selector is invalid".into()),
"invalid authentication configuration: credential selector is invalid"
)]
#[case::acquisition(
Error::CredentialAcquisition("token expired".into()),
"credential acquisition failed: token expired"
)]
#[case::caller(
Error::EmptyCallerCredential("empty token"),
"credential caller failed: empty token"
)]
#[case::provider(
Error::ProviderAuthentication("provider rejected credentials".into()),
"provider rejected credentials"
)]
#[case::chain(
Error::CredentialChain(vec![
Error::CredentialAcquisition("token expired".into()),
Error::EmptyCallerCredential("empty token"),
]),
"credential acquisition failed: credential acquisition failed: token expired; credential caller failed: empty token"
)]
fn display_preserves_failure_phase_and_caller_context(
#[case] error: Error,
#[case] expected: &str,
) {
assert_eq!(error.to_string(), expected);
}
#[rstest]
#[case::configuration(true)]
#[case::acquisition(false)]
fn contextual_failures_keep_the_original_source(#[case] configuration: bool) {
use litellm_auth_types::ErrorDetail;
let detail = ErrorDetail::failed(
"test credential",
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
);
let error = if configuration {
Error::InvalidConfiguration(detail)
} else {
Error::CredentialAcquisition(detail)
};
let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| {
error.source()
})
.find_map(|error| error.downcast_ref::<std::io::Error>())
.expect("the original credential error remains available");
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
assert!(error.to_string().contains("test credential failed:"));
assert!(error.to_string().ends_with(&source.to_string()));
}

3
litellm-rust/crates/auth-types/tests/params.rs
View file

@ -1,8 +1,9 @@
use litellm_auth_types::{AwsParams, ParamSpec};
use litellm_auth_types::{AwsParams, ParamKind, ParamSpec};
use rstest::rstest;
const LOCATION: ParamSpec = ParamSpec {
setting: "location",
kind: ParamKind::Setting,
wire: &["vertex_location", "vertex_ai_location"],
module_global: Some("vertex_location"),
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],

19
litellm-rust/crates/auth-types/tests/token.rs
View file

@ -1,5 +1,4 @@
use std::{
error::Error as StdError,
future::{Future, poll_fn},
sync::{
Arc,
@ -53,7 +52,7 @@ async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: boo
#[rstest]
#[tokio::test]
async fn callback_errors_preserve_the_original_source() {
async fn callback_errors_keep_their_typed_cause() {
let provider = TokenProviderHandle::from_callback(|| async {
Err(Error::CredentialAcquisition(ErrorDetail::failed(
"caller credential",
@ -62,13 +61,15 @@ async fn callback_errors_preserve_the_original_source() {
});
let error = provider.acquire().await.unwrap_err();
assert!(matches!(error, Error::CredentialAcquisition(_)));
let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| {
(*error).source()
})
.find_map(|error| error.downcast_ref::<std::io::Error>())
.unwrap();
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
let Error::CredentialAcquisition(ErrorDetail::Failed { cause, .. }) = &error else {
panic!("the callback's error passes through unchanged: {error:?}");
};
assert_eq!(
cause
.downcast_ref::<std::io::Error>()
.map(std::io::Error::kind),
Some(std::io::ErrorKind::PermissionDenied)
);
}
struct Release(Arc<AtomicBool>);

1
litellm-rust/crates/cache-azure-blob/Cargo.toml
View file

@ -16,6 +16,7 @@ azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true

30
litellm-rust/crates/cache-azure-blob/src/cache.rs
View file

@ -18,6 +18,10 @@ use url::Url;
use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
#[derive(Debug, thiserror::Error)]
#[error("account URL cannot hold a container path")]
struct NotABaseUrl;
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
codec: C,
@ -57,12 +61,12 @@ impl<C: CacheCodec> AzureBlobCache<C> {
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?;
let parsed = Url::parse(account_url).map_err(Error::unavailable)?;
let account_url = parsed.as_str().trim_end_matches('/').to_string();
let container_url = {
let mut url = parsed;
url.path_segments_mut()
.map_err(|()| Error::Unavailable)?
.map_err(|()| Error::unavailable(NotABaseUrl))?
.pop_if_empty()
.push(container);
url
@ -75,7 +79,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
..BlobContainerClientOptions::default()
}),
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let cache = Self {
container: client,
codec,
@ -101,7 +105,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => {
Ok(())
}
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -116,7 +120,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
{
Ok(_) => Ok(()),
Err(error) if !overwrite && is_already_present(&error) => Ok(()),
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -126,13 +130,9 @@ impl<C: CacheCodec> AzureBlobCache<C> {
Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => {
return Ok(None);
}
Err(_) => return Err(Error::Unavailable),
Err(error) => return Err(Error::unavailable(error)),
};
let bytes = response
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?;
let bytes = response.body.collect().await.map_err(Error::unavailable)?;
self.codec.decode(&bytes).map(Some)
}
@ -140,16 +140,16 @@ impl<C: CacheCodec> AzureBlobCache<C> {
let mut pages = self
.container
.list_blobs(None)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.into_pages();
while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? {
let page = page.into_model().map_err(|_| Error::Unavailable)?;
while let Some(page) = pages.try_next().await.map_err(Error::unavailable)? {
let page = page.into_model().map_err(Error::unavailable)?;
for name in page.blob_items.into_iter().filter_map(|item| item.name) {
self.container
.blob_client(&name)
.delete(None)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
Ok(())

15
litellm-rust/crates/cache-azure-blob/src/credential.rs
View file

@ -17,6 +17,10 @@ const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
type EnvLookup = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
#[derive(Debug, thiserror::Error)]
#[error("no Azure credential is available for blob storage")]
struct NoCredential;
pub struct AzureBlobCredential {
service: AzureAuthService,
env_lookup: EnvLookup,
@ -62,15 +66,8 @@ impl TokenCredential for AzureBlobCredential {
&lookup,
)
.await
.map_err(|error| {
azure_core::Error::with_message(ErrorKind::Credential, error.to_string())
})?
.ok_or_else(|| {
azure_core::Error::with_message(
ErrorKind::Credential,
"no Azure credential is available for blob storage",
)
})?;
.map_err(|error| azure_core::Error::new(ErrorKind::Credential, error))?
.ok_or_else(|| azure_core::Error::new(ErrorKind::Credential, NoCredential))?;
let (token, expires_on) = match credential.into_value() {
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
ResolvedCredential::Static(token) => (token, None),

7
litellm-rust/crates/cache-azure-blob/src/transport.rs
View file

@ -24,6 +24,7 @@ impl HttpClient for ReqwestTransport {
Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)),
};
let response = outgoing.send().await.map_err(|error| {
let error = litellm_http::Error::from(error);
let kind = if error.is_connect() {
ErrorKind::Connection
} else {
@ -41,9 +42,9 @@ impl HttpClient for ReqwestTransport {
);
}
}
let body = response
.bytes_stream()
.map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
let body = response.bytes_stream().map_err(|error| {
azure_core::Error::new(ErrorKind::Io, litellm_http::Error::from(error))
});
Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
}
}

12
litellm-rust/crates/cache-azure-blob/tests/cache.rs
View file

@ -128,7 +128,7 @@ fn connect_surfaces_service_failures() {
ResponseCacheCodec,
runtime.handle().clone(),
));
assert!(matches!(result, Err(Error::Unavailable)));
assert!(matches!(result, Err(Error::Unavailable(_))));
}
#[rstest]
@ -263,7 +263,7 @@ fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[ca
fixture.service.seed_blob(key, bytes);
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(_))
));
}
@ -378,15 +378,15 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.runtime.block_on(
@ -394,7 +394,7 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
}

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