mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
Merge remote-tracking branch 'upstream/main' into openrouter-decisions-images
This commit is contained in:
commit
a6c6bfb526
1702 changed files with 116520 additions and 20896 deletions
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
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
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Database access
|
||||
|
|
|
|||
|
|
@ -1,14 +1,14 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [ "${GITHUB_ACTIONS:-}" = true ]; then
|
||||
echo "Integration contracts are owned by CircleCI" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
suite="${1:?integration suite required}"
|
||||
mode="${2:-standard}"
|
||||
side="${3:-}"
|
||||
requested_workers="${INTEGRATION_WORKERS:-}"
|
||||
license_env=()
|
||||
if [ -n "${LITELLM_LICENSE:-}" ]; then
|
||||
license_env=("LITELLM_LICENSE=$LITELLM_LICENSE")
|
||||
fi
|
||||
if [ "$mode" = replica ]; then
|
||||
results="test-results/integration-${suite}-replica"
|
||||
elif [ "$mode" = parity ]; then
|
||||
|
|
@ -84,7 +84,7 @@ if [ "$suite" = browser ]; then
|
|||
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
|
||||
test -f "$LITELLM_UI_PATH/index.html"
|
||||
fi
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256(((os.environ.get("CIRCLE_SHA1") or os.environ.get("GITHUB_SHA") or "local") + (os.environ.get("CIRCLE_WORKFLOW_ID") or os.environ.get("GITHUB_RUN_ID") or "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
|
||||
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
|
||||
|
|
@ -107,7 +107,7 @@ sudo iptables -N integration_only
|
|||
guard_created=true
|
||||
sudo iptables -A integration_only -o lo -j ACCEPT
|
||||
sudo iptables -A integration_only -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
|
||||
for service in postgres-db redis-cache; do
|
||||
for service in "${INTEGRATION_POSTGRES_CONTAINER:-postgres-db}" "${INTEGRATION_REDIS_CONTAINER:-redis-cache}"; do
|
||||
address="$(docker inspect --format '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$service")"
|
||||
sudo iptables -A integration_only -d "$address" -j ACCEPT
|
||||
done
|
||||
|
|
@ -145,6 +145,9 @@ fi
|
|||
if [ "$suite" = mcp ]; then
|
||||
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
|
||||
fi
|
||||
if [ -n "$requested_workers" ]; then
|
||||
export INTEGRATION_WORKERS="$requested_workers"
|
||||
fi
|
||||
isolated_workers=false
|
||||
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
|
||||
isolated_workers=true
|
||||
|
|
@ -185,7 +188,7 @@ start_proxy() {
|
|||
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
|
||||
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
|
||||
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
"${license_env[@]}" \
|
||||
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=true "${cost_map_env[@]}" \
|
||||
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
|
||||
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
|
||||
|
|
@ -257,7 +260,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
|||
"${shared_services[@]}" "${worker_coverage[@]}" \
|
||||
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
|
||||
LITELLM_MODE=PRODUCTION \
|
||||
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
|
||||
"${license_env[@]}" \
|
||||
INTEGRATION_SEED="$INTEGRATION_SEED" \
|
||||
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
|
||||
INTEGRATION_CASE_SHARDS="${CIRCLE_NODE_TOTAL:-1}" INTEGRATION_CASE_SHARD_INDEX="${CIRCLE_NODE_INDEX:-0}" \
|
||||
|
|
|
|||
2
.github/CODEOWNERS
vendored
2
.github/CODEOWNERS
vendored
|
|
@ -1,4 +1,6 @@
|
|||
/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
|
||||
/ci_cd/cost_map_pins.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
|
||||
/helm/ @yassin-berriai
|
||||
/terraform/ @yassin-berriai
|
||||
/.github/workflows/zizmor.yml @yuneng-berri
|
||||
|
|
|
|||
7
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
7
.github/PULL_REQUEST_TEMPLATE/general.md
vendored
|
|
@ -84,9 +84,10 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
## Screenshots / Proof of Fix
|
||||
|
||||
<!-- Include screenshots, screen recordings, or command (e.g., curl) + output demonstrating that your changes work as expected
|
||||
The proof must be completely e2e with no mocks, using actual LLM calls costing real $$$ if applicable. `pytest` commands are not enough
|
||||
Show ONLY the latest run: capture Before at the merge base and After at the PR's current tip, and when new commits change behavior, replace this whole section with the fresh run instead of stacking it on top of older ones. The run must be up to date. As soon as a new commit is made and it makes this PR description's after sha stale (it's no longer tip of PR), you must re-run the QA
|
||||
Structure the section exactly as below: Before and After one heading level below this section, each naming the commit hash it was captured at, one lower-level heading per case inside each, the same case names in the same order on both sides, and numbered steps (command, observed output) under every case, never loose prose; shared setup (config, payloads) goes above Before, and with a single case, drop the case headings and number the steps directly
|
||||
Simple PR exception (2026-10-28, #49030). A mechanical change fully proven by focused checks may skip before/after e2e. Follow AGENTS.md's eligibility rule, trace callers, and verify bug-fix regression sensitivity. Record E2E exception at <tip SHA>, scope, rationale, commands and results here. Tests/static checks suffice only for this exception. Keep CI and bot gates, reassess changed diffs, and honor explicit e2e requests. Omit the Before/After scaffold for an exception
|
||||
Otherwise proof must be fully e2e, no mocks, with real provider calls where applicable. Pytest alone is insufficient. The following format applies to live proof
|
||||
Show only the latest proof. Before uses the merge base and After the tested tip. Rerun affected QA after behavior changes. Verified neutral diffs may carry proof with the tested SHA, current tip, equivalence reason and focused checks, following AGENTS.md and /qa
|
||||
Use matching Before/After case headings with numbered commands and observed output. Put shared setup above them. For one case, omit case headings
|
||||
|
||||
### Before (<hash>)
|
||||
|
||||
|
|
|
|||
119
.github/scripts/assert_secret_absent.py
vendored
Normal file
119
.github/scripts/assert_secret_absent.py
vendored
Normal file
|
|
@ -0,0 +1,119 @@
|
|||
import base64
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import urllib.parse
|
||||
import zipfile
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TextIO
|
||||
|
||||
SECRET_VARIABLE: Final = "SCANNED_SECRET"
|
||||
MIN_ENCODED_LENGTH: Final = 12
|
||||
MAX_ARCHIVE_DEPTH: Final = 3
|
||||
MASK: Final = "***"
|
||||
|
||||
Forms = tuple[tuple[str, str], ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Hit:
|
||||
location: str
|
||||
form: str
|
||||
|
||||
|
||||
def _aligned_encodings(raw: bytes, name: str, encode: Callable[[bytes], bytes]) -> Iterator[tuple[str, str]]:
|
||||
for offset in range(3):
|
||||
chunk: Final = raw[offset:]
|
||||
yield f"{name} offset {offset}", encode(chunk[: len(chunk) - len(chunk) % 3]).decode()
|
||||
|
||||
|
||||
def secret_forms(secret: str) -> Forms:
|
||||
encoded: Final = (
|
||||
("url-encoded", urllib.parse.quote(secret, safe="")),
|
||||
*_aligned_encodings(secret.encode(), "base64", base64.b64encode),
|
||||
*_aligned_encodings(secret.encode(), "base64url", base64.urlsafe_b64encode),
|
||||
)
|
||||
by_value: Final = MappingProxyType(
|
||||
{value: name for name, value in encoded if value != secret and len(value) >= MIN_ENCODED_LENGTH}
|
||||
)
|
||||
return (("raw", secret), *((name, value) for value, name in by_value.items()))
|
||||
|
||||
|
||||
def mask(text: str, forms: Forms) -> str:
|
||||
return reduce(lambda masked, form: masked.replace(form[1], MASK), forms, text)
|
||||
|
||||
|
||||
def _name_hits(location: str, name: str, forms: Forms) -> Iterator[Hit]:
|
||||
for form, value in forms:
|
||||
if value in name:
|
||||
yield Hit(mask(location, forms), f"{form} in name")
|
||||
|
||||
|
||||
def _content_hits(location: str, data: bytes, forms: Forms, depth: int) -> Iterator[Hit]:
|
||||
for form, value in forms:
|
||||
if value.encode() in data:
|
||||
yield Hit(mask(location, forms), form)
|
||||
if depth >= MAX_ARCHIVE_DEPTH or not zipfile.is_zipfile(io.BytesIO(data)):
|
||||
return
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as archive:
|
||||
for member in archive.infolist():
|
||||
if member.is_dir():
|
||||
continue
|
||||
yield from _content_hits(f"{location}!{member.filename}", archive.read(member), forms, depth + 1)
|
||||
|
||||
|
||||
def _walk(root: Path) -> Iterator[tuple[Path, bool]]:
|
||||
for directory, subdirectories, files in os.walk(root):
|
||||
yield from ((Path(directory, name), True) for name in subdirectories if Path(directory, name).is_symlink())
|
||||
yield from ((Path(directory, name), False) for name in files)
|
||||
|
||||
|
||||
def _path_hits(root: Path, path: Path, linked_directory: bool, forms: Forms) -> Iterator[Hit]:
|
||||
location: Final = str(path.relative_to(root))
|
||||
if linked_directory:
|
||||
yield Hit(mask(location, forms), "directory symlink, not scanned")
|
||||
return
|
||||
yield from _name_hits(location, location, forms)
|
||||
yield from _content_hits(location, path.read_bytes(), forms, 0)
|
||||
|
||||
|
||||
def scan(root: Path, secret: str) -> tuple[int, tuple[Hit, ...]]:
|
||||
forms: Final = secret_forms(secret)
|
||||
entries: Final = tuple(_walk(root))
|
||||
hits: Final = tuple(chain.from_iterable(_path_hits(root, path, linked, forms) for path, linked in entries))
|
||||
return len(entries), hits
|
||||
|
||||
|
||||
def main(argv: tuple[str, ...], environ: Mapping[str, str], out: TextIO) -> int:
|
||||
if len(argv) != 1:
|
||||
print(f"usage: assert_secret_absent.py <directory>, with the secret in ${SECRET_VARIABLE}", file=out)
|
||||
return 2
|
||||
secret: Final = environ.get(SECRET_VARIABLE, "")
|
||||
root: Final = Path(argv[0])
|
||||
if not secret:
|
||||
print(f"{SECRET_VARIABLE} is empty, nothing to scan for", file=out)
|
||||
return 0
|
||||
if not root.is_dir():
|
||||
print(f"{root} is not a directory, nothing to scan", file=out)
|
||||
return 0
|
||||
try:
|
||||
scanned, hits = scan(root, secret)
|
||||
except Exception as error:
|
||||
print(f"scan failed with {type(error).__name__}, treating the results as unsafe", file=out)
|
||||
return 1
|
||||
for hit in hits:
|
||||
print(f"secret found: {hit.location} ({hit.form})", file=out)
|
||||
if hits:
|
||||
print(f"{len(hits)} copies of the secret under {root}, refusing to upload", file=out)
|
||||
return 1
|
||||
print(f"no copy of the secret in {scanned} files under {root}", file=out)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main(tuple(sys.argv[1:]), os.environ, sys.stdout))
|
||||
9
.github/scripts/smoke_test_native_wheel.py
vendored
9
.github/scripts/smoke_test_native_wheel.py
vendored
|
|
@ -7,6 +7,8 @@ import zipfile
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json"
|
||||
|
||||
CHILD_SCRIPT: Final = """
|
||||
from importlib.util import module_from_spec, spec_from_file_location
|
||||
from pathlib import Path
|
||||
|
|
@ -34,6 +36,9 @@ else:
|
|||
after = module.gil_stats()
|
||||
if not isinstance(after.get("releases"), int):
|
||||
raise AssertionError(f"native module unusable after panic: {after!r}")
|
||||
|
||||
if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes():
|
||||
raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout")
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -55,7 +60,9 @@ def main() -> int:
|
|||
|
||||
native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
|
||||
native_path.write_bytes(archive.read(native_members[0]))
|
||||
result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False)
|
||||
result: Final = subprocess.run(
|
||||
(sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")
|
||||
|
|
|
|||
38
.github/workflows/cost-map-guard.yml
vendored
38
.github/workflows/cost-map-guard.yml
vendored
|
|
@ -1,7 +1,7 @@
|
|||
name: Cost map guard
|
||||
|
||||
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed
|
||||
pull_request_target:
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- "litellm_**"
|
||||
|
|
@ -19,17 +19,20 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.base.sha }}
|
||||
persist-credentials: false
|
||||
- name: Fetch the pull request head and its merge base
|
||||
- name: Fetch the pull request head, its merge base, and the base branch tip
|
||||
id: revisions
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
BASE_REF: ${{ github.event.pull_request.base.ref }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
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"
|
||||
target="$(gh api "repos/${GITHUB_REPOSITORY}/branches/${BASE_REF}" --jq '.commit.sha')"
|
||||
merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${target}...${HEAD_SHA}" --jq '.merge_base_commit.sha')"
|
||||
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA" "$target"
|
||||
echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT"
|
||||
echo "target=$target" >> "$GITHUB_OUTPUT"
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
|
|
@ -37,7 +40,30 @@ jobs:
|
|||
- name: Run the guard
|
||||
env:
|
||||
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
|
||||
COST_MAP_GUARD_TARGET: ${{ steps.revisions.outputs.target }}
|
||||
COST_MAP_GUARD_AUTHOR_TYPE: ${{ github.event.pull_request.user.type }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
HEAD_REF: ${{ github.event.pull_request.head.ref }}
|
||||
run: |
|
||||
uv run --no-project python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"
|
||||
- name: Check whether the pull request changes the cost map
|
||||
id: changed
|
||||
env:
|
||||
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
if git diff --quiet "$MERGE_BASE" "$HEAD_SHA" -- model_prices_and_context_window.json; then
|
||||
echo "cost_map=false" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "cost_map=true" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
- if: steps.changed.outputs.cost_map == 'true'
|
||||
run: rustup toolchain install --no-self-update
|
||||
- name: Load the cost map through model-catalog
|
||||
if: steps.changed.outputs.cost_map == 'true'
|
||||
env:
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
git show "${HEAD_SHA}:model_prices_and_context_window.json" > "$RUNNER_TEMP/model_prices_and_context_window.json"
|
||||
cargo run --locked --quiet --manifest-path litellm-rust/Cargo.toml -p litellm-model-catalog --bin check-registry \
|
||||
-- "$RUNNER_TEMP/model_prices_and_context_window.json"
|
||||
|
|
|
|||
91
.github/workflows/cost-map-pins-recheck.yml
vendored
Normal file
91
.github/workflows/cost-map-pins-recheck.yml
vendored
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
name: Cost map pins recheck
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- model_prices_and_context_window.json
|
||||
- litellm/model_prices_and_context_window_backup.json
|
||||
- ci_cd/cost_map_pins.json
|
||||
|
||||
permissions: {}
|
||||
|
||||
jobs:
|
||||
rerun-cost-map-guard:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
actions: write
|
||||
contents: read
|
||||
pull-requests: read
|
||||
steps:
|
||||
- name: Re-run the cost map guard on open pull requests the push could break
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_REF: ${{ github.ref_name }}
|
||||
BEFORE_SHA: ${{ github.event.before }}
|
||||
AFTER_SHA: ${{ github.event.after }}
|
||||
FORCED: ${{ github.event.forced }}
|
||||
COST_MAP_PATHS: '["model_prices_and_context_window.json", "litellm/model_prices_and_context_window_backup.json", "ci_cd/cost_map_pins.json"]'
|
||||
PINS_PATHS: '["ci_cd/cost_map_pins.json"]'
|
||||
run: |
|
||||
if [[ "$BEFORE_SHA" =~ ^0+$ || "$FORCED" == true ]] \
|
||||
|| ! gh api "repos/${GITHUB_REPOSITORY}/compare/${BEFORE_SHA}...${AFTER_SHA}?per_page=1" \
|
||||
--jq '(env.PINS_PATHS | fromjson) as $pins | (.files | length) < 300 and all(.files[]; [.filename, .previous_filename] | any(IN($pins[])) | not)' \
|
||||
| grep -qx true; then
|
||||
export RECHECK_PATHS="$COST_MAP_PATHS"
|
||||
else
|
||||
export RECHECK_PATHS="$PINS_PATHS"
|
||||
fi
|
||||
rerun_cutoff="$(date -u -d '30 days ago' +%Y-%m-%dT%H:%M:%SZ)"
|
||||
candidates="$(gh api graphql --paginate -F owner="${GITHUB_REPOSITORY_OWNER}" -F name="${GITHUB_REPOSITORY#*/}" -F base="$BASE_REF" \
|
||||
-f query='query($owner: String!, $name: String!, $base: String!, $endCursor: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequests(states: OPEN, baseRefName: $base, first: 25, after: $endCursor) {
|
||||
nodes { number headRefOid changedFiles files(first: 100) { nodes { path } } }
|
||||
pageInfo { hasNextPage endCursor }
|
||||
}
|
||||
}
|
||||
}' \
|
||||
--jq '(env.RECHECK_PATHS | fromjson) as $paths | .data.repository.pullRequests.nodes[]
|
||||
| {number, headRefOid, listed: any(.files.nodes[].path; IN($paths[])), unlisted: (.changedFiles > 100)}
|
||||
| select(.listed or .unlisted) | "\(.number) \(.headRefOid) \(.listed)"')"
|
||||
failed=0
|
||||
declare -A rerun_heads
|
||||
while read -r number head_sha listed; do
|
||||
[ -n "$number" ] || continue
|
||||
if [ "$listed" = false ]; then
|
||||
if ! changed="$(gh api --paginate "repos/${GITHUB_REPOSITORY}/pulls/${number}/files?per_page=100" \
|
||||
--jq '(env.RECHECK_PATHS | fromjson) as $paths | .[].filename | select(IN($paths[]))')"; then
|
||||
echo "::error::could not list the files of #$number"
|
||||
failed=1
|
||||
continue
|
||||
fi
|
||||
[ -n "$changed" ] || continue
|
||||
fi
|
||||
[ -z "${rerun_heads[$head_sha]:-}" ] || continue
|
||||
rerun_heads[$head_sha]=1
|
||||
if ! runs="$(gh run list --repo "$GITHUB_REPOSITORY" --workflow cost-map-guard.yml --commit "$head_sha" --limit 100 --json databaseId,createdAt --jq '.[] | "\(.databaseId) \(.createdAt)"')"; then
|
||||
echo "::error::could not list the cost map guard runs of #$number"
|
||||
failed=1
|
||||
continue
|
||||
fi
|
||||
if [ -z "$runs" ]; then
|
||||
echo "::warning::no cost map guard run found for #$number"
|
||||
continue
|
||||
fi
|
||||
while read -r run_id created_at; do
|
||||
if [[ "$created_at" < "$rerun_cutoff" ]]; then
|
||||
echo "::warning::cost map guard run $run_id for #$number is over 30 days old, so GitHub cannot re-run it; #$number gets the pin checks on its next push"
|
||||
continue
|
||||
fi
|
||||
gh run watch "$run_id" --repo "$GITHUB_REPOSITORY" > /dev/null || true
|
||||
if ! gh run rerun "$run_id" --repo "$GITHUB_REPOSITORY"; then
|
||||
echo "::error::could not re-run cost map guard run $run_id for #$number"
|
||||
failed=1
|
||||
fi
|
||||
done <<< "$runs"
|
||||
done <<< "$candidates"
|
||||
exit "$failed"
|
||||
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -13,8 +13,10 @@ on:
|
|||
- migrations/Dockerfile
|
||||
- migrations/run.py
|
||||
- gateway/Dockerfile
|
||||
- gateway/launch.py
|
||||
- gateway/main.py
|
||||
- gateway/routes/allowlist.py
|
||||
- litellm_cli/**
|
||||
- backend/Dockerfile
|
||||
- backend/main.py
|
||||
- litellm-rust/**
|
||||
|
|
|
|||
|
|
@ -19,7 +19,7 @@ jobs:
|
|||
name: interactions-spec-canary
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
outputs:
|
||||
|
|
@ -39,8 +39,11 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up the Rust bridge
|
||||
uses: ./.github/actions/rust-bridge
|
||||
|
||||
- name: Install dependencies
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Run the compliance tests against Google's live spec
|
||||
id: compliance
|
||||
|
|
|
|||
10
.github/workflows/spend-logs-guard.yml
vendored
10
.github/workflows/spend-logs-guard.yml
vendored
|
|
@ -1,13 +1,9 @@
|
|||
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:
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
pull_request_review:
|
||||
types:
|
||||
- submitted
|
||||
- dismissed
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -23,7 +19,7 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.base.ref }}
|
||||
ref: ${{ github.event.pull_request.base.sha }}
|
||||
persist-credentials: false
|
||||
- name: Fetch the pull request head, its merge base and its reviews
|
||||
id: revisions
|
||||
|
|
|
|||
199
.github/workflows/test-integration.yml
vendored
Normal file
199
.github/workflows/test-integration.yml
vendored
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
name: integration
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [labeled]
|
||||
schedule:
|
||||
- cron: "5 8-20/3 * * 1-6"
|
||||
timezone: America/Los_Angeles
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
suite:
|
||||
description: Integration group to run
|
||||
type: choice
|
||||
options:
|
||||
- all
|
||||
- providers
|
||||
- extensions
|
||||
- accounting
|
||||
- management
|
||||
- management-replica
|
||||
- database
|
||||
- database-replica
|
||||
- mcp
|
||||
- security
|
||||
- cost
|
||||
- sdk
|
||||
- browser
|
||||
default: all
|
||||
workers:
|
||||
description: Workers per group, blank keeps each group's default
|
||||
type: string
|
||||
default: ""
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name }}-${{ inputs.suite }}-${{ inputs.workers }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
integration:
|
||||
name: integration ${{ matrix.suite }}
|
||||
if: github.event_name != 'pull_request' || github.event.label.name == 'run-ci'
|
||||
environment: integration-license
|
||||
runs-on: ${{ fromJSON('{"providers":"ubuntu-latest-64-cores","extensions":"ubuntu-latest-64-cores","accounting":"ubuntu-latest-32-cores","management":"ubuntu-latest-32-cores","management-replica":"ubuntu-latest-32-cores"}')[matrix.suite] || 'ubuntu-latest-16-cores' }}
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
suite: ${{ fromJSON((inputs.suite || 'all') == 'all' && '["providers","extensions","accounting","management","management-replica","database","database-replica","mcp","security","cost","sdk","browser"]' || format('["{0}"]', inputs.suite)) }}
|
||||
env:
|
||||
UV_PYTHON: "3.12"
|
||||
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"32","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
command: -c max_connections=1000 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000
|
||||
env:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: circle_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--shm-size 1g
|
||||
--health-cmd "pg_isready -U postgres"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 20
|
||||
redis:
|
||||
image: redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6
|
||||
ports:
|
||||
- 6379:6379
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 20
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: ${{ env.UV_PYTHON }}
|
||||
|
||||
- name: Set up uv
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Install dependencies
|
||||
timeout-minutes: 20
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
|
||||
|
||||
- name: Install redis-server 7.2.16
|
||||
timeout-minutes: 3
|
||||
run: |
|
||||
mkdir -p "$RUNNER_TEMP/redis-bin"
|
||||
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
|
||||
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
|
||||
docker rm "$cid"
|
||||
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
|
||||
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
|
||||
|
||||
- name: Install MCP SDK1 peer
|
||||
if: matrix.suite == 'mcp'
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
uv venv --python 3.12 .venv-mcp-peer
|
||||
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
|
||||
echo "MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Node.js
|
||||
if: matrix.suite == 'browser'
|
||||
timeout-minutes: 5
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
cache: "npm"
|
||||
cache-dependency-path: |
|
||||
ui/litellm-dashboard/package-lock.json
|
||||
tests/e2e/ui/package-lock.json
|
||||
|
||||
- name: Install browser dependencies and build the dashboard
|
||||
if: matrix.suite == 'browser'
|
||||
timeout-minutes: 20
|
||||
run: |
|
||||
(cd ui/litellm-dashboard && npm ci)
|
||||
(cd tests/e2e/ui && npm ci)
|
||||
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive \
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js install-deps chromium
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js install chromium
|
||||
(cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build)
|
||||
|
||||
- name: Run integration contracts
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
SUITE: ${{ matrix.suite }}
|
||||
INTEGRATION_POSTGRES_CONTAINER: ${{ job.services.postgres.id }}
|
||||
INTEGRATION_REDIS_CONTAINER: ${{ job.services.redis.id }}
|
||||
LITELLM_LICENSE: ${{ secrets.LITELLM_LICENSE }}
|
||||
run: |
|
||||
case "$SUITE" in
|
||||
*-replica) bash .circleci/scripts/run_integration.sh "${SUITE%-replica}" replica ;;
|
||||
*) bash .circleci/scripts/run_integration.sh "$SUITE" ;;
|
||||
esac
|
||||
|
||||
- name: Verify every case ran exactly once
|
||||
if: ${{ !cancelled() && matrix.suite != 'browser' }}
|
||||
timeout-minutes: 2
|
||||
run: |
|
||||
record="$(find test-results -mindepth 2 -maxdepth 2 -name execution.json -print -quit)"
|
||||
test -n "$record"
|
||||
mkdir -p "$RUNNER_TEMP/shards/0"
|
||||
cp "$record" "$RUNNER_TEMP/shards/0/execution.json"
|
||||
.venv/bin/python -m tests.integration._support.sharding "$RUNNER_TEMP/shards" --shards 1
|
||||
|
||||
- name: Report MCP coverage
|
||||
if: ${{ !cancelled() && matrix.suite == 'mcp' }}
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
coverage_dir=test-results/integration-mcp/coverage
|
||||
test -f "$coverage_dir/coverage.txt" && exit 0
|
||||
export COVERAGE_FILE="$PWD/$coverage_dir/data"
|
||||
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
|
||||
.venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml > "$coverage_dir/coverage.txt"
|
||||
.venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml -d "$coverage_dir/html"
|
||||
tail -n 1 "$coverage_dir/coverage.txt"
|
||||
|
||||
- name: Check results for the license
|
||||
id: secret-scan
|
||||
if: ${{ !cancelled() }}
|
||||
timeout-minutes: 5
|
||||
env:
|
||||
SCANNED_SECRET: ${{ secrets.LITELLM_LICENSE }}
|
||||
run: python .github/scripts/assert_secret_absent.py test-results
|
||||
|
||||
- name: Upload results
|
||||
if: ${{ !cancelled() && steps.secret-scan.outcome == 'success' }}
|
||||
timeout-minutes: 10
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: integration-${{ matrix.suite }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: test-results/
|
||||
retention-days: 14
|
||||
7
.github/workflows/test-linting.yml
vendored
7
.github/workflows/test-linting.yml
vendored
|
|
@ -376,10 +376,10 @@ jobs:
|
|||
exit 0
|
||||
fi
|
||||
files="$(git diff --name-only "$base" HEAD)"
|
||||
if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
|
||||
if grep -Eq '^(litellm/|enterprise/|(pyproject\.toml|uv\.lock)$|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
|
||||
echo "relevant=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "No proxy, types or generator changes in this pull request, nothing to verify."
|
||||
echo "No litellm, enterprise, dependency or generator changes in this pull request, nothing to verify."
|
||||
echo "relevant=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
|
|
@ -561,6 +561,9 @@ jobs:
|
|||
- name: router_code_coverage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
|
||||
|
||||
- name: router_surface_check
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/router_surface_check.py
|
||||
|
||||
- name: test_chat_completion_imports
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py
|
||||
|
||||
|
|
|
|||
7
.github/workflows/test-merge-smoke.yml
vendored
7
.github/workflows/test-merge-smoke.yml
vendored
|
|
@ -132,7 +132,7 @@ jobs:
|
|||
windows-sdk:
|
||||
name: windows-sdk
|
||||
runs-on: windows-2025
|
||||
timeout-minutes: 10
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
defaults:
|
||||
|
|
@ -154,8 +154,11 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up the Rust bridge
|
||||
uses: ./.github/actions/rust-bridge
|
||||
|
||||
- name: Install dependencies
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11
|
||||
|
||||
- name: Run the Windows tests
|
||||
run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests
|
||||
|
|
|
|||
2
.github/workflows/test-rust.yml
vendored
2
.github/workflows/test-rust.yml
vendored
|
|
@ -4,6 +4,7 @@ on:
|
|||
push:
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- "model_prices_and_context_window.json"
|
||||
- "litellm/rust_bridge/**"
|
||||
- "scripts/generate_trace_types.py"
|
||||
- "scripts/trace_codegen/**"
|
||||
|
|
@ -38,6 +39,7 @@ on:
|
|||
- "litellm_**"
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- "model_prices_and_context_window.json"
|
||||
- "litellm/rust_bridge/**"
|
||||
- "scripts/generate_trace_types.py"
|
||||
- "scripts/trace_codegen/**"
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ on:
|
|||
paths:
|
||||
- "terraform/provider/**"
|
||||
- "litellm/proxy/**"
|
||||
- "enterprise/**"
|
||||
- ".github/workflows/test-terraform-provider.yml"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
16
.github/workflows/test-unit.yml
vendored
16
.github/workflows/test-unit.yml
vendored
|
|
@ -42,7 +42,7 @@ jobs:
|
|||
id: rust-bridge-key
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'pyproject.toml') }}
|
||||
RUST_BRIDGE_CACHE_KEY: ${{ runner.os }}-rust-bridge-editable-${{ hashFiles('litellm-rust/Cargo.lock', 'litellm-rust/Cargo.toml', 'litellm-rust/crates/**', 'model_prices_and_context_window.json', 'pyproject.toml') }}
|
||||
run: echo "key=$RUST_BRIDGE_CACHE_KEY" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Set up Python
|
||||
|
|
@ -137,6 +137,7 @@ jobs:
|
|||
include:
|
||||
- shard: core-utils
|
||||
test-path: |-
|
||||
tests/unit/accounting
|
||||
tests/unit/batch_completion
|
||||
tests/unit/decisions
|
||||
tests/unit/extraction
|
||||
|
|
@ -235,6 +236,7 @@ jobs:
|
|||
tests/unit/llms/azure_ai
|
||||
tests/unit/llms/base_llm
|
||||
tests/unit/llms/baseten
|
||||
tests/unit/llms/bigmodel
|
||||
tests/unit/llms/black_forest_labs
|
||||
tests/unit/llms/brave
|
||||
tests/unit/llms/bytez
|
||||
|
|
@ -244,6 +246,7 @@ jobs:
|
|||
tests/unit/llms/clarifai
|
||||
tests/unit/llms/claude_code
|
||||
tests/unit/llms/cloudflare
|
||||
tests/unit/llms/codestral
|
||||
tests/unit/llms/codex
|
||||
tests/unit/llms/cohere
|
||||
tests/unit/llms/cometapi
|
||||
|
|
@ -319,6 +322,7 @@ jobs:
|
|||
tests/unit/llms/openai_like
|
||||
tests/unit/llms/opencode
|
||||
tests/unit/llms/opensandbox
|
||||
tests/unit/llms/opensearch
|
||||
tests/unit/llms/openrouter
|
||||
tests/unit/llms/ovhcloud
|
||||
tests/unit/llms/parallel_ai
|
||||
|
|
@ -1306,6 +1310,8 @@ jobs:
|
|||
done
|
||||
openai-sdk-compat:
|
||||
name: OpenAI ${{ matrix.openai-version }} compatibility
|
||||
needs: rust-bridge
|
||||
if: ${{ !cancelled() }}
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
|
@ -1342,10 +1348,18 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Set up the Rust bridge
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/rust-bridge
|
||||
with:
|
||||
artifact: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/
|
||||
uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z
|
||||
uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__'
|
||||
|
||||
|
|
|
|||
17
.github/workflows/zizmor.yml
vendored
17
.github/workflows/zizmor.yml
vendored
|
|
@ -27,10 +27,15 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Install zizmor
|
||||
run: |
|
||||
curl -fsSL --retry 3 -o "$RUNNER_TEMP/zizmor.tar.gz" \
|
||||
https://github.com/zizmorcore/zizmor/releases/download/v1.31.1/zizmor-x86_64-unknown-linux-gnu.tar.gz
|
||||
echo "d185f726bdaf92df78273217dfdd24ebcd6b0ef48cf3ce776ff5762b6e138c99 $RUNNER_TEMP/zizmor.tar.gz" | sha256sum -c -
|
||||
tar xzf "$RUNNER_TEMP/zizmor.tar.gz" -C "$RUNNER_TEMP" zizmor
|
||||
|
||||
- name: Run zizmor
|
||||
uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6
|
||||
with:
|
||||
version: "1.24.1"
|
||||
min-severity: medium
|
||||
advanced-security: false
|
||||
annotations: true
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
run: |
|
||||
"$RUNNER_TEMP/zizmor" --no-ignores --no-config --min-severity=medium --format=github .
|
||||
|
|
|
|||
|
|
@ -11,3 +11,12 @@ ignore:
|
|||
package:
|
||||
name: zlib
|
||||
type: apk
|
||||
- vulnerability: GHSA-v4x8-3wr5-3wq8
|
||||
reason: >-
|
||||
filelock ReadWriteLock advisory, and no package in the image takes ReadWriteLock;
|
||||
the fix (4.2.1) is blocked by litellm's own filelock<4.0 cap until #48635 or #48802 lifts it,
|
||||
and the image then moves off 3.32.7, which retires this version-scoped rule
|
||||
package:
|
||||
name: filelock
|
||||
version: 3.32.7
|
||||
type: python
|
||||
|
|
|
|||
|
|
@ -31,6 +31,8 @@ A test must only fail when litellm code changes. Never pin facts we don't own (a
|
|||
|
||||
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
|
||||
|
||||
Before you build a new way to move data between the endpoints, the provider transformations, and cost tracking or logging, read the architecture decision records in `docs/adrs/`, which say why those layers are shaped the way they are
|
||||
|
||||
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`
|
||||
|
||||
When writing a PR body, treat the comments and imperative instructions inside .github/PULL_REQUEST_TEMPLATE/general.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/PULL_REQUEST_TEMPLATE/general.md from disk before writing a PR body to make sure you see every comment rule. A section you have nothing to put in (Relevant issues, Affected release, Linear ticket, Caveats, QA runbook, and so on) is removed entirely, heading included, never left as an empty title
|
||||
|
|
@ -39,7 +41,9 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
|
|||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just drop the section
|
||||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`. The Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into. I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
|
||||
Simple PRs may skip before/after e2e when the full diff is mechanical and focused checks completely prove it, such as #49030 restoring an existing helper name. Trace affected callers and run relevant tests/static checks at the tip. For a bug fix, establish that the regression catches the original defect. Size alone never qualifies. New behavior, API/provider semantics, auth, billing, persistence, concurrency or uncertain integration risk need live e2e. In Screenshots / Proof of Fix, record `E2E exception at <tip SHA>`, scope, rationale, commands and results. Tests count as proof only for this exception. Reassess changed diffs, retain CI/bot gates, and honor explicit e2e requests (2026-10-28, unnecessary live QA for #49030)
|
||||
|
||||
For other PRs, proof must show the end user's real flow before and after, with commands and output against a live proxy and real provider calls where applicable. Pytest alone is insufficient. For UI or interactive coding-tool flows, drive that surface and embed before/after screenshots, with URLs, clicks and input fields a reviewer can reproduce
|
||||
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
|
|
@ -105,6 +109,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega
|
|||
- No file sprawl: deliberate file and folder structure
|
||||
- Standard over hand-rolled: use the official SDK or a library where one exists. Where none does, follow industry standards instead of inventing local conventions
|
||||
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module. If none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
|
||||
- Current default for new model names: logic that picks behavior from a model-name table (tokenizer, message framing, supported params) gives a provider's model the table does not list yet that provider's current behavior, never its oldest, since providers ship names faster than any table is updated
|
||||
|
||||
Follow conventional commits for commit names and PR titles
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
|
||||
This document helps contributors understand where to make changes in LiteLLM.
|
||||
|
||||
For why a layer is shaped the way it is, read the architecture decision records in [`docs/adrs/`](docs/adrs/AGENTS.md) (Python SDK and proxy) and [`litellm-rust/docs/adrs/`](litellm-rust/docs/adrs/AGENTS.md) (Rust core)
|
||||
|
||||
---
|
||||
|
||||
## How It Works
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Here are the core requirements for any PR submitted to LiteLLM:
|
|||
|
||||
- [ ] **Sign the Contributor License Agreement (CLA)** - [see details](#contributor-license-agreement-cla)
|
||||
- [ ] **Keep scope isolated** - Your changes should address 1 specific problem at a time
|
||||
- [ ] **Prove it works end to end** - If your change affects what the proxy or SDK does, paste the commands you ran against a live proxy and their output in the PR - [see details](#proving-your-change-works-end-to-end)
|
||||
|
||||
#### Proxy (Backend) PRs
|
||||
|
||||
|
|
@ -248,6 +249,41 @@ If `make test-unit` fails:
|
|||
- **Test edge cases**: Don't just test the happy path
|
||||
- **Update documentation**: If you change APIs, update docs
|
||||
|
||||
## Proving Your Change Works End to End
|
||||
|
||||
Tests are required, but they do not prove that your change works for a user. A PR that changes what the proxy or SDK does shows a real run in its "Screenshots / Proof of Fix" section, meaning the exact commands you sent to a live proxy and the output that came back, against real provider APIs and not mocks. `pytest` output does not count as proof, because a passing test only shows that the code does what its mocks expect. A PR that changes only docs or CI shows the rendered page or the CI run instead
|
||||
|
||||
A good proof looks like this:
|
||||
|
||||
1. Install the proxy dependencies with `make install-proxy-dev`, then start the proxy from your branch with a config that has the model you changed, `uv run litellm --config your_config.yaml --detailed_debug`
|
||||
2. Send the request a user would send with `curl`, and paste both the command and the response. When cost or routing is involved, include the response headers, since `x-litellm-response-cost` is what the user sees
|
||||
3. For a bug fix, do this twice, once at the merge base with BerriAI/litellm's main (`git merge-base origin/main HEAD`, with your remote for BerriAI/litellm in place of `origin` in a fork clone) to show the bug and once at your branch's latest commit to show the fix. Name both commit hashes
|
||||
4. If your change is reachable from more than one endpoint (`/v1/chat/completions`, `/v1/responses`, `/v1/messages`), show each one
|
||||
5. If the request can stream, show a `"stream": true` run too, because streaming builds usage and cost on a different path. A stream sends its headers before its usage exists, so show the `usage.cost` on the last chunk, which `/v1/chat/completions` sends when the request has `"stream_options": {"include_usage": true}` and other endpoints send with `include_cost_in_streaming_usage: true` under `litellm_settings`, or, when the proxy has a database, the request's row on the Admin UI's Logs page or `GET /spend/logs`
|
||||
6. For UI changes, add before and after screenshots and name the page
|
||||
7. If users reach your change through a coding tool such as Claude Code or Codex, drive that tool the way a user does and add before and after screenshots of it, because a `curl` or a headless run alone does not show what the user sees
|
||||
|
||||
```bash
|
||||
curl -sD - http://localhost:4000/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "your-model", "messages": [{"role": "user", "content": "hi"}]}'
|
||||
```
|
||||
|
||||
When the provider reports its own usage or spend numbers, compare the proxy's numbers against them in the proof
|
||||
|
||||
## Contributing an LLM Provider Integration
|
||||
|
||||
Provider work (a new provider, a new parameter, a new usage or cost field) is the most common kind of contribution, and it is where PRs get reworked most often. Follow these rules
|
||||
|
||||
**Find the existing machinery before you add plumbing.** LiteLLM already has generic paths that move data between the endpoints, the provider transformations, and cost tracking. A second, hand-rolled path for the same data is a common reason a provider PR gets rewritten before merge. [`ARCHITECTURE.md`](ARCHITECTURE.md) shows where the layers live, and the decision records in [`docs/adrs/`](docs/adrs/AGENTS.md) explain why they are shaped that way. For example, [the ADR on provider usage fields](docs/adrs/adr_001_provider_usage_fields_ride_as_usage_extras.md) covers how provider-specific usage fields reach cost tracking. If you cannot find a mechanism for what you need, say so in the PR or ask in [#pr-review on Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA) before you build one
|
||||
|
||||
**Bill a new provider usage field from the usage object.** Leave the field on the usage object as an extra and read it in that provider's `cost_calculator.py`. A provider without one gets a new `cost_calculator.py` and an arm for it in `cost_per_token` (`litellm/cost_calculator.py`), which picks the calculator by provider. If the provider's transform builds `Usage` from named fields, as the Anthropic, Gemini, and Bedrock Converse ones do, pass the new field there too. Streamed requests rebuild usage in `ChunkProcessor.calculate_usage` (`litellm/litellm_core_utils/streaming_chunk_builder_utils.py`), which does not copy extras yet the way the Responses-to-chat bridge does (minus a short exclusion list), so until it does, bill the field on streams by naming it there, as Anthropic's `inference_geo` is. A provider that reports its own charge for the request, as OpenRouter does in `usage.cost`, is billed at that charge from `_hidden_params` before any calculator runs, unless the deployment sets its own prices. When the field counts a built-in tool such as web search, add the provider to `get_cost_for_web_search_request` in `litellm/llms/__init__.py`, and add a gate to `litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py` only when no existing gate fires. Then show the field's cost in your proof on every endpoint it reaches, streaming included
|
||||
|
||||
**Do not change what a caller sees to make internals easier.** Each endpoint keeps the schema of the API it emulates. A `/v1/responses` caller gets `input_tokens` and a `/v1/chat/completions` caller gets `prompt_tokens`, whatever the provider sent and whatever cost tracking needs. A provider usage field that shows up as an extra key is fine, since it only adds a key, but renaming, removing, or reshaping a key the endpoint defines is not. Reshaping a public response to feed an internal consumer breaks every user of that endpoint
|
||||
|
||||
If your change makes an architecture decision, such as a new cross-cutting mechanism or a deliberate difference from a provider's API shape, add an ADR to [`docs/adrs/`](docs/adrs/AGENTS.md) in the same PR
|
||||
|
||||
## Building and Running Locally
|
||||
|
||||
### LiteLLM Proxy Server
|
||||
|
|
|
|||
51
README.md
51
README.md
|
|
@ -141,12 +141,56 @@ response = client.chat.completions.create(
|
|||
)
|
||||
```
|
||||
|
||||
[**Docs: LLM Providers**](https://docs.litellm.ai/docs/providers)
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Auto Routing</b> - Choose a model for each prompt (Python SDK + AI Gateway)</summary>
|
||||
|
||||
### Python SDK
|
||||
|
||||
[**Docs: Autorouter SDK**](https://docs.litellm.ai/docs/autorouter_sdk)
|
||||
|
||||
Install with `uv add litellm` and set `OPENAI_API_KEY` for the models below.
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from litellm import Router
|
||||
|
||||
router = Router(model_list=[
|
||||
{"model_name": "fast", "litellm_params": {"model": "openai/gpt-6-luna"}},
|
||||
{"model_name": "capable", "litellm_params": {"model": "openai/gpt-6.1-sol"}},
|
||||
{"model_name": "auto", "litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"classifier_type": "heuristic",
|
||||
"tiers": {"SIMPLE": "fast", "MEDIUM": "fast",
|
||||
"COMPLEX": "capable", "REASONING": "capable"},
|
||||
"return_raw_model_name": True,
|
||||
},
|
||||
"complexity_router_default_model": "capable",
|
||||
}},
|
||||
])
|
||||
response = asyncio.run(router.acompletion(
|
||||
model="auto", messages=[{"role": "user", "content": "What is 2+2?"}],
|
||||
))
|
||||
print(response.model, response.choices[0].message.content)
|
||||
```
|
||||
|
||||
### AI Gateway (Proxy Server)
|
||||
|
||||
[**Create an Auto Router**](https://docs.litellm.ai/docs/auto_router/setup), then call its model name. Replace `auto` and the URL below with yours, and set `LITELLM_API_KEY` to your virtual key.
|
||||
|
||||
```shell
|
||||
curl http://localhost:4000/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LITELLM_API_KEY" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"auto","messages":[{"role":"user","content":"What is 2+2?"}]}'
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Agents</b> - Invoke A2A Agents (Python SDK + AI Gateway)</summary>
|
||||
<summary><b>A2A</b> - Invoke A2A Agents (Python SDK + AI Gateway)</summary>
|
||||
|
||||
[**Supported Providers**](https://docs.litellm.ai/docs/a2a#add-a2a-agents) - LangGraph, Vertex AI Agent Engine, Azure AI Foundry, Bedrock AgentCore, Pydantic AI
|
||||
|
||||
|
|
@ -427,6 +471,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
|
|||
| [Watsonx Text (`watsonx_text`)](https://docs.litellm.ai/docs/providers/watsonx) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [xAI (`xai`)](https://docs.litellm.ai/docs/providers/xai) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Xinference (`xinference`)](https://docs.litellm.ai/docs/providers/xinference) | | | | ✅ | | | | | | |
|
||||
| [ZeroGPU (`zerogpu`)](https://docs.litellm.ai/docs/providers/zerogpu) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
||||
[**Read the Docs**](https://docs.litellm.ai/docs/)
|
||||
|
||||
|
|
@ -441,7 +486,7 @@ You can use LiteLLM through either the Proxy Server or Python SDK. Both give you
|
|||
<tr>
|
||||
<th style={{width: '14%'}}></th>
|
||||
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/simple_proxy">LiteLLM AI Gateway</a></strong></th>
|
||||
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/">LiteLLM Python SDK</a></strong></th>
|
||||
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/python_sdk">LiteLLM Python SDK</a></strong></th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ ENV UV_PYTHON_DOWNLOADS=0
|
|||
COPY pyproject.toml README.md rust-toolchain.toml ./
|
||||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
COPY model_prices_and_context_window.json ./
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
|
|
|
|||
|
|
@ -1,37 +1,67 @@
|
|||
"""Guard the cost map on pull requests.
|
||||
|
||||
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.
|
||||
Every pull request whose diff against its merge base touches one of the two cost map files or the pins file gets the
|
||||
file checks: the files parse, the backup copy matches the root file, and every capability pinned in
|
||||
ci_cd/cost_map_pins.json (a value verified against a live provider call where the provider's model listing says
|
||||
otherwise) still holds in the tree that merging the pull request into its base branch produces, so a branch cut before
|
||||
a pin was added cannot lower it, and a pull request that deletes the pins file fails too. 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 the two cost map files and may
|
||||
only add or update models, so a sync can never lower a pinned capability or edit the pins. Price sync pull requests
|
||||
(branches named litellm-providers/price-sync*) and pull requests a bot account opened (COST_MAP_GUARD_AUTHOR_TYPE is
|
||||
Bot) get the human checks and may not touch the pins file either.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
|
||||
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
|
||||
PINS_PATH: Final = "ci_cd/cost_map_pins.json"
|
||||
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH)
|
||||
CHECKED_PATHS: Final = (*GUARDED_PATHS, PINS_PATH)
|
||||
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
|
||||
PRICE_SYNC_BRANCH_PREFIX: Final = "litellm-providers/price-sync"
|
||||
TARGET_ENV: Final = "COST_MAP_GUARD_TARGET"
|
||||
AUTHOR_TYPE_ENV: Final = "COST_MAP_GUARD_AUTHOR_TYPE"
|
||||
BOT_AUTHOR_TYPE: Final = "Bot"
|
||||
MERGE_CONFLICT_EXIT_CODE: Final = 1
|
||||
SPECIAL_ROOT_KEYS: Final = frozenset({"sample_spec", "fallback_generalizations"})
|
||||
EVALUATION_OUTPUT_PRICE_SOURCES: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
CostMap = dict[str, object]
|
||||
_MISSING: Final = object()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Snapshot:
|
||||
cost_map: str
|
||||
backup: str
|
||||
pins: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MergeFailed:
|
||||
stderr: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Pin:
|
||||
key: str
|
||||
field: str
|
||||
value: object
|
||||
checked: str
|
||||
source: str
|
||||
|
||||
|
||||
def _parse_object(text: str, path: str) -> CostMap | str:
|
||||
|
|
@ -91,6 +121,48 @@ def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]:
|
|||
return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)}
|
||||
|
||||
|
||||
def _pin(key: str, field: str, spec: object) -> Pin | str:
|
||||
malformed: Final = f"{PINS_PATH}: {key}.{field} needs value, checked, and source"
|
||||
if not isinstance(spec, dict) or "value" not in spec:
|
||||
return malformed
|
||||
checked: Final = spec.get("checked")
|
||||
source: Final = spec.get("source")
|
||||
if not isinstance(checked, str) or not isinstance(source, str) or not checked or not source:
|
||||
return malformed
|
||||
return Pin(key, field, spec["value"], checked, source)
|
||||
|
||||
|
||||
def _pins(parsed: CostMap) -> Iterator[Pin | str]:
|
||||
for key, fields in parsed.items():
|
||||
if not isinstance(fields, dict):
|
||||
yield f"{PINS_PATH}: {key} must map field names to pins"
|
||||
continue
|
||||
for field, spec in fields.items():
|
||||
yield _pin(key, field, spec)
|
||||
|
||||
|
||||
def _pin_violations(parsed: CostMap, entries: Mapping[str, Mapping[str, object]]) -> Iterator[str]:
|
||||
for pin in _pins(parsed):
|
||||
if isinstance(pin, str):
|
||||
yield pin
|
||||
elif pin.key not in entries:
|
||||
yield f"{PINS_PATH} pins {pin.key}.{pin.field} but {COST_MAP_PATH} has no {pin.key} entry"
|
||||
elif entries[pin.key].get(pin.field, _MISSING) != pin.value:
|
||||
yield (
|
||||
f"{COST_MAP_PATH}: {pin.key}.{pin.field} must stay {json.dumps(pin.value)} "
|
||||
f"(verified {pin.checked}: {pin.source}); change {PINS_PATH} with new evidence first"
|
||||
)
|
||||
|
||||
|
||||
def pin_failures(pins_text: str, head_map: CostMap) -> tuple[str, ...]:
|
||||
if not pins_text:
|
||||
return (f"{PINS_PATH} is missing; restore it, its pins were verified against live provider calls",)
|
||||
parsed: Final = _parse_object(pins_text, PINS_PATH)
|
||||
if isinstance(parsed, str):
|
||||
return (parsed,)
|
||||
return tuple(_pin_violations(parsed, _entries(head_map)))
|
||||
|
||||
|
||||
def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]:
|
||||
base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH)
|
||||
if isinstance(base_map, str):
|
||||
|
|
@ -120,8 +192,18 @@ def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str
|
|||
)
|
||||
|
||||
|
||||
def pin_author_failures(head_ref: str, author_type: str, changed_files: Sequence[str]) -> tuple[str, ...]:
|
||||
if PINS_PATH not in changed_files:
|
||||
return ()
|
||||
if head_ref.startswith(PRICE_SYNC_BRANCH_PREFIX):
|
||||
return (f"price sync PRs may not change {PINS_PATH}; a pin moves only in a human PR with new evidence",)
|
||||
if author_type == BOT_AUTHOR_TYPE:
|
||||
return (f"PRs opened by a bot may not change {PINS_PATH}; a pin moves only in a human PR with new evidence",)
|
||||
return ()
|
||||
|
||||
|
||||
def touches_cost_map(changed_files: Sequence[str]) -> bool:
|
||||
return any(path in GUARDED_PATHS for path in changed_files)
|
||||
return any(path in CHECKED_PATHS for path in changed_files)
|
||||
|
||||
|
||||
def contract_for(bot: bool, changed_files: Sequence[str]) -> str:
|
||||
|
|
@ -130,9 +212,21 @@ def contract_for(bot: bool, changed_files: Sequence[str]) -> str:
|
|||
return "human PR, file checks only" if touches_cost_map(changed_files) else "human PR, cost map untouched"
|
||||
|
||||
|
||||
def _merged_pin_failures(base: Snapshot, merged: Snapshot | None) -> tuple[str, ...]:
|
||||
if merged is None:
|
||||
return ("the pull request does not merge cleanly into its base branch; merge the base branch in first",)
|
||||
if not merged.pins and not base.pins:
|
||||
return ()
|
||||
merged_map: Final = _parse_object(merged.cost_map, COST_MAP_PATH)
|
||||
if isinstance(merged_map, str):
|
||||
return (f"after merging into the base branch, {merged_map}",)
|
||||
return pin_failures(merged.pins, merged_map)
|
||||
|
||||
|
||||
def guard_failures(
|
||||
base: Snapshot,
|
||||
head: Snapshot,
|
||||
merged: Snapshot | None,
|
||||
changed_files: Sequence[str],
|
||||
bot: bool,
|
||||
*,
|
||||
|
|
@ -147,6 +241,7 @@ def guard_failures(
|
|||
*_file_failures(
|
||||
head, head_map, evaluation_output_price_sources=evaluation_output_price_sources
|
||||
),
|
||||
*_merged_pin_failures(base, merged),
|
||||
*(_bot_failures(base, head_map, changed_files) if bot else ()),
|
||||
)
|
||||
|
||||
|
|
@ -157,7 +252,22 @@ def _git(*args: str) -> str | None:
|
|||
|
||||
|
||||
def snapshot(revision: str) -> Snapshot:
|
||||
return Snapshot(*(_git("show", f"{revision}:{path}") or "" for path in GUARDED_PATHS))
|
||||
return Snapshot(*(_git("show", f"{revision}:{path}") or "" for path in CHECKED_PATHS))
|
||||
|
||||
|
||||
def merged_tree(merge_base: str, target: str, head: str) -> str | MergeFailed | None:
|
||||
result: Final = subprocess.run(
|
||||
("git", "merge-tree", "--write-tree", f"--merge-base={merge_base}", target, head),
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
tree: Final = result.stdout.split("\n", 1)[0]
|
||||
if result.returncode == 0:
|
||||
return tree
|
||||
if result.returncode == MERGE_CONFLICT_EXIT_CODE and tree:
|
||||
return None
|
||||
return MergeFailed(result.stderr.strip())
|
||||
|
||||
|
||||
def changed_files(base: str, head: str) -> tuple[str, ...] | None:
|
||||
|
|
@ -169,6 +279,13 @@ 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 commit of the pull request")
|
||||
env_target: Final = os.environ.get(TARGET_ENV)
|
||||
parser.add_argument(
|
||||
"--target",
|
||||
default=env_target,
|
||||
required=env_target is None,
|
||||
help=f"tip of the base branch the pull request merges into, defaulting to ${TARGET_ENV}",
|
||||
)
|
||||
parser.add_argument("--head-ref", required=True, help="head branch name of the pull request")
|
||||
args: Final = parser.parse_args(argv)
|
||||
bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX)
|
||||
|
|
@ -176,7 +293,19 @@ def main(argv: Sequence[str]) -> int:
|
|||
if changed is None:
|
||||
print(f"cost map guard failed: git diff {args.base} {args.head} failed, so the changed files are unknown")
|
||||
return 1
|
||||
failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed, bot)
|
||||
if not bot and not touches_cost_map(changed):
|
||||
print(f"cost map guard passed ({contract_for(bot, changed)})")
|
||||
return 0
|
||||
merged: Final = merged_tree(args.base, args.target, args.head)
|
||||
if isinstance(merged, MergeFailed):
|
||||
print(f"cost map guard failed: git merge-tree into {args.target} failed: {merged.stderr}")
|
||||
return 1
|
||||
failures: Final = (
|
||||
*guard_failures(
|
||||
snapshot(args.base), snapshot(args.head), None if merged is None else snapshot(merged), changed, bot
|
||||
),
|
||||
*pin_author_failures(args.head_ref, os.environ.get(AUTHOR_TYPE_ENV, ""), changed),
|
||||
)
|
||||
contract: Final = contract_for(bot, changed)
|
||||
if failures:
|
||||
print(f"cost map guard failed ({contract}):")
|
||||
|
|
|
|||
16
ci_cd/cost_map_pins.json
Normal file
16
ci_cd/cost_map_pins.json
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"fireworks_ai/accounts/fireworks/models/minimax-m3": {
|
||||
"supports_vision": {
|
||||
"value": true,
|
||||
"checked": "2026-09-26",
|
||||
"source": "POST https://api.fireworks.ai/inference/v1/chat/completions read a 512x512 PNG back as digit=7, shape=circle, color=red while GET https://api.fireworks.ai/inference/v1/models still said supports_image_input false and GET https://api.fireworks.ai/v1/serverless/models?format=nested carried no supportsImageInput on the row"
|
||||
}
|
||||
},
|
||||
"fireworks_ai/minimax-m3": {
|
||||
"supports_vision": {
|
||||
"value": true,
|
||||
"checked": "2026-09-26",
|
||||
"source": "POST https://api.fireworks.ai/inference/v1/chat/completions read a 512x512 PNG back as digit=7, shape=circle, color=red while GET https://api.fireworks.ai/inference/v1/models still said supports_image_input false and GET https://api.fireworks.ai/v1/serverless/models?format=nested carried no supportsImageInput on the row"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -2,9 +2,12 @@
|
|||
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.
|
||||
so a path rule would put every migration on the owner. This gate reads the diff instead: a schema.prisma
|
||||
hunk counts when it adds or removes a line inside the `model LiteLLM_SpendLogs { ... }` block, a migration
|
||||
hunk counts when an added line names the table. Without such a hunk it passes at once, otherwise it passes
|
||||
only once the owner's latest review on the pull request is an approval that still covers every spend-log line
|
||||
in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never saw).
|
||||
Reviews do not trigger the workflow, so the check is re-run by hand after approving.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -13,8 +16,10 @@ import argparse
|
|||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
|
@ -25,7 +30,10 @@ WATCHED: Final = (
|
|||
"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)
|
||||
TABLE: Final = re.compile(r"(?<![A-Za-z0-9])LiteLLM_SpendLogs(?![A-Za-z0-9])", re.IGNORECASE)
|
||||
DECIDING_STATES: Final = frozenset({"APPROVED", "CHANGES_REQUESTED", "DISMISSED"})
|
||||
MODEL_START: Final = re.compile(r"^model\s+LiteLLM_SpendLogs\s*\{")
|
||||
HUNK: Final = re.compile(r"^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@[^\n]*\n", re.MULTILINE)
|
||||
|
||||
|
||||
class Reviewer(BaseModel):
|
||||
|
|
@ -37,31 +45,77 @@ class Review(BaseModel):
|
|||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
state: str
|
||||
user: Reviewer
|
||||
commit_id: str | None = None
|
||||
|
||||
|
||||
_REVIEWS: Final = TypeAdapter(list[Review])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Block:
|
||||
"""Lines of `model LiteLLM_SpendLogs { ... }` in one schema.prisma copy at the merge base and at the head."""
|
||||
|
||||
base: range
|
||||
head: range
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Hunk:
|
||||
removed: tuple[tuple[int, str], ...]
|
||||
added: tuple[tuple[int, str], ...]
|
||||
|
||||
|
||||
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 model_block(schema: str) -> range:
|
||||
"""1-based line range of `model LiteLLM_SpendLogs { ... }` including both braces, empty when the schema has none."""
|
||||
lines: Final = schema.splitlines()
|
||||
starts: Final = tuple(number for number, line in enumerate(lines, 1) if MODEL_START.match(line))
|
||||
if not starts:
|
||||
return range(0)
|
||||
ends: Final = tuple(number for number, line in enumerate(lines, 1) if number > starts[0] and line.startswith("}"))
|
||||
return range(starts[0], (ends[0] if ends else len(lines)) + 1)
|
||||
|
||||
|
||||
def spend_log_hits(diff: str) -> tuple[str, ...]:
|
||||
"""`path: added line` for every added line in a watched file that names the table."""
|
||||
def _hunks(section: str) -> tuple[Hunk, ...]:
|
||||
headers: Final = tuple(HUNK.finditer(section))
|
||||
ends: Final = tuple(header.start() for header in headers[1:]) + (len(section),)
|
||||
return tuple(_hunk(header, section[header.end() : end]) for header, end in zip(headers, ends))
|
||||
|
||||
|
||||
def _hunk(header: re.Match[str], body: str) -> Hunk:
|
||||
lines: Final = body.splitlines()
|
||||
return Hunk(
|
||||
removed=tuple(enumerate((line[1:] for line in lines if line.startswith("-")), int(header.group(1)))),
|
||||
added=tuple(enumerate((line[1:] for line in lines if line.startswith("+")), int(header.group(2)))),
|
||||
)
|
||||
|
||||
|
||||
def spend_log_hits(diff: str, blocks: Mapping[str, Block]) -> tuple[str, ...]:
|
||||
"""`path: +line` or `path: -line` for every changed line that touches the table.
|
||||
|
||||
A schema.prisma line counts when its number falls inside that copy's `blocks` entry (head for added lines,
|
||||
merge base for removed ones), a migration line when it is added and 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)
|
||||
paths: 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
|
||||
hit for path, section in paths if watched(path) for hit in _section_hits(path, _hunks(section), blocks)
|
||||
) # comprehension-ok: path then its hits
|
||||
|
||||
|
||||
def _section_hits(path: str, hunks: Sequence[Hunk], blocks: Mapping[str, Block]) -> tuple[str, ...]:
|
||||
if path not in blocks:
|
||||
return tuple(f"{path}: +{line.strip()}" for hunk in hunks for _, line in hunk.added if TABLE.search(line))
|
||||
block: Final = blocks[path]
|
||||
removed: Final = tuple(
|
||||
f"{path}: -{line.strip()}" for hunk in hunks for number, line in hunk.removed if number in block.base
|
||||
) # comprehension-ok: hunk then its lines
|
||||
added: Final = tuple(
|
||||
f"{path}: +{line.strip()}" for hunk in hunks for number, line in hunk.added if number in block.head
|
||||
) # comprehension-ok: hunk then its lines
|
||||
return removed + added
|
||||
|
||||
|
||||
def _path_of(section: str) -> str:
|
||||
|
|
@ -70,11 +124,33 @@ def _path_of(section: str) -> str:
|
|||
|
||||
|
||||
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)
|
||||
"""True when `login`'s most recent review that changes GitHub's approval state is an approval.
|
||||
|
||||
COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too.
|
||||
"""
|
||||
own: Final = tuple(review.state for review in reviews if review.user.login == login)
|
||||
states: Final = tuple(state for state in own if state in DECIDING_STATES)
|
||||
return bool(states) and states[-1] == "APPROVED"
|
||||
|
||||
|
||||
def latest_approval(reviews: Sequence[Review], login: str) -> Review | None:
|
||||
"""`login`'s most recent APPROVED review on the pull request, the one GitHub keeps in force after pushes."""
|
||||
own: Final = tuple(review for review in reviews if review.user.login == login and review.state == "APPROVED")
|
||||
return own[-1] if own else None
|
||||
|
||||
|
||||
def approval_covers(
|
||||
approval: Review, head: str, hits: Sequence[str], hits_at: Callable[[str], Sequence[str] | None]
|
||||
) -> bool:
|
||||
"""True when the approval is on the head, or the head adds no spend-log line the approved commit did not."""
|
||||
if approval.commit_id is None:
|
||||
return False
|
||||
if approval.commit_id == head:
|
||||
return True
|
||||
reviewed: Final = hits_at(approval.commit_id)
|
||||
return reviewed is not None and set(hits) <= set(reviewed)
|
||||
|
||||
|
||||
def added_diff(base: str, head: str) -> str:
|
||||
result: Final = subprocess.run(
|
||||
("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX),
|
||||
|
|
@ -85,6 +161,27 @@ def added_diff(base: str, head: str) -> str:
|
|||
return result.stdout
|
||||
|
||||
|
||||
def file_at(revision: str, path: str) -> str:
|
||||
"""Contents of `path` at `revision`, empty when the file does not exist there."""
|
||||
result: Final = subprocess.run(("git", "show", f"{revision}:{path}"), check=False, capture_output=True, text=True)
|
||||
return result.stdout if result.returncode == 0 else ""
|
||||
|
||||
|
||||
def schema_blocks(base: str, head: str) -> Mapping[str, Block]:
|
||||
return MappingProxyType(
|
||||
{path: Block(base=model_block(file_at(base, path)), head=model_block(file_at(head, path))) for path in WATCHED}
|
||||
)
|
||||
|
||||
|
||||
def fetched(commit: str) -> bool:
|
||||
result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False)
|
||||
return result.returncode == 0
|
||||
|
||||
|
||||
def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None:
|
||||
return spend_log_hits(added_diff(base, commit), schema_blocks(base, commit)) if fetched(commit) else None
|
||||
|
||||
|
||||
class Arguments(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
base: str
|
||||
|
|
@ -100,18 +197,28 @@ def main(argv: Sequence[str]) -> int:
|
|||
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))
|
||||
hits: Final = spend_log_hits(added_diff(args.base, args.head), schema_blocks(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")
|
||||
sys.stdout.write(
|
||||
"No change inside model LiteLLM_SpendLogs in schema.prisma and no added migration line names it; "
|
||||
"nothing to review.\n"
|
||||
)
|
||||
return 0
|
||||
if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver):
|
||||
reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes())
|
||||
approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None
|
||||
if approval is not None and approval_covers(
|
||||
approval, args.head, hits, lambda commit: reviewed_hits(args.base, commit)
|
||||
):
|
||||
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(
|
||||
reason: Final = (
|
||||
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"
|
||||
if approval is None
|
||||
else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, "
|
||||
f"a new approving review is needed:\n"
|
||||
)
|
||||
sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n")
|
||||
return 1
|
||||
|
||||
|
||||
|
|
|
|||
11
codecov.yaml
11
codecov.yaml
|
|
@ -15,9 +15,10 @@ ignore:
|
|||
flag_management:
|
||||
default_rules:
|
||||
carryforward: true
|
||||
# Dead flags no CI job uploads anymore: their carried-forward sessions were
|
||||
# measured against old revisions, and the stale line maps mark comment lines
|
||||
# of since-edited files as missed, sinking patch coverage on unrelated PRs.
|
||||
# Dead flags no CI job uploads anymore, and flags whose job is path-gated on
|
||||
# PRs: their carried-forward sessions were measured against old revisions, and
|
||||
# the stale line maps mark comment lines of since-edited files as missed,
|
||||
# sinking patch coverage on unrelated PRs.
|
||||
individual_flags:
|
||||
- name: proxy-mgmt-behavior
|
||||
carryforward: false
|
||||
|
|
@ -99,6 +100,10 @@ flag_management:
|
|||
carryforward: false
|
||||
- name: proxy-db-endpoints-and-responses
|
||||
carryforward: false
|
||||
- name: roi-postgres
|
||||
carryforward: false
|
||||
- name: redis-compat
|
||||
carryforward: false
|
||||
|
||||
component_management:
|
||||
individual_components:
|
||||
|
|
|
|||
5
docs/adrs/AGENTS.md
Normal file
5
docs/adrs/AGENTS.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
# Architecture decision records for the Python SDK and proxy
|
||||
|
||||
This folder records decisions about the Python SDK and proxy the way [`litellm-rust/docs/adrs/`](../../litellm-rust/docs/adrs/AGENTS.md) records them for the Rust core, and that file's process applies here unchanged: one short `adr_NNN_short_title.md` per decision that is hard to reverse or that a new contributor would reasonably question, with Context, Decision, Alternatives Considered, and Consequences, superseded by a new ADR instead of edited once accepted. Each folder numbers its own ADRs, so a reference to an ADR in the other folder names its path
|
||||
|
||||
Read the ADRs here before you build a new way to move data between the endpoints, the provider transformations, and cost tracking or logging. A PR that makes a decision of that kind adds its ADR in the same PR
|
||||
|
|
@ -0,0 +1,39 @@
|
|||
# ADR 001: Provider usage fields ride as extras on the normalized usage object
|
||||
|
||||
Status: Accepted
|
||||
|
||||
## Context
|
||||
|
||||
Providers report usage fields that no other provider has, such as xAI's `usage.server_side_tool_usage_details.web_search_calls` or Anthropic's `usage.inference_geo`, and cost tracking has to read them or that spend never reaches `LiteLLM_SpendLogs`. One provider response can be served on `/v1/chat/completions`, `/v1/responses`, and `/v1/messages`, streaming or not, so a field that only one of those paths carries is billed wrong on the others
|
||||
|
||||
[#30817](https://github.com/BerriAI/litellm/pull/30817) fixed unbilled xAI web searches. One draft did the copy in the xAI Responses transform to keep provider code out of shared files, and that swapped `/v1/responses` usage to the chat shape, so callers of that endpoint lost the Responses schema. The merged version keeps the schema, copies unknown usage fields in the shared Responses-to-chat bridge for every provider, and gates on the reported field without naming the provider
|
||||
|
||||
In October 2026, `Usage` and `ResponseAPIUsage` kept unknown usage keys as pydantic extras, and that bridge copied them. Paths that build usage from named fields kept only the fields they name: the Anthropic, Gemini, and Bedrock Converse transforms, the stream builder, and the bridges from chat usage to `/v1/responses` and `/v1/messages`
|
||||
|
||||
## Decision
|
||||
|
||||
Provider-only usage fields stay on the usage object as pydantic extras, and cost tracking reads them there, in that provider's cost calculator or in a shared gate that checks the field instead of the provider. An extra reaches callers as an added usage key. No transform renames, removes, or reshapes a usage field the endpoint's schema defines to feed cost tracking, though it may fill one LiteLLM's usage object already defines, such as `server_tool_use`
|
||||
|
||||
A shared path that rebuilds usage, such as the stream builder or a bridge between endpoint shapes, copies the extras instead of naming provider fields. A provider transform that builds usage from named fields names the new field. Neither is bypassed with a side channel. Older code that differs, such as the Azure branches in the shared tracker and the Gemini pricing tier read from `_hidden_params`, predates this ADR
|
||||
|
||||
## Alternatives Considered
|
||||
|
||||
**Reshape usage in the provider transform**, as a draft of #30817 did. It changes the schema that callers of that endpoint rely on, and every provider and endpoint needs its own copy
|
||||
|
||||
**A typed field on `ResponseAPIUsage` and `Usage` for each provider's quirk**. The types grow without bound and every bridge needs an update for each one. A field several providers report, like `server_tool_use`, does get one
|
||||
|
||||
**A side channel such as `_hidden_params`**. Cost calculators receive `Usage`, so a side channel makes a second source of truth for the same number
|
||||
|
||||
**Provider branches or provider field names in shared code**. They put one provider's logic in a file every provider shares, which is why #30817 replaced its xAI-only gate with one that checks the reported field
|
||||
|
||||
## Consequences
|
||||
|
||||
Benefits:
|
||||
|
||||
- Billing a new provider usage field touches that provider's calculator (plus a `cost_per_token` arm when the calculator is new), its transform when that names fields, and for a built-in tool at most one gate and one `get_cost_for_web_search_request` arm
|
||||
|
||||
Costs:
|
||||
|
||||
- Extras are untyped, so a reader validates them before use, as `_usage_reports_server_side_web_search_calls` does
|
||||
- The stream builder and the bridges from chat usage did not copy extras when this was written, so until they do, a field billed on those paths is named in them, and tests and proof runs cover each endpoint and `"stream": true`
|
||||
- This covers usage fields LiteLLM prices, not request-side params or a provider's own charge for the request, which OpenRouter and opencode Go report and `response_cost_calculator` reads from `_hidden_params` first
|
||||
|
|
@ -15,7 +15,10 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails._content_utils import iter_message_text
|
||||
from litellm.proxy.guardrails._content_utils import (
|
||||
is_non_conversational_request,
|
||||
iter_message_text,
|
||||
)
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
|
||||
_THRESHOLD_VALUE: Final = TypeAdapter(object)
|
||||
|
|
@ -94,7 +97,7 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger):
|
|||
|
||||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
data: dict[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: CallTypesLiteral,
|
||||
):
|
||||
|
|
@ -102,6 +105,9 @@ class ENTERPRISE_GoogleTextModeration(CustomLogger):
|
|||
- Calls Google's Text Moderation API
|
||||
- Rejects request if it fails safety check
|
||||
"""
|
||||
if is_non_conversational_request(call_type, data):
|
||||
verbose_proxy_logger.debug("Google text moderation: skipping non-conversational call type %s", call_type)
|
||||
return
|
||||
# Covers multimodal list content + Responses-API input.
|
||||
text = "".join(iter_message_text(data))
|
||||
if text:
|
||||
|
|
|
|||
|
|
@ -33,9 +33,8 @@ class EnterpriseCallbackControls:
|
|||
|
||||
try:
|
||||
disabled_callbacks = EnterpriseCallbackControls.get_disabled_callbacks(litellm_params, standard_callback_dynamic_params)
|
||||
verbose_logger.debug(f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}")
|
||||
verbose_logger.debug(f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}")
|
||||
if disabled_callbacks is not None:
|
||||
verbose_logger.debug(f"Checking if {callback} is in dynamically disabled callbacks: {disabled_callbacks}")
|
||||
#########################################################
|
||||
# premium user check
|
||||
#########################################################
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import effective_skip_tool_message_for_guardrail
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
|
@ -753,7 +754,12 @@ 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_tool_payload) + self._redact_prompt(data)
|
||||
inspected = walk_user_text(
|
||||
data,
|
||||
self.redact_text,
|
||||
self.redact_tool_payload,
|
||||
skip_tool_message=effective_skip_tool_message_for_guardrail(self),
|
||||
) + self._redact_prompt(data)
|
||||
|
||||
if inspected == 0:
|
||||
# Image-only, empty-text, and unsupported payloads inspected
|
||||
|
|
|
|||
|
|
@ -5,9 +5,9 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
|
|||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from datetime import datetime, timedelta
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ from litellm.constants import (
|
|||
from litellm.repositories.table_repositories import ManagedObjectRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.verification_token_repository import VerificationTokenRepository
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
|
@ -42,11 +43,14 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
def _error_names_resource(error: Exception, resource_id: str) -> bool:
|
||||
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
|
||||
tokens: Final[list[str]] = re.findall(r"""[^\s"'`<>(]+""", str(error))
|
||||
return any(token.rstrip(".,:;)") == resource_id for token in tokens)
|
||||
|
||||
|
||||
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
|
||||
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
|
||||
GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
|
||||
|
||||
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
|
||||
"completed",
|
||||
|
|
@ -91,6 +95,10 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
|
|||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _unique_jobs(jobs: Sequence[ManagedObjectRow]) -> tuple[ManagedObjectRow, ...]:
|
||||
return tuple({job.id: job for job in jobs}.values())
|
||||
|
||||
|
||||
class CheckBatchCost:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -98,6 +106,7 @@ class CheckBatchCost:
|
|||
prisma_client: "PrismaClient",
|
||||
llm_router: "Router",
|
||||
track_unmanaged_batch_cost: bool = False,
|
||||
clock: Callable[[], datetime] = get_utc_datetime,
|
||||
):
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
|
@ -106,13 +115,14 @@ class CheckBatchCost:
|
|||
self.prisma_client: PrismaClient = prisma_client
|
||||
self.llm_router: Router = llm_router
|
||||
self._track_unmanaged_batch_cost = track_unmanaged_batch_cost
|
||||
self._clock: Final = clock
|
||||
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost")
|
||||
# Cached after the first poll cycle. Once we know the column is absent we skip
|
||||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
self.batch_processed_support_confirmed: bool = False
|
||||
self._poll_cursor: tuple[datetime, str] | None = None
|
||||
self._poll_before: datetime = datetime.now(timezone.utc)
|
||||
self._poll_before: datetime = clock()
|
||||
self._arrival_cursor: tuple[datetime, str] | None = None
|
||||
self._arrival_before: datetime = self._poll_before
|
||||
self._poll_arrivals: bool = False
|
||||
|
|
@ -147,27 +157,6 @@ class CheckBatchCost:
|
|||
return
|
||||
self.batch_processed_support_confirmed = True
|
||||
|
||||
async def _cleanup_stale_managed_objects(self) -> None:
|
||||
"""
|
||||
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
|
||||
in non-terminal states as 'stale_expired'. These will never complete and
|
||||
should not be polled.
|
||||
"""
|
||||
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
result: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
)
|
||||
if result > 0:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: marked {result} stale managed objects "
|
||||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
|
||||
"""Query batch jobs without the batch_processed filter (for older schemas)."""
|
||||
return await self._find_jobs(
|
||||
|
|
@ -187,20 +176,43 @@ class CheckBatchCost:
|
|||
)
|
||||
|
||||
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
|
||||
arrivals: Final = self._poll_arrivals
|
||||
page: Final = await self._find_job_page(where, arrivals=arrivals)
|
||||
self._poll_arrivals = not arrivals if page else arrivals
|
||||
if page:
|
||||
return page
|
||||
return await self._find_job_page(where, arrivals=not arrivals)
|
||||
lead_arrivals: Final = self._poll_arrivals
|
||||
self._poll_arrivals = not lead_arrivals
|
||||
lead_share: Final = MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2
|
||||
lead: Final = await self._find_job_page(where, arrivals=lead_arrivals, take=lead_share)
|
||||
trail: Final = await self._find_job_page(
|
||||
where, arrivals=not lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead)
|
||||
)
|
||||
spill: Final = (
|
||||
await self._find_job_page(
|
||||
where, arrivals=lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) - len(trail)
|
||||
)
|
||||
if len(lead) == lead_share
|
||||
else ()
|
||||
)
|
||||
return _unique_jobs((*lead, *trail, *spill))
|
||||
|
||||
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
|
||||
async def _find_job_page(
|
||||
self, where: Mapping[str, object], *, arrivals: bool, take: int
|
||||
) -> Sequence[ManagedObjectRow]:
|
||||
if take == 0:
|
||||
return ()
|
||||
resumed_mid_lane: Final = (self._arrival_cursor if arrivals else self._poll_cursor) is not None
|
||||
page: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take)
|
||||
if len(page) == take or not resumed_mid_lane:
|
||||
return page
|
||||
wrapped: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take - len(page))
|
||||
return _unique_jobs((*page, *wrapped))
|
||||
|
||||
async def _find_job_page_from_cursor(
|
||||
self, where: Mapping[str, object], *, arrivals: bool, take: int
|
||||
) -> Sequence[ManagedObjectRow]:
|
||||
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
|
||||
if cursor is None:
|
||||
if arrivals:
|
||||
self._arrival_before = datetime.now(timezone.utc)
|
||||
self._arrival_before = self._clock()
|
||||
else:
|
||||
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
|
||||
self._poll_before = self._clock() - RECENT_BATCH_POLL_WINDOW
|
||||
before: Final = self._arrival_before if arrivals else self._poll_before
|
||||
page: Final = await _managed_object_table(self.prisma_client).find_many(
|
||||
where={
|
||||
|
|
@ -220,10 +232,10 @@ class CheckBatchCost:
|
|||
else {}
|
||||
),
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
take=take,
|
||||
order=[{"created_at": "asc"}, {"id": "asc"}],
|
||||
)
|
||||
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
|
||||
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == take else None
|
||||
if arrivals:
|
||||
self._arrival_cursor = next_cursor
|
||||
else:
|
||||
|
|
@ -231,11 +243,47 @@ class CheckBatchCost:
|
|||
return page
|
||||
|
||||
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
|
||||
return {
|
||||
"id": job.id,
|
||||
"updated_at": job.updated_at,
|
||||
**({"batch_processed": False} if self._has_batch_processed_column else {}),
|
||||
}
|
||||
if self._has_batch_processed_column:
|
||||
return {"id": job.id, "batch_processed": False}
|
||||
return {"id": job.id, "updated_at": job.updated_at}
|
||||
|
||||
def _is_past_give_up_age(self, job: ManagedObjectRow) -> bool:
|
||||
return job.created_at < self._clock() - GIVE_UP_POLLING_AFTER
|
||||
|
||||
async def _expire_if_old_and_unfinished(self, job: ManagedObjectRow) -> None:
|
||||
if job.status in TERMINAL_MANAGED_OBJECT_STATUSES or not self._is_past_give_up_age(job):
|
||||
return
|
||||
await self._expire_unfinished(job)
|
||||
|
||||
@staticmethod
|
||||
def _is_completed_with_output_id_pending(response: "LiteLLMBatch") -> bool:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_completed_batch_safe_to_retire,
|
||||
)
|
||||
|
||||
return (
|
||||
response.status in ("completed", "complete")
|
||||
and response.output_file_id is None
|
||||
and not _completed_batch_safe_to_retire(response)
|
||||
)
|
||||
|
||||
async def _expire_unfinished(self, job: ManagedObjectRow) -> None:
|
||||
try:
|
||||
expired: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where={
|
||||
**self._unprocessed_snapshot(job),
|
||||
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: failed to mark stale job {job.id} as stale_expired: {db_err}")
|
||||
return
|
||||
if expired > 0:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: job {job.id} is older than {GIVE_UP_POLLING_AFTER.days} days "
|
||||
"and still not finished, so it is marked stale_expired and no longer polled"
|
||||
)
|
||||
|
||||
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
|
||||
"""
|
||||
|
|
@ -350,21 +398,27 @@ class CheckBatchCost:
|
|||
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 a potentially recoverable row."""
|
||||
return self.llm_router.get_deployment(model_id=model_id) is not None
|
||||
return self.llm_router.registry().deployment(model_id) is not None
|
||||
|
||||
@staticmethod
|
||||
def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
|
||||
"""A 404 naming the output file means there is nothing to fetch on this or any
|
||||
later poll: providers like Vertex AI advertise an output path for every batch,
|
||||
including terminal ones that never wrote it. Any other failure may be
|
||||
transient, so it keeps retrying until the staleness sweep bounds it."""
|
||||
including terminal ones that never wrote it. GCS names a gs:// object without its scheme,
|
||||
and S3 names an s3:// object by its key alone. Any other failure keeps retrying."""
|
||||
import openai
|
||||
|
||||
from litellm.exceptions import NotFoundError
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
if not output_file_id:
|
||||
return False
|
||||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
|
||||
is_not_found: Final = isinstance(error, (NotFoundError, openai.NotFoundError)) or (
|
||||
isinstance(error, BaseLLMException) and error.status_code == 404
|
||||
)
|
||||
object_path: Final = output_file_id.split("://", 1)[-1]
|
||||
names: Final = (output_file_id, object_path, object_path.split("/", 1)[-1])
|
||||
return is_not_found and any(_error_names_resource(error, name) for name in names)
|
||||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
|
||||
|
|
@ -510,7 +564,7 @@ class CheckBatchCost:
|
|||
def _get_deployment_id_for_bare_model(
|
||||
self, bare_model_name: str, llm_provider: str
|
||||
) -> Optional[str]:
|
||||
model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
|
||||
model_group = self.llm_router.registry().resolve(bare_model_name, kind="from_id")
|
||||
deployment_id = (
|
||||
self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None
|
||||
)
|
||||
|
|
@ -526,7 +580,7 @@ class CheckBatchCost:
|
|||
) -> Optional[str]:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
for deployment in self.llm_router.get_model_list(model_name=None) or []:
|
||||
for deployment in self.llm_router.registry().deployments() or []:
|
||||
litellm_params = deployment.get("litellm_params") or {}
|
||||
actual_model = litellm_params.get("model")
|
||||
if not isinstance(actual_model, str):
|
||||
|
|
@ -569,8 +623,9 @@ class CheckBatchCost:
|
|||
"""
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
||||
for deployment_id in self.llm_router.get_model_ids(model_name=model_group):
|
||||
deployment_info = self.llm_router.get_deployment(model_id=deployment_id)
|
||||
registry = self.llm_router.registry()
|
||||
for deployment_id in registry.names(model_group, kind="ids"):
|
||||
deployment_info = registry.deployment(deployment_id)
|
||||
if deployment_info is None:
|
||||
continue
|
||||
try:
|
||||
|
|
@ -689,7 +744,7 @@ class CheckBatchCost:
|
|||
except (IndexError, AttributeError):
|
||||
pass
|
||||
|
||||
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
_file_content = await afile_content(
|
||||
file_id=cast(str, raw_output_file_id),
|
||||
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
|
||||
|
|
@ -712,7 +767,7 @@ class CheckBatchCost:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
deployment_info = self.llm_router.get_deployment(model_id=model_id)
|
||||
deployment_info = self.llm_router.registry().deployment(model_id)
|
||||
if deployment_info is None:
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping job {job.unified_object_id} because it is not a valid deployment info"
|
||||
|
|
@ -942,7 +997,7 @@ class CheckBatchCost:
|
|||
)
|
||||
|
||||
try:
|
||||
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
|
||||
deployment_info: Final = self.llm_router.registry().deployment(model_id)
|
||||
if deployment_info is None:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
|
||||
|
|
@ -955,7 +1010,7 @@ class CheckBatchCost:
|
|||
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
|
||||
)
|
||||
cost_provider: Final = cast(BatchFileProvider, llm_provider)
|
||||
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
|
||||
file_access: Final = {
|
||||
**credentials,
|
||||
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
|
||||
|
|
@ -1094,13 +1149,6 @@ class CheckBatchCost:
|
|||
|
||||
processed_models: List[Tuple[Optional[str], Optional[str]]] = []
|
||||
|
||||
try:
|
||||
await self._cleanup_stale_managed_objects()
|
||||
except Exception as cleanup_err:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
|
||||
)
|
||||
|
||||
# Look for all batches that have not yet been processed by CheckBatchCost.
|
||||
# self._has_batch_processed_column is cached after the first probe so that
|
||||
# older schemas don't pay a guaranteed-failing primary query + warning on
|
||||
|
|
@ -1142,6 +1190,8 @@ class CheckBatchCost:
|
|||
if routing is None:
|
||||
if self._has_unified_id_without_model(job):
|
||||
await self._retire_job(job, "unified object id has no model id")
|
||||
else:
|
||||
await self._expire_if_old_and_unfinished(job)
|
||||
continue
|
||||
model_id, batch_id = routing
|
||||
|
||||
|
|
@ -1166,6 +1216,14 @@ class CheckBatchCost:
|
|||
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
|
||||
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
|
||||
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
|
||||
else:
|
||||
await self._expire_if_old_and_unfinished(job)
|
||||
continue
|
||||
if self._is_completed_with_output_id_pending(response):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
|
||||
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
|
||||
)
|
||||
continue
|
||||
|
||||
## RETRIEVE THE BATCH JOB OUTPUT FILE
|
||||
|
|
@ -1217,17 +1275,9 @@ class CheckBatchCost:
|
|||
self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True)
|
||||
|
||||
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_completed_batch_safe_to_retire,
|
||||
)
|
||||
|
||||
if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
|
||||
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
|
||||
)
|
||||
continue
|
||||
await self._retire_unbilled_terminal_job(job, response, model_id, batch_id)
|
||||
else:
|
||||
await self._expire_if_old_and_unfinished(job)
|
||||
|
||||
# Record polling run metrics (always, even if nothing was processed)
|
||||
if prom_logger:
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
|
||||
ManagedObjectCreatorAttribution,
|
||||
|
|
@ -273,7 +273,7 @@ class CheckResponsesCost:
|
|||
|
||||
def _resolve_deployment(self, response_id: str) -> Deployment | None:
|
||||
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
|
||||
return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
|
||||
return self.llm_router.registry().deployment(model_id) if model_id is not None else None
|
||||
|
||||
async def _get_response(
|
||||
self,
|
||||
|
|
@ -290,12 +290,14 @@ class CheckResponsesCost:
|
|||
sees provider env vars, so it fails for every deployment whose credentials
|
||||
live in the config; the row then never leaves ``queued``.
|
||||
"""
|
||||
if not via_router:
|
||||
return await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
|
||||
router_response = await self.llm_router.aget_responses(
|
||||
response_id=response_id, litellm_metadata=litellm_metadata
|
||||
retrieved: Final[object] = (
|
||||
await self.llm_router.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
|
||||
if via_router
|
||||
else await litellm.aget_responses(response_id=response_id, litellm_metadata=litellm_metadata)
|
||||
)
|
||||
return cast(ResponsesAPIResponse, router_response)
|
||||
if not isinstance(retrieved, ResponsesAPIResponse):
|
||||
raise TypeError(f"Retrieving {response_id} did not return a response object")
|
||||
return retrieved
|
||||
|
||||
async def _fetch_unbilled(self, target: _PollTarget) -> ResponsesAPIResponse:
|
||||
return await self._get_response(
|
||||
|
|
@ -413,7 +415,7 @@ class CheckResponsesCost:
|
|||
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 ()
|
||||
self.llm_router.registry().deployments(requested, team_id=team_id) or ()
|
||||
)
|
||||
serving_id: Final = target.deployment.model_info.id
|
||||
return MappingProxyType(
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import base64
|
|||
import json
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from datetime import datetime, timedelta
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
|
|
@ -27,7 +28,7 @@ import httpx
|
|||
from fastapi import HTTPException
|
||||
from openai import APIConnectionError, AsyncOpenAI
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, Unpack
|
||||
|
||||
import litellm
|
||||
|
|
@ -39,6 +40,7 @@ from litellm.constants import (
|
|||
BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS,
|
||||
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
|
||||
MANAGED_BACKGROUND_RESPONSE_METADATA_KEY,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_FILE_LIST_LIMIT,
|
||||
PENDING_RESPONSE_STATUSES,
|
||||
)
|
||||
|
|
@ -51,7 +53,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
extract_file_metadata,
|
||||
)
|
||||
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
|
||||
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
build_list_page,
|
||||
build_owner_filter,
|
||||
|
|
@ -113,6 +115,7 @@ from litellm.types.utils import (
|
|||
LLMResponseTypes,
|
||||
SpecialEnums,
|
||||
)
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
|
||||
class _ManagedFileRetrieve(Protocol):
|
||||
|
|
@ -242,7 +245,7 @@ def _provider_file_retrieve_credentials(
|
|||
if llm_router is None or model_id is None:
|
||||
return None
|
||||
try:
|
||||
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id)
|
||||
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
except Exception as error:
|
||||
verbose_logger.warning(
|
||||
"Failed to retrieve credentials for provider file "
|
||||
|
|
@ -306,6 +309,7 @@ def _is_xai_batch_output(
|
|||
class _ManagedFileRow(Protocol):
|
||||
unified_file_id: str
|
||||
file_object: OpenAIFileObject
|
||||
model_mappings: Mapping[str, str]
|
||||
flat_model_file_ids: Sequence[str]
|
||||
storage_backend: Optional[str]
|
||||
storage_url: Optional[str]
|
||||
|
|
@ -359,6 +363,8 @@ class _ManagedResourceDatabase(Protocol):
|
|||
@property
|
||||
def litellm_managedfiletable(self) -> _ManagedFileTableActions: ...
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
@property
|
||||
def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ...
|
||||
|
||||
|
|
@ -382,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False):
|
|||
custom_llm_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase:
|
||||
return prisma_client.db
|
||||
|
||||
|
||||
def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions:
|
||||
return prisma_client.db.litellm_managedfiletable
|
||||
|
||||
|
||||
class _ListedFileRow(Protocol):
|
||||
@property
|
||||
def unified_file_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def file_object(self) -> object: ...
|
||||
|
||||
@property
|
||||
def model_mappings(self) -> Mapping[str, object]: ...
|
||||
|
||||
@property
|
||||
def storage_backend(self) -> str | None: ...
|
||||
|
||||
|
||||
class _ScopedFileRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
unified_file_id: str
|
||||
file_object: JsonValue = None
|
||||
model_mappings: Mapping[str, JsonValue]
|
||||
storage_backend: str | None = None
|
||||
|
||||
|
||||
_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...])
|
||||
|
||||
|
||||
def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]:
|
||||
match owner_filter:
|
||||
case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}:
|
||||
return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id)
|
||||
case {"created_by": str() as user_id}:
|
||||
return '"created_by" = $1', (user_id,)
|
||||
case {"team_id": str() as team_id}:
|
||||
return '"team_id" = $1', (team_id,)
|
||||
case _:
|
||||
return ("FALSE", ()) if owner_filter else ("TRUE", ())
|
||||
|
||||
|
||||
async def _scoped_file_rows(
|
||||
database: _ManagedResourceDatabase,
|
||||
owner_filter: Mapping[str, object],
|
||||
model_mapping_keys: frozenset[str],
|
||||
cursor_id: str | None,
|
||||
direction: Literal["asc", "desc"],
|
||||
take: int,
|
||||
) -> tuple[_ScopedFileRow, ...]:
|
||||
owner_condition, owner_params = _owner_condition(owner_filter)
|
||||
keys_param: Final = len(owner_params) + 1
|
||||
comparison: Final = "<" if direction == "desc" else ">"
|
||||
cursor_condition: Final = (
|
||||
f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" '
|
||||
f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) '
|
||||
if cursor_id
|
||||
else ""
|
||||
)
|
||||
order: Final = direction.upper()
|
||||
rows: Final = await database.query_raw(
|
||||
'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" '
|
||||
f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) '
|
||||
f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} '
|
||||
f"LIMIT ${keys_param + (2 if cursor_id else 1)}",
|
||||
*owner_params,
|
||||
sorted(model_mapping_keys),
|
||||
*((cursor_id,) if cursor_id else ()),
|
||||
take,
|
||||
)
|
||||
return _SCOPED_FILE_ROWS.validate_python(rows)
|
||||
|
||||
|
||||
def _iter_provider_file_id_pairs(
|
||||
rows: Sequence[_ManagedFileRow],
|
||||
requested_provider_file_ids: frozenset[str],
|
||||
|
|
@ -507,10 +586,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
prisma_client: PrismaClient,
|
||||
*,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
clock: Callable[[], datetime] = get_utc_datetime,
|
||||
):
|
||||
self.internal_usage_cache = internal_usage_cache
|
||||
self.prisma_client = prisma_client
|
||||
self._sleep = sleep
|
||||
self._clock = clock
|
||||
|
||||
@staticmethod
|
||||
def _get_prometheus_logger():
|
||||
|
|
@ -2260,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
limit: Optional[int] = None,
|
||||
after: Optional[str] = None,
|
||||
model_scope: Optional[FileListModelScope] = None,
|
||||
before: Optional[str] = None,
|
||||
**data: Dict,
|
||||
) -> FileListPage:
|
||||
"""List the managed files the caller owns, newest first.
|
||||
|
||||
Pagination is keyset based on ``unified_file_id`` so a key that owns
|
||||
every file on the proxy still reads one bounded page at a time.
|
||||
``purpose`` is applied after parsing, because the managed file table
|
||||
keeps it inside the ``file_object`` blob instead of a column, and rows
|
||||
``after`` pages toward older files and ``before`` toward newer ones,
|
||||
each page still ordered newest first.
|
||||
``model_scope`` is applied in the query, so only rows whose
|
||||
``model_mappings`` name one of its deployment ids or model names, and
|
||||
storage-backed rows, leave the database; a storage-backed row then
|
||||
stays only when its model name is one the router sends to the listed
|
||||
provider. ``purpose`` is applied after parsing, because the managed
|
||||
file table keeps it inside the ``file_object`` blob instead of a
|
||||
column. Rows
|
||||
whose blob will not parse drop out there too, so a chunk of rows can
|
||||
yield fewer matches than the page holds. Successive chunks are read
|
||||
until the page is full or the caller's rows run out, which keeps
|
||||
|
|
@ -2276,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
scan has to continue past it, the chunk widens to
|
||||
``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per
|
||||
that many rows instead of one per page. That bound is per query, not
|
||||
per request: the work is still linear in the rows the caller owns, and
|
||||
a filter matching nothing reads every one of them, with no index
|
||||
covering either the owner filter or the sort.
|
||||
per request: a ``purpose`` filter matching nothing still reads every
|
||||
row that passes the query.
|
||||
"""
|
||||
validate_file_list_limit(limit)
|
||||
validate_file_list_purpose(purpose)
|
||||
|
|
@ -2287,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if owner_filter is None:
|
||||
return FileListPage.model_validate(build_list_page([]))
|
||||
|
||||
if after:
|
||||
cursor_row = await _managed_file_table(self.prisma_client).find_first(
|
||||
where={**owner_filter, "unified_file_id": after}
|
||||
if after and before:
|
||||
raise ProxyException(
|
||||
message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.",
|
||||
type="invalid_request_error",
|
||||
param="before_id",
|
||||
code=400,
|
||||
openai_code="invalid_value",
|
||||
)
|
||||
cursor_param: Final = "before_id" if before else "after"
|
||||
cursor: Final = before or after
|
||||
database: Final = _managed_database(self.prisma_client)
|
||||
if cursor:
|
||||
cursor_row = await database.litellm_managedfiletable.find_first(
|
||||
where={**owner_filter, "unified_file_id": cursor}
|
||||
)
|
||||
if cursor_row is None:
|
||||
raise ProxyException(
|
||||
message=f"Invalid 'after' cursor: no file found with id '{after}'.",
|
||||
message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.",
|
||||
type="invalid_request_error",
|
||||
param="after",
|
||||
param=cursor_param,
|
||||
code=400,
|
||||
openai_code="invalid_value",
|
||||
)
|
||||
|
||||
if model_scope is not None and not model_scope.model_mapping_keys:
|
||||
return FileListPage.model_validate(build_list_page([]))
|
||||
|
||||
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
|
||||
matches: Final[List[OpenAIFileObject]] = []
|
||||
cursor_id = after
|
||||
cursor_id = cursor
|
||||
chunk_size = page_size + 1
|
||||
direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc"
|
||||
|
||||
while len(matches) <= page_size:
|
||||
cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {}
|
||||
chunk = await _managed_file_table(self.prisma_client).find_many(
|
||||
where=owner_filter,
|
||||
take=chunk_size,
|
||||
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
|
||||
**cursor_args,
|
||||
chunk: Sequence[_ListedFileRow] = (
|
||||
await _scoped_file_rows(
|
||||
database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size
|
||||
)
|
||||
if model_scope is not None
|
||||
else await database.litellm_managedfiletable.find_many(
|
||||
where=owner_filter,
|
||||
take=chunk_size,
|
||||
order=[{"created_at": direction}, {"unified_file_id": direction}],
|
||||
**cursor_args,
|
||||
)
|
||||
)
|
||||
matches.extend(
|
||||
_public_file_object(parsed_file_object, row.unified_file_id)
|
||||
for row in chunk
|
||||
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
|
||||
if (
|
||||
model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None)
|
||||
)
|
||||
and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
|
||||
and (purpose is None or parsed_file_object.purpose == purpose)
|
||||
)
|
||||
if len(chunk) < chunk_size:
|
||||
|
|
@ -2324,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
cursor_id = chunk[-1].unified_file_id
|
||||
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
|
||||
|
||||
return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size))
|
||||
page: Final = matches[:page_size]
|
||||
return FileListPage.model_validate(
|
||||
build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size)
|
||||
)
|
||||
|
||||
def _is_batch_polling_enabled(self) -> bool:
|
||||
"""
|
||||
|
|
@ -2387,12 +2503,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
"status": {"not_in": ["failed", "expired", "cancelled"]},
|
||||
"status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]},
|
||||
},
|
||||
take=MAX_MATCHES_TO_RETURN,
|
||||
order={"created_at": "desc"},
|
||||
)
|
||||
|
||||
input_file_cutoff: Final = self._clock() - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
referencing_batches: Final[list[dict[str, object]]] = []
|
||||
for batch in batches:
|
||||
try:
|
||||
|
|
@ -2405,7 +2522,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# Extract file IDs from batch
|
||||
# Batches typically reference the unified file ID in input_file_id
|
||||
# Output and error files are generated by the provider
|
||||
input_file_id = batch_data.get("input_file_id")
|
||||
input_file_id = batch_data.get("input_file_id") if batch.created_at >= input_file_cutoff else None
|
||||
output_file_id = batch_data.get("output_file_id")
|
||||
error_file_id = batch_data.get("error_file_id")
|
||||
|
||||
|
|
@ -2533,7 +2650,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
|
||||
}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
delete_data = {
|
||||
**filtered_data,
|
||||
**(
|
||||
|
|
@ -2577,7 +2694,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# against the deployment's configured bucket, which they only
|
||||
# trust from this immutable server-side snapshot, never from
|
||||
# request params.
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
credentials = llm_router.registry().credentials(model_id, form="with_provider")
|
||||
if credentials is not None:
|
||||
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -290,7 +290,7 @@ def _check_team_project_limits(
|
|||
|
||||
|
||||
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
|
||||
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
|
||||
return frozenset(llm_router.registry().access_groups()) if llm_router is not None else frozenset()
|
||||
|
||||
|
||||
def _project_models_expanding_at_request_time(
|
||||
|
|
|
|||
63
examples/bigmodel.md
Normal file
63
examples/bigmodel.md
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
# BigModel (Zhipu AI)
|
||||
|
||||
Use `bigmodel/` for the domestic BigModel platform and `zai/` for international Z.AI. Each platform uses its own credentials and default URL
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: glm
|
||||
litellm_params:
|
||||
model: bigmodel/glm-5.3
|
||||
api_key: os.environ/BIGMODEL_API_KEY
|
||||
api_base: https://open.bigmodel.cn
|
||||
```
|
||||
|
||||
One model configuration supports all three native APIs. The provider selects the upstream endpoint from the requested API
|
||||
|
||||
| Proxy endpoint | BigModel endpoint |
|
||||
| --- | --- |
|
||||
| `/v1/chat/completions` | `/api/paas/v4/chat/completions` |
|
||||
| `/v1/responses` | `/api/v1/responses` |
|
||||
| `/v1/messages` | `/api/anthropic/v1/messages` |
|
||||
|
||||
`api_base` is optional and defaults to `https://open.bigmodel.cn`. Set `BIGMODEL_API_BASE` to override the default, or pass an explicit `api_base` to override the environment value. Existing API-specific base URLs and complete endpoint URLs are normalized, preserving a custom proxy's path prefix
|
||||
|
||||
A GLM Coding Plan key uses `https://open.bigmodel.cn/api/coding/paas/v4` as its `api_base`, since Coding Plan and standard keys do not work on each other's endpoints. Chat Completions then calls the Coding Plan endpoint, and Responses and Messages keep the endpoints in the table above
|
||||
|
||||
The Admin UI lists BigModel (Zhipu AI), displays its official logo, and provides a required password API Key field and an optional API Base field with the default URL. Authentication uses an explicit `api_key` or `BIGMODEL_API_KEY`
|
||||
|
||||
```python
|
||||
import litellm
|
||||
|
||||
chat = litellm.completion(
|
||||
model="bigmodel/glm-5.3",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
thinking={"type": "enabled"},
|
||||
reasoning_effort="low",
|
||||
)
|
||||
|
||||
response = litellm.responses(
|
||||
model="bigmodel/glm-5.3",
|
||||
input="Hello",
|
||||
reasoning={"effort": "low"},
|
||||
store=True,
|
||||
)
|
||||
|
||||
async def messages():
|
||||
return await litellm.anthropic.messages.acreate(
|
||||
model="bigmodel/glm-5.3",
|
||||
max_tokens=1024,
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
thinking={"type": "adaptive"},
|
||||
output_config={"effort": "low"},
|
||||
)
|
||||
```
|
||||
|
||||
Each API retains its native response format and streaming events. The native Messages SDK entry point is asynchronous, as is the proxy's `/v1/messages` route
|
||||
|
||||
GLM-5.3, GLM-5.3-Flash, and GLM-5.3-FlashX use always-on thinking. GLM-5.3 supports `low`, `high`, and `max` reasoning effort. Messages preserves native `thinking`, `output_config`, and `reasoning_effort` parameters
|
||||
|
||||
Set `store=True` when creating a Responses request if you need continuation, retrieval, or deletion. Use the response ID returned by LiteLLM. BigModel documents a seven-day retention period and does not support cancel, compact, native WebSocket, or file_search
|
||||
|
||||
The model catalog and packaged backup include GLM-5.3, GLM-5.3-Flash, and GLM-5.3-FlashX. Prices are expressed in USD/token using a static 7.0 CNY/USD conversion, with the source and verification date in each entry's metadata. Use existing `litellm_params` cost fields for contract-specific deployment prices
|
||||
|
||||
API references: [Chat Completions](https://docs.bigmodel.cn/cn/guide/develop/openai/introduction), [Responses](https://docs.bigmodel.cn/cn/guide/develop/responses/introduction), [Messages](https://docs.bigmodel.cn/cn/guide/develop/claude/introduction), and [pricing](https://docs.bigmodel.cn/cn/guide/start/pricing), checked on 2026-10-27
|
||||
|
|
@ -40,6 +40,7 @@ ENV UV_PYTHON_DOWNLOADS=0
|
|||
COPY pyproject.toml README.md rust-toolchain.toml ./
|
||||
COPY .cargo ./.cargo
|
||||
COPY litellm-rust ./litellm-rust
|
||||
COPY model_prices_and_context_window.json ./
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p litellm/rust_bridge /wheels && \
|
||||
|
|
@ -144,6 +145,7 @@ 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/litellm_cli /app/litellm_cli
|
||||
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
|
||||
|
|
@ -153,7 +155,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
|||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
|
||||
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
|
||||
python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__"
|
||||
|
||||
USER nonroot
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,11 @@ handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
|
|||
wins under token auth too, so exporting it here is enough for every worker to
|
||||
pick the pooled URL up unchanged.
|
||||
|
||||
LiteLLM imports wait for ``main``: every uvicorn worker imports this module again
|
||||
before it can answer its first 5 s health ping, so a slow host would kill each one.
|
||||
Workers then load the gateway app through ``worker_app`` without garbage collection
|
||||
pauses, which hold the GIL long enough to miss a ping.
|
||||
|
||||
Run with:
|
||||
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
|
||||
"""
|
||||
|
|
@ -15,19 +20,19 @@ Run with:
|
|||
import os
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from uvicorn.main import main as uvicorn_main
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import (
|
||||
PgBouncerError,
|
||||
PgBouncerSettings,
|
||||
export_pooled_database_url,
|
||||
start_in_container_pgbouncer,
|
||||
)
|
||||
from litellm_cli.main import without_gc_pauses
|
||||
|
||||
GATEWAY_APP: Final = "gateway.main:app"
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
|
||||
|
||||
GATEWAY_APP: Final = "gateway.launch:worker_app"
|
||||
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
|
||||
|
||||
|
||||
|
|
@ -35,21 +40,23 @@ def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str,
|
|||
"""Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
|
||||
keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
|
||||
if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
|
||||
return (GATEWAY_APP, *argv)
|
||||
return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive)
|
||||
return (GATEWAY_APP, "--factory", *argv)
|
||||
return (GATEWAY_APP, "--factory", *argv, KEEPALIVE_FLAG, keepalive)
|
||||
|
||||
|
||||
def pool_database_url(
|
||||
settings: DatabaseURLSettings,
|
||||
pgbouncer: PgBouncerSettings,
|
||||
settings: "DatabaseURLSettings",
|
||||
pgbouncer: "PgBouncerSettings",
|
||||
environ: Mapping[str, str],
|
||||
) -> str | PgBouncerError | None:
|
||||
) -> "str | PgBouncerError | None":
|
||||
"""Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
|
||||
|
||||
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
|
||||
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
|
||||
auth the pooler mints and renews the upstream token itself.
|
||||
"""
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, start_in_container_pgbouncer
|
||||
|
||||
if not pgbouncer.enabled:
|
||||
return None
|
||||
upstream_url: Final = environ.get("DATABASE_URL")
|
||||
|
|
@ -58,11 +65,20 @@ def pool_database_url(
|
|||
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())
|
||||
|
||||
|
||||
def worker_app() -> "FastAPI":
|
||||
with without_gc_pauses():
|
||||
from gateway.main import app
|
||||
return app
|
||||
|
||||
|
||||
def _serve(argv: Sequence[str]) -> None:
|
||||
uvicorn_main(tuple(argv), prog_name="uvicorn")
|
||||
|
||||
|
||||
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url
|
||||
|
||||
settings: Final = DatabaseURLSettings.from_env()
|
||||
settings.apply_to_env()
|
||||
pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
|
||||
|
|
|
|||
|
|
@ -259,3 +259,9 @@ itself.
|
|||
`lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
|
||||
|
||||
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
|
||||
|
||||
## Shutdown accounting
|
||||
|
||||
`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged
|
||||
|
||||
Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml`
|
||||
|
|
|
|||
|
|
@ -14,11 +14,9 @@ either, so the migration keeps today's behavior for it.
|
|||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable, Iterator
|
||||
from collections.abc import Callable
|
||||
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
|
||||
|
|
@ -43,7 +41,7 @@ MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIME
|
|||
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"]
|
||||
Kind = Literal["too_long"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -54,9 +52,6 @@ class Check:
|
|||
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)
|
||||
|
|
@ -100,33 +95,6 @@ class ColumnType:
|
|||
|
||||
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."""
|
||||
|
|
@ -145,117 +113,6 @@ def migration_data_check_timeout_seconds() -> float:
|
|||
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),)
|
||||
|
|
@ -343,11 +200,16 @@ def btree_max_entry_bytes(block_size: int) -> int:
|
|||
|
||||
|
||||
def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed":
|
||||
"""The bytes a value takes in the index entry: its length, or its stored size when
|
||||
Postgres compressed it, since `index_form_tuple` keeps the compressed bytes."""
|
||||
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)
|
||||
length: Final = (
|
||||
sql.SQL("octet_length({})").format(identifier)
|
||||
if column.type_name in _OCTET_LENGTH_TYPES
|
||||
else sql.SQL("octet_length({}::text)").format(identifier)
|
||||
)
|
||||
return sql.SQL("least({}, pg_column_size({}))").format(length, identifier)
|
||||
|
||||
|
||||
def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed":
|
||||
|
|
@ -419,42 +281,7 @@ def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...],
|
|||
|
||||
|
||||
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
|
||||
)
|
||||
return _entry_size_query(check, schema, types, limit_bytes)
|
||||
|
||||
|
||||
def _applies(
|
||||
|
|
@ -469,15 +296,7 @@ def _applies(
|
|||
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
|
||||
return types is not None and any(column.can_overflow_btree for column in types)
|
||||
|
||||
|
||||
def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
|
||||
|
|
@ -510,21 +329,7 @@ def _run_check(
|
|||
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})")
|
||||
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
|
||||
|
||||
|
||||
def find_blockers(
|
||||
|
|
@ -566,19 +371,7 @@ def find_blockers(
|
|||
|
||||
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.'
|
||||
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
|
||||
|
||||
|
||||
def _describe(finding: Finding) -> str:
|
||||
|
|
@ -602,30 +395,4 @@ def _describe(finding: Finding) -> str:
|
|||
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
|
||||
return f"{REFUSAL}\n{lines}\nNo index was built. Fix the rows and the next index build picks it up."
|
||||
|
|
|
|||
184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
Normal file
184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real, each
|
||||
inside its own transaction that is always rolled back, so Postgres itself says whether this
|
||||
database's rows and schema take them. Each migration holds the locks the real upgrade will,
|
||||
and releases them before the next one runs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from functools import reduce
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm_proxy_extras.migration_data_checks import pending_migrations
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import psycopg
|
||||
|
||||
_CONCURRENTLY: Final = re.compile(
|
||||
r"\b((?:CREATE\s+(?:UNIQUE\s+)?INDEX|DROP\s+INDEX|REINDEX\s+\w+)\s+)CONCURRENTLY\s+", re.IGNORECASE
|
||||
)
|
||||
UNDEFINED_SQLSTATES: Final = frozenset({"42703", "42P01", "42704", "42883"})
|
||||
|
||||
Migration = tuple[str, str]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Applied:
|
||||
name: str
|
||||
seconds: float
|
||||
replayed: tuple[str, ...] = ()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rejected:
|
||||
name: str
|
||||
error: str
|
||||
seconds: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Baseline:
|
||||
"""A database with tables but no `_prisma_migrations`: the boot baselines it instead of replaying
|
||||
history, so the dry run asks the same question the boot does. `refusal` is None when the boot
|
||||
will baseline, otherwise its reason (a custom migration folder, or the schema diff)."""
|
||||
|
||||
refusal: str | None
|
||||
|
||||
@property
|
||||
def matches(self) -> bool:
|
||||
return self.refusal is None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DryRun:
|
||||
applied: tuple[Applied, ...]
|
||||
rejected: Rejected | None
|
||||
baseline: Baseline | None = None
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.rejected is None and (self.baseline is None or self.baseline.matches)
|
||||
|
||||
|
||||
def transactional(migration_sql: str) -> str:
|
||||
"""The migration as it can run inside a transaction: `CONCURRENTLY` index builds become plain ones."""
|
||||
return _CONCURRENTLY.sub(r"\1", migration_sql)
|
||||
|
||||
|
||||
def _attempt(
|
||||
connection: psycopg.Connection[tuple[object, ...]],
|
||||
schema: str,
|
||||
chain: tuple[Migration, ...],
|
||||
clock: Callable[[], float],
|
||||
) -> tuple[float, psycopg.Error | None]:
|
||||
import psycopg
|
||||
from psycopg import sql
|
||||
|
||||
started: Final = clock()
|
||||
try:
|
||||
with connection.transaction(force_rollback=True):
|
||||
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
|
||||
for _, migration_sql in chain:
|
||||
connection.execute(transactional(migration_sql))
|
||||
except psycopg.Error as exc:
|
||||
return clock() - started, exc
|
||||
return clock() - started, None
|
||||
|
||||
|
||||
def _run(
|
||||
connection: psycopg.Connection[tuple[object, ...]],
|
||||
schema: str,
|
||||
pending: tuple[Migration, ...],
|
||||
position: int,
|
||||
done: DryRun,
|
||||
clock: Callable[[], float],
|
||||
) -> DryRun:
|
||||
"""One migration in its own rolled-back transaction. A migration that fails because something
|
||||
an earlier pending migration creates does not exist yet is rerun with those replayed in front."""
|
||||
if done.rejected is not None:
|
||||
return done
|
||||
name: Final = pending[position][0]
|
||||
seconds, error = _attempt(connection, schema, (pending[position],), clock)
|
||||
if error is None:
|
||||
return DryRun(done.applied + (Applied(name, seconds),), None, done.baseline)
|
||||
if error.sqlstate not in UNDEFINED_SQLSTATES or position == 0:
|
||||
return DryRun(done.applied, Rejected(name, str(error).strip(), seconds), done.baseline)
|
||||
earlier: Final = pending[:position]
|
||||
chained_seconds, chained_error = _attempt(connection, schema, earlier + (pending[position],), clock)
|
||||
if chained_error is None:
|
||||
replayed: Final = tuple(earlier_name for earlier_name, _ in earlier)
|
||||
return DryRun(done.applied + (Applied(name, chained_seconds, replayed),), None, done.baseline)
|
||||
return DryRun(done.applied, Rejected(name, str(chained_error).strip(), chained_seconds), done.baseline)
|
||||
|
||||
|
||||
def _has_ledger(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
|
||||
from psycopg import sql
|
||||
|
||||
row: Final = connection.execute(
|
||||
"SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(connection),)
|
||||
).fetchone()
|
||||
return row is not None and row[0] is not None
|
||||
|
||||
|
||||
def _has_litellm_tables(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
|
||||
row: Final = connection.execute(
|
||||
"SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name LIKE 'LiteLLM\\_%%' LIMIT 1",
|
||||
(schema,),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
|
||||
def dry_run(
|
||||
connection: psycopg.Connection[tuple[object, ...]],
|
||||
schema: str,
|
||||
migrations_dir: Path,
|
||||
baseline_check: Callable[[], str | None] | None = None,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> DryRun:
|
||||
"""Apply every pending migration in order, each in its own transaction that is rolled back
|
||||
before the next one starts. The first migration Postgres rejects ends the run with its error.
|
||||
`baseline_check` (None when the boot would baseline, otherwise why it refuses) stands in for
|
||||
the replay on a database that has tables but no migration history, which the boot baselines."""
|
||||
if baseline_check is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema):
|
||||
return DryRun((), None, Baseline(baseline_check()))
|
||||
pending: Final = pending_migrations(connection, schema, migrations_dir)
|
||||
return reduce(
|
||||
lambda done, position: _run(connection, schema, pending, position, done, clock),
|
||||
range(len(pending)),
|
||||
DryRun((), None),
|
||||
)
|
||||
|
||||
|
||||
def _timing(applied: Applied) -> str:
|
||||
replay: Final = f" (depends on {', '.join(applied.replayed)}, replayed in front of it)" if applied.replayed else ""
|
||||
return f"\n {applied.name}: {applied.seconds:.1f} s{replay}"
|
||||
|
||||
|
||||
def report(result: DryRun) -> str:
|
||||
"""What the operator reads: the per-migration time the real upgrade will hold its locks,
|
||||
or the migration Postgres rejected with Postgres's own words. Nothing was applied either way."""
|
||||
if result.baseline is not None:
|
||||
if result.baseline.matches:
|
||||
return (
|
||||
"No migration history and the schema matches this build, "
|
||||
"the proxy baselines it at boot, nothing to dry-run."
|
||||
)
|
||||
return f"No migration history and the proxy cannot baseline it at boot: {result.baseline.refusal}"
|
||||
timings: Final = "".join(_timing(applied) for applied in result.applied)
|
||||
if result.rejected is None:
|
||||
count: Final = len(result.applied)
|
||||
if count == 0:
|
||||
return "No pending migrations, this database is on this version's schema."
|
||||
return (
|
||||
f"{count} pending migration(s) applied one at a time in a dry run, each rolled back before the next "
|
||||
f"so its locks are released, nothing changed:{timings}"
|
||||
)
|
||||
return (
|
||||
f"Migration {result.rejected.name} fails on this database after {result.rejected.seconds:.1f} s, "
|
||||
f"nothing was applied.\n{result.rejected.error}"
|
||||
+ (f"\nApplied before it in the dry run:{timings}" if timings else "")
|
||||
)
|
||||
|
|
@ -86,6 +86,18 @@ def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
|
|||
)
|
||||
|
||||
|
||||
def has_packaged_history(migrations_dir: Path) -> bool:
|
||||
"""Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships,
|
||||
the only history the boot baselines an existing database onto."""
|
||||
packaged_dir: Final = Path(__file__).parent
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
return (
|
||||
bool(migrations)
|
||||
and migrations == migration_files(packaged_dir)
|
||||
and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes()
|
||||
)
|
||||
|
||||
|
||||
def baseline_current_schema(
|
||||
coordinator: MigrationCoordinator,
|
||||
schema: str,
|
||||
|
|
@ -95,14 +107,9 @@ def baseline_current_schema(
|
|||
) -> None:
|
||||
from psycopg import sql
|
||||
|
||||
packaged_dir: Final = Path(__file__).parent
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
if (
|
||||
not migrations
|
||||
or migrations != migration_files(packaged_dir)
|
||||
or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
|
||||
):
|
||||
if not has_packaged_history(migrations_dir):
|
||||
raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
|
||||
migrations: Final = migration_files(migrations_dir)
|
||||
|
||||
coordinator.acquire_prisma_lock()
|
||||
existing: Final = coordinator.connection.execute(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx"
|
||||
ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC);
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE TYPE "AccountingEventKind" AS ENUM ('inbound_usage', 'outbound_usage', 'cost');
|
||||
|
||||
CREATE TYPE "AccountingDeliveryState" AS ENUM ('pending', 'applied', 'parked');
|
||||
|
||||
CREATE TYPE "AccountingDeliveryTarget" AS ENUM ('LiteLLM_VerificationToken', 'LiteLLM_UserTable', 'LiteLLM_TeamTable', 'LiteLLM_TeamMembership', 'LiteLLM_OrganizationTable', 'LiteLLM_OrganizationMembership', 'LiteLLM_EndUserTable', 'LiteLLM_AgentsTable', 'LiteLLM_TagTable', 'LiteLLM_ProjectTable', 'LiteLLM_ModelAccessGroupBudgetTable', 'LiteLLM_BudgetTable', 'LiteLLM_BudgetWindowSpend', 'LiteLLM_DailyUserSpend', 'LiteLLM_DailyTeamSpend', 'LiteLLM_DailyOrganizationSpend', 'LiteLLM_DailyEndUserSpend', 'LiteLLM_DailyAgentSpend', 'LiteLLM_DailyTagSpend', 'LiteLLM_DailyModelUsage', 'LiteLLM_DailyToolSpend', 'LiteLLM_AutoRouterDailySpend');
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingEvent" (
|
||||
"event_id" TEXT NOT NULL,
|
||||
"kind" "AccountingEventKind" NOT NULL,
|
||||
"source_event_id" TEXT,
|
||||
"schema_version" INTEGER NOT NULL,
|
||||
"payload_hash" TEXT NOT NULL,
|
||||
"occurred_at" TIMESTAMP(3) NOT NULL,
|
||||
"recorded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"payload" JSONB NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_AccountingEvent_schema_version_check" CHECK ("schema_version" > 0),
|
||||
CONSTRAINT "LiteLLM_AccountingEvent_pkey" PRIMARY KEY ("event_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingDelivery" (
|
||||
"event_id" TEXT NOT NULL,
|
||||
"target" "AccountingDeliveryTarget" NOT NULL,
|
||||
"version" INTEGER NOT NULL,
|
||||
"state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending',
|
||||
"attempts" INTEGER NOT NULL DEFAULT 0,
|
||||
"next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"applied_at" TIMESTAMP(3),
|
||||
"last_error" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_AccountingDelivery_version_check" CHECK ("version" > 0),
|
||||
CONSTRAINT "LiteLLM_AccountingDelivery_attempts_check" CHECK ("attempts" >= 0),
|
||||
CONSTRAINT "LiteLLM_AccountingDelivery_pkey" PRIMARY KEY ("event_id","target","version")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_recorded_at_idx" ON "LiteLLM_AccountingEvent"("recorded_at");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_kind_source_event_id_idx" ON "LiteLLM_AccountingEvent"("kind", "source_event_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingDelivery_target_state_next_attempt_at_idx" ON "LiteLLM_AccountingDelivery"("target", "state", "next_attempt_at");
|
||||
|
|
@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable {
|
|||
@@index([unified_file_id])
|
||||
@@index([flat_model_file_ids], type: Gin)
|
||||
@@index([team_id, created_at(sort: Desc)])
|
||||
@@index([created_by, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the
|
||||
|
|
@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport {
|
|||
@@index([window_end])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
||||
enum AccountingEventKind {
|
||||
inbound_usage
|
||||
outbound_usage
|
||||
cost
|
||||
}
|
||||
|
||||
enum AccountingDeliveryState {
|
||||
pending
|
||||
applied
|
||||
parked
|
||||
}
|
||||
|
||||
enum AccountingDeliveryTarget {
|
||||
LiteLLM_VerificationToken
|
||||
LiteLLM_UserTable
|
||||
LiteLLM_TeamTable
|
||||
LiteLLM_TeamMembership
|
||||
LiteLLM_OrganizationTable
|
||||
LiteLLM_OrganizationMembership
|
||||
LiteLLM_EndUserTable
|
||||
LiteLLM_AgentsTable
|
||||
LiteLLM_TagTable
|
||||
LiteLLM_ProjectTable
|
||||
LiteLLM_ModelAccessGroupBudgetTable
|
||||
LiteLLM_BudgetTable
|
||||
LiteLLM_BudgetWindowSpend
|
||||
LiteLLM_DailyUserSpend
|
||||
LiteLLM_DailyTeamSpend
|
||||
LiteLLM_DailyOrganizationSpend
|
||||
LiteLLM_DailyEndUserSpend
|
||||
LiteLLM_DailyAgentSpend
|
||||
LiteLLM_DailyTagSpend
|
||||
LiteLLM_DailyModelUsage
|
||||
LiteLLM_DailyToolSpend
|
||||
LiteLLM_AutoRouterDailySpend
|
||||
}
|
||||
|
||||
// Durable accounting ledger. One row per usage or cost event; payload holds the
|
||||
// private event envelope versioned by schema_version. Usage rows retain request-scoped
|
||||
// pricing context; cost rows include resolved rates and the finalized pricing snapshot.
|
||||
model LiteLLM_AccountingEvent {
|
||||
event_id String @id
|
||||
kind AccountingEventKind
|
||||
source_event_id String?
|
||||
schema_version Int
|
||||
payload_hash String
|
||||
occurred_at DateTime
|
||||
recorded_at DateTime @default(now())
|
||||
payload Json
|
||||
|
||||
@@index([recorded_at])
|
||||
@@index([kind, source_event_id])
|
||||
}
|
||||
|
||||
// One row per (event, target table, version): pending work for the projector and,
|
||||
// once applied in the same transaction as the target write, the receipt that stops a retry.
|
||||
model LiteLLM_AccountingDelivery {
|
||||
event_id String
|
||||
target AccountingDeliveryTarget
|
||||
version Int
|
||||
state AccountingDeliveryState @default(pending)
|
||||
attempts Int @default(0)
|
||||
next_attempt_at DateTime @default(now())
|
||||
applied_at DateTime?
|
||||
last_error String?
|
||||
|
||||
@@id([event_id, target, version])
|
||||
@@index([target, state, next_attempt_at])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,18 +11,13 @@ from collections.abc import Callable
|
|||
from contextlib import AbstractContextManager, nullcontext
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
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_dry_run import dry_run
|
||||
from litellm_proxy_extras.migration_dry_run import report as dry_run_report
|
||||
from litellm_proxy_extras.migration_lock import held_migration_lock
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -36,7 +31,6 @@ from litellm_proxy_extras.replica_identity import (
|
|||
apply_replica_identity_full,
|
||||
)
|
||||
from litellm_proxy_extras.request_log_indexes import (
|
||||
REQUEST_LOG_INDEXES,
|
||||
ensure_request_log_indexes,
|
||||
filter_request_log_index_diff,
|
||||
)
|
||||
|
|
@ -52,7 +46,7 @@ class LensCheckConnect(Protocol):
|
|||
) -> "psycopg.Connection[tuple[object, ...]]": ...
|
||||
|
||||
|
||||
def str_to_bool(value: Optional[str]) -> bool:
|
||||
def str_to_bool(value: str | None) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
return value.lower() in ("true", "1", "t", "y", "yes")
|
||||
|
|
@ -169,7 +163,7 @@ def _without_sql_comments(statement: str) -> str:
|
|||
).strip()
|
||||
|
||||
|
||||
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
|
||||
def _without_spend_logs_pk_clauses(statement: str) -> str | None:
|
||||
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
|
||||
if not prefix_match:
|
||||
return statement
|
||||
|
|
@ -217,7 +211,7 @@ def _max_migration_timestamp(names) -> int:
|
|||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
|
||||
def _redacted_command(command: object) -> str | tuple[str, ...] | list[str]:
|
||||
if isinstance(command, tuple):
|
||||
return tuple(_redact_credentials(str(argument)) for argument in command)
|
||||
if isinstance(command, list):
|
||||
|
|
@ -404,7 +398,7 @@ class ProxyExtrasDBManager:
|
|||
return tuple(row) if row is not None else ()
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
|
||||
def _failed_migration_logs(migration_name: str, started_at: str) -> str | None:
|
||||
row: Final = ProxyExtrasDBManager._read_migration_ledger(
|
||||
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
|
||||
"AND finished_at IS NULL AND rolled_back_at IS NULL",
|
||||
|
|
@ -721,10 +715,10 @@ class ProxyExtrasDBManager:
|
|||
return row is not None
|
||||
|
||||
@staticmethod
|
||||
def _prisma_schema_param(url: str) -> Optional[str]:
|
||||
def _prisma_schema_param(url: str) -> str | None:
|
||||
"""The `schema` query param Prisma uses to pick its target schema,
|
||||
or None when the URL does not set one."""
|
||||
from urllib.parse import urlparse, parse_qsl
|
||||
from urllib.parse import parse_qsl, urlparse
|
||||
|
||||
return next(
|
||||
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
|
||||
|
|
@ -1424,7 +1418,6 @@ class ProxyExtrasDBManager:
|
|||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
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
|
||||
|
|
@ -1433,46 +1426,69 @@ class ProxyExtrasDBManager:
|
|||
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."""
|
||||
def _baseline_refusal(prisma_dir: Path) -> str | None:
|
||||
"""Why `baseline_current_schema` would refuse this database at boot, None when it would baseline:
|
||||
the migrations folder must be the packaged one, then `prisma migrate diff` must find the live
|
||||
database and this build's schema.prisma identical."""
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
from litellm_proxy_extras.migration_recovery import has_packaged_history
|
||||
|
||||
if not has_packaged_history(prisma_dir):
|
||||
return (
|
||||
f"{prisma_dir} is a custom migration history (LITELLM_MIGRATION_DIR), it differs from the "
|
||||
"migrations this build ships, and the proxy only baselines onto the packaged history."
|
||||
)
|
||||
try:
|
||||
prisma_toolchain.run_prisma(
|
||||
(
|
||||
_get_prisma_command(),
|
||||
"migrate",
|
||||
"diff",
|
||||
"--from-schema-datasource",
|
||||
str(prisma_dir / "schema.prisma"),
|
||||
"--to-schema-datamodel",
|
||||
str(prisma_dir / "schema.prisma"),
|
||||
"--exit-code",
|
||||
),
|
||||
timeout=prisma_toolchain.prisma_command_timeout(),
|
||||
env=migration_environment(_get_prisma_env()),
|
||||
)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
diff: Final = str(exc.stdout or exc.stderr or exc).strip()
|
||||
return f"the schema differs from this build. prisma migrate diff:\n{diff}"
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
return f"prisma migrate diff did not finish within {exc.timeout} s"
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
|
||||
"""`--check_migrations_against_existing_data`: the pending migrations applied for real on
|
||||
DIRECT_URL or DATABASE_URL inside one transaction that is rolled back, so Postgres reports
|
||||
what the upgrade would hit. (True, the per-migration timings) when they all apply,
|
||||
(False, Postgres's own error) when one does not."""
|
||||
import psycopg
|
||||
|
||||
if not use_migrate:
|
||||
return (True, "The schema is pushed rather than migrated, there is nothing to dry-run.")
|
||||
database_url: Final = os.environ.get("DATABASE_URL")
|
||||
if not database_url:
|
||||
return ()
|
||||
return (False, "DATABASE_URL is not set.")
|
||||
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)
|
||||
prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir())
|
||||
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())
|
||||
result: Final = dry_run(
|
||||
connection,
|
||||
schema,
|
||||
prisma_dir / "migrations",
|
||||
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir),
|
||||
)
|
||||
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)
|
||||
return (False, f"Could not reach the database: {_redact_credentials(str(exc))}")
|
||||
return (result.ok, dry_run_report(result))
|
||||
|
||||
@staticmethod
|
||||
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
|
||||
|
|
@ -1724,5 +1740,4 @@ class ProxyExtrasDBManager:
|
|||
time.sleep(random.randrange(5, 15))
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
pass
|
||||
return False
|
||||
|
|
|
|||
38
litellm-rust/.agents/skills/rust-error/SKILL.md
Normal file
38
litellm-rust/.agents/skills/rust-error/SKILL.md
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
---
|
||||
name: rust-error
|
||||
description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust errors
|
||||
|
||||
`crates/http/src/error.rs` is the reference implementation. Read the reference for the part you are touching before editing, each one carries the reasoning and worked examples behind its rules
|
||||
|
||||
- [shape](references/shape.md)
|
||||
- One top-level `Error` per crate in `src/error.rs`. A public enum by default, an opaque struct over a private `Repr` enum when the variants must stay private. No `#[non_exhaustive]`
|
||||
- A variant is a failure mode a caller handles differently. Provider names and guidance are data, never variants
|
||||
- [messages and causes](references/messages-and-causes.md)
|
||||
- Every message is written once, in `#[error(...)]`, and describes its own layer only. Call sites pass typed facts, never a string
|
||||
- The cause is the `source()`, never part of the message. Keep it typed, never stringify or drop it
|
||||
- `reqwest::Error` enters our types only through `litellm_http::Error`, which strips its URL
|
||||
- Anything shown to someone else is rendered with `litellm_error::Report`, never `to_string()`
|
||||
- [conversions](references/conversions.md)
|
||||
- One `From` per source and target pair, only when the conversion needs no context. It classifies and sanitizes once
|
||||
- Nest a lower crate's error in one variant. Translate when the meaning changes, never mirror its variants
|
||||
- [classification](references/classification.md)
|
||||
- Callers ask predicates named after their decision, they do not match nested variants. Predicates match exhaustively
|
||||
- Classification states facts. Retry and fallback policy lives with the caller that knows idempotency
|
||||
- [boundaries](references/boundaries.md)
|
||||
- An error becomes text for someone else only in one mapping function per error type, in the gateway crates, the Python bridge and host-python. The inference gateway and the bridge decide by the same predicates
|
||||
- A failure before the provider was called is a decline. Both ask `declined()` first and hand it back to the host unchanged
|
||||
- [gateway](references/gateway.md)
|
||||
- One `Error` with `IntoResponse` per gateway crate that serves a router. Every failure flows through the same `From` chain, is classified once in `status()` and rendered by surface and framing
|
||||
- [pyo3](references/pyo3.md)
|
||||
- Two `create_exception!` classes with the payload in `args`, `ValueError` for the caller's fault and `RuntimeError` otherwise, decoded in `failures.py`
|
||||
- [testing](references/testing.md)
|
||||
- Test predicates, chain shape and redaction on errors built the way production builds them
|
||||
- [recipes](references/recipes.md)
|
||||
- Common tasks walked through end to end
|
||||
|
||||
The workspace uses thiserror 2 for every error type, and no library crate uses `anyhow` or `Box<dyn Error>` in a public signature. The [derive docs](https://docs.rs/thiserror/latest/thiserror/) are the authority on macro behavior, and the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html) on what `Display` and `source()` mean
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
# Boundaries
|
||||
|
||||
A boundary is where an error becomes something another program reads: an HTTP response, a Python exception, a stored or logged message. Each boundary has one mapping function per error type it receives, and nothing outside those functions formats an error for someone else
|
||||
|
||||
| Boundary | Mapping | Reference |
|
||||
|---|---|---|
|
||||
| Inference HTTP response | `Error::status()` and its renderers in `crates/gateway-inference/src/error.rs` | [gateway](gateway.md) |
|
||||
| Other gateway surfaces | the `IntoResponse` of each `crates/gateway-*` crate's `Error` | [gateway](gateway.md) |
|
||||
| Python bridge, inference and OCR routes | `route_failure` in `crates/python-bridge/src/errors.rs` and in `crates/python-bridge/src/routes/ocr/errors.rs` | [pyo3](pyo3.md) |
|
||||
| Python bridge, everything else | one `From<Error> for PyErr` or mapping function per crate error, all in `crates/python-bridge` | [pyo3](pyo3.md) |
|
||||
| Python host runtime | `From<Error> for PyErr` in `crates/host-python/src/error.rs` | [pyo3](pyo3.md) |
|
||||
|
||||
## Asking the nested error
|
||||
|
||||
The inference gateway and the Python bridge's route mapping both receive a nested route error and decide by the same predicates in the same order, so a new failure lands correctly on both once its predicates answer right
|
||||
|
||||
| Predicate | Gateway | Bridge |
|
||||
|---|---|---|
|
||||
| `declined()` | 501 | `RustBridgeDeclined`, without the route's public mapping |
|
||||
| `upstream_response()` | the provider's status and body | `RustUpstreamError` |
|
||||
| `is_missing_credential()` | 401 | no class of its own, a `RuntimeError` because `is_request()` answers false |
|
||||
| `is_request()` | 400 | `ValueError` |
|
||||
| none of these | 500 | `RuntimeError` |
|
||||
|
||||
If you find yourself adding an arm for your variant at a boundary, check whether a predicate is missing first
|
||||
|
||||
The two boundaries are not identical. The gateway has variants of its own for body parsing, an unknown model and an unsupported path, and the OCR bridge has Python parity arms such as `FileNotFoundError` for a missing file. Those arms match the boundary's own variants or the top level of a public enum, never a nested error's variants. A missing API key is 401 at the gateway and `RuntimeError` at both bridges, pinned by tests to what Python's own path does
|
||||
|
||||
## Declines
|
||||
|
||||
A failure before anything reached the provider is a decline, and `declined()` on the route error names its reason. It means another host may serve the same request as it was received, so each boundary hands it to the host unchanged: the gateway answers 501 and the Python bridge raises `RustBridgeDeclined`, never a public exception or a `ValueError`. See [gateway](gateway.md) and [pyo3](pyo3.md) for each host's side
|
||||
|
||||
Anything after the send is terminal, never a decline. `declined()` answers `None` for a response fault and for a post-call hook failure even when the wrapped error was a decline, and every route that transforms the provider's answer collapses its errors with `Error::after_send`, which keeps a provider answer and a response fault as they are and nests anything else as the source of an invalid response, so a decline can never describe a call that was billed. Ask `declined()`, do not match the variant behind it, because that after-send cut lives in the predicate
|
||||
|
||||
## What reaches the client
|
||||
|
||||
A boundary renders anything that is not a provider answer with `litellm_error::Report`, so the reader sees every layer of the chain once. That is why messages are written for a reader, each layer leaves its cause to `source()`, and reqwest URLs are stripped on entry. A boundary that must not show the chain, such as the bridge's TLS settings error that names the setting but not the file, renders fixed text of its own and says so in its test
|
||||
|
||||
Provider bodies are passed through intact for gateway and Python error mapping. `litellm_http::Error::upstream(status, body)` retains the complete body, but its `Display`, `Debug` and `Report` diagnostics show only the first 256 characters and mark the cut. Do not log the raw body or `RustUpstreamError` args. OCR's `Provider` error also keeps the whole body for Python's `get_error_class`, bounded by the connection's response size limit
|
||||
|
||||
Never format `Debug` of an error into a response or a Python exception. A type whose sources may hold secrets derives `veil::Redact`, so its `Debug` is safe to log
|
||||
|
||||
## Logging
|
||||
|
||||
Propagate with context, and let the terminal handler own the one log or trace event for a failed call. Do not log at every `?`, and do not log an error you also return. Diagnostic tracing excludes raw error strings and bodies, see [`rust-tracing`](../../rust-tracing/SKILL.md)
|
||||
|
|
@ -0,0 +1,66 @@
|
|||
# Classification
|
||||
|
||||
## Ask, don't match
|
||||
|
||||
Downstream crates ask a lower error questions instead of matching its nested variants. `litellm_http::Error` answers these
|
||||
|
||||
| Predicate | Caller decides |
|
||||
|---|---|
|
||||
| `dispatched()` | whether retrying elsewhere could repeat a billed call |
|
||||
| `is_request()` | the caller's input is at fault, so 400 and no reroute |
|
||||
| `is_response()` | the provider answered with something we cannot use |
|
||||
| `upstream_response()` | the provider's own status and body to pass through |
|
||||
| `status()` | an HTTP status to report |
|
||||
| `is_connect()` | the connection was never made, for a host's own connect-error mapping |
|
||||
| `tls_source()` | which TLS setting to blame in a config error |
|
||||
| `is_computed_header()` | the caller forwarded a header the signer owns, which `litellm_inference::Error` turns into a decline |
|
||||
|
||||
Upper types add the questions their callers ask. `litellm_inference::Error` and the OCR error answer `is_missing_credential()`, and `litellm_llms::Error` and `litellm_inference::Error` answer `declined()`, which the boundaries ask first instead of reaching into `Auth` or `Declined`, see [boundaries](boundaries.md)
|
||||
|
||||
```rust
|
||||
// reaches through two crates' variants and breaks when either changes
|
||||
Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => StatusCode::UNAUTHORIZED,
|
||||
|
||||
// asks the question it cares about
|
||||
Self::Route(error) if error.is_missing_credential() => StatusCode::UNAUTHORIZED,
|
||||
```
|
||||
|
||||
A predicate is an exhaustive `match` with no `_` arm, never `matches!`, so a new variant does not compile until it has answered every question on purpose
|
||||
|
||||
```rust
|
||||
pub fn is_request(&self) -> bool {
|
||||
match &*self.0 {
|
||||
Repr::Request(_) | Repr::BlockedUrl | Repr::DownloadTooLarge => true,
|
||||
Repr::Connect(_) | Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => false,
|
||||
// ... every other variant, listed by name
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Add a predicate when a real caller needs a new distinction, and name it after the fact the caller needs, in the caller's words, not after the variant's spelling. An upper type's predicates call through to the nested error, as in `Self::Http(error) => error.is_request()` and `Self::Auth(error) => error.is_missing_credential()`
|
||||
|
||||
A public enum may be matched one level down by the type that owns the variant, such as the gateway matching its own `Unsupported` or the OCR bridge matching `FileRead`. Going a second level down, into a variant of a variant, is what to avoid, the gateway used to do that for `MissingApiKey` and the old `Unsupported` route variant
|
||||
|
||||
## Facts, not policy
|
||||
|
||||
A predicate reports what happened. The decision to retry, reroute or fall back depends on things the error cannot know, such as whether the request is idempotent, whether its body can be replayed, the deadline and the attempt budget. That decision lives with the caller. `dispatched()` exists for the router's reroute decision, so every transport failure answers it even before that caller is wired
|
||||
|
||||
A timeout is the main trap. A request that timed out may already have reached the provider and been billed, so `Timeout` answers `dispatched() == true` and is not safe to retry as if it were a connect failure. Only a failure that provably happened before any byte reached the provider answers `false`: connecting, building or signing the request, and media downloads, which fetch the caller's document before the provider call is made
|
||||
|
||||
```rust
|
||||
pub fn dispatched(&self) -> bool {
|
||||
match &*self.0 {
|
||||
Repr::Network(_) | Repr::Timeout(_) | Repr::Status { .. } => true,
|
||||
Repr::Connect(_) | Repr::RequestBody(_) | Repr::Download(_) | Repr::BlockedUrl => false,
|
||||
// ... every other variant, listed by name
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
When the transport cannot tell whether a request reached the provider, it must answer `true`. Wrongly answering `false` risks a duplicate billed call, wrongly answering `true` only costs a reroute
|
||||
|
||||
Keep one predicate per question. `status()` currently mixes the provider's status, the media download's status and a synthesized 408 for timeouts, so a caller that wants "did the provider answer" must use `upstream_response()` instead. Do not add more meanings to an existing predicate, add a new one
|
||||
|
||||
## Equality
|
||||
|
||||
`Arc<Repr>` errors compare with `Arc::ptr_eq`, so a clone equals its original and two separately built errors never compare equal. That exists only so a wrapping enum such as `litellm_inference::Error` can keep its `PartialEq` derive. Do not give a new type pointer equality for any other reason, and do not rely on it in tests. Assert predicates instead, see [testing](testing.md)
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
# Conversions
|
||||
|
||||
## One From per source and target pair
|
||||
|
||||
A `From<Foreign> for Ours` exists when the conversion is obvious, infallible and needs no context from the call site. It does all classification and sanitizing for that pair, so every call site is `?` or `.map_err(Ours::from)`
|
||||
|
||||
```rust
|
||||
impl From<reqwest::Error> for Error {
|
||||
fn from(error: reqwest::Error) -> Self {
|
||||
let error = error.without_url();
|
||||
if error.is_timeout() {
|
||||
Repr::Timeout(error).into()
|
||||
} else if error.is_connect() || error.is_builder() {
|
||||
Repr::Connect(error).into()
|
||||
} else {
|
||||
Repr::Network(error).into()
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
A call site never re-derives that classification
|
||||
|
||||
```rust
|
||||
// every route used to repeat this, and each drifted a little
|
||||
.map_err(|err| {
|
||||
if err.is_connect() || err.is_builder() {
|
||||
Error::Transport(Connect(err.to_string()))
|
||||
} else {
|
||||
Error::Transport(Network(err.to_string()))
|
||||
}
|
||||
})?
|
||||
|
||||
// now
|
||||
.map_err(litellm_http::Error::from)?
|
||||
```
|
||||
|
||||
The same foreign type may convert into several of our types, one `From` each. What is not allowed is two places converting the same pair differently
|
||||
|
||||
## When the call site knows something the From cannot
|
||||
|
||||
One `io::Error` can mean "reading the CA bundle failed" or "resolving the media host failed". A global `From<io::Error>` cannot know which, so there is none. Name the operation at the call site instead
|
||||
|
||||
```rust
|
||||
.map_err(Repr::Resolve)?
|
||||
.map_err(|source| Repr::Read { path, tls_source, source })?
|
||||
```
|
||||
|
||||
Do not drop meaningful context just to use `?`
|
||||
|
||||
## Reclassify once, in the shared From
|
||||
|
||||
When a lower error needs a different meaning for every caller, do it in the upper type's `From`, not per route
|
||||
|
||||
```rust
|
||||
impl From<litellm_http::Error> for Error {
|
||||
fn from(error: litellm_http::Error) -> Self {
|
||||
if error.is_computed_header() {
|
||||
return Self::Declined(Decline::new("request forwards a header AWS SigV4 computes"));
|
||||
}
|
||||
Self::Http(error)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
All four inference routes classify a header the signer computes the same way, because none of them does it themselves
|
||||
|
||||
## Nest, translate or add context, never mirror
|
||||
|
||||
| The upper layer | Do | Example |
|
||||
|---|---|---|
|
||||
| has nothing to add | nest in one transparent variant and delegate predicates | `Self::Http(error) => error.is_request()` |
|
||||
| changes what the failure means | translate to an upper variant, keep the cause attached when it helps | computed header -> `Declined` |
|
||||
| knows something the lower one does not | a variant with the extra fields plus the cause as its source | `Read { path, tls_source, source }` |
|
||||
|
||||
Never re-declare a lower crate's variants in the upper type with an upper-layer name and a 1:1 mapping function. OCR once did that for every media download failure, `MediaError::BlockedUrl => Error::BlockedDocumentUrl` and so on, which doubled every change and let the two lists drift.
|
||||
|
||||
## Third-party and host faults
|
||||
|
||||
Converting a fault from a crate we do not own into a string variant is still flattening. Prefer a variant that carries the fault, and when an existing `ErrorDetail` constructor fits, use it
|
||||
|
||||
```rust
|
||||
// flattens
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))
|
||||
|
||||
// keeps the cause typed
|
||||
.map_err(|error| Error::InvalidRequest(ErrorDetail::invalid("header", error)))
|
||||
```
|
||||
|
||||
## Through a foreign error type
|
||||
|
||||
A trait we implement for another crate sometimes forces its error type on us, such as `io::Error` from a body stream or `sqlx::Error` from a driver. Wrap ours as that type's source (`io::Error::other(error)`, `azure_core::Error::new(kind, error)`) and recover it with `downcast_ref` on the other side, never `with_message(error.to_string())` or `from_msg`. The error must come back out as the same value it went in as, so one failure classifies the same way on every path
|
||||
56
litellm-rust/.agents/skills/rust-error/references/gateway.md
Normal file
56
litellm-rust/.agents/skills/rust-error/references/gateway.md
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
# Gateway responses
|
||||
|
||||
The gateway is where a Rust error becomes an HTTP response. This page is about the path an error takes to get there and the one place it is classified. The gateway crates' own conventions, such as which envelope each API surface uses, which extractors exist and how the router is composed, live in the [`rust-gateway`](../../rust-gateway/SKILL.md) skill
|
||||
|
||||
## The flow
|
||||
|
||||
```text
|
||||
reqwest / provider failure
|
||||
-> litellm_http::Error From<reqwest::Error>, strips the URL, classifies once
|
||||
-> litellm_inference::Error::Http(_) From, one transparent variant, predicates delegate
|
||||
-> gateway_inference::Error::Route(_) From, one variant per nested crate
|
||||
-> handler returns Result<_, Error> every `?` ends here
|
||||
-> IntoResponse status() picks the code, a renderer builds the body
|
||||
```
|
||||
|
||||
Each hop is one `From` that nests the lower error in one variant, see [conversions](conversions.md). No hop re-derives what the lower crate already decided, and nothing between the provider and `IntoResponse` formats an error for the client. That is what makes "one `Error` with `IntoResponse` per gateway crate that serves a router" work: a handler can `?` any failure from any depth and the response comes out right
|
||||
|
||||
## Classify once
|
||||
|
||||
`status()` is pure and runs once per error. Variants the gateway declares itself, such as a bad body, an unknown model or an oversized body, are matched at the top of its own enum. A nested route error goes to `route_status` and an OCR error to `ocr_status`, which ask the predicates [boundaries](boundaries.md) lists
|
||||
|
||||
| Answer | Status |
|
||||
|---|---|
|
||||
| `declined()` | 501, for a route error, the gateway matched and parsed the request but core cannot serve it, nothing reached the provider |
|
||||
| `upstream_response()` | the provider's own status, 502 when it is not a valid code |
|
||||
| `is_missing_credential()` | 401 |
|
||||
| `is_request()` | 400, for a route error |
|
||||
| `http_status_code()` | the transport's status, for an OCR error, such as 408 for a timeout |
|
||||
| otherwise | 500 |
|
||||
|
||||
Do not reach into a nested enum's variants from the gateway. When a new nested failure needs a different status, add a predicate on the nested error, see [classification](classification.md)
|
||||
|
||||
## Render by surface and framing
|
||||
|
||||
Rendering is separate from classification. The surface picks the envelope and the moment picks the framing, a JSON body before the headers went out and an `event: error` SSE frame after. Both framings of one surface render the same status and `Report` message, so a streaming failure and a JSON failure of the same error agree on everything but framing
|
||||
|
||||
| Surface | JSON | SSE |
|
||||
|---|---|---|
|
||||
| Anthropic, `/v1/messages` | `body(request_id)` | `sse_frame()` |
|
||||
| OpenAI, chat, responses and the rest | `openai_response()` | `responses_sse_frame()`, the Responses API `error` event |
|
||||
|
||||
The message of anything that is not a provider answer is the `litellm_error::Report` of the error, and in the OpenAI and Anthropic JSON envelopes the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Responses `error` event has the fixed type `error` and the status as its `code`. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else. The OpenAI envelope always wraps, because that is what the Python proxy returns there
|
||||
|
||||
## Every way in uses the same Error
|
||||
|
||||
A request can fail before the handler runs, in the handler, or after the headers went out. All three end in the same `Error`, so they get the same classification and envelope
|
||||
|
||||
- An extractor sets `type Rejection = Error`, so a bad body is the surface's envelope rather than axum's plain text
|
||||
- A path the gateway knows but does not serve returns `Error::Unsupported(path)` from one handler, so the client gets a 501 in the surface's envelope rather than a bare 404
|
||||
- A stream that fails after its headers went out hands the surface's SSE frame of `Error::from(error)` to `litellm_host_http::Sse::new`. `litellm_host_http::Error<E>` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the `From` maps the second to `Internal`
|
||||
|
||||
The axum hooks behind each line are in [`rust-gateway`](../../rust-gateway/SKILL.md)
|
||||
|
||||
## Testing
|
||||
|
||||
`status()` and `body()` are pure, so unit tests call them on errors built the way production builds them, one `#[case]` per fault, as `status_follows_who_is_at_fault`, `an_anthropic_provider_body_passes_through` and `other_failures_get_the_anthropic_envelope_with_their_message` in `crates/gateway-inference/src/error.rs` do. Assert the status and the fields the surface promises, not the whole body, except where passing the body through untouched is the contract. Driving a whole router belongs to the gateway crate's own tests, see [`rust-gateway`](../../rust-gateway/references/testing.md)
|
||||
|
|
@ -0,0 +1,112 @@
|
|||
# Messages and causes
|
||||
|
||||
## Messages belong to the variant
|
||||
|
||||
Every message template is written once, in the variant's `#[error(...)]`. A caller passes the typed facts the template needs and nothing else
|
||||
|
||||
```rust
|
||||
// none of these may appear at a call site
|
||||
Error::Network("WebSocket handshake returned a wrong Sec-WebSocket-Accept".into())
|
||||
Error::Request(format!("bad header {name}"))
|
||||
Error::Signature(error.to_string())
|
||||
|
||||
// the call site names what happened, the variant owns the words
|
||||
Repr::Handshake(HandshakeFault::Accept).into()
|
||||
Error::signature(error)
|
||||
Error::NonAsciiHeader { name }
|
||||
```
|
||||
|
||||
When a check fails for a distinct reason, that reason is a variant or a field, never a string inside a generic variant. A `Message(String)` variant, or a `From<String>` / `From<&str>` into an error, is the hole every stringly call site goes through, so neither exists
|
||||
|
||||
Format arguments can call functions and reach into fields, so a hand-written `Display` is rarely needed
|
||||
|
||||
```rust
|
||||
#[error("could not read {}", path.display())]
|
||||
#[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))]
|
||||
```
|
||||
|
||||
Keep messages short and lowercase, with no trailing period, unless Python parity pins the exact text. A reporter appends `: {cause}` after them. Never format `Debug` (`{:?}`) of an error, or of a type with no `Display`, into a message
|
||||
|
||||
## The cause is the source, never the message
|
||||
|
||||
This is the rule from the [`std::error::Error` docs](https://doc.rust-lang.org/std/error/trait.Error.html): an error's `Display` describes that layer only, and the error underneath is returned by `source()`. thiserror wires `source()` from a field marked `#[source]`, a `#[from]` field, or a field named `source`
|
||||
|
||||
```rust
|
||||
// prints the cause twice under any reporter
|
||||
#[error("could not read {}: {source}", path.display())]
|
||||
Read { path: PathBuf, source: io::Error },
|
||||
#[error("invalid JSON: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
|
||||
// this layer's words, the cause is the source
|
||||
#[error("could not read {}", path.display())]
|
||||
Read { path: PathBuf, source: io::Error },
|
||||
#[error("invalid JSON")]
|
||||
Json(#[from] serde_json::Error),
|
||||
```
|
||||
|
||||
| Field | thiserror does | Use for |
|
||||
|---|---|---|
|
||||
| `source: E` or `#[source] cause: E` | exposes it as `source()` | a cause with context fields next to it |
|
||||
| `#[from] E` | implies `#[source]` and adds `From` | context-free wrapping, see [conversions](conversions.md) |
|
||||
| `#[error(transparent)]` | forwards both `Display` and `source()` to the inner value | nesting another of our errors that adds no text |
|
||||
| no attribute, any other name | nothing, the field is plain data | facts the message renders, such as a path or a status |
|
||||
|
||||
`#[error(transparent)]` is an invisible layer, not an extra chain node, so a transparent wrapper's `source()` is the inner value's `source()`
|
||||
|
||||
A field that holds an error but is neither rendered nor a source hides that error from every reader. Every error-typed field is a source. A fieldless reason enum, such as `HandshakeFault` or `PagesFault`, is data and not a cause, so the variant renders it in its message
|
||||
|
||||
`source()` is one chain, so an error with several causes, such as `CredentialChain` trying each credential in turn, makes the most relevant one its source and renders the others in its message, each with `Report`
|
||||
|
||||
## Keep causes typed
|
||||
|
||||
Store the lower-level error itself, never its string, and never drop it with `map_err(|_| ...)`
|
||||
|
||||
| Flattening or dropping loses | Because |
|
||||
|---|---|
|
||||
| the chain | a `String` has no `source()` |
|
||||
| downcasting | the type is gone |
|
||||
| redaction | `reqwest::Error`'s `Display` includes the request URL, which can carry credentials in the path or query |
|
||||
|
||||
If the cause is not `Clone` and the error must be, hold it as `#[source] litellm_error::Source`, a cloneable `Arc<dyn Error + Send + Sync>` that is transparent to `Display` and `source()`, or make the crate's error opaque. `litellm_error::find::<T>(&error)` walks the chain and looks through `Source`, so it is how anyone reaches a typed cause. A third-party `Box<dyn Error + Send + Sync>` converts with `Source::from`, and an `anyhow::Error` a dependency returns goes through `Box` first, `Source::from(Box::<dyn Error + Send + Sync>::from(error))`, and a dependency's own boxed alias such as `tokenizers::Error` is held as it is. A source boxed to keep the variant small is a `Box<E>` node of the chain, so it is reached as `find::<Box<E>>`. Never put `#[error(transparent)]` on a `Source`, the erased error would then be no node of the chain and `find` could not reach it. Do not stringify or drop a cause to get `Clone` or `PartialEq`
|
||||
|
||||
```rust
|
||||
// loses the type and leaks the URL
|
||||
.map_err(|error| Error::Network(error.to_string()))
|
||||
// loses everything
|
||||
.map_err(|_| Error::Transport)
|
||||
|
||||
// keeps the cause, the From impl strips the URL
|
||||
.map_err(litellm_http::Error::from)
|
||||
```
|
||||
|
||||
A cause that must not be shown at all, such as a secret manager's response, is still a source. The type derives [`veil::Redact`](https://docs.rs/veil) with `#[redact]` on the field, so `Debug` hides it, and the boundary decides how much of the chain to show, see [boundaries](boundaries.md)
|
||||
|
||||
## Strip reqwest URLs on entry
|
||||
|
||||
Every `reqwest::Error` enters our types through `From<reqwest::Error> for litellm_http::Error`, which calls `without_url()`. A crate that talks HTTP holds `litellm_http::Error`, never `reqwest::Error`. When a dependency's own error wraps a `reqwest::Error`, as vaultrs and rustify do, one `From` takes the reqwest error out and converts it, as `secrets-hashicorp` does. That way so every later `Display`, `Debug` or chain walk is safe from that URL
|
||||
|
||||
`without_url()` only removes the URL reqwest attached. It does not scrub headers, response bodies, provider messages, file paths or anything a variant stores next to it. Those need their own policy, see [boundaries](boundaries.md)
|
||||
|
||||
## Rendering
|
||||
|
||||
Anything that shows an error to someone else renders the whole chain with `litellm_error::Report`, which writes each layer once joined by `: `, the single-line form of the unstable [`std::error::Report`](https://doc.rust-lang.org/std/error/struct.Report.html). `to_string()` of an error is only its top layer and is never what a client, a Python exception or a log line shows
|
||||
|
||||
```rust
|
||||
let message = litellm_error::Report::new(&error).to_string();
|
||||
// "could not reach the provider: error sending request: client error (Connect): tcp connect error: Connection refused"
|
||||
```
|
||||
|
||||
## #[from] limits
|
||||
|
||||
`#[from]` only works on a variant whose only fields are the source and an optional backtrace. A variant that needs context such as a path, a phase or a `TlsSource` takes its cause through an explicit `map_err` or a named constructor
|
||||
|
||||
```rust
|
||||
std::fs::read(&path).map_err(|source| Repr::Read {
|
||||
path: path.clone(),
|
||||
tls_source: TlsSource::CaBundle,
|
||||
source,
|
||||
})?;
|
||||
```
|
||||
|
||||
`#[backtrace]` and automatic `provide()` need a nightly compiler and are not used here
|
||||
76
litellm-rust/.agents/skills/rust-error/references/pyo3.md
Normal file
76
litellm-rust/.agents/skills/rust-error/references/pyo3.md
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
# Python bridge
|
||||
|
||||
`crates/python-bridge` turns a Rust error into a Python exception through one mapping function per error type. A `#[pyfunction]` returns `Err(PyErr)` and pyo3 raises it, so those functions are the whole boundary
|
||||
|
||||
| Error | Mapping |
|
||||
|---|---|
|
||||
| `litellm_inference::Error` from a route | `errors::route_failure`, with `route_error_to_pyerr` as its shorthand for routes that set no marker |
|
||||
| the OCR error | `route_failure` and `to_pyerr` in `routes/ocr/errors.rs` |
|
||||
| token counter, cache, spend, secrets, argument checks | one `From<Error> for PyErr` or mapping function next to the code that registers the functions |
|
||||
|
||||
A route host never builds a `PyErr` from a route error itself. It calls the mapping, and the mapping is the only place that knows the exception classes and their attributes
|
||||
|
||||
## Exception classes
|
||||
|
||||
| Class | Meaning | Payload |
|
||||
|---|---|---|
|
||||
| `RustUpstreamError` | the provider call was issued and the provider answered with a failure | `args == (status, body)` and a `headers` attribute of `(name, value)` pairs, always set |
|
||||
| `RustBridgeDeclined` | the route stopped before calling the provider, so the host may run its own path | `args == (message,)` |
|
||||
| a standard class from `pyo3::exceptions` | whichever Python class means the same thing, such as `ValueError` or `TypeError` for the caller's input, `FileNotFoundError` or `OSError` for a file, `KeyError` or `OverflowError` where Python raises them, and `RuntimeError` for anything else | the rendered chain, or a parity text where Python pins one |
|
||||
|
||||
The two native classes come from `create_exception!` in `crates/python-bridge/src/errors.rs` and are exported with `#[pymodule_export]`, which is how the [pyo3 guide](https://github.com/PyO3/pyo3/blob/main/guide/src/exception.md) makes an exception importable. Their docstrings are the contract Python reads. `errors::upstream_error(py, status, body, headers)` is the one constructor of `RustUpstreamError`, so the payload cannot differ between routes
|
||||
|
||||
Python decodes a `RustUpstreamError` in `litellm/rust_bridge/failures.py`, which validates `args` and `headers` with a `TypeAdapter` and rebuilds the provider's `httpx.Response`. `litellm/rust_bridge/runtime.py` catches `RustBridgeDeclined` and reads its reason. `_raise_upstream` in `runtime.py` still reads `args` a second time by hand for the `attempt` path, which is the one decode to fold into `failures.py`. A new fact goes into `args` or `headers`, and the decoder learns to read it in the same change
|
||||
|
||||
Two routes also flag a request fault with a boolean attribute, read by their Python route module: `ocr_request_format_error` in `host/routes/ocr.py` and `messages_request_error` in `host/routes/messages.py`. OCR additionally sets `status_code` and `message` for Python parity, and `native_settings_error`, which only the OCR host reads back to tell a settings failure from the caller's exception. Do not add another attribute, nothing on the Python side types them
|
||||
|
||||
## Why args, not a typed subclass
|
||||
|
||||
The bridge builds with pyo3's `abi3-py310` feature. Under abi3, `#[pyclass(extends = PyException)]` is only allowed on Python 3.12 and newer, so a class with typed fields is not available. Even where it is allowed, pyo3 passes an empty tuple to the native base's `tp_new`, so `.args` and `str()` would need manual overrides. The payload therefore travels in `args`, which a `create_exception!` class fills from `new_err`
|
||||
|
||||
## Mapping order
|
||||
|
||||
A core decline is returned first, as `RustBridgeDeclined` with the decline's reason as its only arg, without calling the route's public mapping. Then a Python exception raised inside the call, such as a secret manager callback's, is returned as it is. Then the provider's answer passes through as `RustUpstreamError`. On `/v1/messages` a request the route rejected becomes a `ValueError` flagged with `messages_request_error`. Everything else is split by `is_request()` into `ValueError` or `RuntimeError`, with the message rendered by `litellm_error::Report`. A missing API key is a `RuntimeError` because `is_request()` answers false for it, and `missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors` pins that
|
||||
|
||||
OCR puts Python parity arms in front of that order, each matching a top-level variant of its own public enum
|
||||
|
||||
| Arm | Exception |
|
||||
|---|---|
|
||||
| the provider answered | `RustUpstreamError` with the provider's headers |
|
||||
| `RequestFormat` | `ValueError` flagged with `ocr_request_format_error` |
|
||||
| `FileRead` whose cause is `NotFound` | `FileNotFoundError` naming the path |
|
||||
| other `FileRead` | `OSError` with the io message |
|
||||
|
||||
Both routes ask `litellm_auth::Error::is_request()` through their own `is_request()`, so a missing key or API base is a `RuntimeError` on both, and a missing setting the caller can pass, such as a region, is a `ValueError`
|
||||
|
||||
Use the standard class whose Python meaning matches. `PyFileNotFoundError` for a missing file is right, a `RuntimeError` whose text says the file is missing is not. `import_exception!` exists for an exception defined in Python, none is used today because `failures.py` maps to LiteLLM's public exceptions through `litellm.exception_type`
|
||||
|
||||
## Declining
|
||||
|
||||
`RustBridgeDeclined` means no provider call was made and the host may run its own path for the same call. The tokenizer, embeddings and token counter routes raise it from their own pre-checks, and `route_failure` raises it for a route error whose `declined()` answers, so a decline from deep inside a transformation reaches Python the same way. Never raise it for a failure after the request went out, the host would repeat a billed call. `a_core_decline_hands_the_call_back_to_python_before_the_public_mapping` and `route_failures_are_terminal` in `errors.rs` pin both sides
|
||||
|
||||
Route hosts pass a terminal failure through the Python route module's `map_failure`, which runs `litellm.exception_type` and returns the public exception a Python caller expects. A decline must skip that step. Once wrapped it looks like a provider failure to the dispatcher and the host no longer falls back, so `route_failure` returns it before it calls `public`
|
||||
|
||||
## Rules
|
||||
|
||||
- Never format `Debug` of an error into a `PyErr`
|
||||
- Never build a `PyErr` from a string at a call site. A failed argument check is a variant of a typed error, as `ProjectionError` in `coercion.rs` is, mapped once
|
||||
- A Python exception raised inside the call stays that exception. Return it as is, never wrap it in a Rust message or stringify it into a new `ValueError`. The one policy exception is a caller-supplied callable, whose failure `callable::wrap_failure` reports as a `RuntimeError` with the original as `__cause__`
|
||||
- A `panic!`, `unwrap` or `expect` inside a `#[pyfunction]` or `#[getter]` reaches Python as `pyo3_runtime.PanicException`, which is not an `Exception` subclass. Model the failure as a value instead
|
||||
- When the mapping itself calls Python, as `setattr` does, propagate its failure with `?` rather than discarding it with `.ok()`
|
||||
|
||||
## Testing
|
||||
|
||||
Build the Rust error the way production builds it, map it, and assert the contract Python sees. `Python::initialize()` once, then `Python::attach`
|
||||
|
||||
```rust
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let failure = route_error_to_pyerr(litellm_http::Error::upstream(429, "slow down".into()).into());
|
||||
assert!(failure.is_instance_of::<RustUpstreamError>(py));
|
||||
let args: (u16, String) = failure.value(py).getattr("args").unwrap().extract().unwrap();
|
||||
assert_eq!(args, (429, "slow down".into()));
|
||||
});
|
||||
```
|
||||
|
||||
One `#[case]` per fault class, asserting `is_instance_of` for the class and `extract` for the payload. The Python side of the contract is tested in `tests/unit/rust_bridge`, through `failures.py`
|
||||
64
litellm-rust/.agents/skills/rust-error/references/recipes.md
Normal file
64
litellm-rust/.agents/skills/rust-error/references/recipes.md
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
# Recipes
|
||||
|
||||
## A new failure in litellm_http
|
||||
|
||||
Say media downloads must reject a `Content-Type` outside an allow-list
|
||||
|
||||
1. Add a `Repr` variant with typed facts and the message, `#[error("media type {0:?} is not allowed")] DisallowedMediaType(String)`
|
||||
2. Return it where the check fails, `return Err(Repr::DisallowedMediaType(content_type.to_owned()).into())`
|
||||
3. Add an arm to every predicate's exhaustive `match`, the compiler lists them, the way its neighbor `DownloadTooLarge` answers. A media download happens before the provider call, so it is not in `dispatched()`. The caller chose the URL, so it is in `is_request()`
|
||||
4. Leave both boundaries alone. With those answers the gateway returns 400 and Python raises `ValueError`
|
||||
5. Test it with a mock server answering `Content-Type: application/x-msdownload` and assert `error.is_request()`
|
||||
|
||||
## A provider-specific failure
|
||||
|
||||
A second OCR provider needs a missing-key error whose text names its environment variable and how to set it
|
||||
|
||||
1. Look for an existing provider-neutral variant whose template fits. `litellm_auth::Error::MissingApiKey { provider, environment_variable }` renders `Missing {provider} API Key - Set ...`
|
||||
2. If none fits, find what varies between providers with the same failure, here the environment variable and the guidance after the dash, and add one neutral variant with those as fields, `#[error("Missing {environment_variable} - {guidance}")] MissingCredential { environment_variable: &'static str, guidance: &'static str }`
|
||||
3. Pass the provider's values from its transformation. The next provider passes its own and adds no variant, and `is_missing_credential()` lists `MissingCredential { .. }` once
|
||||
|
||||
## Wrapping a new lower crate in litellm_inference
|
||||
|
||||
Say routes start calling `litellm_cache` and its errors must surface
|
||||
|
||||
1. Add one variant, `#[error(transparent)] Cache(#[from] litellm_cache::Error)`
|
||||
2. Extend every predicate's `match` with an arm that calls through to the cache error's own predicates, adding a predicate there if it has none
|
||||
3. Do not copy the cache error's variants into `litellm_inference::Error`
|
||||
|
||||
## A foreign error that needs context
|
||||
|
||||
Reading a file fails with `io::Error`, and the message must name the file
|
||||
|
||||
1. A variant with the context and the cause, `Read { path: PathBuf, source: io::Error }` with `#[error("could not read {}", path.display())]`
|
||||
2. Build it with `map_err(|source| Repr::Read { path: path.clone(), source })`, not a `From<io::Error>`
|
||||
3. If the error must be `Clone`, the crate's error is opaque or the field is `litellm_error::Source`
|
||||
|
||||
## Migrating a stringly error
|
||||
|
||||
Call sites look like `Error::Network(error.to_string())`, `Error::Request(format!(...))` or `.map_err(|_| Error::Unavailable)`
|
||||
|
||||
1. List every call site of the string variant and group them by what actually failed
|
||||
2. Each group becomes a variant holding the typed cause or facts. Groups no caller tells apart share one
|
||||
3. Where a foreign error was stringified with classification logic around it, move that logic into the `From` and turn the call sites into `?`
|
||||
4. Delete the string variant. If it cannot be deleted yet, nothing new may construct it
|
||||
5. Re-run the boundary tests, since messages a client sees may change
|
||||
|
||||
## Moving a crate to source chains
|
||||
|
||||
A variant reads `#[error("invalid JSON: {0}")] Json(#[from] serde_json::Error)` or stores `cause` and renders `{cause}`
|
||||
|
||||
1. Drop the cause from the template, `#[error("invalid JSON")]`
|
||||
2. Make the cause the source, `#[from]`, `#[source]` or a field named `source`
|
||||
3. Find what renders this error for someone else and make it use `litellm_error::Report`
|
||||
4. Update tests that asserted the cause in `to_string()` to assert it in the report
|
||||
|
||||
## Making a public enum opaque
|
||||
|
||||
Downstream crates match nested variants, or a variant needs a non-`Clone` cause
|
||||
|
||||
1. Rename the enum to `pub(crate) enum Repr` and add `pub struct Error(Arc<Repr>)` with `#[error(transparent)]` and `From<Repr>`
|
||||
2. For each downstream `match`, find the question it asks and add a predicate. Replace the `match` with the predicate
|
||||
3. For each downstream construction, add a constructor taking typed data, or move the failure into this crate
|
||||
4. Swap test assertions to predicates and add the `#[cfg(test)] repr()` accessor for in-crate tests
|
||||
5. Implement `PartialEq` with `Arc::ptr_eq` only if a wrapping enum's derive needs it
|
||||
105
litellm-rust/.agents/skills/rust-error/references/shape.md
Normal file
105
litellm-rust/.agents/skills/rust-error/references/shape.md
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
# Error shape
|
||||
|
||||
## One top-level Error per crate
|
||||
|
||||
A crate's errors live in `src/error.rs` and are re-exported from `lib.rs` as `Error`. A caller of any public function in the crate should only ever see `Result<_, crate::Error>`, so one `?` and one set of predicates cover the whole crate
|
||||
|
||||
This deliberately departs from the advice to put an error type next to each unit of fallibility ([Sabrina Jewson, "Modular Errors"](https://sabrinajewson.org/blog/errors)). Our callers do not branch on which function failed, they ask the predicates in [classification](classification.md), and the boundaries map one type per crate. Splitting by function would multiply those predicates and mappings without giving any caller a decision it lacks. Layering, the other half of that advice, is kept: each crate nests the one below it, see [conversions](conversions.md)
|
||||
|
||||
That one type takes one of two forms
|
||||
|
||||
| Form | Use when | Example |
|
||||
|---|---|---|
|
||||
| Public enum | Callers outside the crate match variants, and every variant holds our own data or another crate's `Error` that is already `Clone` and comparable | `litellm_inference::Error`, `litellm_auth::Error` |
|
||||
| Opaque struct over a private `Repr` | The variants hold third-party causes, or the variants must be free to change without breaking callers | `litellm_http::Error` |
|
||||
|
||||
Both forms are still one top-level `Error`. The opaque form just moves the enum behind `Repr`, where it is `pub(crate)`
|
||||
|
||||
```rust
|
||||
// public enum
|
||||
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("missing required field: {0}")]
|
||||
MissingField(&'static str),
|
||||
#[error(transparent)]
|
||||
Http(litellm_http::Error),
|
||||
}
|
||||
|
||||
// opaque struct
|
||||
#[derive(Clone, Debug, thiserror::Error)]
|
||||
#[error(transparent)]
|
||||
pub struct Error(Arc<Repr>);
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum Repr {
|
||||
#[error("could not reach the provider")]
|
||||
Connect(#[source] reqwest::Error),
|
||||
#[error("upstream request failed with status {status}: {body}")]
|
||||
Status { status: u16, body: String },
|
||||
}
|
||||
```
|
||||
|
||||
Start with the public enum. Move to opaque when a variant needs a non-`Clone` third-party cause, or when downstream crates keep matching variants in ways that duplicate classification. `reqwest::Error`, `std::io::Error` and `serde_json::Error` are built the opaque way
|
||||
|
||||
In the opaque form other crates cannot name a variant, so they cannot build an error with their own message. Public constructors exist only for errors another crate genuinely produces, take typed data and never a message, as in `Error::upstream(status, body)`, `Error::computed_header(name)`, `Error::unsignable_body()` and `Error::signature(cause)`
|
||||
|
||||
`Arc` makes cloning a refcount bump, so `Repr` never has to be `Clone`. `Box<Repr>` is enough when the goal is only a small `Result`. `Arc` is for when callers clone
|
||||
|
||||
Every crate here is internal to the workspace, so do not mark errors or variants `#[non_exhaustive]`. It only forces a wildcard arm on the other workspace crates and turns off the exhaustiveness check that makes a new variant visit every `match`. When the variants must be free to change, use the opaque form
|
||||
|
||||
## What else may sit next to Error
|
||||
|
||||
These are not a second top-level error and need no folding in
|
||||
|
||||
- Leaf data types whose fields are typed and whose message is in `#[error(...)]`, such as `litellm_http::request::HeaderError`. A wrapping crate nests it as one variant
|
||||
- Plain data used by predicates, such as `litellm_http::TlsSource`
|
||||
- A private fault enum one variant carries, such as `PemFault` inside `Repr::InvalidPem`
|
||||
|
||||
## Variants are failure modes
|
||||
|
||||
A public variant is something a caller handles differently: a phase, a status code, retry versus no retry, fault on the request versus the environment, or a message Python parity pins exactly. Failures no caller tells apart share one variant and differ in their data
|
||||
|
||||
```rust
|
||||
// one variant per provider, so every new provider adds a variant and a match arm
|
||||
#[error("Missing Azure API Base - set AZURE_API_BASE")]
|
||||
MissingAzureApiBase,
|
||||
#[error("Missing Vertex API Base - set VERTEXAI_API_BASE")]
|
||||
MissingVertexApiBase,
|
||||
|
||||
// one failure mode, the provider and its guidance are data
|
||||
#[error("Missing {provider} API Base - {guidance}")]
|
||||
MissingApiBase { provider: &'static str, guidance: &'static str },
|
||||
```
|
||||
|
||||
An exact parity message does not justify a provider-specific variant when caller-supplied data can produce the same text
|
||||
|
||||
Inside a private `Repr` a variant costs no API, so one per distinct reason is fine when each one carries a different message. When several only differ in a word and share every predicate answer, prefer one variant with a fieldless reason enum, so predicates list one arm instead of four, as `Repr::Handshake(HandshakeFault)` does
|
||||
|
||||
```rust
|
||||
#[error("WebSocket handshake returned {0}")]
|
||||
Handshake(HandshakeFault),
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub(crate) enum HandshakeFault {
|
||||
#[error("a wrong Sec-WebSocket-Accept")]
|
||||
Accept,
|
||||
#[error("an invalid Upgrade header")]
|
||||
Upgrade,
|
||||
#[error("an invalid Connection header")]
|
||||
Connection,
|
||||
#[error("an invalid subprotocol")]
|
||||
Subprotocol,
|
||||
}
|
||||
```
|
||||
|
||||
## When to split
|
||||
|
||||
Split into a separate type only when different functions fail in disjoint ways and different callers see them, so no caller matches variants its function can never return. A type per function overdoes it
|
||||
|
||||
Name a split type after what went wrong, not after the function that returns it. `ConnectError` reads better than `OpenSessionError`, and a unit struct is fine for a single failure mode
|
||||
|
||||
A split type still converts into the crate's `Error` when a caller further up needs one `Result` type
|
||||
|
||||
`litellm-llms` has one recorded exception, the OCR error in `src/base_llm/ocr/error.rs`, kept apart until OCR folds into `litellm_llms::Error`, see `crates/inference-ocr/AGENTS.md`. Do not add another
|
||||
|
||||
A utility crate whose modules are unrelated, such as `core-utils`, has no crate-wide failure, so each module that can fail owns one error next to its code instead of a `src/error.rs`. The rest of this skill applies to each of them
|
||||
37
litellm-rust/.agents/skills/rust-error/references/testing.md
Normal file
37
litellm-rust/.agents/skills/rust-error/references/testing.md
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
# Testing errors
|
||||
|
||||
Assert behavior through predicates and accessors, on errors produced the way production produces them
|
||||
|
||||
```rust
|
||||
// guesses the message and the variant, and cannot be written once the type is opaque
|
||||
assert_eq!(error, Error::Transport(TransportError::Http { status: 429, body: "slow down".into() }));
|
||||
|
||||
// asserts what callers rely on
|
||||
assert_eq!(error.upstream_response(), Some((429, "slow down")));
|
||||
assert!(error.dispatched());
|
||||
assert!(!error.is_request());
|
||||
```
|
||||
|
||||
## Producing the error
|
||||
|
||||
In order of preference
|
||||
|
||||
1. A real failing call. Bind a `TcpListener`, drop it, then connect to its address for a connect failure. Accept and never answer for a timeout
|
||||
2. A real cause, such as `tungstenite::Error::ConnectionClosed` or a `serde_json::from_str("{")` error, passed through our `From`
|
||||
3. A public constructor, such as `Error::upstream(500, "boom".into())`
|
||||
|
||||
Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, reach the cause the test expects with `litellm_error::find::<T>(&error)`. A transparent layer is no node of the chain, so `find` reaches foreign causes underneath it but never the nested error of ours itself. Ask that error's predicates through the outer type instead, which is what callers do too
|
||||
|
||||
## What to cover
|
||||
|
||||
| Property | Test shape |
|
||||
|---|---|
|
||||
| classification | one `#[case]` per input class, asserting the predicate |
|
||||
| redaction | build from a URL holding `api_key=secret`, assert the message lacks `secret` and the path |
|
||||
| cause kept | `litellm_error::find::<T>(&error)` returns the expected cause |
|
||||
| rendered once | the cause's text appears exactly once in `litellm_error::Report::new(&error).to_string()` and not in `to_string()` |
|
||||
| boundary mapping | the gateway status and body, or the Python exception type and its `status` attribute |
|
||||
|
||||
Assert exact message text only where it is a contract, such as a message Python parity pins. Everywhere else assert that the message contains the facts it must carry. Matching a fieldless variant with `matches!` is fine, full `assert_eq!` on an error that holds a cause is not, since causes compare by pointer
|
||||
|
||||
A test that returns `Result` reports only `Debug` of the error on failure. Prefer `expect` with a short reason, or `unwrap_err` followed by assertions on the predicates
|
||||
22
litellm-rust/.agents/skills/rust-gateway/SKILL.md
Normal file
22
litellm-rust/.agents/skills/rust-gateway/SKILL.md
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
---
|
||||
name: rust-gateway
|
||||
description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust gateway
|
||||
|
||||
`crates/gateway` mounts endpoint routers and serves the listener. `crates/gateway-inference`, `gateway-mcp` and `gateway-ui` own their endpoints, `gateway-auth` is the inbound authentication middleware and `gateway-management` holds management logic that has no router yet. A crate's `AGENTS.md`, where it has one, says what belongs in it. Read the reference for the part you are touching
|
||||
|
||||
- [surfaces](references/surfaces.md)
|
||||
- Each API surface mirrors the envelope the Python proxy returns there, so the shapes differ on purpose
|
||||
- The error type string is derived from the status and provider bodies pass through bounded
|
||||
- [requests](references/requests.md)
|
||||
- Extractors reject with the crate's `Error`, handlers take the body as a `Result` when the error needs the request id
|
||||
- Every alias of a path is registered together, a known but unserved path goes to one handler
|
||||
- Streams relay through `litellm_host_http::Sse`, with the error frame rendered by the crate's `Error`
|
||||
- [testing](references/testing.md)
|
||||
- Build the router through `tests/support` and drive it with `oneshot`, assert the fields the surface promises
|
||||
|
||||
How an error is classified before it reaches a response is the [`rust-error`](../rust-error/SKILL.md) skill, in particular its [gateway](../rust-error/references/gateway.md) page. Axum's own conventions are in its [docs](https://docs.rs/axum/latest/axum/)
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
# Requests
|
||||
|
||||
## Extractors reject with the crate's Error
|
||||
|
||||
A custom extractor sets `type Rejection = Error`, so a bad body gets the surface's envelope rather than axum's plain-text rejection. Never take axum's own `Json`, `Query` or `Path` in a handler, their rejection is plain text. The rejection is the variant's source, never its string. `JsonObject` and `InferenceBody` in `crates/gateway-inference/src/request.rs` are the pattern
|
||||
|
||||
```rust
|
||||
impl<S: Send + Sync> FromRequest<S> for JsonObject {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
|
||||
let body = Bytes::from_request(request, state).await.map_err(|error| {
|
||||
if error.status() == StatusCode::PAYLOAD_TOO_LARGE {
|
||||
Error::BodyTooLarge
|
||||
} else {
|
||||
BodyError::Read(error).into()
|
||||
}
|
||||
})?;
|
||||
Ok(Self(object(&body)?))
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
When the handler needs something from the request in the error body, such as the request id on `/v1/messages`, take the extractor as `Result<JsonObject, Error>` so the handler maps the rejection itself, as `messages::create` does. `Infallible` is the right rejection for an extractor that only reads an optional header, as `RequestId` does
|
||||
|
||||
## Paths
|
||||
|
||||
Every alias the Python proxy registers for a route is registered together in `crates/gateway-inference/src/lib.rs`, such as the bare path, its `/v1` form and the `/engines/{model}` and `/openai/deployments/{model}` forms, otherwise one alias silently behaves differently from the others. A path the gateway knows but does not serve, such as `/v1/completions`, is routed to `request::unsupported`, which returns `Error::Unsupported(path)` so the client gets a 501 in the surface's envelope rather than a bare 404
|
||||
|
||||
`crates/gateway` merges the mounted routers and adds the shared inbound middleware as `route_layer`s, so authentication runs only for matched routes. Endpoint paths and parsing stay in the mounted crate, see the two crates' `AGENTS.md`
|
||||
|
||||
## Streaming
|
||||
|
||||
Once the response headers are sent the status cannot change. A streaming handler builds `litellm_host_http::Sse::new(Json, |error| Bytes::from(Error::from(error).sse_frame()))`, where the closure renders a mid-stream failure as an `event: error` frame with the same status and message a JSON failure on that surface carries, `sse_frame()` for Anthropic and `responses_sse_frame()` for the Responses API's `error` event. `litellm_host_http::Error<E>` separates `Call(E)`, the route's own error, from `Protocol`, a host bug, and the crate's `From` maps the second to `Internal`
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
# Envelopes follow the API surface
|
||||
|
||||
Each surface mirrors what the Python proxy returns there, so the shapes differ on purpose
|
||||
|
||||
| Surface | Envelope | Where |
|
||||
|---|---|---|
|
||||
| `/v1/messages` | Anthropic, `{"type": "error", "error": {"type", "message"}, "request_id"}` | `Error::body()` and `sse_frame()` in `crates/gateway-inference/src/error.rs` |
|
||||
| OpenAI-shaped routes | `{"error": {"message", "type", "param", "code"}}` | `Error::openai_response()` in the same file |
|
||||
| `/v1/responses` stream | the Responses `error` event, `{"type": "error", "code", "message", "param"}` | `Error::responses_sse_frame()` in the same file |
|
||||
| MCP | `{"detail": message}` | `crates/gateway-mcp/src/error.rs` |
|
||||
| UI | `{"error": {"message"}}`, plus `Retry-After` on 429 | `crates/gateway-ui/src/error.rs` |
|
||||
| auth rejections | plain text, the top-level message only | `crates/gateway-auth/src/error.rs` |
|
||||
|
||||
Every surface takes its status from one `status()` on its crate's `Error` and renders the message of anything that is not a provider answer with `litellm_error::Report`. In `gateway-inference` the error `type` string is derived from the status by `error_type()`, so a new failure never names its own type. The Anthropic envelope passes a provider body already in its shape through untouched and extracts the message from anything else, and the OpenAI envelope always wraps, as the Python proxy does. A provider body is bounded by `litellm_http::Error::upstream` when the error is built
|
||||
|
||||
Two surfaces show less than the chain on purpose, and a test in each pins it. An auth rejection goes to a caller who has not authenticated, so it never shows what a secret manager or key store said. An MCP upstream or authorization failure shows only its own layer, because rmcp renders the transport error, URL included, into its own `Display`, and an authorizer's cause is the same backend detail auth rejections hide. The inference surfaces show only the top layer of an auth rejection for the same reason
|
||||
|
||||
`/v1/messages` adds the request id to the envelope when the provider did not set one, which is why its handler takes the body as `Result<JsonObject, Error>` and renders the rejection itself, see [requests](requests.md)
|
||||
|
||||
A new surface gets its own renderer next to the others, named after the shape it produces, and its status still comes from `Error::status()`. Picking the status is the [`rust-error`](../../rust-error/references/gateway.md) skill's concern, not the surface's
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
# Testing
|
||||
|
||||
Test HTTP contracts in the crate that owns the endpoint: status codes, forwarded headers, error envelopes and streaming behavior. Inference and provider behavior is tested in the `inference-<fmt>` and `llms` crates, and mounting and server lifecycle in `crates/gateway`
|
||||
|
||||
`crates/gateway-inference/tests/support/mod.rs` builds a `Router` with injected deployments, pools and settings through `app()` and its variants, and `post()` drives it with `tower::ServiceExt::oneshot`. `json()` reads the response body back. Use these instead of building a router or a client by hand in a test
|
||||
|
||||
Assert the status and the fields the surface promises, not the whole body. A test that pins a whole envelope breaks on every unrelated field, one that checks `status`, `error.type` and the presence of `request_id` catches the regressions that matter
|
||||
|
||||
`status()` and `body()` of the error type are pure, so a classification case belongs in the error module's unit tests, see the [`rust-error`](../../rust-error/references/gateway.md) skill. Router tests cover what only the full path shows, such as an extractor rejection arriving in the right envelope or an SSE error frame after the headers went out
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-string-enums
|
||||
description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust string enums
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
---
|
||||
name: rust-tracing
|
||||
description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery
|
||||
metadata:
|
||||
internal: true
|
||||
---
|
||||
|
||||
# Rust tracing
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.a
|
|||
|
||||
For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md)
|
||||
|
||||
For error types, error variants and conversions between crate errors, follow [.agents/skills/rust-error/SKILL.md](.agents/skills/rust-error/SKILL.md)
|
||||
|
||||
For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of a hand-listed `ALL` array. Derive Strum string conversions instead of hand-written variant-to-string matches
|
||||
|
||||
Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all`
|
||||
|
|
@ -24,13 +26,3 @@ Use [`#[rstest]`](https://docs.rs/rstest/latest/rstest/attr.rstest.html) for new
|
|||
|
||||
- Never loop over inputs (`for`, `.iter().for_each`, `.all`) inside a test body. Give each input its own `#[case::name(...)]`, or use `#[values(...)]` for a cross product
|
||||
- Exception: a test pinning a Rust table against a repo-owned data file (for example `include_str!` of a JSON config) may iterate that file's entries
|
||||
|
||||
## Error definitions
|
||||
|
||||
- A crate's errors live in `src/error.rs`, defined with `thiserror`, and re-exported from `lib.rs`
|
||||
- Put message templates in the variant's `#[error(...)]` declaration. Callers pass only the small typed arguments needed to fill them, never `Error::Variant(format!(...))` or a preformatted message. Keep the smallest set of neutral variants that callers need to distinguish. Different wording or providers do not justify new variants
|
||||
- Default to one top-level `Error` enum per crate, with one variant per failure mode and a `#[error(...)]` message on each. A failure mode is something a caller handles differently (phase, status code, retry, a message Python parity pins exactly). Failures no caller tells apart share one variant and differ only in its message
|
||||
- Keep shared error enums minimal and provider-neutral. Provider names, credential types, configuration fields, and setup guidance belong in caller-supplied data, not dedicated variants or hardcoded shared messages. Reuse a variant for the same failure mode across providers, such as `MissingApiBase { provider: "Azure", guidance: "..." }`. An exact parity message does not justify a provider-specific variant when caller-supplied context can preserve it
|
||||
- Wrap a lower-level error as a variant with `#[from]` or `#[source]` instead of flattening it to a string
|
||||
- Exception: split into separate types when different functions fail in disjoint ways, especially when different callers see them. A shared enum would force every caller to match variants its function can never return
|
||||
- Name a split type after what went wrong (a unit struct is fine for a single failure mode), not after the function that returns it
|
||||
|
|
|
|||
107
litellm-rust/Cargo.lock
generated
107
litellm-rust/Cargo.lock
generated
|
|
@ -3150,6 +3150,7 @@ dependencies = [
|
|||
"defmt",
|
||||
"jiff-core",
|
||||
"jiff-static",
|
||||
"jiff-tzdb",
|
||||
"log",
|
||||
"portable-atomic",
|
||||
"portable-atomic-util",
|
||||
|
|
@ -3178,6 +3179,12 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jiff-tzdb"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fa8377070c6bae868759445e5a77f66d84f0b72f3a054bfb00e6d038b8282da7"
|
||||
|
||||
[[package]]
|
||||
name = "jni"
|
||||
version = "0.22.4"
|
||||
|
|
@ -3385,7 +3392,10 @@ dependencies = [
|
|||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-types",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"moka",
|
||||
"reqwest",
|
||||
|
|
@ -3405,6 +3415,7 @@ dependencies = [
|
|||
"azure_core",
|
||||
"azure_identity",
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"moka",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
|
|
@ -3423,6 +3434,7 @@ dependencies = [
|
|||
"google-cloud-auth",
|
||||
"http 1.4.2",
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"moka",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -3437,6 +3449,7 @@ dependencies = [
|
|||
name = "litellm-auth-types"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-error",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -3450,6 +3463,7 @@ dependencies = [
|
|||
name = "litellm-cache"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-error",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -3475,6 +3489,7 @@ dependencies = [
|
|||
"reqwest",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"wiremock",
|
||||
|
|
@ -3493,6 +3508,7 @@ dependencies = [
|
|||
"serde-pickle",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
|
|
@ -3505,11 +3521,13 @@ dependencies = [
|
|||
"litellm-auth-types",
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"percent-encoding",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
@ -3532,6 +3550,7 @@ dependencies = [
|
|||
"futures-util",
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"qdrant-client",
|
||||
"reqwest",
|
||||
|
|
@ -3552,11 +3571,13 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"r2d2",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
|
|
@ -3609,6 +3630,7 @@ dependencies = [
|
|||
"litellm-auth-aws",
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
|
|
@ -3631,11 +3653,13 @@ dependencies = [
|
|||
"litellm-cache",
|
||||
"litellm-cache-redis",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"uuid",
|
||||
]
|
||||
|
|
@ -3645,6 +3669,7 @@ name = "litellm-callbacks-legacy-python"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-auth",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
"litellm-host-python",
|
||||
"proptest",
|
||||
|
|
@ -3652,6 +3677,7 @@ dependencies = [
|
|||
"rstest",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3667,6 +3693,7 @@ name = "litellm-config"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"litellm-router-types",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -3684,6 +3711,7 @@ name = "litellm-core-utils"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"fancy-regex 0.19.2",
|
||||
"litellm-error",
|
||||
"litellm-llms-types",
|
||||
"litellm-tracing",
|
||||
"rstest",
|
||||
|
|
@ -3708,9 +3736,13 @@ dependencies = [
|
|||
name = "litellm-cost"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"criterion",
|
||||
"proptest",
|
||||
"jiff",
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3740,6 +3772,14 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-error"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"rstest",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-framer"
|
||||
version = "0.1.0"
|
||||
|
|
@ -3767,6 +3807,7 @@ dependencies = [
|
|||
"litellm-auth-types",
|
||||
"litellm-circuit-breaker",
|
||||
"litellm-config",
|
||||
"litellm-error",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-gateway-inference",
|
||||
"litellm-gateway-mcp",
|
||||
|
|
@ -3800,6 +3841,7 @@ dependencies = [
|
|||
"futures-util",
|
||||
"litellm-auth-types",
|
||||
"litellm-config",
|
||||
"litellm-error",
|
||||
"litellm-secrets",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -3823,6 +3865,7 @@ dependencies = [
|
|||
"litellm-auth",
|
||||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-error",
|
||||
"litellm-gateway-auth",
|
||||
"litellm-host",
|
||||
"litellm-host-http",
|
||||
|
|
@ -3842,6 +3885,7 @@ dependencies = [
|
|||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tower",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
@ -3852,6 +3896,7 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"litellm-gateway-auth",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
|
|
@ -3869,6 +3914,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"litellm-auth-types",
|
||||
"litellm-config",
|
||||
"litellm-error",
|
||||
"litellm-secrets",
|
||||
"moka",
|
||||
"rmcp",
|
||||
|
|
@ -3883,6 +3929,7 @@ dependencies = [
|
|||
"tower",
|
||||
"url",
|
||||
"uuid",
|
||||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3892,9 +3939,11 @@ dependencies = [
|
|||
"axum",
|
||||
"axum-login",
|
||||
"base64 0.22.1",
|
||||
"futures-util",
|
||||
"governor",
|
||||
"jsonwebtoken",
|
||||
"litellm-auth-types",
|
||||
"litellm-error",
|
||||
"litellm-gateway-auth",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
|
|
@ -3919,6 +3968,7 @@ dependencies = [
|
|||
"litellm-coroutine",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
|
|
@ -3956,6 +4006,7 @@ dependencies = [
|
|||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-coroutine",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
"parking_lot",
|
||||
"pyo3",
|
||||
|
|
@ -3976,6 +4027,7 @@ dependencies = [
|
|||
"http 1.4.2",
|
||||
"hyper-util",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"rcgen",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
|
|
@ -3990,6 +4042,7 @@ dependencies = [
|
|||
"tokio-rustls 0.26.4",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
"url",
|
||||
"veil",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
|
@ -4005,6 +4058,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-framer",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
|
|
@ -4032,6 +4086,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
|
|
@ -4059,6 +4114,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-framer",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
|
|
@ -4090,6 +4146,7 @@ dependencies = [
|
|||
"litellm-auth-aws",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
|
|
@ -4107,6 +4164,7 @@ dependencies = [
|
|||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
|
|
@ -4122,6 +4180,7 @@ dependencies = [
|
|||
"litellm-cache-memory",
|
||||
"litellm-cache-response",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
"litellm-host-native",
|
||||
"litellm-http",
|
||||
|
|
@ -4159,6 +4218,7 @@ dependencies = [
|
|||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-inference",
|
||||
"litellm-inference-testing",
|
||||
|
|
@ -4187,6 +4247,7 @@ dependencies = [
|
|||
"litellm-auth-azure",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-framer",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
|
|
@ -4226,16 +4287,35 @@ name = "litellm-model-catalog"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"indexmap 2.14.0",
|
||||
"jiff",
|
||||
"jsonschema",
|
||||
"litellm-llms-types",
|
||||
"litellm-cost",
|
||||
"litellm-error",
|
||||
"macro_rules_attribute",
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_with",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-model-catalog-refresh"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-model-catalog",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
|
|
@ -4253,6 +4333,7 @@ dependencies = [
|
|||
"litellm-cache-response",
|
||||
"litellm-callbacks-legacy-python",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-host",
|
||||
"litellm-host-python",
|
||||
"litellm-http",
|
||||
|
|
@ -4264,6 +4345,7 @@ dependencies = [
|
|||
"litellm-inference-transcription",
|
||||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-model-catalog",
|
||||
"litellm-router-types",
|
||||
"litellm-secrets",
|
||||
"litellm-secrets-aws",
|
||||
|
|
@ -4283,6 +4365,7 @@ dependencies = [
|
|||
"serde_with",
|
||||
"sqlx",
|
||||
"strum",
|
||||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
|
|
@ -4333,12 +4416,17 @@ name = "litellm-router"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-config",
|
||||
"litellm-core-utils",
|
||||
"litellm-inference",
|
||||
"litellm-inference-messages",
|
||||
"litellm-router-types",
|
||||
"rstest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-router-core"
|
||||
version = "0.1.0"
|
||||
|
||||
[[package]]
|
||||
name = "litellm-router-types"
|
||||
version = "0.1.0"
|
||||
|
|
@ -4360,6 +4448,7 @@ dependencies = [
|
|||
"google-cloud-auth",
|
||||
"google-cloud-kms-v1",
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-python-compat",
|
||||
"litellm-secrets-aws",
|
||||
|
|
@ -4418,6 +4507,7 @@ dependencies = [
|
|||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"veil",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
@ -4441,6 +4531,7 @@ dependencies = [
|
|||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"veil",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
@ -4467,6 +4558,7 @@ dependencies = [
|
|||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"veil",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
@ -4476,7 +4568,10 @@ name = "litellm-secrets-hashicorp"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-core-utils",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-secrets-types",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"rustify",
|
||||
"rustify_derive",
|
||||
|
|
@ -4510,6 +4605,7 @@ name = "litellm-spend-clickhouse"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"flate2",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-storage-clickhouse",
|
||||
"rstest",
|
||||
|
|
@ -4530,6 +4626,7 @@ name = "litellm-storage-clickhouse"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"flate2",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"rstest",
|
||||
"serde",
|
||||
|
|
@ -4549,6 +4646,7 @@ version = "0.1.0"
|
|||
dependencies = [
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"litellm-http",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"semver",
|
||||
|
|
@ -4572,6 +4670,7 @@ dependencies = [
|
|||
"criterion",
|
||||
"indexmap 2.14.0",
|
||||
"itoa",
|
||||
"litellm-error",
|
||||
"litellm-token-counter-fast",
|
||||
"litellm-token-counter-huggingface",
|
||||
"litellm-token-counter-tiktoken",
|
||||
|
|
@ -4588,6 +4687,7 @@ name = "litellm-token-counter-fast"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"litellm-error",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"rustc-hash",
|
||||
|
|
@ -4613,6 +4713,7 @@ name = "litellm-token-counter-tiktoken"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"litellm-error",
|
||||
"once_cell",
|
||||
"rstest",
|
||||
"rustc-hash",
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ repository = "https://github.com/BerriAI/litellm"
|
|||
|
||||
[workspace.dependencies]
|
||||
litellm-config = { path = "crates/config" }
|
||||
litellm-error = { path = "crates/error" }
|
||||
litellm-router = { path = "crates/router" }
|
||||
litellm-router-types = { path = "crates/router-types" }
|
||||
litellm-tracing = { path = "crates/tracing" }
|
||||
|
|
@ -72,6 +73,9 @@ 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" }
|
||||
litellm-cost = { path = "crates/cost" }
|
||||
litellm-model-catalog = { path = "crates/model-catalog" }
|
||||
litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" }
|
||||
|
||||
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
|
||||
tracing = "0.1"
|
||||
|
|
@ -108,6 +112,7 @@ serde = { version = "1.0", features = ["derive"] }
|
|||
serde_json = { version = "1.0", features = ["float_roundtrip"] }
|
||||
serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] }
|
||||
sha2 = "0.10"
|
||||
sha1 = "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"] }
|
||||
|
|
|
|||
|
|
@ -7,7 +7,9 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-auth-types.workspace = true
|
||||
litellm-error.workspace = true
|
||||
litellm-http.workspace = true
|
||||
reqwest.workspace = true
|
||||
|
||||
moka = { workspace = true, features = ["sync"] }
|
||||
serde.workspace = true
|
||||
|
|
@ -26,5 +28,6 @@ aws-smithy-runtime-api = "1.13.0"
|
|||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
reqwest.workspace = true
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::time::Duration;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
|
|
@ -12,16 +11,17 @@ use aws_sigv4::http_request::{
|
|||
};
|
||||
use aws_sigv4::sign::v4;
|
||||
use aws_smithy_runtime_api::client::identity::Identity;
|
||||
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
|
||||
|
||||
use litellm_auth_types::AwsParams;
|
||||
|
||||
use super::Error;
|
||||
use super::constants::{
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
|
||||
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
|
||||
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
|
||||
DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES,
|
||||
};
|
||||
use crate::error::{CredentialFlow, Error, SigningStage};
|
||||
|
||||
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
|
||||
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
|
||||
|
|
@ -271,7 +271,7 @@ impl AwsAuthService {
|
|||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsProfile(error.to_string()))
|
||||
.map_err(Error::credentials(CredentialFlow::Profile))
|
||||
}
|
||||
AwsAuthFlow::AssumeRole { role, session_name } => {
|
||||
if is_already_running_as_role(&role, &resolved).await? {
|
||||
|
|
@ -287,7 +287,7 @@ impl AwsAuthService {
|
|||
let credentials = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
|
||||
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
|
||||
self.set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
|
|
@ -327,7 +327,7 @@ impl AwsAuthService {
|
|||
provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
|
||||
.map_err(Error::credentials(CredentialFlow::AssumeRole))
|
||||
}
|
||||
AwsAuthFlow::WebIdentity {
|
||||
token,
|
||||
|
|
@ -350,12 +350,14 @@ impl AwsAuthService {
|
|||
.web_identity_token(token)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
|
||||
.map_err(Error::credentials(CredentialFlow::WebIdentity))?;
|
||||
let credentials = response
|
||||
.credentials()
|
||||
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
|
||||
let expiration = SystemTime::try_from(*credentials.expiration())
|
||||
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
|
||||
.ok_or(Error::MissingWebIdentityCredentials)?;
|
||||
let expiration =
|
||||
SystemTime::try_from(*credentials.expiration()).map_err(|source| {
|
||||
Error::WebIdentityExpiration(litellm_error::Source::new(source))
|
||||
})?;
|
||||
Ok(Credentials::new(
|
||||
credentials.access_key_id(),
|
||||
credentials.secret_access_key(),
|
||||
|
|
@ -376,7 +378,7 @@ impl AwsAuthService {
|
|||
let credentials = provider
|
||||
.provide_credentials()
|
||||
.await
|
||||
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
|
||||
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
|
||||
self.set_cached_credentials(
|
||||
key,
|
||||
credentials.clone(),
|
||||
|
|
@ -435,12 +437,12 @@ fn default_session_name() -> String {
|
|||
/// client header cannot change the canonical request and invalidate the
|
||||
/// signature. Signing everything instead makes the request 403 on a header the
|
||||
/// caller supplied, on a deployment that works on the Python path.
|
||||
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
|
||||
pub fn aws_signature_headers(headers: &HeaderMap) -> HeaderMap {
|
||||
headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
let name = name.to_ascii_lowercase();
|
||||
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|
||||
let name = name.as_str();
|
||||
AWS_SIGNED_HEADER_NAMES.contains(&name)
|
||||
|| name.starts_with("x-amz-")
|
||||
|| name.starts_with("x-amzn-")
|
||||
})
|
||||
|
|
@ -456,45 +458,62 @@ pub fn is_sigv4_computed_header(name: &str) -> bool {
|
|||
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
|
||||
}
|
||||
|
||||
pub fn sign_post(
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct SigningScope<'a> {
|
||||
pub region: &'a str,
|
||||
pub service: &'a str,
|
||||
pub credentials: &'a Credentials,
|
||||
pub time: SystemTime,
|
||||
}
|
||||
|
||||
pub fn sign_headers(
|
||||
method: &str,
|
||||
url: &str,
|
||||
body: &[u8],
|
||||
headers: &BTreeMap<String, String>,
|
||||
region: &str,
|
||||
service: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
let identity: Identity = credentials.clone().into();
|
||||
headers: &HeaderMap,
|
||||
scope: SigningScope<'_>,
|
||||
) -> Result<HeaderMap, Error> {
|
||||
let identity: Identity = scope.credentials.clone().into();
|
||||
let params = v4::SigningParams::builder()
|
||||
.identity(&identity)
|
||||
.region(region)
|
||||
.name(service)
|
||||
.time(signing_time)
|
||||
.region(scope.region)
|
||||
.name(scope.service)
|
||||
.time(scope.time)
|
||||
.settings(SigningSettings::default())
|
||||
.build()
|
||||
.map(SigningParams::from)
|
||||
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
|
||||
.map_err(Error::signing(SigningStage::Parameters))?;
|
||||
let header_refs = headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()));
|
||||
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
|
||||
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
|
||||
.map(|(name, value)| {
|
||||
value
|
||||
.to_str()
|
||||
.map(|value| (name.as_str(), value))
|
||||
.map_err(|_| Error::NonAsciiHeader {
|
||||
name: name.to_string(),
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let request = SignableRequest::new(
|
||||
method,
|
||||
url,
|
||||
header_refs.into_iter(),
|
||||
SignableBody::Bytes(body),
|
||||
)
|
||||
.map_err(Error::signing(SigningStage::Request))?;
|
||||
let (instructions, _) = sign(request, ¶ms)
|
||||
.map_err(|error| Error::AwsSigning(error.to_string()))?
|
||||
.map_err(Error::signing(SigningStage::Signature))?
|
||||
.into_parts();
|
||||
Ok(instructions
|
||||
instructions
|
||||
.headers()
|
||||
.map(|(name, value)| {
|
||||
let normalized_name = match name {
|
||||
"authorization" => "Authorization",
|
||||
"x-amz-date" => "X-Amz-Date",
|
||||
"x-amz-security-token" => "X-Amz-Security-Token",
|
||||
_ => name,
|
||||
};
|
||||
(normalized_name.to_string(), value.to_string())
|
||||
let name = HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(Error::signing(SigningStage::Output))?;
|
||||
let value =
|
||||
HeaderValue::from_str(value).map_err(Error::signing(SigningStage::Output))?;
|
||||
Ok((name, value))
|
||||
})
|
||||
.collect())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Model-id and region parsing shared by every Bedrock route.
|
||||
|
|
@ -639,6 +658,7 @@ pub fn host_supplied_credentials(params: &AwsParams) -> Option<Credentials> {
|
|||
mod tests {
|
||||
use super::*;
|
||||
use crate::constants::{AWS_REGION, BEDROCK_SERVICE};
|
||||
use reqwest::header::CONTENT_TYPE;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
|
|
@ -731,12 +751,12 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
|
||||
fn parity_inputs() -> (String, Vec<u8>, HeaderMap) {
|
||||
(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
|
||||
.to_string(),
|
||||
br#"{"input":"hello"}"#.to_vec(),
|
||||
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
|
||||
HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -961,9 +981,9 @@ mod tests {
|
|||
// cannot change the canonical request. Signing it instead makes the
|
||||
// request 403 the moment anything on the wire rewrites or drops it.
|
||||
let (url, body, mut headers) = parity_inputs();
|
||||
headers.insert("x-request-id".to_string(), "abc-123".to_string());
|
||||
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
|
||||
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
|
||||
headers.insert("x-request-id", HeaderValue::from_static("abc-123"));
|
||||
headers.insert("Accept-Encoding", HeaderValue::from_static("gzip"));
|
||||
headers.insert("x-amzn-trace-id", HeaderValue::from_static("Root=1-abc"));
|
||||
let signable = aws_signature_headers(&headers);
|
||||
|
||||
assert!(!signable.contains_key("x-request-id"));
|
||||
|
|
@ -979,18 +999,21 @@ mod tests {
|
|||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_post(
|
||||
let signed = sign_headers(
|
||||
"POST",
|
||||
&url,
|
||||
&body,
|
||||
&signable,
|
||||
"us-east-1",
|
||||
BEDROCK_SERVICE,
|
||||
&credentials,
|
||||
SystemTime::UNIX_EPOCH,
|
||||
SigningScope {
|
||||
region: "us-east-1",
|
||||
service: BEDROCK_SERVICE,
|
||||
credentials: &credentials,
|
||||
time: SystemTime::UNIX_EPOCH,
|
||||
},
|
||||
)
|
||||
.expect("signs");
|
||||
let authorization = signed
|
||||
.get("Authorization")
|
||||
let authorization = signed["Authorization"]
|
||||
.to_str()
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
|
|
@ -1012,29 +1035,24 @@ mod tests {
|
|||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_post(
|
||||
let signed = sign_headers(
|
||||
"POST",
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
"us-east-1",
|
||||
BEDROCK_SERVICE,
|
||||
&credentials,
|
||||
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
SigningScope {
|
||||
region: "us-east-1",
|
||||
service: BEDROCK_SERVICE,
|
||||
credentials: &credentials,
|
||||
time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
},
|
||||
)
|
||||
.expect("golden signature");
|
||||
assert_eq!(signed["X-Amz-Date"], "20240102T030405Z");
|
||||
assert_eq!(signed["X-Amz-Security-Token"], "session-token");
|
||||
assert_eq!(
|
||||
signed.get("X-Amz-Date").map(String::as_str),
|
||||
Some("20240102T030405Z")
|
||||
);
|
||||
assert_eq!(
|
||||
signed.get("X-Amz-Security-Token").map(String::as_str),
|
||||
Some("session-token")
|
||||
);
|
||||
assert_eq!(
|
||||
signed.get("Authorization").map(String::as_str),
|
||||
Some(
|
||||
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
|
||||
)
|
||||
signed["Authorization"],
|
||||
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -1048,19 +1066,71 @@ mod tests {
|
|||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_post(
|
||||
let signed = sign_headers(
|
||||
"POST",
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
"us-east-1",
|
||||
BEDROCK_SERVICE,
|
||||
&credentials,
|
||||
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
SigningScope {
|
||||
region: "us-east-1",
|
||||
service: BEDROCK_SERVICE,
|
||||
credentials: &credentials,
|
||||
time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
|
||||
},
|
||||
)
|
||||
.expect("signature");
|
||||
assert!(!signed.contains_key("X-Amz-Security-Token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_header_value_outside_visible_ascii_is_named_instead_of_signed() {
|
||||
let (url, body, mut headers) = parity_inputs();
|
||||
headers.insert("x-amz-meta-note", HeaderValue::from_str("café").unwrap());
|
||||
let credentials = Credentials::new("AKIDEXAMPLE", "secret", None, None, "test");
|
||||
|
||||
let error = sign_headers(
|
||||
"POST",
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
SigningScope {
|
||||
region: "us-east-1",
|
||||
service: BEDROCK_SERVICE,
|
||||
credentials: &credentials,
|
||||
time: UNIX_EPOCH,
|
||||
},
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::NonAsciiHeader {
|
||||
name: "x-amz-meta-note".into()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_shared_auth_error_nests_the_aws_failure_and_its_cause() {
|
||||
let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "profile not found");
|
||||
let error =
|
||||
litellm_auth_types::Error::from(Error::credentials(CredentialFlow::Profile)(cause));
|
||||
|
||||
assert!(matches!(
|
||||
litellm_error::find::<Error>(&error),
|
||||
Some(Error::Credentials {
|
||||
flow: CredentialFlow::Profile,
|
||||
..
|
||||
})
|
||||
));
|
||||
let report = litellm_error::Report::new(&error).to_string();
|
||||
assert!(
|
||||
report.contains("could not load AWS profile credentials"),
|
||||
"{report}"
|
||||
);
|
||||
assert_eq!(report.matches("profile not found").count(), 1, "{report}");
|
||||
}
|
||||
|
||||
#[ignore]
|
||||
#[tokio::test]
|
||||
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
|
@ -1068,7 +1138,7 @@ mod tests {
|
|||
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
|
||||
let headers =
|
||||
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
|
||||
HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]);
|
||||
let credentials = AwsAuthService::default()
|
||||
.resolve_credentials(
|
||||
AwsAuthConfig {
|
||||
|
|
@ -1087,23 +1157,25 @@ mod tests {
|
|||
let url = format!(
|
||||
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
|
||||
);
|
||||
let signed_headers = sign_post(
|
||||
let signed_headers = sign_headers(
|
||||
"POST",
|
||||
&url,
|
||||
&body,
|
||||
&headers,
|
||||
region,
|
||||
BEDROCK_SERVICE,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
SigningScope {
|
||||
region,
|
||||
service: BEDROCK_SERVICE,
|
||||
credentials: &credentials,
|
||||
time: SystemTime::now(),
|
||||
},
|
||||
)?;
|
||||
let mut request = client.post(&url).body(body.clone());
|
||||
for (name, value) in &headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
for (name, value) in signed_headers {
|
||||
request = request.header(name, value);
|
||||
}
|
||||
let response = request.send().await?;
|
||||
let response = client
|
||||
.post(&url)
|
||||
.headers(headers.clone())
|
||||
.headers(signed_headers)
|
||||
.body(body.clone())
|
||||
.send()
|
||||
.await?;
|
||||
let status = response.status();
|
||||
let response_body = response.text().await?;
|
||||
let snippet: String = response_body.chars().take(240).collect();
|
||||
|
|
|
|||
|
|
@ -1,46 +1,73 @@
|
|||
use thiserror::Error as ThisError;
|
||||
use litellm_error::Source;
|
||||
|
||||
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
|
||||
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("AWS profile credentials failed: {0}")]
|
||||
AwsProfile(String),
|
||||
#[error("AWS default credentials failed: {0}")]
|
||||
AwsDefaultChain(String),
|
||||
#[error("AWS role credentials failed: {0}")]
|
||||
AwsAssumeRole(String),
|
||||
#[error("AWS web identity credentials failed: {0}")]
|
||||
AwsWebIdentity(String),
|
||||
#[error("AWS web identity expiration was invalid: {0}")]
|
||||
AwsWebIdentityExpiration(String),
|
||||
#[error("AWS signing parameters failed: {0}")]
|
||||
AwsSigningParameters(String),
|
||||
#[error("AWS signable request failed: {0}")]
|
||||
AwsSignableRequest(String),
|
||||
#[error("AWS request signing failed: {0}")]
|
||||
AwsSigning(String),
|
||||
#[error("could not load AWS {flow} credentials")]
|
||||
Credentials {
|
||||
flow: CredentialFlow,
|
||||
#[source]
|
||||
source: Source,
|
||||
},
|
||||
#[error("AWS web identity response had no credentials")]
|
||||
AwsMissingWebIdentityCredentials,
|
||||
MissingWebIdentityCredentials,
|
||||
#[error("AWS web identity response had an invalid expiration")]
|
||||
WebIdentityExpiration(#[source] Source),
|
||||
#[error("header {name} has a value that is not visible ASCII")]
|
||||
NonAsciiHeader { name: String },
|
||||
#[error("{stage}")]
|
||||
Signing {
|
||||
stage: SigningStage,
|
||||
#[source]
|
||||
source: Source,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum CredentialFlow {
|
||||
#[error("profile")]
|
||||
Profile,
|
||||
#[error("default chain")]
|
||||
DefaultChain,
|
||||
#[error("assumed role")]
|
||||
AssumeRole,
|
||||
#[error("web identity")]
|
||||
WebIdentity,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum SigningStage {
|
||||
#[error("SigV4 signing parameters are invalid")]
|
||||
Parameters,
|
||||
#[error("SigV4 cannot sign this request")]
|
||||
Request,
|
||||
#[error("SigV4 could not compute the signature")]
|
||||
Signature,
|
||||
#[error("SigV4 produced an invalid header")]
|
||||
Output,
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub(crate) fn credentials<E: std::error::Error + Send + Sync + 'static>(
|
||||
flow: CredentialFlow,
|
||||
) -> impl FnOnce(E) -> Self {
|
||||
move |source| Self::Credentials {
|
||||
flow,
|
||||
source: Source::new(source),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn signing<E: std::error::Error + Send + Sync + 'static>(
|
||||
stage: SigningStage,
|
||||
) -> impl FnOnce(E) -> Self {
|
||||
move |source| Self::Signing {
|
||||
stage,
|
||||
source: Source::new(source),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Error> for litellm_auth_types::Error {
|
||||
fn from(error: Error) -> Self {
|
||||
Self::ProviderAuthentication(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::Error;
|
||||
|
||||
#[test]
|
||||
fn converts_to_shared_auth_error_without_losing_context() {
|
||||
let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into()));
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
litellm_auth_types::Error::ProviderAuthentication(
|
||||
"AWS profile credentials failed: profile not found".into()
|
||||
)
|
||||
);
|
||||
Self::provider(error)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,5 +5,5 @@ mod signer;
|
|||
|
||||
pub use aws::*;
|
||||
pub use aws_credential_types::Credentials;
|
||||
pub use error::Error;
|
||||
pub use error::{CredentialFlow, Error, SigningStage};
|
||||
pub use signer::SigV4Signer;
|
||||
|
|
|
|||
|
|
@ -1,11 +1,13 @@
|
|||
use std::{collections::BTreeMap, time::SystemTime};
|
||||
use std::time::SystemTime;
|
||||
|
||||
use aws_credential_types::Credentials;
|
||||
use reqwest::{Request, header::HeaderMap};
|
||||
|
||||
use crate::{
|
||||
AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header,
|
||||
sign_post,
|
||||
AwsAuthService, AwsCredentialSource, Error, SigningScope, aws_signature_headers,
|
||||
is_sigv4_computed_header, sign_headers,
|
||||
};
|
||||
use aws_credential_types::Credentials;
|
||||
use litellm_http::outbound::{RequestSigner, UnsignedRequest};
|
||||
use litellm_http::outbound::RequestSigner;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct SigV4Signer {
|
||||
|
|
@ -45,41 +47,54 @@ impl SigV4Signer {
|
|||
}
|
||||
|
||||
impl RequestSigner for SigV4Signer {
|
||||
fn sign(
|
||||
&self,
|
||||
request: UnsignedRequest<'_>,
|
||||
) -> Result<Vec<(String, String)>, litellm_http::Error> {
|
||||
if let Some((name, _)) = request
|
||||
.headers
|
||||
.iter()
|
||||
.find(|(name, _)| is_sigv4_computed_header(name))
|
||||
fn sign(&self, request: &Request) -> Result<HeaderMap, litellm_http::Error> {
|
||||
if let Some(name) = request
|
||||
.headers()
|
||||
.keys()
|
||||
.find(|name| is_sigv4_computed_header(name.as_str()))
|
||||
{
|
||||
return Err(litellm_http::Error::ComputedHeader(name.clone()));
|
||||
return Err(litellm_http::Error::computed_header(name.as_str()));
|
||||
}
|
||||
let headers: BTreeMap<String, String> = request.headers.iter().cloned().collect();
|
||||
sign_post(
|
||||
request.url,
|
||||
request.body,
|
||||
&aws_signature_headers(&headers),
|
||||
&self.region,
|
||||
self.service,
|
||||
&self.credentials,
|
||||
(self.clock)(),
|
||||
let body = match request.body() {
|
||||
None => &[],
|
||||
Some(body) => body
|
||||
.as_bytes()
|
||||
.ok_or_else(litellm_http::Error::unsignable_body)?,
|
||||
};
|
||||
sign_headers(
|
||||
request.method().as_str(),
|
||||
request.url().as_str(),
|
||||
body,
|
||||
&aws_signature_headers(request.headers()),
|
||||
SigningScope {
|
||||
region: &self.region,
|
||||
service: self.service,
|
||||
credentials: &self.credentials,
|
||||
time: (self.clock)(),
|
||||
},
|
||||
)
|
||||
.map(|signature| signature.into_iter().collect())
|
||||
.map_err(|error| litellm_http::Error::Signature(error.to_string()))
|
||||
.map_err(|error| match error {
|
||||
Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
|
||||
Error::Credentials { .. }
|
||||
| Error::MissingWebIdentityCredentials
|
||||
| Error::WebIdentityExpiration(_)
|
||||
| Error::Signing { .. } => litellm_http::Error::signature(error),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use reqwest::{Body, Method, Url};
|
||||
use serde_json::{Value, json as json_value};
|
||||
use std::time::{Duration, UNIX_EPOCH};
|
||||
|
||||
use litellm_http::outbound::OutboundRequest;
|
||||
use serde_json::{Value, json};
|
||||
use litellm_http::outbound::{json, signed};
|
||||
|
||||
use super::*;
|
||||
|
||||
const URL: &str = "https://textract.us-east-1.amazonaws.com/";
|
||||
|
||||
fn fixed_clock() -> SystemTime {
|
||||
UNIX_EPOCH + Duration::from_secs(1_700_000_000)
|
||||
}
|
||||
|
|
@ -93,82 +108,121 @@ mod tests {
|
|||
.with_clock(fixed_clock)
|
||||
}
|
||||
|
||||
fn target() -> (String, String) {
|
||||
("X-Amz-Target".into(), "Textract.DetectDocumentText".into())
|
||||
}
|
||||
|
||||
fn authorization(body: &Value, service: &'static str) -> String {
|
||||
OutboundRequest::signed_json(
|
||||
"https://textract.us-east-1.amazonaws.com/".into(),
|
||||
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
|
||||
body,
|
||||
None,
|
||||
signed(
|
||||
json(URL, &[target()], body, None).unwrap(),
|
||||
&signer(service),
|
||||
)
|
||||
.unwrap()
|
||||
.header("Authorization")
|
||||
.unwrap()
|
||||
.to_string()
|
||||
.headers()["authorization"]
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_signature_verifies_against_the_bytes_that_are_sent() {
|
||||
let sent = OutboundRequest::signed_json(
|
||||
"https://textract.us-east-1.amazonaws.com/".into(),
|
||||
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
|
||||
&json!({"Document": {"Bytes": "aGk="}}),
|
||||
None,
|
||||
let sent = signed(
|
||||
json(
|
||||
URL,
|
||||
&[target()],
|
||||
&json_value!({"Document": {"Bytes": "aGk="}}),
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
&signer("textract"),
|
||||
)
|
||||
.unwrap();
|
||||
let unsigned: BTreeMap<String, String> = sent
|
||||
let unsigned: HeaderMap = sent
|
||||
.headers()
|
||||
.iter()
|
||||
.filter(|(name, _)| !is_sigv4_computed_header(name))
|
||||
.cloned()
|
||||
.filter(|(name, _)| !is_sigv4_computed_header(name.as_str()))
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect();
|
||||
let recomputed = sign_post(
|
||||
sent.url(),
|
||||
sent.body(),
|
||||
let recomputed = sign_headers(
|
||||
"POST",
|
||||
sent.url().as_str(),
|
||||
sent.body().unwrap().as_bytes().unwrap(),
|
||||
&aws_signature_headers(&unsigned),
|
||||
"us-east-1",
|
||||
"textract",
|
||||
&Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
|
||||
fixed_clock(),
|
||||
SigningScope {
|
||||
region: "us-east-1",
|
||||
service: "textract",
|
||||
credentials: &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
|
||||
time: fixed_clock(),
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
sent.header("Authorization"),
|
||||
Some(recomputed["Authorization"].as_str())
|
||||
);
|
||||
assert_eq!(sent.headers()["authorization"], recomputed["authorization"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_signature_depends_on_the_body_and_the_service() {
|
||||
let original = authorization(&json!({"text": "card 4111"}), "textract");
|
||||
let original = authorization(&json_value!({"text": "card 4111"}), "textract");
|
||||
|
||||
assert_ne!(
|
||||
original,
|
||||
authorization(&json!({"text": "card [REDACTED]"}), "textract")
|
||||
authorization(&json_value!({"text": "card [REDACTED]"}), "textract")
|
||||
);
|
||||
assert_ne!(
|
||||
original,
|
||||
authorization(&json!({"text": "card 4111"}), "bedrock")
|
||||
authorization(&json_value!({"text": "card 4111"}), "bedrock")
|
||||
);
|
||||
assert!(original.contains("/us-east-1/textract/aws4_request"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() {
|
||||
let error = OutboundRequest::signed_json(
|
||||
"https://textract.us-east-1.amazonaws.com/".into(),
|
||||
vec![("authorization".into(), "Bearer caller".into())],
|
||||
&json!({}),
|
||||
None,
|
||||
let error = signed(
|
||||
json(
|
||||
URL,
|
||||
&[("authorization".into(), "Bearer caller".into())],
|
||||
&json_value!({}),
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
&signer("textract"),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
litellm_http::Error::ComputedHeader("authorization".into())
|
||||
);
|
||||
assert!(error.is_computed_header());
|
||||
assert!(error.to_string().contains("authorization"), "{error}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_header_that_cannot_be_signed_is_the_callers_fault() {
|
||||
let error = signed(
|
||||
json(
|
||||
URL,
|
||||
&[("x-amz-meta-note".into(), "café".into())],
|
||||
&json_value!({}),
|
||||
None,
|
||||
)
|
||||
.unwrap(),
|
||||
&signer("textract"),
|
||||
)
|
||||
.unwrap_err();
|
||||
|
||||
let report = litellm_error::Report::new(&error).to_string();
|
||||
assert!(report.contains("x-amz-meta-note"), "{report}");
|
||||
assert!(error.is_request());
|
||||
assert!(!error.dispatched());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_streaming_body_cannot_be_signed() {
|
||||
let mut request = Request::new(Method::POST, Url::parse(URL).unwrap());
|
||||
*request.body_mut() = Some(Body::wrap_stream(futures_util::stream::once(async {
|
||||
Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{}"))
|
||||
})));
|
||||
|
||||
assert!(matches!(
|
||||
signed(request, &signer("textract")),
|
||||
Err(error) if !error.dispatched() && !error.is_request()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,5 +18,6 @@ azure_core = "1.0.0"
|
|||
azure_identity = { version = "1.0.0", features = ["tokio"] }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-error.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ use azure_identity::{
|
|||
};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use litellm_auth_types::Error;
|
||||
use litellm_auth_types::{Error, ErrorDetail};
|
||||
use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
|
||||
|
||||
use super::credential_provider_cache::{
|
||||
|
|
@ -133,7 +133,9 @@ impl NativeAzureTokenAcquirer {
|
|||
let token = credential
|
||||
.get_token(&[scope.as_str()], None)
|
||||
.await
|
||||
.map_err(|error| Error::CredentialAcquisition(error.to_string().into()))?;
|
||||
.map_err(|error| {
|
||||
Error::CredentialAcquisition(ErrorDetail::failed("Azure token request", error))
|
||||
})?;
|
||||
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
|
||||
.ok()
|
||||
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
|
||||
|
|
@ -250,11 +252,8 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
|
|||
let Some(authority) = authority else {
|
||||
return Ok(());
|
||||
};
|
||||
let url = url::Url::parse(authority.value()).map_err(|_| {
|
||||
Error::InvalidConfiguration(
|
||||
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
|
||||
.into(),
|
||||
)
|
||||
let url = url::Url::parse(authority.value()).map_err(|error| {
|
||||
Error::InvalidConfiguration(ErrorDetail::failed("Azure authority parsing", error))
|
||||
})?;
|
||||
if url.scheme() != "https"
|
||||
|| url.host_str().is_none()
|
||||
|
|
@ -264,10 +263,10 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
|
|||
|| url.fragment().is_some()
|
||||
|| !matches!(url.path(), "" | "/")
|
||||
{
|
||||
return Err(Error::InvalidConfiguration(
|
||||
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
|
||||
.into(),
|
||||
));
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
|
||||
field: "Azure authority".into(),
|
||||
expected: "an HTTPS origin without credentials, query, or fragment",
|
||||
}));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -377,7 +376,7 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
|
|||
|
||||
fn mixed_sources<T>() -> Result<T, Error> {
|
||||
Err(Error::InvalidConfiguration(
|
||||
"request-controlled Azure auth inputs cannot be combined with host credentials".into(),
|
||||
ErrorDetail::MixedInputSources { provider: "Azure" },
|
||||
))
|
||||
}
|
||||
|
||||
|
|
@ -444,7 +443,7 @@ fn build_credential(
|
|||
.map(|credential| credential as Arc<dyn TokenCredential>),
|
||||
}
|
||||
.map_err(|error| {
|
||||
Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed(
|
||||
Error::InvalidConfiguration(ErrorDetail::failed(
|
||||
"Azure credential initialization",
|
||||
error,
|
||||
))
|
||||
|
|
@ -662,12 +661,11 @@ mod tests {
|
|||
))
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
litellm_auth_types::Error::InvalidConfiguration(
|
||||
assert!(
|
||||
litellm_error::Report::new(&error).to_string().contains(
|
||||
"request-controlled Azure auth inputs cannot be combined with host credentials"
|
||||
.into()
|
||||
)
|
||||
),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -695,12 +693,11 @@ mod tests {
|
|||
authority,
|
||||
))
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
error,
|
||||
litellm_auth_types::Error::InvalidConfiguration(
|
||||
assert!(
|
||||
litellm_error::Report::new(&error).to_string().contains(
|
||||
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
|
||||
.into()
|
||||
)
|
||||
),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
use litellm_auth_types::Error;
|
||||
use litellm_auth_types::{
|
||||
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
|
||||
SecretValue, Sourced, TokenProviderHandle,
|
||||
};
|
||||
use litellm_auth_types::{Error, ErrorDetail};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
|
@ -91,9 +91,9 @@ impl AzureAuthService {
|
|||
AzureCredentialPlan::Caller(caller) => {
|
||||
let credential = caller.acquire().await?;
|
||||
if credential.secret().expose().is_empty() {
|
||||
return Err(Error::EmptyCallerCredential(
|
||||
"Azure AD token provider returned an empty token",
|
||||
));
|
||||
return Err(Error::EmptyCallerCredential {
|
||||
caller: "Azure AD token provider",
|
||||
});
|
||||
}
|
||||
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
|
||||
}
|
||||
|
|
@ -106,11 +106,9 @@ impl AzureAuthService {
|
|||
} => {
|
||||
let assertion = resolve_reference(inputs, env_lookup, reference.value())
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
Error::CredentialAcquisition(
|
||||
"Azure OIDC reference did not resolve to a value".into(),
|
||||
)
|
||||
})?;
|
||||
.ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
|
||||
subject: "Azure OIDC reference",
|
||||
}))?;
|
||||
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
|
||||
tenant_id,
|
||||
client_id,
|
||||
|
|
@ -142,7 +140,13 @@ impl AzureAuthService {
|
|||
Err(error) => failures.push(error),
|
||||
}
|
||||
}
|
||||
Err(Error::CredentialChain(failures))
|
||||
match failures.pop() {
|
||||
Some(last) => Err(Error::CredentialChain {
|
||||
last: Box::new(last),
|
||||
earlier: failures,
|
||||
}),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
AzureCredentialPlan::Missing => Ok(None),
|
||||
}
|
||||
|
|
@ -173,7 +177,12 @@ pub(crate) fn select_auth_plan(
|
|||
.map(|selector| Sourced::new(selector, value.source()))
|
||||
})
|
||||
.transpose()
|
||||
.map_err(|_| Error::InvalidConfiguration("invalid Azure credential selector".into()))?;
|
||||
.map_err(|error| {
|
||||
Error::InvalidConfiguration(ErrorDetail::failed(
|
||||
"Azure credential selector parsing",
|
||||
error,
|
||||
))
|
||||
})?;
|
||||
let federated_token_file = configured_string(
|
||||
&inputs.federated_token_file,
|
||||
AZURE_FEDERATED_TOKEN_FILE_ENV,
|
||||
|
|
@ -263,9 +272,12 @@ fn select_native_plan(
|
|||
let selection_source = selected.source();
|
||||
|
||||
match selected.into_value() {
|
||||
AzureCredentialType::ClientSecretCredential => Err(Error::InvalidConfiguration(
|
||||
"ClientSecretCredential requires tenant_id, client_id, and client_secret".into(),
|
||||
)),
|
||||
AzureCredentialType::ClientSecretCredential => {
|
||||
Err(Error::InvalidConfiguration(ErrorDetail::MissingSettings {
|
||||
credential: "ClientSecretCredential",
|
||||
settings: "tenant_id, client_id, and client_secret",
|
||||
}))
|
||||
}
|
||||
AzureCredentialType::WorkloadIdentityCredential => {
|
||||
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
|
||||
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
|
||||
|
|
@ -349,22 +361,21 @@ fn workload_request(
|
|||
authority: Option<Sourced<String>>,
|
||||
) -> Result<NativeAzureRequest, Error> {
|
||||
Ok(NativeAzureRequest::WorkloadIdentity {
|
||||
tenant_id: tenant_id.ok_or_else(|| {
|
||||
Error::InvalidConfiguration("WorkloadIdentityCredential requires tenant_id".into())
|
||||
})?,
|
||||
client_id: client_id.ok_or_else(|| {
|
||||
Error::InvalidConfiguration("WorkloadIdentityCredential requires client_id".into())
|
||||
})?,
|
||||
token_file_path: token_file_path.ok_or_else(|| {
|
||||
Error::InvalidConfiguration(
|
||||
"WorkloadIdentityCredential requires azure_federated_token_file".into(),
|
||||
)
|
||||
})?,
|
||||
tenant_id: tenant_id.ok_or(workload_requires("tenant_id"))?,
|
||||
client_id: client_id.ok_or(workload_requires("client_id"))?,
|
||||
token_file_path: token_file_path.ok_or(workload_requires("azure_federated_token_file"))?,
|
||||
scope,
|
||||
authority,
|
||||
})
|
||||
}
|
||||
|
||||
fn workload_requires(settings: &'static str) -> Error {
|
||||
Error::InvalidConfiguration(ErrorDetail::MissingSettings {
|
||||
credential: "WorkloadIdentityCredential",
|
||||
settings,
|
||||
})
|
||||
}
|
||||
|
||||
fn configured_string(
|
||||
configured: &ConfigValue<String>,
|
||||
environment_name: &str,
|
||||
|
|
@ -410,11 +421,10 @@ async fn resolve_reference(
|
|||
.map_or(CredentialLookup::Missing, CredentialLookup::Found),
|
||||
CredentialRef::None => return Ok(None),
|
||||
CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => {
|
||||
let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| {
|
||||
Error::InvalidConfiguration(
|
||||
"credential reference requires a host credential resolver".into(),
|
||||
)
|
||||
})?;
|
||||
let resolver = inputs
|
||||
.credential_resolver
|
||||
.as_ref()
|
||||
.ok_or(Error::InvalidConfiguration(ErrorDetail::MissingResolver))?;
|
||||
resolver.resolve(reference).await?
|
||||
}
|
||||
};
|
||||
|
|
@ -433,7 +443,7 @@ fn oidc_reference(
|
|||
let value = token.value().expose();
|
||||
if token.source() == InputSource::Request && value.starts_with("oidc/") {
|
||||
return Err(Error::InvalidConfiguration(
|
||||
"request-controlled Azure credential references are not allowed".into(),
|
||||
ErrorDetail::RequestControlledReference { provider: "Azure" },
|
||||
));
|
||||
}
|
||||
if let Some(name) = value.strip_prefix("oidc/env/") {
|
||||
|
|
@ -457,19 +467,17 @@ fn oidc_reference(
|
|||
}
|
||||
if value.starts_with("oidc/") {
|
||||
return Err(Error::InvalidConfiguration(
|
||||
"unsupported OIDC reference".into(),
|
||||
ErrorDetail::UnsupportedReference,
|
||||
));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
|
||||
fn non_empty_reference(value: &str, kind: &'static str) -> Result<String, Error> {
|
||||
if value.is_empty() {
|
||||
return Err(Error::InvalidConfiguration(
|
||||
litellm_auth_types::ErrorDetail::Empty {
|
||||
subject: kind.into(),
|
||||
},
|
||||
));
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::Empty {
|
||||
subject: kind,
|
||||
}));
|
||||
}
|
||||
Ok(value.to_string())
|
||||
}
|
||||
|
|
@ -518,7 +526,9 @@ mod tests {
|
|||
expires_on: None,
|
||||
})
|
||||
} else {
|
||||
Err(Error::CredentialAcquisition(kind.into()))
|
||||
Err(Error::CredentialAcquisition(
|
||||
litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
|
||||
))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -624,7 +634,11 @@ mod tests {
|
|||
)))
|
||||
.expect_err("unsupported backend must fail validation");
|
||||
|
||||
assert!(error.to_string().contains("unsupported OIDC reference"));
|
||||
assert!(
|
||||
litellm_error::Report::new(&error)
|
||||
.to_string()
|
||||
.contains("unsupported OIDC reference")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
|
|
@ -649,11 +663,11 @@ mod tests {
|
|||
})
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::InvalidConfiguration(
|
||||
"request-controlled Azure credential references are not allowed".into()
|
||||
)
|
||||
assert!(
|
||||
litellm_error::Report::new(&error)
|
||||
.to_string()
|
||||
.contains("request-controlled Azure credential references are not allowed"),
|
||||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -715,7 +729,16 @@ mod tests {
|
|||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
|
||||
assert!(matches!(&error, Error::CredentialChain { earlier, .. } if earlier.len() == 1));
|
||||
let report = litellm_error::Report::new(&error).to_string();
|
||||
assert!(
|
||||
report.contains("managed-identity did not resolve"),
|
||||
"{report}"
|
||||
);
|
||||
assert!(
|
||||
report.contains("developer-tools did not resolve"),
|
||||
"{report}"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
|
|
@ -763,7 +786,9 @@ mod tests {
|
|||
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::EmptyCallerCredential("Azure AD token provider returned an empty token")
|
||||
Error::EmptyCallerCredential {
|
||||
caller: "Azure AD token provider"
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -24,5 +24,6 @@ google-cloud-auth = { workspace = true, optional = true }
|
|||
http = { workspace = true, optional = true }
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-error.workspace = true
|
||||
rstest.workspace = true
|
||||
tempfile.workspace = true
|
||||
|
|
|
|||
|
|
@ -73,13 +73,17 @@ impl VertexAuth {
|
|||
.or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV));
|
||||
let project_id = get_vertex_ai_project(config, env_lookup);
|
||||
|
||||
if !has_authorization && static_token.is_none() {
|
||||
let access = self.get_access_token(config, env_lookup).await?;
|
||||
return Ok(VertexEnvironment {
|
||||
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
|
||||
project_id: project_id.unwrap_or(access.project_id),
|
||||
});
|
||||
}
|
||||
let static_token = match (has_authorization, static_token) {
|
||||
(true, _) => None,
|
||||
(false, Some(token)) => Some(token),
|
||||
(false, None) => {
|
||||
let access = self.get_access_token(config, env_lookup).await?;
|
||||
return Ok(VertexEnvironment {
|
||||
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
|
||||
project_id: project_id.unwrap_or(access.project_id),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let project_id = match project_id {
|
||||
Some(project_id) => project_id,
|
||||
|
|
@ -90,14 +94,9 @@ impl VertexAuth {
|
|||
.await?
|
||||
}
|
||||
};
|
||||
let headers = if has_authorization {
|
||||
headers
|
||||
} else {
|
||||
apply_credential(
|
||||
headers,
|
||||
static_token.as_deref().expect("static token was checked"),
|
||||
CredentialPlacement::Bearer,
|
||||
)?
|
||||
let headers = match static_token {
|
||||
Some(token) => apply_credential(headers, &token, CredentialPlacement::Bearer)?,
|
||||
None => headers,
|
||||
};
|
||||
Ok(VertexEnvironment {
|
||||
headers,
|
||||
|
|
@ -234,14 +233,13 @@ enum CredentialCacheKey {
|
|||
Adc,
|
||||
}
|
||||
|
||||
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
|
||||
Set it to a readable file path, or to the credentials JSON itself.";
|
||||
|
||||
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
|
||||
if is_base64_encoded_json(configured) {
|
||||
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
|
||||
BASE64_CREDENTIALS.into(),
|
||||
)));
|
||||
return Err(Error::InvalidConfiguration(
|
||||
ErrorDetail::Base64EncodedJson {
|
||||
field: "vertex_credentials",
|
||||
},
|
||||
));
|
||||
}
|
||||
if names_a_credentials_file(configured) {
|
||||
CustomServiceAccount::from_file(configured)
|
||||
|
|
@ -299,7 +297,8 @@ mod tests {
|
|||
else {
|
||||
panic!("{configured:?} should not load");
|
||||
};
|
||||
assert!(error.to_string().contains(expected), "{error}");
|
||||
let report = litellm_error::Report::new(&error).to_string();
|
||||
assert!(report.contains(expected), "{report}");
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
|
|
@ -321,7 +320,7 @@ mod tests {
|
|||
.await
|
||||
{
|
||||
Ok(_) => panic!("{configured:?} should not load"),
|
||||
Err(error) => error.to_string(),
|
||||
Err(error) => litellm_error::Report::new(&error).to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -55,7 +55,7 @@ impl CredentialsProvider for GoogleCredentials {
|
|||
entity_tag: EntityTag::new(),
|
||||
data,
|
||||
})
|
||||
.map_err(|_| CredentialsError::from_msg(false, "Google authentication failed"))
|
||||
.map_err(|error| CredentialsError::from_source(false, error))
|
||||
}
|
||||
|
||||
async fn universe_domain(&self) -> Option<String> {
|
||||
|
|
@ -102,5 +102,9 @@ mod tests {
|
|||
);
|
||||
let error = credentials.headers(Extensions::new()).await.unwrap_err();
|
||||
assert!(!format!("{error:?}").contains("private"));
|
||||
assert_eq!(
|
||||
std::error::Error::source(&error).and_then(|cause| cause.downcast_ref::<Error>()),
|
||||
Some(&Error::InvalidHeader)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-error.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
subtle.workspace = true
|
||||
|
|
|
|||
|
|
@ -86,9 +86,9 @@ impl CredentialPlan {
|
|||
Self::Caller(caller) => {
|
||||
let credential = caller.acquire().await?;
|
||||
if credential.secret().expose().is_empty() {
|
||||
return Err(Error::EmptyCallerCredential(
|
||||
"credential caller returned an empty credential",
|
||||
));
|
||||
return Err(Error::EmptyCallerCredential {
|
||||
caller: "credential caller",
|
||||
});
|
||||
}
|
||||
Ok(CredentialPlanResolution::Resolved(credential))
|
||||
}
|
||||
|
|
@ -104,8 +104,8 @@ mod tests {
|
|||
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
|
||||
CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
};
|
||||
use crate::Error;
|
||||
use crate::SecretValue;
|
||||
use crate::{Error, ErrorDetail};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct HostResolver;
|
||||
|
|
@ -150,9 +150,10 @@ mod tests {
|
|||
impl CredentialResolver for FailingResolver {
|
||||
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
|
||||
Box::pin(async {
|
||||
Err(Error::CredentialAcquisition(
|
||||
"host credential lookup failed".into(),
|
||||
))
|
||||
Err(Error::CredentialAcquisition(ErrorDetail::failed(
|
||||
"host credential lookup",
|
||||
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
|
||||
)))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -168,9 +169,14 @@ mod tests {
|
|||
.await
|
||||
.expect_err("acquisition errors cannot become fallback");
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
Error::CredentialAcquisition("host credential lookup failed".into())
|
||||
assert!(
|
||||
matches!(error, Error::CredentialAcquisition(_)),
|
||||
"{error:?}"
|
||||
);
|
||||
assert!(
|
||||
litellm_error::Report::new(&error)
|
||||
.to_string()
|
||||
.contains("host credential lookup failed")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,17 +1,23 @@
|
|||
use thiserror::Error as ThisError;
|
||||
use litellm_error::{Report, Source};
|
||||
|
||||
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
|
||||
use crate::{ParamKind, ParamSpec};
|
||||
|
||||
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("invalid authentication configuration: {0}")]
|
||||
#[error("invalid authentication configuration")]
|
||||
InvalidConfiguration(#[source] ErrorDetail),
|
||||
#[error("credential acquisition failed: {0}")]
|
||||
#[error("credential acquisition failed")]
|
||||
CredentialAcquisition(#[source] ErrorDetail),
|
||||
#[error("credential caller failed: {0}")]
|
||||
EmptyCallerCredential(&'static str),
|
||||
#[error("{0}")]
|
||||
ProviderAuthentication(String),
|
||||
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
|
||||
CredentialChain(Vec<Error>),
|
||||
#[error("{caller} returned an empty credential")]
|
||||
EmptyCallerCredential { caller: &'static str },
|
||||
#[error("provider authentication failed")]
|
||||
Provider(#[source] Source),
|
||||
#[error("no credential in the chain succeeded{}", earlier_attempts(earlier))]
|
||||
CredentialChain {
|
||||
#[source]
|
||||
last: Box<Error>,
|
||||
earlier: Vec<Error>,
|
||||
},
|
||||
#[error(
|
||||
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
|
||||
)]
|
||||
|
|
@ -22,7 +28,7 @@ pub enum Error {
|
|||
#[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())]
|
||||
MissingParam {
|
||||
provider: &'static str,
|
||||
spec: &'static crate::ParamSpec,
|
||||
spec: &'static ParamSpec,
|
||||
},
|
||||
#[error("Missing {provider} API Base - {guidance}")]
|
||||
MissingApiBase {
|
||||
|
|
@ -33,78 +39,282 @@ pub enum Error {
|
|||
InvalidHeader,
|
||||
}
|
||||
|
||||
fn earlier_attempts(earlier: &[Error]) -> String {
|
||||
if earlier.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
let attempts = earlier
|
||||
.iter()
|
||||
.map(|failure| Report::new(failure).to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("; ");
|
||||
format!(" (earlier attempts: {attempts})")
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub fn provider(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
|
||||
Self::Provider(Source::new(cause))
|
||||
}
|
||||
|
||||
pub fn is_missing_credential(&self) -> bool {
|
||||
match self {
|
||||
Self::MissingApiKey { .. } => true,
|
||||
Self::MissingParam { spec, .. } => match spec.kind {
|
||||
ParamKind::Credential => true,
|
||||
ParamKind::Setting => false,
|
||||
},
|
||||
Self::InvalidConfiguration(_)
|
||||
| Self::CredentialAcquisition(_)
|
||||
| Self::EmptyCallerCredential { .. }
|
||||
| Self::Provider(_)
|
||||
| Self::CredentialChain { .. }
|
||||
| Self::MissingApiBase { .. }
|
||||
| Self::InvalidHeader => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_request(&self) -> bool {
|
||||
match self {
|
||||
Self::InvalidConfiguration(_)
|
||||
| Self::EmptyCallerCredential { .. }
|
||||
| Self::InvalidHeader => true,
|
||||
Self::CredentialAcquisition(_) | Self::Provider(_) | Self::CredentialChain { .. } => {
|
||||
false
|
||||
}
|
||||
Self::MissingParam { spec, .. } => match spec.kind {
|
||||
ParamKind::Setting => true,
|
||||
ParamKind::Credential => false,
|
||||
},
|
||||
Self::MissingApiKey { .. } | Self::MissingApiBase { .. } => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum ErrorDetail {
|
||||
#[error("{0}")]
|
||||
Message(String),
|
||||
#[error("{field} must be {expected}")]
|
||||
InvalidType {
|
||||
field: String,
|
||||
expected: &'static str,
|
||||
},
|
||||
#[error("{subject} cannot be empty")]
|
||||
Empty { subject: String },
|
||||
Empty { subject: &'static str },
|
||||
#[error("credential header {0} already exists")]
|
||||
DuplicateHeader(&'static str),
|
||||
#[error("{operation} failed: {source}")]
|
||||
#[error("credential plan is not allowed by the provider auth policy")]
|
||||
PlanNotAllowed,
|
||||
#[error("{credential} requires {settings}")]
|
||||
MissingSettings {
|
||||
credential: &'static str,
|
||||
settings: &'static str,
|
||||
},
|
||||
#[error("credentials cannot be sent to a request-controlled {provider} endpoint")]
|
||||
RequestControlledEndpoint { provider: &'static str },
|
||||
#[error("request-controlled {provider} credential references are not allowed")]
|
||||
RequestControlledReference { provider: &'static str },
|
||||
#[error("request-controlled {provider} auth inputs cannot be combined with host credentials")]
|
||||
MixedInputSources { provider: &'static str },
|
||||
#[error("unsupported OIDC reference")]
|
||||
UnsupportedReference,
|
||||
#[error("credential reference requires a host credential resolver")]
|
||||
MissingResolver,
|
||||
#[error("{subject} did not resolve to a value")]
|
||||
Unresolved { subject: &'static str },
|
||||
#[error(
|
||||
"The `{field}` value is base64-encoded JSON, not the JSON itself. Set it to a readable file path, or to the credentials JSON itself."
|
||||
)]
|
||||
Base64EncodedJson { field: &'static str },
|
||||
#[error("{operation} failed")]
|
||||
Failed {
|
||||
operation: &'static str,
|
||||
#[source]
|
||||
source: ErrorSource,
|
||||
cause: Source,
|
||||
},
|
||||
}
|
||||
|
||||
impl ErrorDetail {
|
||||
pub fn failed(
|
||||
operation: &'static str,
|
||||
source: impl std::error::Error + Send + Sync + 'static,
|
||||
cause: impl std::error::Error + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
Self::Failed {
|
||||
operation,
|
||||
source: ErrorSource::new(source),
|
||||
cause: Source::new(cause),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for ErrorDetail {
|
||||
fn from(message: String) -> Self {
|
||||
Self::Message(message)
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io;
|
||||
|
||||
use litellm_error::{Report, Source};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Error, ErrorDetail};
|
||||
use crate::{AwsParams, VertexParams};
|
||||
|
||||
#[rstest]
|
||||
#[case::api_key(
|
||||
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
|
||||
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
|
||||
)]
|
||||
#[case::another_api_key(
|
||||
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
|
||||
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
|
||||
)]
|
||||
#[case::api_base(
|
||||
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
|
||||
"Missing Example API Base - Pass api_base"
|
||||
)]
|
||||
#[case::another_api_base(
|
||||
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
|
||||
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
|
||||
)]
|
||||
fn missing_inputs_keep_the_python_text(#[case] error: Error, #[case] expected: &str) {
|
||||
assert_eq!(error.to_string(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configuration(
|
||||
Error::InvalidConfiguration(ErrorDetail::Empty { subject: "credential selector" }),
|
||||
"invalid authentication configuration: credential selector cannot be empty"
|
||||
)]
|
||||
#[case::acquisition(
|
||||
Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }),
|
||||
"credential acquisition failed: OIDC reference did not resolve to a value"
|
||||
)]
|
||||
#[case::empty_caller(
|
||||
Error::EmptyCallerCredential { caller: "token provider" },
|
||||
"token provider returned an empty credential"
|
||||
)]
|
||||
fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) {
|
||||
assert!(
|
||||
Report::new(&error).to_string().contains(expected),
|
||||
"{}",
|
||||
Report::new(&error)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
|
||||
"test credential",
|
||||
denied()
|
||||
)))]
|
||||
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
|
||||
"test credential",
|
||||
denied()
|
||||
)))]
|
||||
fn a_failure_keeps_its_cause_as_the_source_and_renders_it_once(#[case] error: Error) {
|
||||
let cause =
|
||||
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
|
||||
.find_map(|cause| cause.downcast_ref::<Source>()?.downcast_ref::<io::Error>())
|
||||
.expect("the io cause is in the source chain");
|
||||
assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied);
|
||||
let report = Report::new(&error).to_string();
|
||||
assert_eq!(report.matches("access denied").count(), 1, "{report}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
|
||||
"test credential",
|
||||
denied()
|
||||
)))]
|
||||
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
|
||||
"test credential",
|
||||
denied()
|
||||
)))]
|
||||
fn a_contextual_failure_leaves_its_cause_out_of_its_own_message(#[case] error: Error) {
|
||||
assert!(!error.to_string().contains("access denied"), "{error}");
|
||||
assert!(
|
||||
Report::new(&error)
|
||||
.to_string()
|
||||
.contains("test credential failed: access denied")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_provider_failure_keeps_the_provider_error_as_its_source() {
|
||||
let error = Error::provider(denied());
|
||||
assert_eq!(
|
||||
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
|
||||
Some(io::ErrorKind::PermissionDenied)
|
||||
);
|
||||
assert!(!error.to_string().contains("access denied"), "{error}");
|
||||
let report = Report::new(&error).to_string();
|
||||
assert_eq!(report.matches("access denied").count(), 1, "{report}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_chain_keeps_its_last_failure_as_the_source_and_renders_every_attempt_once() {
|
||||
let error = Error::CredentialChain {
|
||||
last: Box::new(Error::provider(denied())),
|
||||
earlier: vec![
|
||||
Error::CredentialAcquisition(ErrorDetail::Unresolved {
|
||||
subject: "OIDC reference",
|
||||
}),
|
||||
Error::EmptyCallerCredential {
|
||||
caller: "token provider",
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
|
||||
Some(io::ErrorKind::PermissionDenied)
|
||||
);
|
||||
let report = Report::new(&error).to_string();
|
||||
assert_eq!(report.matches("access denied").count(), 1, "{report}");
|
||||
assert_eq!(
|
||||
report.matches("OIDC reference did not resolve").count(),
|
||||
1,
|
||||
"{report}"
|
||||
);
|
||||
assert_eq!(
|
||||
report
|
||||
.matches("token provider returned an empty credential")
|
||||
.count(),
|
||||
1,
|
||||
"{report}"
|
||||
);
|
||||
}
|
||||
|
||||
fn denied() -> io::Error {
|
||||
io::Error::new(io::ErrorKind::PermissionDenied, "access denied")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, true)]
|
||||
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), true)]
|
||||
#[case::aws_secret_access_key(AwsParams::SECRET_ACCESS_KEY.missing("AWS"), true)]
|
||||
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), false)]
|
||||
#[case::aws_region(AwsParams::REGION.missing("AWS"), false)]
|
||||
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
|
||||
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)]
|
||||
fn only_an_absent_credential_is_a_missing_credential(
|
||||
#[case] error: Error,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(error.is_missing_credential(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), true)]
|
||||
#[case::aws_region(AwsParams::REGION.missing("AWS"), true)]
|
||||
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
|
||||
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), false)]
|
||||
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, false)]
|
||||
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), true)]
|
||||
#[case::invalid_header(Error::InvalidHeader, true)]
|
||||
#[case::identity_provider_outage(
|
||||
Error::CredentialAcquisition(ErrorDetail::MissingResolver),
|
||||
false
|
||||
)]
|
||||
#[case::provider_sdk_failure(Error::provider(denied()), false)]
|
||||
fn a_setting_the_caller_can_pass_is_theirs_to_fix_but_credentials_and_api_bases_are_not(
|
||||
#[case] error: Error,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(error.is_request(), expected);
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for ErrorDetail {
|
||||
fn from(message: &str) -> Self {
|
||||
Self::Message(message.into())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ErrorSource(std::sync::Arc<dyn std::error::Error + Send + Sync>);
|
||||
|
||||
impl std::ops::Deref for ErrorSource {
|
||||
type Target = dyn std::error::Error + Send + Sync;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
self.0.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ErrorSource {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
std::fmt::Display::fmt(&self.0, formatter)
|
||||
}
|
||||
}
|
||||
|
||||
impl ErrorSource {
|
||||
pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self {
|
||||
Self(std::sync::Arc::new(error))
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for ErrorSource {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
std::sync::Arc::ptr_eq(&self.0, &other.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for ErrorSource {}
|
||||
|
|
|
|||
|
|
@ -21,9 +21,9 @@ pub fn apply_credential(
|
|||
placement: CredentialPlacement,
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
if credential.trim().is_empty() {
|
||||
return Err(Error::InvalidConfiguration(
|
||||
"credential cannot be empty".into(),
|
||||
));
|
||||
return Err(Error::InvalidConfiguration(crate::ErrorDetail::Empty {
|
||||
subject: "credential",
|
||||
}));
|
||||
}
|
||||
if headers
|
||||
.iter()
|
||||
|
|
@ -71,6 +71,10 @@ mod tests {
|
|||
)
|
||||
.expect_err("provider policy must handle existing credentials");
|
||||
|
||||
assert!(error.to_string().contains("already exists"));
|
||||
assert!(
|
||||
litellm_error::Report::new(&error)
|
||||
.to_string()
|
||||
.contains("already exists")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -51,9 +51,9 @@ pub use credential::{
|
|||
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
|
||||
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
|
||||
};
|
||||
pub use error::{Error, ErrorDetail, ErrorSource};
|
||||
pub use error::{Error, ErrorDetail};
|
||||
pub use http::CredentialPlacement;
|
||||
pub use params::{AwsParams, ParamSpec, VertexParams};
|
||||
pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams};
|
||||
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
|
||||
pub use secret::SecretValue;
|
||||
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};
|
||||
|
|
|
|||
|
|
@ -4,15 +4,17 @@ use serde::{Deserialize, Serialize};
|
|||
use serde_json::{Map, Value};
|
||||
use veil::Redact;
|
||||
|
||||
use super::ParamSpec;
|
||||
use super::{ParamKind, ParamSpec};
|
||||
|
||||
const fn spec(
|
||||
setting: &'static str,
|
||||
kind: ParamKind,
|
||||
wire: &'static [&'static str],
|
||||
env: &'static [&'static str],
|
||||
) -> ParamSpec {
|
||||
ParamSpec {
|
||||
setting,
|
||||
kind,
|
||||
wire,
|
||||
module_global: None,
|
||||
env,
|
||||
|
|
@ -53,40 +55,67 @@ pub struct AwsParams {
|
|||
impl AwsParams {
|
||||
pub const ACCESS_KEY_ID: ParamSpec = spec(
|
||||
"access key id",
|
||||
ParamKind::Credential,
|
||||
&["aws_access_key_id"],
|
||||
&["AWS_ACCESS_KEY_ID"],
|
||||
);
|
||||
pub const SECRET_ACCESS_KEY: ParamSpec = spec(
|
||||
"secret access key",
|
||||
ParamKind::Credential,
|
||||
&["aws_secret_access_key"],
|
||||
&["AWS_SECRET_ACCESS_KEY"],
|
||||
);
|
||||
pub const SESSION_TOKEN: ParamSpec = spec(
|
||||
"session token",
|
||||
ParamKind::Credential,
|
||||
&["aws_session_token"],
|
||||
&["AWS_SESSION_TOKEN"],
|
||||
);
|
||||
pub const REGION: ParamSpec = spec(
|
||||
"region",
|
||||
ParamKind::Setting,
|
||||
&["aws_region_name"],
|
||||
&["AWS_REGION_NAME", "AWS_REGION"],
|
||||
);
|
||||
pub const SESSION_NAME: ParamSpec =
|
||||
spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]);
|
||||
pub const PROFILE_NAME: ParamSpec =
|
||||
spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]);
|
||||
pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]);
|
||||
pub const SESSION_NAME: ParamSpec = spec(
|
||||
"session name",
|
||||
ParamKind::Setting,
|
||||
&["aws_session_name"],
|
||||
&["AWS_SESSION_NAME"],
|
||||
);
|
||||
pub const PROFILE_NAME: ParamSpec = spec(
|
||||
"profile name",
|
||||
ParamKind::Setting,
|
||||
&["aws_profile_name"],
|
||||
&["AWS_PROFILE_NAME"],
|
||||
);
|
||||
pub const ROLE_NAME: ParamSpec = spec(
|
||||
"role name",
|
||||
ParamKind::Setting,
|
||||
&["aws_role_name"],
|
||||
&["AWS_ROLE_NAME"],
|
||||
);
|
||||
pub const WEB_IDENTITY_TOKEN: ParamSpec = spec(
|
||||
"web identity token",
|
||||
ParamKind::Credential,
|
||||
&["aws_web_identity_token"],
|
||||
&["AWS_WEB_IDENTITY_TOKEN"],
|
||||
);
|
||||
pub const STS_ENDPOINT: ParamSpec =
|
||||
spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]);
|
||||
pub const EXTERNAL_ID: ParamSpec =
|
||||
spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]);
|
||||
pub const STS_ENDPOINT: ParamSpec = spec(
|
||||
"STS endpoint",
|
||||
ParamKind::Setting,
|
||||
&["aws_sts_endpoint"],
|
||||
&["AWS_STS_ENDPOINT"],
|
||||
);
|
||||
pub const EXTERNAL_ID: ParamSpec = spec(
|
||||
"external id",
|
||||
ParamKind::Setting,
|
||||
&["aws_external_id"],
|
||||
&["AWS_EXTERNAL_ID"],
|
||||
);
|
||||
pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec(
|
||||
"Bedrock runtime endpoint",
|
||||
ParamKind::Setting,
|
||||
&["aws_bedrock_runtime_endpoint"],
|
||||
&["AWS_BEDROCK_RUNTIME_ENDPOINT"],
|
||||
);
|
||||
|
|
|
|||
|
|
@ -10,11 +10,18 @@ pub use vertex::VertexParams;
|
|||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct ParamSpec {
|
||||
pub setting: &'static str,
|
||||
pub kind: ParamKind,
|
||||
pub wire: &'static [&'static str],
|
||||
pub module_global: Option<&'static str>,
|
||||
pub env: &'static [&'static str],
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ParamKind {
|
||||
Credential,
|
||||
Setting,
|
||||
}
|
||||
|
||||
impl ParamSpec {
|
||||
/// The first non-blank value: each wire name through `params`, then each environment
|
||||
/// name through `env`.
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::Unexpected};
|
|||
use serde_json::Value;
|
||||
use veil::Redact;
|
||||
|
||||
use super::ParamSpec;
|
||||
use super::{ParamKind, ParamSpec};
|
||||
|
||||
/// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the
|
||||
/// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read.
|
||||
|
|
@ -27,18 +27,21 @@ pub struct VertexParams {
|
|||
impl VertexParams {
|
||||
pub const CREDENTIALS: ParamSpec = ParamSpec {
|
||||
setting: "credentials",
|
||||
kind: ParamKind::Credential,
|
||||
wire: &["vertex_credentials", "vertex_ai_credentials"],
|
||||
module_global: None,
|
||||
env: &["VERTEXAI_CREDENTIALS"],
|
||||
};
|
||||
pub const PROJECT: ParamSpec = ParamSpec {
|
||||
setting: "project",
|
||||
kind: ParamKind::Setting,
|
||||
wire: &["vertex_project", "vertex_ai_project"],
|
||||
module_global: Some("vertex_project"),
|
||||
env: &["VERTEXAI_PROJECT"],
|
||||
};
|
||||
pub const LOCATION: ParamSpec = ParamSpec {
|
||||
setting: "location",
|
||||
kind: ParamKind::Setting,
|
||||
wire: &["vertex_location", "vertex_ai_location"],
|
||||
module_global: Some("vertex_location"),
|
||||
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use crate::http::apply_credential;
|
||||
use crate::{CredentialPlacement, Error, ResolvedCredential};
|
||||
use crate::{CredentialPlacement, Error, ErrorDetail, ResolvedCredential};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum CredentialPlanKind {
|
||||
|
|
@ -31,11 +31,18 @@ pub struct ProviderAuthPolicy {
|
|||
|
||||
impl ProviderAuthPolicy {
|
||||
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, _)| {
|
||||
self.accepted_existing_headers
|
||||
.iter()
|
||||
.any(|accepted| name.eq_ignore_ascii_case(accepted))
|
||||
})
|
||||
self.existing_credential_header(headers).is_some()
|
||||
}
|
||||
|
||||
fn existing_credential_header(&self, headers: &[(String, String)]) -> Option<&'static str> {
|
||||
self.accepted_existing_headers
|
||||
.iter()
|
||||
.copied()
|
||||
.find(|accepted| {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case(accepted))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn apply(
|
||||
|
|
@ -44,11 +51,11 @@ impl ProviderAuthPolicy {
|
|||
kind: CredentialPlanKind,
|
||||
credential: &ResolvedCredential,
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
if self.has_existing_credential(&headers) {
|
||||
if let Some(existing) = self.existing_credential_header(&headers) {
|
||||
return match self.existing_header_behavior {
|
||||
ExistingHeaderBehavior::Preserve => Ok(headers),
|
||||
ExistingHeaderBehavior::Reject => Err(Error::InvalidConfiguration(
|
||||
"credential header already exists".into(),
|
||||
ErrorDetail::DuplicateHeader(existing),
|
||||
)),
|
||||
};
|
||||
}
|
||||
|
|
@ -56,11 +63,7 @@ impl ProviderAuthPolicy {
|
|||
.rules
|
||||
.iter()
|
||||
.find(|rule| rule.kind == kind)
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidConfiguration(
|
||||
"credential plan is not allowed by the provider auth policy".into(),
|
||||
)
|
||||
})?;
|
||||
.ok_or(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed))?;
|
||||
apply_credential(headers, credential.secret().expose(), rule.placement)
|
||||
}
|
||||
}
|
||||
|
|
@ -108,6 +111,10 @@ mod tests {
|
|||
)
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.to_string().contains("not allowed"));
|
||||
assert!(
|
||||
litellm_error::Report::new(&error)
|
||||
.to_string()
|
||||
.contains("not allowed")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,74 +0,0 @@
|
|||
use litellm_auth_types::Error;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::api_key(
|
||||
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
|
||||
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
|
||||
)]
|
||||
#[case::another_api_key(
|
||||
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
|
||||
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
|
||||
)]
|
||||
#[case::api_base(
|
||||
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
|
||||
"Missing Example API Base - Pass api_base"
|
||||
)]
|
||||
#[case::another_api_base(
|
||||
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
|
||||
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
|
||||
)]
|
||||
#[case::configuration(
|
||||
Error::InvalidConfiguration("credential selector is invalid".into()),
|
||||
"invalid authentication configuration: credential selector is invalid"
|
||||
)]
|
||||
#[case::acquisition(
|
||||
Error::CredentialAcquisition("token expired".into()),
|
||||
"credential acquisition failed: token expired"
|
||||
)]
|
||||
#[case::caller(
|
||||
Error::EmptyCallerCredential("empty token"),
|
||||
"credential caller failed: empty token"
|
||||
)]
|
||||
#[case::provider(
|
||||
Error::ProviderAuthentication("provider rejected credentials".into()),
|
||||
"provider rejected credentials"
|
||||
)]
|
||||
#[case::chain(
|
||||
Error::CredentialChain(vec![
|
||||
Error::CredentialAcquisition("token expired".into()),
|
||||
Error::EmptyCallerCredential("empty token"),
|
||||
]),
|
||||
"credential acquisition failed: credential acquisition failed: token expired; credential caller failed: empty token"
|
||||
)]
|
||||
fn display_preserves_failure_phase_and_caller_context(
|
||||
#[case] error: Error,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(error.to_string(), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::configuration(true)]
|
||||
#[case::acquisition(false)]
|
||||
fn contextual_failures_keep_the_original_source(#[case] configuration: bool) {
|
||||
use litellm_auth_types::ErrorDetail;
|
||||
|
||||
let detail = ErrorDetail::failed(
|
||||
"test credential",
|
||||
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
|
||||
);
|
||||
let error = if configuration {
|
||||
Error::InvalidConfiguration(detail)
|
||||
} else {
|
||||
Error::CredentialAcquisition(detail)
|
||||
};
|
||||
let source = std::iter::successors(Some(&error as &dyn std::error::Error), |error| {
|
||||
error.source()
|
||||
})
|
||||
.find_map(|error| error.downcast_ref::<std::io::Error>())
|
||||
.expect("the original credential error remains available");
|
||||
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
|
||||
assert!(error.to_string().contains("test credential failed:"));
|
||||
assert!(error.to_string().ends_with(&source.to_string()));
|
||||
}
|
||||
|
|
@ -1,8 +1,9 @@
|
|||
use litellm_auth_types::{AwsParams, ParamSpec};
|
||||
use litellm_auth_types::{AwsParams, ParamKind, ParamSpec};
|
||||
use rstest::rstest;
|
||||
|
||||
const LOCATION: ParamSpec = ParamSpec {
|
||||
setting: "location",
|
||||
kind: ParamKind::Setting,
|
||||
wire: &["vertex_location", "vertex_ai_location"],
|
||||
module_global: Some("vertex_location"),
|
||||
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
use std::{
|
||||
error::Error as StdError,
|
||||
future::{Future, poll_fn},
|
||||
sync::{
|
||||
Arc,
|
||||
|
|
@ -53,7 +52,7 @@ async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: boo
|
|||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn callback_errors_preserve_the_original_source() {
|
||||
async fn callback_errors_keep_their_typed_cause() {
|
||||
let provider = TokenProviderHandle::from_callback(|| async {
|
||||
Err(Error::CredentialAcquisition(ErrorDetail::failed(
|
||||
"caller credential",
|
||||
|
|
@ -62,13 +61,15 @@ async fn callback_errors_preserve_the_original_source() {
|
|||
});
|
||||
|
||||
let error = provider.acquire().await.unwrap_err();
|
||||
assert!(matches!(error, Error::CredentialAcquisition(_)));
|
||||
let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| {
|
||||
(*error).source()
|
||||
})
|
||||
.find_map(|error| error.downcast_ref::<std::io::Error>())
|
||||
.unwrap();
|
||||
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
|
||||
let Error::CredentialAcquisition(ErrorDetail::Failed { cause, .. }) = &error else {
|
||||
panic!("the callback's error passes through unchanged: {error:?}");
|
||||
};
|
||||
assert_eq!(
|
||||
cause
|
||||
.downcast_ref::<std::io::Error>()
|
||||
.map(std::io::Error::kind),
|
||||
Some(std::io::ErrorKind::PermissionDenied)
|
||||
);
|
||||
}
|
||||
|
||||
struct Release(Arc<AtomicBool>);
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ azure_core = "1.1.0"
|
|||
azure_storage_blob = "1.1.0"
|
||||
futures-util.workspace = true
|
||||
reqwest.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,10 @@ use url::Url;
|
|||
|
||||
use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("account URL cannot hold a container path")]
|
||||
struct NotABaseUrl;
|
||||
|
||||
pub struct AzureBlobCache<C> {
|
||||
container: BlobContainerClient,
|
||||
codec: C,
|
||||
|
|
@ -57,12 +61,12 @@ impl<C: CacheCodec> AzureBlobCache<C> {
|
|||
codec: C,
|
||||
runtime: Handle,
|
||||
) -> Result<Self, Error> {
|
||||
let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?;
|
||||
let parsed = Url::parse(account_url).map_err(Error::unavailable)?;
|
||||
let account_url = parsed.as_str().trim_end_matches('/').to_string();
|
||||
let container_url = {
|
||||
let mut url = parsed;
|
||||
url.path_segments_mut()
|
||||
.map_err(|()| Error::Unavailable)?
|
||||
.map_err(|()| Error::unavailable(NotABaseUrl))?
|
||||
.pop_if_empty()
|
||||
.push(container);
|
||||
url
|
||||
|
|
@ -75,7 +79,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
|
|||
..BlobContainerClientOptions::default()
|
||||
}),
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
let cache = Self {
|
||||
container: client,
|
||||
codec,
|
||||
|
|
@ -101,7 +105,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
|
|||
Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => {
|
||||
Ok(())
|
||||
}
|
||||
Err(_) => Err(Error::Unavailable),
|
||||
Err(error) => Err(Error::unavailable(error)),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -116,7 +120,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
|
|||
{
|
||||
Ok(_) => Ok(()),
|
||||
Err(error) if !overwrite && is_already_present(&error) => Ok(()),
|
||||
Err(_) => Err(Error::Unavailable),
|
||||
Err(error) => Err(Error::unavailable(error)),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -126,13 +130,9 @@ impl<C: CacheCodec> AzureBlobCache<C> {
|
|||
Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => {
|
||||
return Ok(None);
|
||||
}
|
||||
Err(_) => return Err(Error::Unavailable),
|
||||
Err(error) => return Err(Error::unavailable(error)),
|
||||
};
|
||||
let bytes = response
|
||||
.body
|
||||
.collect()
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
let bytes = response.body.collect().await.map_err(Error::unavailable)?;
|
||||
self.codec.decode(&bytes).map(Some)
|
||||
}
|
||||
|
||||
|
|
@ -140,16 +140,16 @@ impl<C: CacheCodec> AzureBlobCache<C> {
|
|||
let mut pages = self
|
||||
.container
|
||||
.list_blobs(None)
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.into_pages();
|
||||
while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? {
|
||||
let page = page.into_model().map_err(|_| Error::Unavailable)?;
|
||||
while let Some(page) = pages.try_next().await.map_err(Error::unavailable)? {
|
||||
let page = page.into_model().map_err(Error::unavailable)?;
|
||||
for name in page.blob_items.into_iter().filter_map(|item| item.name) {
|
||||
self.container
|
||||
.blob_client(&name)
|
||||
.delete(None)
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
|
|
|
|||
|
|
@ -17,6 +17,10 @@ const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
|
|||
|
||||
type EnvLookup = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("no Azure credential is available for blob storage")]
|
||||
struct NoCredential;
|
||||
|
||||
pub struct AzureBlobCredential {
|
||||
service: AzureAuthService,
|
||||
env_lookup: EnvLookup,
|
||||
|
|
@ -62,15 +66,8 @@ impl TokenCredential for AzureBlobCredential {
|
|||
&lookup,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
azure_core::Error::with_message(ErrorKind::Credential, error.to_string())
|
||||
})?
|
||||
.ok_or_else(|| {
|
||||
azure_core::Error::with_message(
|
||||
ErrorKind::Credential,
|
||||
"no Azure credential is available for blob storage",
|
||||
)
|
||||
})?;
|
||||
.map_err(|error| azure_core::Error::new(ErrorKind::Credential, error))?
|
||||
.ok_or_else(|| azure_core::Error::new(ErrorKind::Credential, NoCredential))?;
|
||||
let (token, expires_on) = match credential.into_value() {
|
||||
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
|
||||
ResolvedCredential::Static(token) => (token, None),
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ impl HttpClient for ReqwestTransport {
|
|||
Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)),
|
||||
};
|
||||
let response = outgoing.send().await.map_err(|error| {
|
||||
let error = litellm_http::Error::from(error);
|
||||
let kind = if error.is_connect() {
|
||||
ErrorKind::Connection
|
||||
} else {
|
||||
|
|
@ -41,9 +42,9 @@ impl HttpClient for ReqwestTransport {
|
|||
);
|
||||
}
|
||||
}
|
||||
let body = response
|
||||
.bytes_stream()
|
||||
.map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
|
||||
let body = response.bytes_stream().map_err(|error| {
|
||||
azure_core::Error::new(ErrorKind::Io, litellm_http::Error::from(error))
|
||||
});
|
||||
Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -128,7 +128,7 @@ fn connect_surfaces_service_failures() {
|
|||
ResponseCacheCodec,
|
||||
runtime.handle().clone(),
|
||||
));
|
||||
assert!(matches!(result, Err(Error::Unavailable)));
|
||||
assert!(matches!(result, Err(Error::Unavailable(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -263,7 +263,7 @@ fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[ca
|
|||
fixture.service.seed_blob(key, bytes);
|
||||
assert!(matches!(
|
||||
fixture.cache.get_cache(key, &no_ttl()),
|
||||
Err(Error::InvalidEntry)
|
||||
Err(Error::InvalidEntry(_))
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -378,15 +378,15 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
|
|||
fixture.service.set_failing(true);
|
||||
assert!(matches!(
|
||||
fixture.cache.get_cache("key", &no_ttl()),
|
||||
Err(Error::Unavailable)
|
||||
Err(Error::Unavailable(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
|
||||
Err(Error::Unavailable)
|
||||
Err(Error::Unavailable(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
fixture.cache.flush_cache(),
|
||||
Err(Error::Unavailable)
|
||||
Err(Error::Unavailable(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
fixture.runtime.block_on(
|
||||
|
|
@ -394,7 +394,7 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
|
|||
.cache
|
||||
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
|
||||
),
|
||||
Err(Error::Unavailable)
|
||||
Err(Error::Unavailable(_))
|
||||
));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ rand.workspace = true
|
|||
rusqlite = { version = "0.39", features = ["bundled"] }
|
||||
serde-pickle = "1.2"
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
|
|
|||
|
|
@ -64,7 +64,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
|
|||
{
|
||||
tokio::task::spawn_blocking(move || operation(&store))
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -162,7 +162,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D
|
|||
.map(|key| match self.get_cache(key, context) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect()
|
||||
|
|
@ -186,7 +186,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D
|
|||
Some(value) => match self.decode_stored(value) {
|
||||
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
|
||||
Ok(None) => Ok(BatchEntry::Miss),
|
||||
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
|
||||
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
|
||||
Err(error) => Err(error),
|
||||
},
|
||||
})
|
||||
|
|
@ -272,7 +272,7 @@ fn increment<A: ValueAdapter, D: DiskStore>(
|
|||
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
|
||||
};
|
||||
store.update(key, unix_now(), &mut apply)?;
|
||||
result.ok_or(Error::InvalidEntry)
|
||||
result.ok_or(Error::InvalidEntry(None))
|
||||
}
|
||||
|
||||
fn unix_now() -> f64 {
|
||||
|
|
|
|||
|
|
@ -3,8 +3,8 @@ use py_literal::Value;
|
|||
use serde_json::{Map, Number};
|
||||
|
||||
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
|
||||
let value = serde_pickle::value_from_slice(bytes, Default::default())
|
||||
.map_err(|_| Error::InvalidEntry)?;
|
||||
let value =
|
||||
serde_pickle::value_from_slice(bytes, Default::default()).map_err(Error::invalid_entry)?;
|
||||
from_pickle_value(value)
|
||||
}
|
||||
|
||||
|
|
@ -70,7 +70,7 @@ fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Err
|
|||
}
|
||||
|
||||
fn integer(value: String) -> Result<Value, Error> {
|
||||
value.parse().map_err(|_| Error::InvalidEntry)
|
||||
value.parse().map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
pub(crate) fn from_json(value: serde_json::Value) -> Value {
|
||||
|
|
@ -101,7 +101,7 @@ pub(crate) fn from_json(value: serde_json::Value) -> Value {
|
|||
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
|
||||
serde_json::from_str(value)
|
||||
.map(from_json)
|
||||
.map_err(|_| Error::InvalidEntry)
|
||||
.map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
pub(crate) fn is_truthy(value: &Value) -> bool {
|
||||
|
|
@ -131,7 +131,7 @@ pub(crate) fn to_f64(value: &Value) -> Option<f64> {
|
|||
}
|
||||
|
||||
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
|
||||
serde_json::to_vec(&to_json_value(value)?).map_err(Error::invalid_entry)
|
||||
}
|
||||
|
||||
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
|
||||
|
|
@ -142,12 +142,12 @@ fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
|
|||
value
|
||||
.to_string()
|
||||
.parse::<Number>()
|
||||
.map_err(|_| Error::InvalidEntry)?,
|
||||
.map_err(Error::invalid_entry)?,
|
||||
),
|
||||
Value::Float(value) => {
|
||||
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
|
||||
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry(None))?)
|
||||
}
|
||||
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
|
||||
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry(None)),
|
||||
Value::String(value) => serde_json::Value::String(value.clone()),
|
||||
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
|
||||
serde_json::Value::Array(
|
||||
|
|
@ -162,7 +162,7 @@ fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
|
|||
.iter()
|
||||
.map(|(key, value)| {
|
||||
let Value::String(key) = key else {
|
||||
return Err(Error::InvalidEntry);
|
||||
return Err(Error::InvalidEntry(None));
|
||||
};
|
||||
Ok((key.clone(), to_json_value(value)?))
|
||||
})
|
||||
|
|
|
|||
|
|
@ -48,13 +48,13 @@ struct Row {
|
|||
impl DiskcacheSqliteStore {
|
||||
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
|
||||
let directory = directory.as_ref().to_path_buf();
|
||||
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
|
||||
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
|
||||
fs::create_dir_all(&directory).map_err(Error::unavailable)?;
|
||||
let directory = std::path::absolute(&directory).map_err(Error::unavailable)?;
|
||||
let database = directory.join("cache.db");
|
||||
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
|
||||
let connection = Connection::open(database).map_err(Error::unavailable)?;
|
||||
connection
|
||||
.busy_timeout(std::time::Duration::from_secs(60))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
|
||||
let mut settings = read_settings(&connection)?;
|
||||
for (key, value) in default_settings() {
|
||||
|
|
@ -74,7 +74,7 @@ impl DiskcacheSqliteStore {
|
|||
value
|
||||
)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
for (key, value) in &settings {
|
||||
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
|
||||
connection
|
||||
|
|
@ -82,7 +82,7 @@ impl DiskcacheSqliteStore {
|
|||
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
|
||||
params![key, value],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
}
|
||||
for (key, value) in [
|
||||
|
|
@ -96,7 +96,7 @@ impl DiskcacheSqliteStore {
|
|||
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
|
||||
params![key, value],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
connection
|
||||
.execute_batch(
|
||||
|
|
@ -117,7 +117,7 @@ impl DiskcacheSqliteStore {
|
|||
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
|
||||
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
|
||||
let eviction_policy = setting_string(&settings, "eviction_policy")
|
||||
.unwrap_or_else(|| "least-recently-stored".to_string());
|
||||
|
|
@ -128,23 +128,27 @@ impl DiskcacheSqliteStore {
|
|||
.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
"least-recently-used" => {
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
"least-frequently-used" => {
|
||||
connection
|
||||
.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
_ => {
|
||||
return Err(Error::unavailable(SettingFault::EvictionPolicy(
|
||||
eviction_policy,
|
||||
)));
|
||||
}
|
||||
_ => return Err(Error::Unavailable),
|
||||
}
|
||||
connection
|
||||
.execute_batch(
|
||||
|
|
@ -170,12 +174,12 @@ impl DiskcacheSqliteStore {
|
|||
UPDATE Settings SET value = value - OLD.size
|
||||
WHERE key = \"size\"; END;",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
|
||||
let min_file_size = setting_i64(&settings, "disk_min_file_size")
|
||||
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
|
||||
.try_into()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
|
||||
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
|
||||
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
|
||||
|
|
@ -207,7 +211,7 @@ impl DiskcacheSqliteStore {
|
|||
|row| row.get::<_, Option<String>>(0),
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.flatten()
|
||||
{
|
||||
cleanup.push(old_filename);
|
||||
|
|
@ -221,7 +225,7 @@ impl DiskcacheSqliteStore {
|
|||
|row| row.get::<_, i64>(0),
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
if let Some(rowid) = rowid {
|
||||
connection
|
||||
.execute(
|
||||
|
|
@ -230,7 +234,7 @@ impl DiskcacheSqliteStore {
|
|||
WHERE rowid = ?",
|
||||
params![now, expire_time, now, size, mode, filename, value, rowid],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
} else {
|
||||
connection
|
||||
.execute(
|
||||
|
|
@ -240,7 +244,7 @@ impl DiskcacheSqliteStore {
|
|||
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
|
||||
params![key, now, expire_time, now, size, mode, filename, value],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
cleanup.extend(self.cull(connection, now)?);
|
||||
Ok(cleanup)
|
||||
|
|
@ -257,13 +261,13 @@ impl DiskcacheSqliteStore {
|
|||
WHERE expire_time IS NOT NULL AND expire_time < ?
|
||||
ORDER BY expire_time LIMIT ?",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.query_map(params![now, self.cull_limit], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
for (_, filename) in &expired {
|
||||
if let Some(filename) = filename {
|
||||
cleanup.push(filename.clone());
|
||||
|
|
@ -272,7 +276,7 @@ impl DiskcacheSqliteStore {
|
|||
for (rowid, _) in &expired {
|
||||
connection
|
||||
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
|
||||
if remaining <= 0 || self.volume(connection)? < self.size_limit {
|
||||
|
|
@ -283,19 +287,23 @@ impl DiskcacheSqliteStore {
|
|||
"least-recently-stored" => "store_time",
|
||||
"least-recently-used" => "access_time",
|
||||
"least-frequently-used" => "access_count",
|
||||
_ => return Err(Error::Unavailable),
|
||||
policy => {
|
||||
return Err(Error::unavailable(SettingFault::EvictionPolicy(
|
||||
policy.to_owned(),
|
||||
)));
|
||||
}
|
||||
};
|
||||
let rows = connection
|
||||
.prepare(&format!(
|
||||
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
|
||||
))
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.query_map(params![remaining], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
for (_, filename) in &rows {
|
||||
if let Some(filename) = filename {
|
||||
cleanup.push(filename.clone());
|
||||
|
|
@ -304,7 +312,7 @@ impl DiskcacheSqliteStore {
|
|||
for (rowid, _) in rows {
|
||||
connection
|
||||
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
Ok(cleanup)
|
||||
}
|
||||
|
|
@ -312,15 +320,15 @@ impl DiskcacheSqliteStore {
|
|||
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
|
||||
let page_count: i64 = connection
|
||||
.query_row("PRAGMA page_count", [], |row| row.get(0))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
let page_size: i64 = connection
|
||||
.query_row("PRAGMA page_size", [], |row| row.get(0))
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
let size: i64 = connection
|
||||
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
|
||||
row.get(0)
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
Ok(page_count.saturating_mul(page_size).saturating_add(size))
|
||||
}
|
||||
}
|
||||
|
|
@ -331,13 +339,13 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
}
|
||||
|
||||
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let connection = self.connection.lock()?;
|
||||
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
|
||||
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
|
||||
let row = connection
|
||||
.query_row(select, params![key, now], row_from_query)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
if !self.statistics && !has_get_update(&self.eviction_policy) {
|
||||
return row
|
||||
.map(|row| fetch_row(&self.directory, row))
|
||||
|
|
@ -348,7 +356,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
let row = connection
|
||||
.query_row(select, params![key, now], row_from_query)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
let Some(row) = row else {
|
||||
if self.statistics {
|
||||
connection
|
||||
|
|
@ -356,7 +364,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
|
||||
[],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
return Ok(None);
|
||||
};
|
||||
|
|
@ -369,14 +377,14 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
|
||||
[],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
} else if !hit && self.statistics {
|
||||
connection
|
||||
.execute(
|
||||
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
|
||||
[],
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
if has_get_update(&self.eviction_policy) && hit {
|
||||
let update = match self.eviction_policy.as_str() {
|
||||
|
|
@ -384,16 +392,20 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
"least-frequently-used" => {
|
||||
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
|
||||
}
|
||||
_ => return Err(Error::Unavailable),
|
||||
policy => {
|
||||
return Err(Error::unavailable(SettingFault::EvictionPolicy(
|
||||
policy.to_owned(),
|
||||
)));
|
||||
}
|
||||
};
|
||||
if self.eviction_policy == "least-recently-used" {
|
||||
connection
|
||||
.execute(update, params![now, rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
} else {
|
||||
connection
|
||||
.execute(update, params![rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
}
|
||||
}
|
||||
value
|
||||
|
|
@ -409,7 +421,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
) -> Result<(), Error> {
|
||||
let columns = store_value(&self.directory, self.min_file_size, value)?;
|
||||
let new_filename = columns.filename.clone();
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let connection = self.connection.lock()?;
|
||||
let result = transactional(&connection, |connection| {
|
||||
self.set_locked(connection, key, columns, expire_time, now)
|
||||
});
|
||||
|
|
@ -428,7 +440,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
}
|
||||
|
||||
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let connection = self.connection.lock()?;
|
||||
let selected = transactional(&connection, |connection| {
|
||||
let row = connection
|
||||
.query_row(
|
||||
|
|
@ -439,13 +451,13 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
row_from_query,
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
let Some(row) = row else {
|
||||
return Ok(None);
|
||||
};
|
||||
connection
|
||||
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
Ok(Some(row))
|
||||
})?;
|
||||
let Some(row) = selected else {
|
||||
|
|
@ -460,7 +472,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
}
|
||||
|
||||
fn clear(&self) -> Result<(), Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let connection = self.connection.lock()?;
|
||||
let mut last_rowid = 0_i64;
|
||||
loop {
|
||||
let batch = transactional(&connection, |connection| {
|
||||
|
|
@ -469,13 +481,13 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
"SELECT rowid, filename FROM Cache
|
||||
WHERE rowid > ? ORDER BY rowid LIMIT 100",
|
||||
)
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.query_map(params![last_rowid], |row| {
|
||||
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
|
||||
})
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
if rows.is_empty() {
|
||||
return Ok(rows);
|
||||
}
|
||||
|
|
@ -486,7 +498,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
.join(",");
|
||||
connection
|
||||
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
Ok(rows)
|
||||
})?;
|
||||
if batch.is_empty() {
|
||||
|
|
@ -509,7 +521,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
now: f64,
|
||||
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
|
||||
) -> Result<(), Error> {
|
||||
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
let connection = self.connection.lock()?;
|
||||
let mut created_filename = None;
|
||||
let result = transactional(&connection, |connection| {
|
||||
let current = connection
|
||||
|
|
@ -521,7 +533,7 @@ impl DiskStore for DiskcacheSqliteStore {
|
|||
row_from_query,
|
||||
)
|
||||
.optional()
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.map(|row| fetch_row(&self.directory, row))
|
||||
.transpose()?
|
||||
.flatten();
|
||||
|
|
@ -582,21 +594,35 @@ fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Erro
|
|||
};
|
||||
statement
|
||||
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map_err(Error::unavailable)?
|
||||
.collect::<Result<HashMap<_, _>, _>>()
|
||||
.map_err(|_| Error::Unavailable)
|
||||
.map_err(Error::unavailable)
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum SettingFault {
|
||||
#[error("unknown eviction policy {0:?}")]
|
||||
EvictionPolicy(String),
|
||||
#[error("setting {0:?} is not an SQLite pragma")]
|
||||
NotPragma(String),
|
||||
#[error("setting {0:?} has an unsupported value type")]
|
||||
PragmaValue(String),
|
||||
}
|
||||
|
||||
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
|
||||
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
|
||||
let pragma = key
|
||||
.strip_prefix("sqlite_")
|
||||
.ok_or_else(|| Error::unavailable(SettingFault::NotPragma(key.to_owned())))?;
|
||||
match value {
|
||||
Value::Integer(value) => connection
|
||||
.pragma_update(None, pragma, value)
|
||||
.map_err(|_| Error::Unavailable),
|
||||
.map_err(Error::unavailable),
|
||||
Value::Text(value) => connection
|
||||
.pragma_update(None, pragma, value)
|
||||
.map_err(|_| Error::Unavailable),
|
||||
_ => Err(Error::Unavailable),
|
||||
.map_err(Error::unavailable),
|
||||
_ => Err(Error::unavailable(SettingFault::PragmaValue(
|
||||
key.to_owned(),
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -634,7 +660,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
|
|||
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
|
||||
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
|
||||
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
|
||||
Value::Null => Err(Error::InvalidEntry),
|
||||
Value::Null => Err(Error::InvalidEntry(None)),
|
||||
},
|
||||
MODE_BINARY | MODE_PICKLE => {
|
||||
let bytes = match row.value {
|
||||
|
|
@ -645,7 +671,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
|
|||
};
|
||||
value
|
||||
}
|
||||
_ => return Err(Error::InvalidEntry),
|
||||
_ => return Err(Error::InvalidEntry(None)),
|
||||
};
|
||||
Ok(Some(if row.mode == MODE_BINARY {
|
||||
StoredValue::Bytes(bytes)
|
||||
|
|
@ -663,24 +689,24 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
|
|||
}
|
||||
Value::Blob(value) => value,
|
||||
Value::Text(value) => value.into_bytes(),
|
||||
_ => return Err(Error::InvalidEntry),
|
||||
_ => return Err(Error::InvalidEntry(None)),
|
||||
};
|
||||
Ok(Some(StoredValue::Text(
|
||||
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
|
||||
String::from_utf8(bytes).map_err(Error::invalid_entry)?,
|
||||
)))
|
||||
}
|
||||
_ => Err(Error::InvalidEntry),
|
||||
_ => Err(Error::InvalidEntry(None)),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
|
||||
let Some(filename) = filename else {
|
||||
return Err(Error::InvalidEntry);
|
||||
return Err(Error::InvalidEntry(None));
|
||||
};
|
||||
match fs::read(directory.join(filename)) {
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
|
||||
Err(_) => Err(Error::Unavailable),
|
||||
Err(error) => Err(Error::unavailable(error)),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -712,7 +738,7 @@ fn store_value(
|
|||
let bytes = value.into_bytes();
|
||||
let filename = write_file(directory, &bytes)?;
|
||||
Ok(StoredColumns {
|
||||
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
|
||||
size: i64::try_from(bytes.len()).map_err(Error::unavailable)?,
|
||||
mode: MODE_TEXT,
|
||||
filename: Some(filename),
|
||||
value: None,
|
||||
|
|
@ -727,7 +753,7 @@ fn store_value(
|
|||
StoredValue::Bytes(value) => {
|
||||
let filename = write_file(directory, &value)?;
|
||||
Ok(StoredColumns {
|
||||
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
|
||||
size: i64::try_from(value.len()).map_err(Error::unavailable)?,
|
||||
mode: MODE_BINARY,
|
||||
filename: Some(filename),
|
||||
value: None,
|
||||
|
|
@ -742,7 +768,7 @@ fn store_value(
|
|||
StoredValue::Pickle(value) => {
|
||||
let filename = write_file(directory, &value)?;
|
||||
Ok(StoredColumns {
|
||||
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
|
||||
size: i64::try_from(value.len()).map_err(Error::unavailable)?,
|
||||
mode: MODE_PICKLE,
|
||||
filename: Some(filename),
|
||||
value: None,
|
||||
|
|
@ -761,14 +787,14 @@ fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
|
|||
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
|
||||
let path = directory.join(&filename);
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
|
||||
fs::create_dir_all(parent).map_err(Error::unavailable)?;
|
||||
}
|
||||
let mut file = OpenOptions::new()
|
||||
.write(true)
|
||||
.create_new(true)
|
||||
.open(path)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
file.write_all(bytes).map_err(Error::unavailable)?;
|
||||
Ok(filename)
|
||||
}
|
||||
|
||||
|
|
@ -789,12 +815,12 @@ fn transactional<T>(
|
|||
) -> Result<T, Error> {
|
||||
connection
|
||||
.execute_batch("BEGIN IMMEDIATE")
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
match operation(connection) {
|
||||
Ok(value) => {
|
||||
connection
|
||||
.execute_batch("COMMIT")
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
.map_err(Error::unavailable)?;
|
||||
Ok(value)
|
||||
}
|
||||
Err(error) => {
|
||||
|
|
|
|||
|
|
@ -84,7 +84,7 @@ fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation
|
|||
match expected {
|
||||
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
|
||||
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
|
||||
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
|
||||
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry(_)))),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue