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

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

Browse source
This commit is contained in:
mateo-berri 2026-10-28 01:45:13 -07:00
commit f6d5cdd12d
4079 changed files with 331611 additions and 62535 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 ;;

90
.githooks/pre-push
View file

@ -1,90 +0,0 @@
#!/usr/bin/env bash
#
# pre-push — enforce Conventional Branches
# https://conventional-branch.github.io/
#
# Branch format: <type>/<description>
# <type> must be one of: feature, bugfix, hotfix, release, chore
#
# Protected branches (always allowed):
# - main
# - dependabot/*
# - gh-readonly-queue/*
#
# Tag pushes and branch deletions are skipped.
# Bypass: git push --no-verify.
set -eu
ZERO_OID="0000000000000000000000000000000000000000"
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
PROTECTED_NAMES="main"
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
is_protected() {
branch="$1"
for name in $PROTECTED_NAMES; do
if [ "$branch" = "$name" ]; then
return 0
fi
done
for prefix in $PROTECTED_PREFIXES; do
case "$branch" in "$prefix"*) return 0 ;; esac
done
return 1
}
invalid=""
while read -r local_ref local_oid remote_ref remote_oid; do
# Branch deletion (no local commit being pushed).
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
continue
fi
# Only validate branch pushes; ignore tags and other ref namespaces.
case "$remote_ref" in
refs/heads/*) ;;
*) continue ;;
esac
branch="${remote_ref#refs/heads/}"
if is_protected "$branch"; then
continue
fi
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
invalid="$invalid $branch"
fi
done
if [ -n "$invalid" ]; then
cat >&2 <<EOF
✗ Branch name does not follow Conventional Branches.
Invalid:$invalid
Expected: <type>/<description>
Allowed types: feature, bugfix, hotfix, release, chore
Examples:
feature/weighted-round-robin
bugfix/streaming-empty-chunks
chore/bump-deps
hotfix/auth-bypass
Protected (always allowed): main, dependabot/*, gh-readonly-queue/*.
See https://conventional-branch.github.io/
Rename with: git branch -m <new-name>
To bypass (use sparingly): git push --no-verify
EOF
exit 1
fi
exit 0

6
.github/CODEOWNERS vendored
View file

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

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

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

2
.github/ci-coverage-allowlist.yml vendored
View file

@ -5,7 +5,7 @@ description: >-
test_paths:
- reason: >-
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
litellm.agent() end-to-end suite. It drives the real claude, codex, opencode and pi CLIs and
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and

69
.github/e2e-stack/assert_tests_ran.py vendored
View file

@ -1,69 +0,0 @@
import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
from coverage_registry.management_cases import MANAGEMENT_CASES
def main() -> int:
selected: Final = tuple(sys.argv[2:])
try:
report: Final = ET.parse(Path(sys.argv[1])).getroot()
except (ET.ParseError, OSError):
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
missing: Final = tuple(path for path in selected if path not in passed)
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
passed_nodes: Final = frozenset(
prop.get("value")
for case in cases
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
for prop in case.findall("./properties/property")
if prop.get("name") == "management_node"
)
missing_nodes: Final = required_nodes - passed_nodes
for node in sorted(missing_nodes):
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
for path in selected:
collected: Final = sum(case.get("file") == path for case in cases)
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
for case in cases:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
continue
outcome = "skipped" if case.find("skipped") is not None else "failed"
_ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
for prop in case.findall("./properties/property"):
name = prop.get("name", "")
value = prop.get("value", "")
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
r"[A-Za-z0-9_.:<>-]{1,240}", value
):
_ = sys.stdout.write(f" {name}: {value}\n")
if expected_count is not None and (
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
):
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
return 1
if (
selected
and not missing
and not missing_nodes
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
):
return 0
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
return 1
if __name__ == "__main__":
sys.exit(main())

58
.github/scripts/assert_ci_coverage.py vendored
View file

@ -32,6 +32,7 @@ IGNORE_ARG_RE: Final = re.compile(r"--ignore(?:-glob)?[= ](\S+)")
DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
GLOB_CHARS = frozenset("*?")
DEDICATED_INTEGRATION_WORKFLOWS: Final = MappingProxyType({"image-scan.yml": "tests/integration/database/image/"})
# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
# tests has to be named by some shard or it runs nowhere. A child listed here is
@ -560,6 +561,17 @@ def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
return frozenset(ast.literal_eval(values))
def _owns(workflow: str, relative_path: str) -> bool:
owned_root: Final = DEDICATED_INTEGRATION_WORKFLOWS.get(workflow)
return owned_root is not None and relative_path.startswith(owned_root)
def _tokens_outside_owner(relative_path: str, workflow_tokens: Mapping[str, frozenset[str]]) -> frozenset[str]:
return frozenset[str]().union(
*(tokens for workflow, tokens in workflow_tokens.items() if not _owns(workflow, relative_path))
)
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
runner: Final = repo_root / "tests/integration/run.py"
if not runner.exists():
@ -611,15 +623,20 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
)
!= 1
)
gha_tokens: Final = _invoked_test_tokens(
scalar
for path in (repo_root / ".github/workflows").glob("*.y*ml")
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
workflow_tokens: Final = MappingProxyType(
{
path.name: _invoked_test_tokens(_scalars(yaml.safe_load(path.read_text()), path.name))
for path in (repo_root / ".github/workflows").glob("*.y*ml")
}
)
gha_tokens: Final = frozenset[str]().union(*workflow_tokens.values())
circle_owned: Final = (
frozenset(str(path.relative_to(repo_root)) for path in integration_root.rglob("test_*.py")) - github_files
)
findings: Final = (
tuple(
Finding(path, "integration contract is also selected by GitHub Actions")
for path in paths
Finding(path, "integration file outside GITHUB_FILES is selected by GitHub Actions")
for path in sorted(circle_owned)
if any(_token_covers(token, path) for token in gha_tokens)
)
+ tuple(
@ -632,6 +649,11 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
for path in sorted(github_files)
if not (repo_root / path).is_file()
)
+ tuple(
Finding(path, "GitHub-owned integration contract runs outside a dedicated integration workflow")
for path in sorted(github_files)
if any(_token_covers(token, path) for token in _tokens_outside_owner(path, workflow_tokens))
)
)
browser_commands: Final = tuple(
scalar.value
@ -688,13 +710,9 @@ def main() -> int:
scalars = _all_scalars()
integration_paths, ownership_findings = _integration_ownership()
test_findings = (
_uncovered_tests(
allowlist,
_invoked_selections(scalars)
+ (Selection(included=integration_paths, ignored=frozenset()),),
)
+ ownership_findings
test_findings = _uncovered_tests(
allowlist,
_invoked_selections(scalars) + (Selection(included=integration_paths, ignored=frozenset()),),
)
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
@ -711,13 +729,25 @@ def main() -> int:
test_findings,
"Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if ownership_findings:
browser_remedy = (
" CircleCI's browser group owns the integrationCritical Playwright contracts, so no GitHub workflow may "
"select them and the default Playwright config must exclude them."
if any("browser" in finding.detail or "Playwright" in finding.detail for finding in ownership_findings)
else ""
)
_report(
"integration contracts with the wrong CI owner",
ownership_findings,
"Follow tests/integration/AGENTS.md, which says which CI job owns each integration file." + browser_remedy,
)
if dockerfile_findings:
_report(
"Dockerfiles that no CI job builds",
dockerfile_findings,
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if stale_findings or test_findings or dockerfile_findings:
if stale_findings or test_findings or ownership_findings or dockerfile_findings:
return 1
_write(

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

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

@ -1,39 +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: Regenerate JSON Schema
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create Pull Request
run: |
git add model_prices_and_context_window.json model_prices_and_context_window.schema.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"

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

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

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"

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

@ -1,70 +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: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- 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 \
model_prices_and_context_window.schema.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 }}

10
.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
@ -539,8 +545,8 @@ jobs:
- name: test_unit_passed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_unit_passed_gate.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py tests/code_coverage_tests/test_e2e_idp_stack.py
- name: test_e2e_idp_stack
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_idp_stack.py
- name: test_e2e_metadata
env:

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:

180
.github/workflows/test-mcp-oauth-e2e.yml vendored
View file

@ -1,180 +0,0 @@
name: MCP OAuth happy path
on:
pull_request:
paths:
- '.github/workflows/test-mcp-oauth-e2e.yml'
- '.github/e2e-stack/**'
- 'tests/e2e/*.py'
- 'tests/e2e/pytest.ini'
- 'tests/e2e/idp_realm.json'
- 'tests/e2e/mcp/**'
- 'litellm/experimental_mcp_client/**'
- 'litellm/proxy/_experimental/mcp_server/**'
- 'litellm/proxy/auth/**'
- 'litellm/proxy/management_endpoints/mcp_management_endpoints.py'
- 'litellm/proxy/_types.py'
- 'litellm/types/mcp_server/mcp_server_manager.py'
- 'litellm/proxy/management_endpoints/*sso*.py'
- 'litellm/proxy/management_endpoints/sso/**'
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
- 'litellm/proxy/proxy_server.py'
- 'litellm/proxy/schema.prisma'
- 'ui/litellm-dashboard/src/app/connect/**'
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
- 'pyproject.toml'
- 'uv.lock'
workflow_dispatch:
permissions: {}
concurrency:
group: mcp-oauth-${{ github.ref }}
cancel-in-progress: true
jobs:
oauth:
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
environment: e2e-changed
timeout-minutes: 45
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_HOST: 127.0.0.1
DATABASE_PORT: '5432'
DATABASE_USER: litellm
DATABASE_PASSWORD: dbpassword9090
DATABASE_NAME: litellm
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
E2E_KEYCLOAK_ADMIN_USER: admin
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
E2E_FIXTURE_MODE: live
E2E_PROVIDER_CACHE: '0'
E2E_MCP_OAUTH_LIVE: '1'
E2E_REQUIRED_TEST_COUNT: '4'
steps:
- name: Checkout the tested source
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
persist-credentials: false
- name: Require and materialize the upstream login
env:
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
run: |
umask 077
python3 - <<'PY'
import base64
import json
import os
import secrets
from pathlib import Path
encoded = os.environ.get("STORAGE_STATE", "")
if not encoded:
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
state = json.loads(base64.b64decode(encoded, validate=True))
if not isinstance(state, dict) or not state.get("cookies"):
raise SystemExit("The captured login must contain browser cookies")
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
directory.mkdir(mode=0o700)
path = directory / "linear-state.json"
path.write_text(json.dumps(state))
with open(os.environ["GITHUB_ENV"], "a") as output:
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
value = "sk-e2e-" + secrets.token_hex(24)
print(f"::add-mask::{value}")
output.write(f"{name}={value}\n")
PY
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: '3.13'
- uses: ./.github/actions/setup-uv-with-retries
with:
version: '0.10.9'
- uses: ./.github/actions/cache-cargo-build
- name: Install the frozen E2E environment
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
uv run --no-sync python scripts/prisma_generate_if_needed.py
uv run --no-sync playwright install --with-deps chromium
- name: Configure license access
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: mcp-oauth-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Load the E2E license
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
test -n "${license}"
echo "::add-mask::${license}"
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
- name: Build the gateway consent UI at the tested commit
run: |
cd ui/litellm-dashboard
../../scripts/with_dashboard_node.sh npm ci
../../scripts/with_dashboard_node.sh npm run build
mkdir -p ../../litellm/proxy/_experimental/out
cp -r out/. ../../litellm/proxy/_experimental/out/
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
mkdir -p "${page%.html}"
mv "${page}" "${page%.html}/index.html"
done
- name: Prepare the isolated database and IdP
run: |
umask 077
bash .github/e2e-stack/start-idp.sh
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
- name: Run every required OAuth variant without retries
run: |
umask 077
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
- name: Report JUnit results and reject skipped or missing cases
if: always()
run: |
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
- name: Remove private login and logs
if: always()
run: |
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"

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

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

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

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

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

@ -139,6 +139,7 @@ jobs:
test-path: |-
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
@ -215,10 +216,12 @@ jobs:
- shard: llms-providers
test-path: |-
tests/unit/llms/test_api_base_trailing_slash_routes.py
tests/unit/llms/test_cache_control_and_reasoning.py
tests/unit/llms/test_custom_llm.py
tests/unit/llms/test_file_content_block.py
tests/unit/llms/test_file_search_responses.py
tests/unit/llms/test_get_complete_url_trailing_slash.py
tests/unit/llms/test_lifecycle_fix.py
tests/unit/llms/test_oss_decision.py
tests/unit/llms/test_polling_url_origin_match.py
@ -266,6 +269,7 @@ jobs:
tests/unit/llms/featherless_ai
tests/unit/llms/firecrawl
tests/unit/llms/fireworks_ai
tests/unit/llms/friendliai
tests/unit/llms/gdc
tests/unit/llms/gemini
tests/unit/llms/gigachat
@ -298,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
@ -321,6 +326,7 @@ jobs:
tests/unit/llms/pass_through
tests/unit/llms/perplexity
tests/unit/llms/pg_vector
tests/unit/llms/pi
tests/unit/llms/publicai
tests/unit/llms/ragflow
tests/unit/llms/recraft
@ -410,6 +416,7 @@ jobs:
tests/unit/embeddings
tests/unit/endpoints
tests/unit/files
tests/unit/fine_tuning
tests/unit/images
tests/unit/interactions
tests/unit/messages
@ -419,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"
@ -430,6 +438,13 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py
tests/test_litellm_rust/embeddings/test_dispatch.py
tests/test_litellm_rust/messages/test_dispatch.py
tests/test_litellm_rust/ocr/test_dispatch.py
tests/test_litellm_rust/responses/test_dispatch.py
tests/unit/compression
tests/unit/harness
tests/unit/rust_bridge
@ -565,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
@ -589,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
@ -658,6 +673,7 @@ jobs:
- shard: proxy-extras
test-path: >-
tests/unit/litellm_proxy_extras
tests/unit/db_scripts
python-version: "3.12"
workers: 2
reruns: 0
@ -796,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
@ -813,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
@ -935,7 +953,7 @@ jobs:
dist: loadscope
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
timeout-minutes: 3
timeout-minutes: ${{ matrix.shard == 'db-and-spend' && 10 || 3 }}
with:
persist-credentials: false
fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}

2
.gitignore vendored
View file

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

4
AGENTS.md
View file

@ -29,7 +29,7 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
`tests/unit/` mirrors `litellm/` in a parallel path (see `tests/unit/AGENTS.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
End-to-end tests belong in `tests/e2e/`. A rule about one directory's files, like its harness or which CI runs it, goes in that directory's `AGENTS.md`, never only in a README or here. Read a directory's `AGENTS.md` before editing files in it
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
@ -72,7 +72,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Internal contributors branch a new PR off the repository's current default branch, with a branch name prefixed with litellm_
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch

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

12
CONTRIBUTING.md
View file

@ -39,12 +39,12 @@ git clone https://github.com/YOUR_USERNAME/litellm.git
cd litellm
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
git checkout -b feature/your-feature
git checkout -b litellm_your_feature
# Install development dependencies
make install-dev
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
# Install git hooks that enforce commit conventions (one-time, opt-in)
make install-hooks
# Verify your setup works
@ -55,7 +55,7 @@ That's it! Your local development environment is ready.
## Commit and Branch Conventions
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/). Internal contributors name branches with the `litellm_` prefix and no `/`, and fork branches can use any name. Run `make install-hooks` once per clone to enable the local hook that enforces the commit convention. See the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for details
### 2. Development Workflow
@ -78,8 +78,8 @@ uv run pytest tests/unit/<your_test_file>.py -v
git add .
git commit -m "feat(scope): your descriptive commit message"
# Push and create a PR (branch must follow Conventional Branches — see above)
git push origin feature/your-feature
# Push and create a PR
git push origin litellm_your_feature
```
## Adding Testing
@ -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

8
Makefile
View file

@ -22,7 +22,7 @@ help:
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
@echo " make install-test-deps - Install the full local test environment"
@echo " make install-helm-unittest - Install helm unittest plugin"
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
@echo " make install-hooks - Install git hooks (Conventional Commits)"
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
@echo " make pre-commit - Legacy alias for make check"
@echo " make format - Apply ruff format code formatting"
@ -120,7 +120,7 @@ install-helm-unittest:
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.8.2; \
}
# Install git hooks that enforce Conventional Commits and Conventional Branches.
# Install git hooks that enforce Conventional Commits.
# Opt-in: not chained into install-dev.
install-hooks:
./scripts/install_git_hooks.sh
@ -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

19
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(
@ -287,7 +287,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
</details>
<details>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode, pi or Deep Agents on any model (Python SDK)</summary>
### Python SDK - Agents
@ -296,7 +296,7 @@ import litellm
from litellm import Harness, sandbox
result = litellm.agent(
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.PI, Harness.DEEPAGENTS
"Find why tests/test_router.py is flaky and fix it.",
sandbox=sandbox.local("./repo"),
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
@ -305,7 +305,7 @@ result = litellm.agent(
print(result.text, result.cost, [f.path for f in result.files])
```
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents.
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
@ -384,6 +384,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | |
| [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | |
| [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | |
| [Nace (`nace`)](https://console.nace.ai/docs/drex) | | | | | | | | | | |
| [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | |
| [NLP Cloud (`nlp_cloud`)](https://docs.litellm.ai/docs/providers/nlp_cloud) | ✅ | ✅ | ✅ | | | | | | | |
| [Novita AI (`novita`)](https://novita.ai/models/llm?utm_source=github_litellm&utm_medium=github_readme&utm_campaign=github_link) | ✅ | ✅ | ✅ | | | | | | | |
@ -615,15 +616,15 @@ The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws
### Verify Docker Image Signatures
All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Every release is signed with the same key introduced in [commit `0112e53`](https://github.com/BerriAI/litellm/commit/0112e53046018d726492c814b3644b7d376029d0).
All LiteLLM Docker images published to GHCR are signed with [cosign](https://docs.sigstore.dev/cosign/overview/). Releases are signed with the key introduced in [commit `e82a3c6`](https://github.com/BerriAI/litellm/commit/e82a3c646908ba0a6ee0beddc4310d7a6cac4065).
**Verify using the pinned commit hash (recommended):**
A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the original signing key:
A commit hash is cryptographically immutable, so this is the strongest way to ensure you are using the published signing key:
```bash
cosign verify \
--key https://raw.githubusercontent.com/BerriAI/litellm/0112e53046018d726492c814b3644b7d376029d0/cosign.pub \
--key https://raw.githubusercontent.com/BerriAI/litellm/e82a3c646908ba0a6ee0beddc4310d7a6cac4065/cosign.pub \
ghcr.io/berriai/litellm:<release-tag>
```
@ -659,6 +660,8 @@ This covers:
We welcome contributions to LiteLLM! Whether you're fixing bugs, adding features, or improving documentation, we appreciate your help.
If you want to work closely with our team, check out [LiteLLM OSS Champions](https://www.litellm.ai/champions), our program for trusted OSS contributors with their own Slack channel and a voice in the roadmap
## Quick Start for Contributors
This requires uv to be installed.

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

47
backend/Dockerfile
View file

@ -1,9 +1,36 @@
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
# ---------- Rust extension ----------
# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
FROM $LITELLM_BUILD_IMAGE AS rust-native
WORKDIR /app
USER root
COPY --from=uvbin /uv /uvx /usr/local/bin/
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
ENV UV_PYTHON_DOWNLOADS=0
COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
python -c "import maturin; maturin.build_wheel('/wheels')"
# ---------- Builder ----------
FROM $LITELLM_BUILD_IMAGE AS builder
@ -16,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
@ -54,7 +81,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
COPY . .
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-default-groups --no-editable \
uv sync --frozen --no-install-project --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -63,6 +90,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra saml \
--python python3.13
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@ -79,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
@ -94,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 && \

1
backend/routes/allowlist.py
View file

@ -165,6 +165,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
BACKEND_MOUNT_PATHS: frozenset[str] = frozenset(
{
"/admin",
"/metrics",
"/swagger", # API documentation static assets belong to the backend
"/mcp", # lazily-mounted MCP sub-app serves on the backend component
}

41
ci_cd/cost_map_guard.py
View file

@ -1,11 +1,10 @@
"""Guard the cost map on pull requests.
Every pull request whose diff against its merge base touches one of the three cost map files gets the file
checks: the files parse, the backup copy matches the root file, and the JSON schema is in sync and validates the
map. A pull request that leaves all three 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 three
files and may only add or update models.
Every pull request whose diff against its merge base touches one of the two cost map files gets the file
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 and may only add or update models.
"""
from __future__ import annotations
@ -19,13 +18,11 @@ from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
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"})
EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
CostMap = dict[str, object]
@ -35,7 +32,6 @@ CostMap = dict[str, object]
class Snapshot:
cost_map: str
backup: str
schema: str
def _parse_object(text: str, path: str) -> CostMap | str:
@ -46,13 +42,6 @@ def _parse_object(text: str, path: str) -> CostMap | str:
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
def _rendered_schema(cost_map: CostMap) -> str:
try:
return render(build_schema(cost_map))
except SystemExit as error:
return str(error)
def _evaluation_output_price_failure(
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
) -> str | None:
@ -78,29 +67,13 @@ def _file_failures(
*,
evaluation_output_price_sources: Mapping[str, str] = EVALUATION_OUTPUT_PRICE_SOURCES,
) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
backup_failure: Final = (
()
if head.backup == head.cost_map
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
)
schema_failure: Final = (
()
if head.schema == schema_text
else (
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
)
)
return (
*backup_failure,
*schema_failure,
*(
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
*(
failure
for key, entry in head_map.items()

418
ci_cd/generate_model_prices_schema.py
View file

@ -1,418 +0,0 @@
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Final, Optional
import jsonschema
REPO_ROOT = Path(__file__).parent.parent
PRICES_PATH = REPO_ROOT / "model_prices_and_context_window.json"
SCHEMA_PATH = REPO_ROOT / "model_prices_and_context_window.schema.json"
SPECIAL_ROOT_KEYS = frozenset({"sample_spec", "fallback_generalizations"})
JsonSchema = dict
NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
WEEKDAY_PATTERN: Final = (
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
"gemini_native_audio",
"gemini_audio_only_live",
"uses_embed_content",
"use_openai_responses_path",
"bedrock_converse_supports_strict_tools",
"thinking_always_on",
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
"properties": {
"search_context_size_low": NONNEG_NUMBER,
"search_context_size_medium": NONNEG_NUMBER,
"search_context_size_high": NONNEG_NUMBER,
},
"additionalProperties": False,
},
"guardrail_cost_per_unit": {
"type": "object",
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
"additionalProperties": NONNEG_NUMBER,
},
"metadata": {
"type": "object",
"description": "Free-form notes about the entry (e.g. pricing derivation).",
},
"provider_specific_entry": {
"type": "object",
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
},
}
ARRAY_KEYS: dict[str, JsonSchema] = {
"supported_audio_formats": {
"type": "array",
"description": "Audio container formats the model can return.",
"items": {"type": "string", "enum": ["mp3", "wav"]},
},
"supported_endpoints": {
"type": "array",
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
"items": STRING,
},
"supported_modalities": {
"type": "array",
"description": "Input modalities the model accepts.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video"]},
},
"supported_output_modalities": {
"type": "array",
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"reasoning_effort_levels": {
"type": "array",
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
"items": STRING,
},
"tiered_pricing": {
"type": "array",
"description": "Context-length or result-count tiered rates; each tier's costs apply within its range.",
"items": {
"type": "object",
"properties": {
"range": {
"type": "array",
"description": "[min, max] prompt-token span this tier applies to.",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"max_results_range": {
"type": "array",
"description": "[min, max] result-count span this tier applies to (search models).",
"items": NONNEG_NUMBER,
"minItems": 2,
"maxItems": 2,
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
"input_cost_per_query": NONNEG_NUMBER,
},
"additionalProperties": False,
},
},
}
INTEGER_KEYS: dict[str, JsonSchema] = {
"max_tokens": {
**NONNEG_INTEGER,
"description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.",
},
"max_input_tokens": {
**NONNEG_INTEGER,
"description": "Maximum prompt/context tokens the model accepts.",
},
"max_output_tokens": {
**NONNEG_INTEGER,
"description": "Maximum tokens the model can generate in one response.",
},
"output_vector_size": {
**NONNEG_INTEGER,
"description": "Embedding dimension for embedding models.",
},
"prompt_cache_min_tokens": {
**NONNEG_INTEGER,
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
},
"token_count_framing": {
**NONNEG_INTEGER,
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
},
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
}
NUMBER_KEYS: dict[str, JsonSchema] = {
"regional_processing_uplift_multiplier_eu": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for EU data residency (e.g. 1.10 = +10%).",
},
"regional_processing_uplift_multiplier_us": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).",
},
"regional_endpoint_uplift_multiplier": {
"type": "number",
"minimum": 1,
"description": "Multiplier applied to all token costs when served from a non-global Vertex AI endpoint (e.g. 1.10 = +10%).",
},
}
COST_DESCRIPTIONS: dict[str, str] = {
"input_cost_per_token": "USD per prompt token.",
"output_cost_per_token": "USD per generated token.",
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
"google_maps_grounding_cost_per_query": (
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
),
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
"output_cost_per_token_batches": "USD per generated token via the provider's batch API.",
}
def cost_description(key: str) -> Optional[str]:
if key in COST_DESCRIPTIONS:
return COST_DESCRIPTIONS[key]
if key.endswith("_flex"):
return "Flex service-tier rate for the same-named base field."
if key.endswith("_priority"):
return "Priority service-tier rate for the same-named base field."
if "_above_" in key:
return "Rate applied once the prompt exceeds the token threshold in the field name."
return None
def cost_schema(key: str) -> JsonSchema:
description = cost_description(key)
return {**NONNEG_NUMBER, "description": description} if description else dict(NONNEG_NUMBER)
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
return {
"litellm_provider": {
"type": "string",
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
},
"mode": {
"type": "string",
"description": "Primary API surface / task type of the model.",
"enum": list(modes),
},
"source": {
"type": "string",
"description": "URL of the provider pricing/model page this entry was taken from.",
},
"deprecation_date": {
"type": "string",
"description": "Date the provider deprecates the model, YYYY-MM-DD.",
"format": "date",
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$",
},
"web_search_billing_unit": {
"type": "string",
"description": "Whether web search is billed per query or per prompt.",
"enum": ["per_query", "per_prompt"],
},
"bedrock_output_config_effort_ceiling": {
"type": "string",
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"default_reasoning_effort": {
"type": "string",
"description": (
"Reasoning effort the provider applies when the request omits reasoning_effort. "
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
"accepted, which hold only when the effort resolves to 'none'."
),
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"token_count_fallback_model": {
"type": "string",
"minLength": 1,
"description": "Model that counts tokens for this one when the provider cannot count it.",
},
"audio_transcription_config": STRING,
"vertex_ai_audio_api": {
"type": "string",
"enum": ["lyria_predict", "lyria_interactions"],
},
}
def classify(key: str, modes: tuple) -> Optional[JsonSchema]:
curated = {**OBJECT_KEYS, **ARRAY_KEYS, **string_key_schemas(modes), **INTEGER_KEYS, **NUMBER_KEYS}
if key in curated:
return curated[key]
if key.startswith("supports_") or key in EXTRA_BOOLEAN_KEYS:
return BOOLEAN
if "cost" in key:
return cost_schema(key)
return None
def build_schema(prices: dict) -> JsonSchema:
entries = {name: entry for name, entry in prices.items() if name not in SPECIAL_ROOT_KEYS}
all_keys = tuple(sorted({key for entry in entries.values() for key in entry}))
modes = tuple(sorted({entry["mode"] for entry in entries.values() if "mode" in entry}))
unclassified = tuple(key for key in all_keys if classify(key, modes) is None)
if unclassified:
raise SystemExit(
f"Unclassified keys in {PRICES_PATH.name}: {', '.join(unclassified)}. "
f"Add them to the key tables in {Path(__file__).name} and rerun it."
)
entry_properties = {key: classify(key, modes) for key in all_keys}
return {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"title": "LiteLLM model_prices_and_context_window.json",
"description": (
"Schema for LiteLLM's model price and context window registry "
"(https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json). "
"Every top-level key except 'sample_spec' and 'fallback_generalizations' is a model id, "
"optionally prefixed with its provider (e.g. 'azure/gpt-5.4'), mapping to a model entry. "
"All costs are USD per unit. New optional fields are added regularly, so consumers should "
"ignore unknown fields rather than reject them."
),
"type": "object",
"properties": {
"sample_spec": {
"type": "object",
"description": (
"Documentation placeholder illustrating the entry shape; not a real model and not "
"schema-conformant (several values are prose)."
),
},
"fallback_generalizations": {
"type": "object",
"description": "Regex rules that generalize unknown model ids to known families; not a model entry.",
"properties": {
"rules": {
"type": "array",
"items": {
"type": "object",
"properties": {
"name": STRING,
"pattern": STRING,
"description": STRING,
},
"required": ["name", "pattern"],
"additionalProperties": True,
},
}
},
"additionalProperties": False,
},
},
"additionalProperties": {"$ref": "#/$defs/modelEntry"},
"$defs": {
"modelEntry": {
"type": "object",
"description": (
"Pricing, limits, and capability flags for one model. Fields other than litellm_provider "
"are optional; boolean capability flags are simply omitted when unknown or false."
),
"required": ["litellm_provider"],
"properties": entry_properties,
"additionalProperties": True,
}
},
}
def render(schema: JsonSchema) -> str:
return json.dumps(schema, indent=2) + "\n"
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)
)
def main() -> int:
check = "--check" in sys.argv[1:]
prices = json.loads(PRICES_PATH.read_text())
rendered = render(build_schema(prices))
errors = validation_errors(prices, json.loads(rendered))
if errors:
print(f"{PRICES_PATH.name} does not validate against the generated schema:")
print("\n".join(errors[:20]))
return 1
if not check:
SCHEMA_PATH.write_text(rendered)
print(f"wrote {SCHEMA_PATH}")
return 0
if not SCHEMA_PATH.exists() or SCHEMA_PATH.read_text() != rendered:
print(
f"{SCHEMA_PATH.name} is out of sync with {PRICES_PATH.name}. "
f"Run `python {Path(__file__).relative_to(REPO_ROOT)}` and commit the result."
)
return 1
print(f"{SCHEMA_PATH.name} is in sync and {PRICES_PATH.name} validates against it")
return 0
if __name__ == "__main__":
sys.exit(main())

119
ci_cd/spend_logs_guard.py Normal file
View file

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

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

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

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

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

4
cosign.pub
View file

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

201
db_scripts/strip_vector_store_internal_metadata.py Normal file
View file

@ -0,0 +1,201 @@
import argparse
import os
import sys
from collections.abc import Iterator, Sequence
from dataclasses import dataclass
from itertools import accumulate, repeat, takewhile
from typing import Final, Literal, TextIO
import httpx
from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
from litellm.vector_stores.utils import is_litellm_written_metadata_key
DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
LIST_PAGE_SIZE: Final = "100"
class VectorStore(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
id: str
metadata: dict[str, JsonValue] | None = None
class VectorStorePage(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
data: tuple[VectorStore, ...]
has_more: bool = False
class OpenAIErrorBody(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
message: str = ""
class OpenAIErrorResponse(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
error: OpenAIErrorBody
@dataclass(frozen=True, slots=True)
class Options:
api_key: str
api_base: str
vector_store_ids: tuple[str, ...]
extra_keys: frozenset[str]
apply: bool
@dataclass(frozen=True, slots=True)
class StoreResult:
store_id: str
matched_keys: tuple[str, ...]
status: Literal["dry_run", "updated", "failed"]
error: str = ""
def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool:
return is_litellm_written_metadata_key(key) or key in extra_keys
def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]:
return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys)))
def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage:
params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after}
return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content)
def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None:
if page is None or not page.has_more or not page.data:
return None
return fetch_page(client, after=page.data[-1].id)
def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]:
pages: Final = accumulate(
repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None)
)
for page in takewhile(lambda page: page is not None, pages):
yield from page.data
def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]:
for store_id in store_ids:
yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content)
def describe_http_error(response: httpx.Response) -> str:
try:
message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message
except ValidationError:
return f"HTTP {response.status_code}"
return f"HTTP {response.status_code}: {message}"
def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult:
try:
response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)})
except httpx.TransportError as error:
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__)
if response.is_success:
return StoreResult(store_id=store_id, matched_keys=keys, status="updated")
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response))
def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None:
keys: Final = matching_keys(store, options.extra_keys)
if not keys:
return None
if not options.apply:
return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run")
return delete_keys(client, store.id, keys)
def format_result(result: StoreResult) -> str:
key_list: Final = ", ".join(result.matched_keys)
count: Final = len(result.matched_keys)
match result.status:
case "dry_run":
return f"{result.store_id}: would delete {count} keys: {key_list}"
case "updated":
return f"{result.store_id}: deleted {count} keys: {key_list}"
case "failed":
return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}"
def clean_and_report(
client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO
) -> Iterator[StoreResult | None]:
for store in stores:
result: Final = clean_store(client, store, options)
if result is not None:
print(format_result(result), file=out)
yield result
def run(options: Options, client: httpx.Client, out: TextIO) -> int:
stores: Final = (
retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client)
)
scanned: Final = tuple(clean_and_report(client, stores, options, out))
results: Final = tuple(result for result in scanned if result is not None)
updated: Final = sum(1 for result in results if result.status == "updated")
failed: Final = sum(1 for result in results if result.status == "failed")
mode: Final = "applied" if options.apply else "dry run, pass --apply to delete"
print(
f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})",
file=out,
)
return 1 if failed else 0
def parse_options(argv: Sequence[str] | None) -> Options:
parser: Final = argparse.ArgumentParser(
description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default."
)
parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY")
parser.add_argument("--api-base", default=DEFAULT_API_BASE)
parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store")
parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete")
parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them")
args: Final = parser.parse_args(argv)
if not args.api_key:
parser.error("--api-key or OPENAI_API_KEY is required")
return Options(
api_key=args.api_key,
api_base=args.api_base.rstrip("/"),
vector_store_ids=tuple(args.vector_store_id),
extra_keys=frozenset(args.extra_key),
apply=args.apply,
)
def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int:
options: Final = parse_options(argv)
with httpx.Client(
base_url=options.api_base,
headers={"Authorization": f"Bearer {options.api_key}"},
timeout=60.0,
transport=transport,
) as client:
try:
return run(options, client, sys.stdout)
except httpx.HTTPStatusError as error:
sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n")
return 2
except httpx.TransportError as error:
sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n")
return 2
except ValidationError:
sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n")
return 2
if __name__ == "__main__":
sys.exit(main())

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

16
enterprise/litellm_enterprise/integrations/custom_guardrail.py
View file

@ -20,9 +20,15 @@ class EnterpriseCustomGuardrailHelper:
the tag's value (the mode for that tag).
- If no request tag matches, fall back to default mode(s).
"""
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
)
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
)
from litellm.router_strategy.tag_based_routing import (
without_baggage_mapped_tags,
)
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import premium_user
@ -36,9 +42,13 @@ class EnterpriseCustomGuardrailHelper:
proxy_server_request = data.get("proxy_server_request", {})
request_tags = StandardLoggingPayloadSetup.get_request_tags(
litellm_params=data,
proxy_server_request=proxy_server_request,
request_tags = without_baggage_mapped_tags(
StandardLoggingPayloadSetup.get_request_tags(
litellm_params=data,
proxy_server_request=proxy_server_request,
),
data.get(get_metadata_variable_name_from_kwargs(data)),
keep_serving_deployment_tags=True,
)
# Check if any request tag matches a Mode tag key

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

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

@ -3,6 +3,8 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
"""
import asyncio
import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -12,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
row_metadata,
)
import litellm
@ -38,7 +41,12 @@ if TYPE_CHECKING:
from litellm.types.utils import LiteLLMBatch
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -83,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _row_metadata(row: object) -> dict[str, object]:
metadata: Final = getattr(row, "metadata", None)
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
class CheckBatchCost:
def __init__(
self,
@ -108,6 +111,11 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
self._line_item_store_turn: Final = asyncio.Lock()
self._poll_idle: Final = asyncio.Event()
@ -160,29 +168,9 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
return await self._find_jobs(
where={
"file_purpose": "batch",
"status": {
@ -196,10 +184,59 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
**where,
"created_at": {
"lte": before,
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
},
**(
{
"OR": [
{"created_at": {"gt": cursor[0]}},
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
]
}
if cursor is not None
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
if arrivals:
self._arrival_cursor = next_cursor
else:
self._poll_cursor = next_cursor
return page
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
@ -207,27 +244,22 @@ class CheckBatchCost:
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
return
if retired == 0:
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool:
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
"""
Atomically flip batch_processed from false to true, returning whether this pod won
the row. Every pod and uvicorn worker schedules its own poller against the shared
@ -241,10 +273,16 @@ class CheckBatchCost:
"""
if not self._has_batch_processed_column:
return True
from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json
try:
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": False},
data={"batch_processed": True},
data={
"batch_processed": True,
"status": response.status if response.status != "completed" else "complete",
"file_object": postgres_storable_json(response),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
@ -253,7 +291,9 @@ class CheckBatchCost:
return False
return claimed > 0
async def _release_job_claim(self, job: "ManagedObjectRow") -> None:
async def _release_job_claim(
self, job: "ManagedObjectRow", original_status: str | None, original_file_object: object
) -> None:
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
Safe to match on batch_processed=True: while this poller is active the retrieve
@ -265,7 +305,11 @@ class CheckBatchCost:
try:
await _managed_object_table(self.prisma_client).update_many(
where={"id": job.id, "batch_processed": True},
data={"batch_processed": False},
data={
"batch_processed": False,
"status": original_status,
"file_object": json.dumps(original_file_object),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
@ -299,13 +343,13 @@ class CheckBatchCost:
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
404 must not retire a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
@ -320,7 +364,7 @@ class CheckBatchCost:
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -342,18 +386,18 @@ class CheckBatchCost:
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
if finalized == 0:
return False
verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
return True
except Exception as db_err:
verbose_proxy_logger.error(
@ -564,8 +608,6 @@ class CheckBatchCost:
@staticmethod
def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
file_object: Final = job.file_object
@ -608,7 +650,11 @@ class CheckBatchCost:
from litellm.files.types import FileContentCallOptions
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
mask_api_base_credentials,
resolve_vertex_location_for_cost,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -718,8 +764,7 @@ class CheckBatchCost:
)
except Exception as _e:
verbose_proxy_logger.warning(
f"CheckBatchCost: failed to create managed file ID for "
f"{_file_attr}={_raw_file_id!r}: {_e}"
f"CheckBatchCost: failed to create managed file ID for {_file_attr}={_raw_file_id!r}: {_e}"
)
# Pass the deployment's router-registered pricing (litellm_params custom
@ -738,6 +783,12 @@ class CheckBatchCost:
custom_llm_provider=batch_file_provider,
model_name=model_name,
model_info=deployment_model_info,
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
@ -755,8 +806,9 @@ class CheckBatchCost:
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
)
)
logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id
logging_obj = LiteLLMLogging(
model=batch_result.models[0],
model=logging_model,
messages=[{"role": "user", "content": "<retrieve_batch>"}],
stream=False,
call_type="aretrieve_batch",
@ -787,27 +839,35 @@ class CheckBatchCost:
custom_llm_provider=str(llm_provider) if llm_provider else None,
)
if not await self._claim_job_for_costing(job):
original_status: Final = job.status
original_file_object: Final = job.file_object
if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
"in this window, so its cost is already being tracked there"
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)
raise
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
finally:
if accounting.retryable_failure:
await self._release_job_claim(job, original_status, original_file_object)
if accounting.retryable_failure:
return None
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:
@ -876,7 +936,10 @@ class CheckBatchCost:
)
from litellm.batches.batch_utils import fetch_batch_managed_file_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
resolve_vertex_location_for_cost,
)
try:
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
@ -920,6 +983,12 @@ class CheckBatchCost:
cost_provider=cost_provider,
model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
input_content: Final = (
b""
@ -971,8 +1040,8 @@ class CheckBatchCost:
from litellm.proxy.proxy_server import proxy_config
from litellm.secret_managers.main import str_to_bool
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
api_key: Final = job.api_key if hasattr(job, "api_key") else None
team_id: Final = job.team_id if hasattr(job, "team_id") else None
has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key)
try:
key_row: Final = (
@ -990,7 +1059,9 @@ class CheckBatchCost:
if (has_key_row and key_row is None) or (team_id and team_row is None):
return True
settings: Final = get_dynamic_logging_metadata(
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
UserAPIKeyAuth(
metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id
),
proxy_config,
)
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
@ -1015,6 +1086,7 @@ class CheckBatchCost:
"""
try:
from litellm.integrations.prometheus import PrometheusLogger
prom_logger = PrometheusLogger.get_instance()
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
@ -1039,7 +1111,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await _managed_object_table(self.prisma_client).find_many(
jobs = await self._find_jobs(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1052,8 +1124,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:
@ -1099,10 +1169,7 @@ class CheckBatchCost:
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE
if (
response.status in PROVIDER_TERMINAL_BATCH_STATUSES
and response.output_file_id is not None
):
if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None:
try:
tracked = await self._track_completed_batch_cost(
job=job,
@ -1123,7 +1190,7 @@ class CheckBatchCost:
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
f"(job {job.id}): {tracking_err}"
)
self._record_error(prom_logger, "cost_tracking_error")
continue
@ -1133,24 +1200,19 @@ class CheckBatchCost:
# Track this job for the final metrics summary
processed_models.append(tracked)
# mark the job as complete
try:
update_data: dict = {
"status": response.status if response.status != "completed" else "complete",
"file_object": response.model_dump_json(),
}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data=update_data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
if not self._has_batch_processed_column:
# nothing claimed the row, so it stays in the poll page and its lines wait for that retry
if not self._has_batch_processed_column:
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
data={
"status": response.status if response.status != "completed" else "complete",
"file_object": response.model_dump_json(),
},
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
)
continue
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)

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

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

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

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

92
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 itertools import chain
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -43,6 +44,7 @@ from litellm.constants import (
)
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import require_attributes
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
from litellm.litellm_core_utils.internal_call_metadata import requests_background
from litellm.litellm_core_utils.prompt_templates.common_utils import (
@ -132,7 +134,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
Span = Union[_Span, Any]
Span = _Span
InternalUsageCache = _InternalUsageCache
PrismaClient = _PrismaClient
else:
@ -416,7 +418,7 @@ def _storable_json_value(value: JsonValue) -> JsonValue:
return value
def _postgres_storable_json(file_object: BaseModel) -> str:
def postgres_storable_json(file_object: BaseModel) -> str:
dumped: Final = file_object.model_dump_json()
if "\\u0000" not in dumped:
return dumped
@ -479,9 +481,24 @@ def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str,
_MANAGED_FILES_TARGET: Final = "managed_files"
_PROVIDER_FILE_ID_KWARGS: Final = MappingProxyType(
{
CallTypes.acreate_batch: ("input_file_id",),
CallTypes.acreate_fine_tuning_job: ("training_file", "validation_file"),
}
)
_PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDER_FILE_ID_KWARGS.values())))
_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__(
@ -650,7 +667,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
# FIX: Update status and file_object on every operation to keep state in sync
update_columns: Final = {
"file_object": _postgres_storable_json(stored_object),
"file_object": postgres_storable_json(stored_object),
"status": file_object.status,
"updated_by": user_api_key_dict.user_id,
}
@ -666,7 +683,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
data={
"create": {
"unified_object_id": unified_object_id,
"file_object": _postgres_storable_json(stored_object),
"file_object": postgres_storable_json(stored_object),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": resolve_resource_owner_id(user_api_key_dict),
@ -702,7 +719,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"unified_object_id": response_id,
"model_object_id": decrypted_response_id(response_id),
"file_purpose": "response",
"file_object": _postgres_storable_json(_without_model_output(response)),
"file_object": postgres_storable_json(_without_model_output(response)),
"status": PENDING_RESPONSE_STATUSES[0],
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
@ -1562,6 +1579,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if isinstance(file_id, str) and file_id
)
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
if fine_tuning_file_ids:
data["model_file_id_mapping"] = await self.get_model_file_id_mapping(
fine_tuning_file_ids, user_api_key_dict.parent_otel_span
)
return data
@ -1576,50 +1597,49 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if request_kwargs is None:
return healthy_deployments
input_file_id = cast(Optional[str], request_kwargs.get("input_file_id"))
model_file_id_mapping = cast(
Optional[Dict[str, Dict[str, str]]],
request_kwargs.get("model_file_id_mapping"),
request: Final = cast(Mapping[str, object], request_kwargs)
model_file_id_mapping: Final = cast(Mapping[str, Mapping[str, str]], request.get("model_file_id_mapping") or {})
requested_file_ids: Final = (request.get(key) for key in _PROVIDER_FILE_ID_KEYS)
model_ids_per_file: Final = tuple(
frozenset(model_file_id_mapping[file_id])
for file_id in requested_file_ids
if isinstance(file_id, str) and file_id in model_file_id_mapping
)
allowed_model_ids: Final[frozenset[str]] = (
model_ids_per_file[0].intersection(*model_ids_per_file[1:]) if model_ids_per_file else frozenset()
)
allowed_model_ids = []
if input_file_id and model_file_id_mapping:
model_id_dict = model_file_id_mapping.get(input_file_id, {})
allowed_model_ids = list(model_id_dict.keys())
if len(allowed_model_ids) == 0:
return healthy_deployments
return [
holders: Final[list[dict[str, object]]] = [
deployment
for deployment in healthy_deployments
if deployment.get("model_info", {}).get("id") in allowed_model_ids
]
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.
"""
accessor_key: Optional[str] = None
if call_type and call_type == CallTypes.acreate_batch:
accessor_key = "input_file_id"
elif call_type and call_type == CallTypes.acreate_fine_tuning_job:
accessor_key = "training_file"
else:
accessor_keys: Final = _PROVIDER_FILE_ID_KWARGS.get(call_type) if call_type else None
if accessor_keys is None:
return kwargs
if accessor_key:
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], _keyed_get(kwargs.get("model_info", {}), "id"))
if model_id is None:
model_id = cast(
Optional[str],
_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))
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))
if model_id is None:
model_id = cast(
Optional[str],
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
)
mapped_file_id: Optional[str] = None
if input_file_id and model_file_id_mapping and model_id:
mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None)
@ -1813,7 +1833,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
)
async def get_model_file_id_mapping(self, file_ids: List[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
@ -1855,7 +1875,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
llm_router: Optional[Router],
_create_file_request: CreateFileRequest,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
) -> List[OpenAIFileObject]:
if llm_router is None:
raise Exception("LLM Router not initialized. Ensure models added to proxy.")
@ -1871,7 +1891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
create_file_request: CreateFileRequest,
llm_router: Router,
target_model_names_list: List[str],
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
user_api_key_dict: UserAPIKeyAuth,
) -> OpenAIFileObject:
responses = await self.create_file_for_each_model(
@ -1938,7 +1958,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_objects: List[OpenAIFileObject],
create_file_request: CreateFileRequest,
internal_usage_cache: InternalUsageCache,
litellm_parent_otel_span: Span,
litellm_parent_otel_span: Span | None,
target_model_names_list: List[str],
) -> OpenAIFileObject:
## GET THE FILE TYPE FROM THE CREATE FILE REQUEST
@ -2641,7 +2661,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Continue with other files even if one fails
continue
def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str:
def _get_content_type_from_file_object(self, file_object: object) -> str:
"""
Determine content type from file object.

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

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

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

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

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

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

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

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

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.77"
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.77"
version = "0.1.80"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

49
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
@ -20,6 +20,33 @@ RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${
make -j"$(nproc)" pgbouncer && \
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
# ---------- Rust extension ----------
# Copies only the native build's inputs, so the fat-LTO compile stays cached across Python-only changes
FROM $LITELLM_BUILD_IMAGE AS rust-native
WORKDIR /app
USER root
COPY --from=uvbin /uv /uvx /usr/local/bin/
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
ENV UV_PYTHON_DOWNLOADS=0
COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \
touch litellm/__init__.py litellm/rust_bridge/__init__.py && \
uvx --python python3.13 --from maturin==1.15.0 --with puccinialin==0.1.14 \
python -c "import maturin; maturin.build_wheel('/wheels')"
# ---------- Builder ----------
FROM $LITELLM_BUILD_IMAGE AS builder
@ -32,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
@ -66,11 +93,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
# Stage 2 — copy source and install the workspace members.
COPY . .
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync --frozen --no-default-groups --no-editable \
uv sync --frozen --no-install-project --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--group admin-mcp \
@ -79,6 +106,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3.13
RUN --mount=type=bind,from=rust-native,source=/wheels,target=/wheels \
uv pip install --python /app/.venv/bin/python --no-deps /wheels/litellm-*.whl
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
# compiled Rust extension must live next to the source or it is never imported.
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
@ -95,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
@ -110,7 +140,12 @@ ENV HOME=/home/nonroot \
PYTHONUNBUFFERED=1 \
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY --from=builder --chown=nonroot:nonroot /app /app
COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer

4
gateway/routes/allowlist.py
View file

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

13
helm/litellm-helm/README.md
View file

@ -70,11 +70,18 @@ proxyConfigMap:
proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
model_list: []
```
The chart ships no models, so add them under `proxy_config.model_list` or from the Admin UI with `store_model_in_db: true`
```
proxy_config:
model_list:
- model_name: gpt-3.5-turbo
- model_name: gpt-5.6
litellm_params:
model: gpt-3.5-turbo
api_key: eXaMpLeOnLy
model: openai/gpt-5.6
api_key: os.environ/OPENAI_API_KEY
```
#### Example using existing `proxyConfigMap` instead of creating it:

26
helm/litellm-helm/tests/deployment_tests.yaml
View file

@ -270,6 +270,32 @@ tests:
- matchRegex:
path: data["config.yaml"]
pattern: '(?m)^\s*health_check_details:\s*false$'
- it: should not render any placeholder models when the user sets proxy_config without model_list
template: configmap-litellm.yaml
set:
proxy_config.general_settings.store_model_in_db: true
asserts:
- notMatchRegex:
path: data["config.yaml"]
pattern: 'model_name:'
- matchRegex:
path: data["config.yaml"]
pattern: 'store_model_in_db: true'
- matchRegex:
path: data["config.yaml"]
pattern: 'master_key: os.environ/PROXY_MASTER_KEY'
- it: should render a user-supplied model_list as the only models
template: configmap-litellm.yaml
set:
proxy_config.model_list:
- model_name: my-model
litellm_params:
model: openai/my-model
api_key: os.environ/OPENAI_API_KEY
asserts:
- matchRegex:
path: data["config.yaml"]
pattern: 'model_list:\n- litellm_params:\n api_key: os.environ/OPENAI_API_KEY\n model: openai/my-model\n model_name: my-model\n$'
- it: should allow overriding liveness, readiness, and startup probes
template: deployment.yaml
set:

12
helm/litellm-helm/values.yaml
View file

@ -180,17 +180,7 @@ proxyConfigMap:
# Examples: https://github.com/BerriAI/litellm/tree/main/litellm/proxy/example_config_yaml
# Reference: https://docs.litellm.ai/docs/proxy/configs
proxy_config:
model_list:
# At least one model must exist for the proxy to start.
- model_name: gpt-3.5-turbo
litellm_params:
model: gpt-3.5-turbo
api_key: eXaMpLeOnLy
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
model_list: []
general_settings:
master_key: os.environ/PROXY_MASTER_KEY

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

23
litellm-proxy-extras/litellm_proxy_extras/migrations/20261015120000_add_environment_variables_revision_trigger/migration.sql Normal file
View file

@ -0,0 +1,23 @@
SET LOCAL lock_timeout = '10s';
CREATE SEQUENCE IF NOT EXISTS "litellm_environment_variables_revision_seq";
GRANT USAGE ON SEQUENCE "litellm_environment_variables_revision_seq" TO PUBLIC;
CREATE OR REPLACE FUNCTION litellm_bump_environment_variables_revision() RETURNS TRIGGER AS $$
DECLARE
next_revision BIGINT := nextval(format('%I.litellm_environment_variables_revision_seq', TG_TABLE_SCHEMA)::regclass);
BEGIN
IF TG_OP = 'UPDATE' THEN
NEW."reload_revision" := GREATEST(next_revision, OLD."reload_revision" + 1);
ELSE
NEW."reload_revision" := next_revision;
END IF;
RETURN NEW;
END;
$$ LANGUAGE plpgsql;
DROP TRIGGER IF EXISTS "litellm_environment_variables_revision" ON "LiteLLM_Config";
CREATE TRIGGER "litellm_environment_variables_revision"
BEFORE INSERT OR UPDATE ON "LiteLLM_Config"
FOR EACH ROW WHEN (NEW."param_name" = 'environment_variables')
EXECUTE PROCEDURE litellm_bump_environment_variables_revision();

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

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

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

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

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

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

80
litellm-rust/Cargo.lock generated
View file

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

9
litellm-rust/Cargo.toml
View file

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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